Fix units JSON units error in notebook (#3206)

* fix error loading various units (counts, nJy, etc) in notebook

CHANGES.rst

@@ -6,7 +6,7 @@ New Features
 
 - Added flux/surface brightness translation and surface brightness
   unit conversion in Cubeviz and Specviz. [#2781, #2940, #3088, #3111, #3113, #3129,
-  #3139, #3149, #3155, #3178, #3185, #3187, #3190, #3156, #3200, #3192]
+  #3139, #3149, #3155, #3178, #3185, #3187, #3190, #3156, #3200, #3192, #3206]
 
 - Plugin tray is now open by default. [#2892]
 

jdaviz/app.py

@@ -53,9 +53,9 @@
 from jdaviz.utils import (SnackbarQueue, alpha_index, data_has_valid_wcs, layer_is_table_data,
                           MultiMaskSubsetState, _wcs_only_label, flux_conversion,
                           spectral_axis_conversion)
+from jdaviz.core.custom_units import SPEC_PHOTON_FLUX_DENSITY_UNITS
 from jdaviz.core.validunits import (check_if_unit_is_per_solid_angle,
                                     combine_flux_and_angle_units,
-                                    locally_defined_flux_units,
                                     supported_sq_angle_units)
 
 __all__ = ['Application', 'ALL_JDAVIZ_CONFIGS', 'UnitConverterWithSpectral']
@@ -75,7 +75,7 @@ class UnitConverterWithSpectral:
     def equivalent_units(self, data, cid, units):
         if cid.label == "flux":
             eqv = u.spectral_density(1 * u.m)  # Value does not matter here.
-            all_flux_units = locally_defined_flux_units() + ['ct']
+            all_flux_units = SPEC_PHOTON_FLUX_DENSITY_UNITS + ['ct']
             angle_units = supported_sq_angle_units()
             all_sb_units = combine_flux_and_angle_units(all_flux_units, angle_units)
 

jdaviz/configs/cubeviz/plugins/tests/test_cubeviz_unit_conversion.py

@@ -5,8 +5,7 @@
 from regions import PixCoord, CirclePixelRegion
 from specutils import Spectrum1D
 
-from jdaviz.core.custom_units import PIX2
-from jdaviz.core.validunits import locally_defined_flux_units
+from jdaviz.core.custom_units import PIX2, SPEC_PHOTON_FLUX_DENSITY_UNITS
 
 
 def cubeviz_wcs_dict():
@@ -36,19 +35,42 @@ def test_basic_unit_conversions(cubeviz_helper, angle_unit):
 
     # load cube with flux units of MJy
     w, wcs_dict = cubeviz_wcs_dict()
-    flux = np.zeros((30, 20, 3001), dtype=np.float32)
+    flux = np.zeros((3, 4, 5), dtype=np.float32)
     cube = Spectrum1D(flux=flux * u.MJy / angle_unit, wcs=w, meta=wcs_dict)
     cubeviz_helper.load_data(cube, data_label="test")
 
     # get all available flux units for translation. Since cube is loaded
-    # in Jy, this will be all items in 'locally_defined_flux_units'
-
-    all_flux_units = locally_defined_flux_units()
+    # in Jy, this will be all items in 'spectral_and_photon_flux_density_units'
 
     uc_plg = cubeviz_helper.plugins['Unit Conversion']
 
-    for flux_unit in all_flux_units:
+    for flux_unit in SPEC_PHOTON_FLUX_DENSITY_UNITS:
         uc_plg.flux_unit = flux_unit
+        assert cubeviz_helper.app._get_display_unit('spectral_y') == flux_unit
+
+
+@pytest.mark.parametrize("flux_unit, expected_choices", [(u.count, ['ct']),
+                                                         (u.Jy, SPEC_PHOTON_FLUX_DENSITY_UNITS),
+                                                         (u.nJy, SPEC_PHOTON_FLUX_DENSITY_UNITS + ['nJy'])])  # noqa
+def test_flux_unit_choices(cubeviz_helper, flux_unit, expected_choices):
+    """
+    Test that cubes loaded with various flux units have the expected default
+    flux unit selection in the unit conversion plugin, and that the list of
+    convertable flux units in the dropdown is correct.
+    """
+
+    w, wcs_dict = cubeviz_wcs_dict()
+    flux = np.zeros((3, 4, 5), dtype=np.float32)
+    # load cube in flux_unit, will become cube in flux_unit / pix2
+    cube = Spectrum1D(flux=flux * flux_unit, wcs=w, meta=wcs_dict)
+    cubeviz_helper.load_data(cube)
+
+    uc_plg = cubeviz_helper.plugins['Unit Conversion']
+
+    assert uc_plg.angle_unit.selected == 'pix2'  # will always be pix2
+
+    assert uc_plg.flux_unit.selected == flux_unit.to_string()
+    assert uc_plg.flux_unit.choices == expected_choices
 
 
 @pytest.mark.parametrize("angle_unit", [u.sr, PIX2])

jdaviz/configs/default/plugins/viewers.py

@@ -736,7 +736,7 @@ def _plot_uncertainties(self):
                 self.figure.marks = list(self.figure.marks) + [error_line_mark]
 
     def set_plot_axes(self):
-        # Set y axes labels for the spectrum viewer
+        # Set x and y axes labels for the spectrum viewer
         y_display_unit = self.state.y_display_unit
         y_unit = (
             u.Unit(y_display_unit) if y_display_unit and y_display_unit != 'None'

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

@@ -4,6 +4,8 @@
 from astropy.nddata import InverseVariance
 from specutils import Spectrum1D
 
+from jdaviz.core.custom_units import SPEC_PHOTON_FLUX_DENSITY_UNITS
+
 
 # On failure, should not crash; essentially a no-op.
 @pytest.mark.parametrize(
@@ -133,3 +135,22 @@ def test_non_stddev_uncertainty(specviz_helper):
         np.abs(viewer.figure.marks[-1].y - viewer.figure.marks[-1].y.mean(0)),
         stddev
     )
+
+
+@pytest.mark.parametrize("flux_unit, expected_choices", [(u.count, ['ct']),
+                                                         (u.Jy, SPEC_PHOTON_FLUX_DENSITY_UNITS),
+                                                         (u.nJy, SPEC_PHOTON_FLUX_DENSITY_UNITS + ['nJy'])])  # noqa
+def test_flux_unit_choices(specviz_helper, flux_unit, expected_choices):
+    """
+    Test that cubes loaded with various flux units have the expected default
+    flux unit selection in the unit conversion plugin, and that the list of
+    convertable flux units in the dropdown is correct.
+    """
+
+    spec = Spectrum1D([1, 2, 3] * flux_unit, [4, 5, 6] * u.um)
+    specviz_helper.load_data(spec)
+
+    uc_plg = specviz_helper.plugins['Unit Conversion']
+
+    assert uc_plg.flux_unit.selected == flux_unit.to_string()
+    assert uc_plg.flux_unit.choices == expected_choices

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

@@ -6,7 +6,6 @@
 from traitlets import List, Unicode, observe, Bool
 
 from jdaviz.configs.default.plugins.viewers import JdavizProfileView
-from jdaviz.core.custom_units import PIX2
 from jdaviz.core.events import GlobalDisplayUnitChanged, AddDataMessage
 from jdaviz.core.registries import tray_registry
 from jdaviz.core.template_mixin import (PluginTemplateMixin, UnitSelectPluginComponent,
@@ -133,14 +132,16 @@ def __init__(self, *args, **kwargs):
                                                    items='flux_unit_items',
                                                    selected='flux_unit_selected')
         # NOTE: will switch to count only if first data loaded into viewer in in counts
-        self.flux_unit.choices = create_flux_equivalencies_list(u.Jy, u.Hz)
+        # initialize flux choices to empty list, will be populated when data is loaded
+        self.flux_unit.choices = []
 
         self.has_angle = self.config in ('cubeviz', 'specviz', 'mosviz')
         self.angle_unit = UnitSelectPluginComponent(self,
                                                     items='angle_unit_items',
                                                     selected='angle_unit_selected')
         # NOTE: will switch to pix2 only if first data loaded into viewer is in pix2 units
-        self.angle_unit.choices = create_angle_equivalencies_list(u.sr)
+        # initialize flux choices to empty list, will be populated when data is loaded
+        self.angle_unit.choices = []
 
         self.has_sb = self.has_angle or self.config in ('imviz',)
         # NOTE: sb_unit is read_only, exposed through sb_unit property
@@ -219,19 +220,15 @@ def _on_add_data_to_viewer(self, msg):
                 flux_unit = data_obj.flux.unit if angle_unit is None else data_obj.flux.unit * angle_unit  # noqa
 
                 if not self.flux_unit_selected:
-                    if flux_unit in (u.count, u.DN, u.electron / u.s):
-                        self.flux_unit.choices = [flux_unit]
-                    elif flux_unit not in self.flux_unit.choices:
-                        # ensure that the native units are in the list of choices
-                        self.flux_unit.choices += [flux_unit]
+                    self.flux_unit.choices = create_flux_equivalencies_list(flux_unit)
                     try:
                         self.flux_unit.selected = str(flux_unit)
                     except ValueError:
                         self.flux_unit.selected = ''
 
                 if not self.angle_unit_selected:
-                    if angle_unit == PIX2:
-                        self.angle_unit.choices = ['pix2']
+                    self.angle_unit.choices = create_angle_equivalencies_list(angle_unit)
+
                     try:
                         if angle_unit is None:
                             # default to sr if input spectrum is not in surface brightness units

jdaviz/core/custom_units.py

@@ -1,4 +1,24 @@
 import astropy.units as u
 
+__all__ = ["PIX2", "SPEC_PHOTON_FLUX_DENSITY_UNITS"]
+
 # define custom composite units here
 PIX2 = u.pix * u.pix
+
+
+def _spectral_and_photon_flux_density_units():
+    """
+    This function returns an alphabetically sorted list of string representations
+    of spectral and photon flux density units. This list represents flux units
+    that the unit conversion plugin supports conversion to and from if the input
+    data unit is compatible with items in the list (i.e is equivalent directly
+    or with u.spectral_density(cube_wave)).
+    """
+    flux_units = ['Jy', 'mJy', 'uJy', 'MJy', 'W / (Hz m2)', 'eV / (Hz s m2)',
+                  'erg / (Hz s cm2)', 'erg / (Angstrom s cm2)',
+                  'ph / (Angstrom s cm2)', 'ph / (Hz s cm2)']
+
+    return sorted(flux_units)
+
+
+SPEC_PHOTON_FLUX_DENSITY_UNITS = _spectral_and_photon_flux_density_units()

jdaviz/core/validunits.py

@@ -1,9 +1,9 @@
 from astropy import units as u
 import itertools
 
-from jdaviz.core.custom_units import PIX2
+from jdaviz.core.custom_units import PIX2, SPEC_PHOTON_FLUX_DENSITY_UNITS
 
-__all__ = ['supported_sq_angle_units', 'locally_defined_flux_units',
+__all__ = ['supported_sq_angle_units',
            'combine_flux_and_angle_units', 'units_to_strings',
            'create_spectral_equivalencies_list',
            'create_flux_equivalencies_list', 'check_if_unit_is_per_solid_angle']
@@ -16,20 +16,6 @@ def supported_sq_angle_units(as_strings=False):
     return units
 
 
-def locally_defined_flux_units():
-    """
-    This function returns a list of string representations of flux units. This
-    list represents flux units that the unit conversion plugin supports
-    conversion to and from if the input data unit is compatible with items in the
-    list (i.e is equivalent directly or with u.spectral_density(cube_wave)).
-
-    """
-    flux_units = ['Jy', 'mJy', 'uJy', 'MJy', 'W / (Hz m2)', 'eV / (Hz s m2)',
-                  'erg / (Hz s cm2)', 'erg / (Angstrom s cm2)',
-                  'ph / (Angstrom s cm2)', 'ph / (Hz s cm2)']
-    return flux_units
-
-
 def combine_flux_and_angle_units(flux_units, angle_units):
     """
     Combine (list of) flux_units and angle_units to create a list of string
@@ -92,36 +78,39 @@ def create_spectral_equivalencies_list(spectral_axis_unit,
     return sorted(units_to_strings(local_units)) + spectral_axis_unit_equivalencies_titles
 
 
-def create_flux_equivalencies_list(flux_unit, spectral_axis_unit):
-    """Get all possible conversions for flux from current flux units."""
-    if ((flux_unit in (u.count, u.dimensionless_unscaled))
-            or (spectral_axis_unit in (u.pix, u.dimensionless_unscaled))):
-        return []
-
-    # Get unit equivalencies. Value passed into u.spectral_density() is irrelevant.
-    try:
-        curr_flux_unit_equivalencies = flux_unit.find_equivalent_units(
-            equivalencies=u.spectral_density(1 * spectral_axis_unit),
-            include_prefix_units=False)
-    except u.core.UnitConversionError:
-        return []
+def create_flux_equivalencies_list(flux_unit):
+    """
+    Get all possible conversions for flux from flux_unit, to populate 'flux'
+    dropdown menu in the unit conversion plugin.
 
-    mag_units = ['bol', 'AB', 'ST']
-    # remove magnitude units from list
-    curr_flux_unit_equivalencies = [unit for unit in curr_flux_unit_equivalencies if not any(mag in unit.name for mag in mag_units)]  # noqa
+    If flux_unit is a spectral or photon density (i.e convertable to units in
+    SPEC_PHOTON_FLUX_DENSITY_UNITS), then the loaded unit and all of the
+    units in SPEC_PHOTON_FLUX_DENSITY_UNITS.
 
-    # Get local flux units.
-    local_units = [u.Unit(unit) for unit in locally_defined_flux_units()]
+    If the loaded flux unit is count, dimensionless_unscaled, DN, e/s, then
+    there will be no additional items available for unit conversion and the
+    only item in the dropdown will be the native unit.
+    """
 
-    # Remove overlap units.
-    curr_flux_unit_equivalencies = list(set(curr_flux_unit_equivalencies)
-                                        - set(local_units))
+    flux_unit_str = flux_unit.to_string()
 
-    # Convert equivalencies into readable versions of the units and sort them alphabetically.
-    flux_unit_equivalencies_titles = sorted(units_to_strings(curr_flux_unit_equivalencies))
+    # if flux_unit is a spectral or photon flux density unit, then the flux unit
+    # dropdown options should be the loaded unit (which may have a different
+    # prefix e.g nJy) in addition to items in SPEC_PHOTON_FLUX_DENSITY_UNITS
+    equiv = u.spectral_density(1 * u.m)  # spec. unit doesn't matter here, we're not evaluating
+    for un in SPEC_PHOTON_FLUX_DENSITY_UNITS:
+        if flux_unit.is_equivalent(un, equiv):
+            if flux_unit_str not in SPEC_PHOTON_FLUX_DENSITY_UNITS:
+                return SPEC_PHOTON_FLUX_DENSITY_UNITS + [flux_unit_str]
+            else:
+                return SPEC_PHOTON_FLUX_DENSITY_UNITS
 
-    # Concatenate both lists with the local units coming first.
-    return sorted(units_to_strings(local_units)) + flux_unit_equivalencies_titles
+    else:
+        # for any other units, including counts, DN, e/s, DN /s, etc,
+        # no other conversions between flux units available as we only support
+        # conversions to and from spectral and photon flux density flux unit.
+        # dropdown will only contain one item (the input unit)
+        return [flux_unit_str]
 
 
 def create_angle_equivalencies_list(solid_angle_unit):
@@ -146,7 +135,7 @@ def create_angle_equivalencies_list(solid_angle_unit):
 
     """
 
-    if solid_angle_unit is None:
+    if solid_angle_unit is None or solid_angle_unit is PIX2:
         # if there was no solid angle in the unit when calling this function
         # can only represent that unit as per square pixel
         return ['pix^2']