From 46addabaa82401660a8425b595a228a18efc1dae Mon Sep 17 00:00:00 2001
From: Thomas Robitaille <thomas.robitaille@gmail.com>
Date: Wed, 22 May 2024 11:58:54 +0100
Subject: [PATCH] Fix unit conversion from flux to surface brightness

---
 jdaviz/app.py            | 52 ++++++++++++++++------------------------
 jdaviz/tests/test_app.py |  7 ++++++
 2 files changed, 28 insertions(+), 31 deletions(-)

diff --git a/jdaviz/app.py b/jdaviz/app.py
index e65de0270e..8e54eb7ea7 100644
--- a/jdaviz/app.py
+++ b/jdaviz/app.py
@@ -111,37 +111,27 @@ def to_unit(self, data, cid, values, original_units, target_units):
             else:
                 # Ensure a spectrum passed through Spectral Extraction plugin
                 if '_pixel_scale_factor' in spec.meta:
-                    # if spectrum data collection item is in Surface Brightness units
-                    if u.sr in spec.unit.bases:
-                        # Data item in data collection does not update from conversion/translation.
-                        # App wide orginal data units are used for conversion, orginal_units and
-                        # target_units dicate the conversion to take place.
-                        if (u.sr in u.Unit(original_units).bases) and \
-                           (u.sr not in u.Unit(target_units).bases):
-                            # Surface Brightness -> Flux
-                            eqv = [(u.MJy / u.sr,
-                                    u.MJy,
-                                    lambda x: (x * spec.meta['_pixel_scale_factor']),
-                                    lambda x: x)]
-                        else:
-                            # Flux -> Surface Brightness
-                            eqv = u.spectral_density(spec.spectral_axis)
-
-                    # if spectrum data collection item is in Flux units
-                    elif u.sr not in spec.unit.bases:
-                        # Data item in data collection does not update from conversion/translation.
-                        # App wide orginal data units are used for conversion, orginal_units and
-                        # target_units dicate the conversion to take place.
-                        if (u.sr not in u.Unit(original_units).bases) and \
-                           (u.sr in u.Unit(target_units).bases):
-                            # Flux -> Surface Brightness
-                            eqv = [(u.MJy,
-                                    u.MJy / u.sr,
-                                    lambda x: (x / spec.meta['_pixel_scale_factor']),
-                                    lambda x: x)]
-                        else:
-                            # Surface Brightness -> Flux
-                            eqv = u.spectral_density(spec.spectral_axis)
+
+                    # Data item in data collection does not update from conversion/translation.
+                    # App wide original data units are used for conversion, original and
+                    # target_units dictate the conversion to take place.
+
+                    if (u.sr in u.Unit(original_units).bases) and \
+                        (u.sr not in u.Unit(target_units).bases):
+                        # Surface Brightness -> Flux
+                        eqv = [(u.MJy / u.sr,
+                                u.MJy,
+                                lambda x: (x * spec.meta['_pixel_scale_factor']),
+                                lambda x: x)]
+                    elif (u.sr not in u.Unit(original_units).bases) and \
+                        (u.sr in u.Unit(target_units).bases):
+                        # Flux -> Surface Brightness
+                        eqv = [(u.MJy,
+                                u.MJy / u.sr,
+                                lambda x: (x / spec.meta['_pixel_scale_factor']),
+                                lambda x: x)]
+                    else:
+                        eqv = u.spectral_density(spec.spectral_axis)
 
                 elif len(values) == 2:
                     # Need this for setting the y-limits
diff --git a/jdaviz/tests/test_app.py b/jdaviz/tests/test_app.py
index 9e923ab301..90677a48ca 100644
--- a/jdaviz/tests/test_app.py
+++ b/jdaviz/tests/test_app.py
@@ -225,6 +225,9 @@ def test_to_unit(cubeviz_helper):
     cid = cubeviz_helper.app.data_collection[0].data.find_component_id('flux')
     data = cubeviz_helper.app.data_collection[-1].data
     values = 1
+
+    # Surface brightness to flux
+
     original_units = u.MJy / u.sr
     target_units = u.MJy
 
@@ -232,5 +235,9 @@ def test_to_unit(cubeviz_helper):
 
     assert np.allclose(value, 4.7945742429049767e-11)
 
+    # Flux to surface brightness
+
     original_units = u.MJy
     target_units = u.MJy / u.sr
+
+    value = uc.to_unit(cubeviz_helper, data, cid, values, original_units, target_units)