From 99f67edcc63c40b0083f57c3a6345e157a1a5d27 Mon Sep 17 00:00:00 2001
From: Thomas Robitaille <thomas.robitaille@gmail.com>
Date: Wed, 22 May 2024 12:40:40 +0100
Subject: [PATCH] Fix case in to_unit where limits are converted from Fnu to
 Flam

---
 jdaviz/app.py            | 23 +++++++++++++----------
 jdaviz/tests/test_app.py | 29 +++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 10 deletions(-)

diff --git a/jdaviz/app.py b/jdaviz/app.py
index 8e54eb7ea7..4376fbf055 100644
--- a/jdaviz/app.py
+++ b/jdaviz/app.py
@@ -109,6 +109,15 @@ def to_unit(self, data, cid, values, original_units, target_units):
             except RuntimeError:
                 eqv = []
             else:
+
+                # If there are only two values, this is likely the limits being converted, so then
+                # in case we need to use the spectral density equivalency, we need to provide only
+                # to spectral axis values. If there is only one value
+                if not np.isscalar(values) and len(values) == 2:
+                    spectral_values = spec.spectral_axis[0]
+                else:
+                    spectral_values = spec.spectral_axis
+
                 # Ensure a spectrum passed through Spectral Extraction plugin
                 if '_pixel_scale_factor' in spec.meta:
 
@@ -117,29 +126,23 @@ def to_unit(self, data, cid, values, original_units, target_units):
                     # 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):
+                           (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):
+                            (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
-                    spec_limits = [spec.spectral_axis[0].value, spec.spectral_axis[-1].value]
-                    eqv = u.spectral_density(spec_limits * spec.spectral_axis.unit)
-
+                        eqv = u.spectral_density(spectral_values)
                 else:
-                    eqv = u.spectral_density(spec.spectral_axis)
+                    eqv = u.spectral_density(spectral_values)
 
         else:  # spectral axis
             eqv = u.spectral() + u.pixel_scale(1*u.pix)
diff --git a/jdaviz/tests/test_app.py b/jdaviz/tests/test_app.py
index 91566e6e70..33d6ad4c0f 100644
--- a/jdaviz/tests/test_app.py
+++ b/jdaviz/tests/test_app.py
@@ -243,3 +243,32 @@ def test_to_unit(cubeviz_helper):
     value = uc.to_unit(cubeviz_helper, data, cid, value, original_units, target_units)
 
     assert np.allclose(value, 1)
+
+    # Change from Fnu to Flam (with values shape matching spectral axis)
+
+    values = np.ones(3001)
+    original_units = u.MJy
+    target_units = u.erg / u.cm**2 / u.s / u.AA
+
+    new_values = uc.to_unit(cubeviz_helper, data, cid, values, original_units, target_units)
+
+    assert np.allclose(new_values,
+                       (values * original_units)
+                       .to_value(target_units,
+                                 equivalencies=u.spectral_density(cube.spectral_axis)))
+
+    # Change from Fnu to Flam (with a shape (2,) array of values indicating we
+    # are probably converting the limits)
+
+    values = [1, 2]
+    original_units = u.MJy
+    target_units = u.erg / u.cm**2 / u.s / u.AA
+
+    new_values = uc.to_unit(cubeviz_helper, data, cid, values, original_units, target_units)
+
+    # In this case we do a regular spectral density conversion, but using the
+    # first value in the spectral axis for the equivalency
+    assert np.allclose(new_values,
+                       ([1, 2] * original_units)
+                       .to_value(target_units,
+                                 equivalencies=u.spectral_density(cube.spectral_axis[0])))