Improve spaxel unit handling for cubes (#3307)

* Add spaxel to the list of recognized solid angle units

* Update test to use a newer file that would fail the test on main

* convert per spaxel to per pixel squared on data load

* Only add spaxel once, on app load.

* Codestyle

* Set bad uncert values to 0 instead of NaN here

* Remove inf uncertainty replacement

* Remove unused import

jdaviz/app.py

@@ -52,7 +52,7 @@
 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_and_equivs import SPEC_PHOTON_FLUX_DENSITY_UNITS
+from jdaviz.core.custom_units_and_equivs import SPEC_PHOTON_FLUX_DENSITY_UNITS, enable_spaxel_unit
 from jdaviz.core.unit_conversion_utils import (check_if_unit_is_per_solid_angle,
                                                combine_flux_and_angle_units,
                                                supported_sq_angle_units)
@@ -62,6 +62,8 @@
 SplitPanes()
 GoldenLayout()
 
+enable_spaxel_unit()
+
 CONTAINER_TYPES = dict(row='gl-row', col='gl-col', stack='gl-stack')
 EXT_TYPES = dict(flux=['flux', 'sci'],
                  uncert=['ivar', 'err', 'var', 'uncert'],

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

@@ -24,9 +24,6 @@
 
 SPECUTILS_LT_1_15_1 = not minversion(specutils, "1.15.1.dev")
 
-spaxel = u.def_unit('spaxel', 1 * u.Unit(""))
-u.add_enabled_units([spaxel])
-
 moment_unit_options = {0: ["Surface Brightness"],
                        1: ["Velocity", "Spectral Unit"],
                        2: ["Velocity", "Velocity^N"]}

jdaviz/configs/cubeviz/plugins/parsers.py

@@ -203,6 +203,10 @@ def _return_spectrum_with_correct_units(flux, wcs, metadata, data_type=None,
     if (apply_pix2 and (data_type != "mask") and
             (not check_if_unit_is_per_solid_angle(flux.unit))):
         target_flux_unit = flux.unit / PIX2
+    elif check_if_unit_is_per_solid_angle(flux.unit, return_unit=True) == "spaxel":
+        # We need to convert spaxel to pixel squared, since spaxel isn't fully supported by astropy
+        # This is horribly ugly but just multiplying by u.Unit("spaxel") doesn't work
+        target_flux_unit = flux.unit * u.Unit('spaxel') / PIX2
 
     if target_wave_unit is None and hdulist is not None:
         found_target = False

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

@@ -16,8 +16,9 @@ def test_data_retrieval(cubeviz_helper):
     return the same spectrum values.
     """
     # This file is originally from
-    # https://data.sdss.org/sas/dr14/manga/spectro/redux/v2_1_2/7495/stack/manga-7495-12704-LOGCUBE.fits.gz
-    URL = 'https://stsci.box.com/shared/static/28a88k1qfipo4yxc4p4d40v4axtlal8y.fits'
+    # https://data.sdss.org/sas/dr17/manga/spectro/redux/v3_1_1/9862/stack/manga-9862-12703-LOGCUBE.fits.gz
+    # (Updated to a newer file 11/19/2024)
+    URL = 'https://stsci.box.com/shared/static/gts87zqt5265msuwi4w5u003b6typ6h0.gz'
 
     spectrum_viewer_reference_name = "spectrum-viewer"
     fn = download_file(URL, cache=True)

jdaviz/configs/default/plugins/gaussian_smooth/gaussian_smooth.py

@@ -1,6 +1,5 @@
 import numpy as np
 
-from astropy import units as u
 from astropy.convolution import convolve, Gaussian2DKernel
 from specutils import Spectrum1D
 from specutils.manipulation import gaussian_smooth
@@ -17,10 +16,6 @@
 __all__ = ['GaussianSmooth']
 
 
-spaxel = u.def_unit('spaxel', 1 * u.Unit(""))
-u.add_enabled_units([spaxel])
-
-
 @tray_registry('g-gaussian-smooth', label="Gaussian Smooth",
                viewer_requirements=['spectrum', 'flux'])
 class GaussianSmooth(PluginTemplateMixin, DatasetSelectMixin, AddResultsMixin):

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

@@ -146,7 +146,7 @@ def __init__(self, *args, **kwargs):
                                                     items='angle_unit_items',
                                                     selected='angle_unit_selected')
         # NOTE: will switch to pix2 only if first data loaded into viewer is in pix2 units
-        # initialize flux choices to empty list, will be populated when data is loaded
+        # initialize angle unit 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',)

jdaviz/core/custom_units_and_equivs.py

@@ -8,6 +8,13 @@
 PIX2 = u.pix * u.pix
 
 
+# Add spaxel to enabled units
+def enable_spaxel_unit():
+    spaxel = u.Unit('spaxel', represents=u.pixel, parse_strict='silent')
+    u.add_enabled_units([spaxel])
+    return
+
+
 def _spectral_and_photon_flux_density_units(freq_only=False, wav_only=False,
                                             as_units=False):
     """
@@ -92,8 +99,15 @@ def _eqv_flux_to_sb_pixel():
                   u.ct,
                   u.DN,
                   u.DN / u.s]
-    return [(flux_unit, flux_unit / PIX2, lambda x: x, lambda x: x)
-            for flux_unit in flux_units]
+
+    equivs = [(flux_unit, flux_unit / PIX2, lambda x: x, lambda x: x)
+              for flux_unit in flux_units]
+
+    # We also need to convert between spaxel and pixel squared
+    equivs += [(flux_unit / u.Unit('spaxel'), flux_unit / PIX2,
+               lambda x: x, lambda x: x) for flux_unit in flux_units]
+
+    return equivs
 
 
 def _eqv_sb_per_pixel_to_per_angle(flux_unit, scale_factor=1):

jdaviz/core/unit_conversion_utils.py

@@ -115,15 +115,11 @@ def check_if_unit_is_per_solid_angle(unit, return_unit=False):
             # 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':
+            if new_unit.physical_type == 'solid angle' or new_unit == PIX2 or new_unit_str == 'spaxel':  # noqa
+                # square pixel and spaxel should be considered square angle units
                 if return_unit:  # area units present and requested to be returned
                     return new_unit
                 return True  # area units present but not requested to be returned
-            # square pixel should be considered a square angle unit
-            if new_unit == PIX2:
-                if return_unit:
-                    return new_unit
-                return True
 
     # in the case there are no area units, but return units were requested
     if return_unit: