update mm plugin to use correct display flux units for moment 0

CHANGES.rst

@@ -71,6 +71,8 @@ Bug Fixes
 Cubeviz
 ^^^^^^^
 
+- Moment map plugin now reflects selected flux / surface brightness unit for moment zero. [#2877]
+
 Imviz
 ^^^^^
 

jdaviz/configs/cubeviz/plugins/moment_maps/moment_maps.py

@@ -19,6 +19,7 @@
                                         SpectralContinuumMixin,
                                         skip_if_no_updates_since_last_active,
                                         with_spinner)
+from jdaviz.core.validunits import check_if_unit_is_per_solid_angle
 from jdaviz.core.user_api import PluginUserApi
 
 __all__ = ['MomentMap']
@@ -27,7 +28,7 @@
 spaxel = u.def_unit('spaxel', 1 * u.Unit(""))
 u.add_enabled_units([spaxel])
 
-moment_unit_options = {0: ["Flux"],
+moment_unit_options = {0: ["Surface Brightness", "Flux"],
                        1: ["Velocity", "Spectral Unit"],
                        2: ["Velocity", "Velocity^N"]}
 
@@ -101,6 +102,10 @@ def __init__(self, *args, **kwargs):
                                                filters=['not_child_layer',
                                                         'layer_in_spectrum_viewer'])
 
+        # when plugin is initialized, there won't be a dataset selected, so
+        # call the output unit 'Flux' for now (rather than surface brightness).
+        # the appropriate label will be chosen once a dataset is selected, and/or
+        # unit conversion is done
         self.output_unit = SelectPluginComponent(self,
                                                  items='output_unit_items',
                                                  selected='output_unit_selected',
@@ -159,14 +164,18 @@ def _set_default_results_label(self, event={}):
 
     @observe("dataset_selected", "n_moment")
     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
         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")
 
-        unit_dict = {"Flux": "",
+        # either 'Flux' or 'Surface Brightness'
+        orig_flux_or_sb = self.output_unit_items[0]['label']
+
+        unit_dict = {orig_flux_or_sb: "",
                      "Spectral Unit": "",
                      "Velocity": "km/s",
                      "Velocity^N": f"km{self.n_moment}/s{self.n_moment}"}
@@ -184,11 +193,36 @@ def _set_data_units(self, event={}):
             self.dataset_spectral_unit = sunit
             unit_dict["Spectral Unit"] = sunit
 
-            unit_dict["Flux"] = data.get_component('flux').units
+            # get flux/SB units
+            if self.spectrum_viewer and hasattr(self.spectrum_viewer.state, 'y_display_unit'):
+                if self.spectrum_viewer.state.y_display_unit is not None:
+                    unit_dict[orig_flux_or_sb] = self.spectrum_viewer.state.y_display_unit
+                else:
+                    # spectrum_viewer.state will only have x/y_display_unit if unit conversion has
+                    # been done if not, get default flux units which should be the units displayed
+                    unit_dict[orig_flux_or_sb] = data.get_component('flux').units
+            else:
+                # spectrum_viewer.state will only have x/y_display_unit if unit conversion has
+                # been done if not, get default flux units which should be the units displayed
+                unit_dict[orig_flux_or_sb] = data.get_component('flux').units
+
+            # figure out if label should say 'Flux' or 'Surface Brightness'
+            sb_or_flux_label = "Flux"
+            is_unit_solid_angle = check_if_unit_is_per_solid_angle(unit_dict[orig_flux_or_sb])
+            if is_unit_solid_angle is True:
+                sb_or_flux_label = "Surface Brightness"
 
         # Update units in selection item dictionary
         for item in self.output_unit_items:
             item["unit_str"] = unit_dict[item["label"]]
+            if item["label"] in ["Flux", "Surface Brightness"]:
+                # change unit label to reflect if unit is flux or SB
+                item["label"] = sb_or_flux_label
+
+        if self.dataset_selected != "":
+            # output_unit_selected might not match (potentially) changed flux/SB label
+            if self.output_unit_selected in ['Flux', 'Surface Brightness']:
+                self.output_unit_selected = sb_or_flux_label
 
         # Filter what we want based on n_moment
         if self.n_moment == 0:
@@ -233,6 +267,7 @@ def calculate_moment(self, add_data=True):
         """
 
         # Check to make sure API use hasn't put us into an invalid state.
+
         try:
             n_moment = int(self.n_moment)
             if n_moment < 0:
@@ -300,6 +335,23 @@ def calculate_moment(self, add_data=True):
         # If n>1 and velocity is desired, need to take nth root of result
         if n_moment > 0 and self.output_unit_selected.lower() == "velocity":
             self.moment = np.power(self.moment, 1/self.n_moment)
+
+        # convert units for moment 0, which is the only currently supported
+        # moment for using converted units.
+        if n_moment == 0:
+            # get flux/SB units
+            if self.spectrum_viewer and hasattr(self.spectrum_viewer.state, 'y_display_unit'):
+                if self.spectrum_viewer.state.y_display_unit is not None:
+                    flux_sb_unit = self.spectrum_viewer.state.y_display_unit
+                else:
+                    flux_sb_unit = data.get_component('flux').units
+            else:
+                flux_sb_unit = data.get_component('flux').units
+
+            # convert unit string to u.Unit so moment map data can be converted
+            flux_or_sb_display_unit = u.Unit(flux_sb_unit)
+            self.moment = self.moment.to(flux_or_sb_display_unit)
+
         # Reattach the WCS so we can load the result
         self.moment = CCDData(self.moment, wcs=data_wcs)
 

jdaviz/configs/cubeviz/plugins/moment_maps/tests/test_moment_maps.py

@@ -7,6 +7,7 @@
 from astropy.nddata import CCDData
 from astropy.wcs import WCS
 from astroquery.mast import Observations
+import astropy.units as u
 import numpy as np
 from numpy.testing import assert_allclose
 from glue.core.roi import XRangeROI
@@ -277,3 +278,46 @@ def test_momentmap_nirspec_prism(cubeviz_helper, tmp_path):
     sky_cube = cubeviz_helper.app.data_collection[0].meta["_orig_spec"].wcs.celestial.pixel_to_world(30, 50)  # noqa: E501
     assert_allclose((sky_moment.ra.deg, sky_moment.dec.deg),
                     (sky_cube.ra.deg, sky_cube.dec.deg))
+
+
+def test_correct_output_flux_or_sb_units(cubeviz_helper, spectrum1d_cube_custom_fluxunit):
+
+    # test that the output unit labels in the moment map plugin reflect any
+    # changes made in the unit conversion plugin.
+    # NOTE: for now this is limited to moment 0, test should be expanded to
+    # test higher-order moments once implemented.
+
+    sb_cube = spectrum1d_cube_custom_fluxunit(fluxunit=u.MJy / u.sr)
+
+    # load surface brigtness cube
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", message="No observer defined on WCS.*")
+        cubeviz_helper.load_data(sb_cube, data_label='test')
+
+    mm = cubeviz_helper.plugins['Moment Maps']._obj
+    mm.open_in_tray()  # plugin has to be open for unit change to take hold
+
+    # check that label is initialized with 'Surface Brightness' since the cube
+    # loaded is in MJy / sr. for the 0th moment, the only item will be the 0th
+    # 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'
+
+    # check that calculated moment has the correct units
+    mm.calculate_moment()
+    assert mm.moment.unit == 'MJy / sr'
+
+    # now change surface brightness units in the unit conversion plugin
+    uc = cubeviz_helper.plugins["Unit Conversion"]
+    uc.open_in_tray()  # plugin has to be open for unit change to take hold
+    uc.flux_or_sb_unit = 'Jy / sr'
+
+    # 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'
+
+    # and that calculated moment has the correct units
+    mm.calculate_moment()
+    assert mm.moment.unit == 'Jy / sr'

jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py

@@ -129,6 +129,7 @@ def _on_glue_y_display_unit_changed(self, y_unit):
             y_unit = _valid_glue_display_unit(y_unit, self.spectrum_viewer, 'y')
             y_u = u.Unit(y_unit)
             choices = create_flux_equivalencies_list(y_u, x_u)
+
             # ensure that original entry is in the list of choices
             if not np.any([y_u == u.Unit(choice) for choice in choices]):
                 choices = [y_unit] + choices
@@ -160,8 +161,11 @@ def _on_spectral_unit_changed(self, *args):
 
     @observe('flux_unit_selected')
     def _on_flux_unit_changed(self, *args):
+
         yunit = _valid_glue_display_unit(self.flux_or_sb_unit.selected, self.spectrum_viewer, 'y')
+
         if self.spectrum_viewer.state.y_display_unit != yunit:
+
             self.spectrum_viewer.state.y_display_unit = yunit
             self.hub.broadcast(GlobalDisplayUnitChanged('flux',
                                                         self.flux_or_sb_unit.selected,

jdaviz/core/tests/test_validunits.py

@@ -0,0 +1,16 @@
+import astropy.units as u
+import pytest
+
+from jdaviz.core.validunits import check_if_unit_is_per_solid_angle
+
+
+@pytest.mark.parametrize("unit, is_solid_angle", [
+    ('erg / sr', True), ('erg / (deg * deg)', True), ('erg / (deg * arcsec)', True),
+    ('erg / (deg**2)', True), ('erg deg**-2', True), ('erg sr^-1', True),
+    ('erg / degree', False), ('erg sr^-2', False)])
+def test_check_if_unit_is_per_solid_angle(unit, is_solid_angle):
+    # test function that tests if a unit string or u.Unit is per solid angle
+    assert check_if_unit_is_per_solid_angle(unit) == is_solid_angle
+
+    # turn string into u.Unit, and make sure it also passes
+    assert check_if_unit_is_per_solid_angle(u.Unit(unit)) == is_solid_angle

jdaviz/core/validunits.py

@@ -1,7 +1,8 @@
 from astropy import units as u
+import itertools
 
 __all__ = ['units_to_strings', 'create_spectral_equivalencies_list',
-           'create_flux_equivalencies_list']
+           'create_flux_equivalencies_list', 'check_if_unit_is_per_solid_angle']
 
 
 def units_to_strings(unit_list):
@@ -98,3 +99,52 @@ def create_flux_equivalencies_list(flux_unit, spectral_axis_unit):
 
     # Concatenate both lists with the local units coming first.
     return sorted(units_to_strings(local_units)) + flux_unit_equivalencies_titles
+
+
+def check_if_unit_is_per_solid_angle(unit):
+    """
+    Check if a given u.Unit or unit string (that can be converted to
+    a u.Unit object) represents some unit per solid angle.
+
+    Parameters
+    ----------
+    unit : str or u.Unit
+        Unit object or string representation of unit.
+
+    Examples
+    --------
+    >>> check_if_unit_is_per_solid_angle('erg / (s cm^2 sr)')
+    True
+    >>> check_if_unit_is_per_solid_angle('erg / s cm^2')
+    False
+    >>> check_if_unit_is_per_solid_angle('Jy * sr^-1')
+    True
+
+    """
+
+    # first, convert string to u.Unit obj.
+    # this will take care of some formatting consistency like
+    # turning something like Jy / (degree*degree) to Jy / deg**2
+    # and erg sr^1 to erg / sr
+    if isinstance(unit, u.core.Unit) or isinstance(unit, u.core.CompositeUnit):
+        unit_str = unit.to_string()
+    elif isinstance(unit, str):
+        unit = u.Unit(unit)
+        unit_str = unit.to_string()
+    else:
+        raise ValueError('Unit must be u.Unit, or string that can be converted into a u.Unit')
+
+    if '/' in unit_str:
+        # might be comprised of several units in denom.
+        denom = unit_str.split('/')[-1].split()
+
+        # find all combos of one or two units, to catch cases where there are two different
+        # units of angle in the denom that might comprise a solid angle when multiplied.
+        for i in [combo for length in (1, 2) for combo in itertools.combinations(denom, length)]:
+            # turn tuple of 1 or 2 units into a string, and turn that into a u.Unit to check type
+            new_unit_str = ' '.join(i).translate(str.maketrans('', '', '()'))
+            new_unit = u.Unit(new_unit_str)
+            if new_unit.physical_type == 'solid angle':
+                return True
+
+    return False