From 8cc6fa5f905b7a3753c65270844426aa70014194 Mon Sep 17 00:00:00 2001
From: Clare Shanahan <cshanahan@stsci.edu>
Date: Wed, 2 Oct 2024 15:03:11 -0400
Subject: [PATCH] add spectral unit to moment map 0 display unit, simplify
 logic

Co-authored by: bmorris3 <bmmorris@stsci.edu>
---
 .../plugins/moment_maps/moment_maps.py        | 45 ++++++++++---------
 .../plugins/moment_maps/moment_maps.vue       |  2 +-
 .../moment_maps/tests/test_moment_maps.py     | 14 ++++--
 jdaviz/core/template_mixin.py                 |  6 +--
 4 files changed, 37 insertions(+), 30 deletions(-)

diff --git a/jdaviz/configs/cubeviz/plugins/moment_maps/moment_maps.py b/jdaviz/configs/cubeviz/plugins/moment_maps/moment_maps.py
index 7748a3511d..6c740f7121 100644
--- a/jdaviz/configs/cubeviz/plugins/moment_maps/moment_maps.py
+++ b/jdaviz/configs/cubeviz/plugins/moment_maps/moment_maps.py
@@ -170,7 +170,7 @@ def _set_data_units(self, event={}):
 
         if isinstance(self.n_moment, str) or self.n_moment < 0:
             return
-        unit_options_index = 2 if self.n_moment > 2 else self.n_moment
+        unit_options_index = min(self.n_moment, 2)
         if self.output_unit_selected not in moment_unit_options[unit_options_index]:
             self.output_unit_selected = moment_unit_options[unit_options_index][0]
         self.send_state("output_unit_selected")
@@ -192,7 +192,7 @@ def _set_data_units(self, event={}):
                 sunit = ""
             self.dataset_spectral_unit = sunit
             unit_dict["Spectral Unit"] = sunit
-            unit_dict["Surface Brightness"] = self.app._get_display_unit('sb')
+            unit_dict["Surface Brightness"] = str(self.moment_zero_unit)
 
         # Update units in selection item dictionary
         for item in self.output_unit_items:
@@ -268,14 +268,12 @@ def calculate_moment(self, add_data=True):
 
         # slice out desired region
         # TODO: should we add a warning for a composite spectral subset?
-        spec_min, spec_max = self.spectral_subset.selected_min_max(cube)
-        display_spectral_axis_unit = self.app._get_display_unit(self.slice_display_unit_name)
+        spec_reg = self.spectral_subset.selected_obj
 
-        # Convert units of min and max if necessary
-        if display_spectral_axis_unit and display_spectral_axis_unit != str(spec_min.unit):
-            spec_min = spec_min.to(display_spectral_axis_unit, equivalencies=u.spectral())
-            spec_max = spec_max.to(display_spectral_axis_unit, equivalencies=u.spectral())
-        slab = manipulation.spectral_slab(cube, spec_min, spec_max)
+        if spec_reg is None:
+            slab = cube
+        else:
+            slab = manipulation.extract_region(cube, spec_reg)
 
         # Calculate the moment and convert to CCDData to add to the viewers
         # Need transpose to align JWST mirror shape: This is because specutils
@@ -313,19 +311,7 @@ def calculate_moment(self, add_data=True):
         # convert units for moment 0, which is the only currently supported
         # moment for using converted units.
         if n_moment == 0:
-            # get display units for moment 0 based on unit conversion plugin selection
-            # the 0th item of this dictionary is the surface brightness unit, kept
-            # up to date with choices from the UC plugin
-            moment_0_display_unit = self.output_unit_items[0]['unit_str']
-
-            # convert unit string to Unit so moment map data can be converted
-            # `display_unit` is the data unit (surface brighntess) translated
-            # to the choice of selected flux and angle unit from the UC plugin.
-            display_unit = u.Unit(moment_0_display_unit)
-
-            moment_new_unit = display_unit * self.app._get_display_unit('spectral')  # noqa: E501
-
-            self.moment = self.moment.to(moment_new_unit)
+            self.moment = self.moment.to(self.moment_zero_unit)
 
         # Reattach the WCS so we can load the result
         self.moment = CCDData(self.moment, wcs=data_wcs)
@@ -344,6 +330,21 @@ def calculate_moment(self, add_data=True):
 
         return self.moment
 
+    @property
+    def moment_zero_unit(self):
+        if self.spectrum_viewer.state.x_display_unit is not None:
+            return (
+                u.Unit(self.app._get_display_unit('sb')) *
+                self.spectrum_viewer.state.x_display_unit
+            )
+        return ''
+
+    @property
+    def spectral_unit_selected(self):
+        # NOTE: should we add this logic to app._get_display_unit?
+        unit_conv = self.app.get_tray_item_from_name('g-unit-conversion')
+        return unit_conv.spectral_unit.selected
+
     def vue_calculate_moment(self, *args):
         self.calculate_moment(add_data=True)
 
diff --git a/jdaviz/configs/cubeviz/plugins/moment_maps/moment_maps.vue b/jdaviz/configs/cubeviz/plugins/moment_maps/moment_maps.vue
index d5f4bfef63..5cf4b7ea7b 100644
--- a/jdaviz/configs/cubeviz/plugins/moment_maps/moment_maps.vue
+++ b/jdaviz/configs/cubeviz/plugins/moment_maps/moment_maps.vue
@@ -127,7 +127,7 @@
           ></v-radio>
         </v-radio-group>
       </v-row>
-      <v-row v-if="output_unit_selected !== 'Spectral Unit' && output_unit_selected !== 'Flux'">
+      <v-row v-if="output_unit_selected !== 'Spectral Unit' && output_unit_selected !== 'Surface Brightness'">
         <v-text-field
         ref="reference_wavelength"
         type="number"
diff --git a/jdaviz/configs/cubeviz/plugins/moment_maps/tests/test_moment_maps.py b/jdaviz/configs/cubeviz/plugins/moment_maps/tests/test_moment_maps.py
index c587ed53d6..06bd361c20 100644
--- a/jdaviz/configs/cubeviz/plugins/moment_maps/tests/test_moment_maps.py
+++ b/jdaviz/configs/cubeviz/plugins/moment_maps/tests/test_moment_maps.py
@@ -347,9 +347,9 @@ def test_correct_output_spectral_y_units(cubeviz_helper, spectrum1d_cube_custom_
     # in this list. also check that the initial unit is correct.
     output_unit_moment_0 = mm.output_unit_items[0]
     assert output_unit_moment_0['label'] == 'Surface Brightness'
-    assert output_unit_moment_0['unit_str'] == 'MJy / sr'
+    assert output_unit_moment_0['unit_str'] == 'MJy m / sr'
 
-    # check that calculated moment has the correct units
+    # check that calculated moment has the correct units of MJy m / sr
     mm.calculate_moment()
     assert mm.moment.unit == f'M{moment_unit}'
 
@@ -360,8 +360,16 @@ def test_correct_output_spectral_y_units(cubeviz_helper, spectrum1d_cube_custom_
     # and make sure this change is propogated
     output_unit_moment_0 = mm.output_unit_items[0]
     assert output_unit_moment_0['label'] == 'Surface Brightness'
-    assert output_unit_moment_0['unit_str'] == 'Jy / sr'
+    assert output_unit_moment_0['unit_str'] == 'Jy m / sr'
 
     # and that calculated moment has the correct units
     mm.calculate_moment()
     assert mm.moment.unit == moment_unit
+
+    # now change the spectral unit and make sure that change is
+    # reflected in MM plugin
+    uc.spectral_unit = 'um'
+    assert output_unit_moment_0['unit_str'] == 'Jy um / sr'
+
+    mm.calculate_moment()
+    assert mm.moment.unit == moment_unit.replace('m', 'um')
diff --git a/jdaviz/core/template_mixin.py b/jdaviz/core/template_mixin.py
index 0eb455691c..f4fed1a278 100644
--- a/jdaviz/core/template_mixin.py
+++ b/jdaviz/core/template_mixin.py
@@ -2953,8 +2953,7 @@ def _get_continuum(self, dataset, spectral_subset, update_marks=False, per_pixel
         if per_pixel:
             if self.app.config != 'cubeviz':
                 raise ValueError("per-pixel only supported for cubeviz")
-            full_spectrum = self.app._jdaviz_helper.get_data(self.dataset.selected,
-                                                             use_display_units=True)
+            full_spectrum = self.app._jdaviz_helper.get_data(self.dataset.selected)
         else:
             full_spectrum = dataset.get_selected_spectrum(use_display_units=True)
 
@@ -2978,8 +2977,7 @@ def _get_continuum(self, dataset, spectral_subset, update_marks=False, per_pixel
             spectrum = full_spectrum
         else:
             sr = self.app.get_subsets(spectral_subset.selected,
-                                      simplify_spectral=True,
-                                      use_display_units=True)
+                                      simplify_spectral=True)
             spectrum = extract_region(full_spectrum, sr, return_single_spectrum=True)
             sr_lower = np.nanmin(spectrum.spectral_axis[spectrum.spectral_axis >= sr.lower])  # noqa
             sr_upper = np.nanmax(spectrum.spectral_axis[spectrum.spectral_axis <= sr.upper])  # noqa