removing unit dictionary, multiply/divide by sr

spacetelescope · Mar 21, 2024 · 374dcaf · 374dcaf
1 parent 42dd35b
commit 374dcaf
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 22 deletions.
diff --git a/jdaviz/configs/cubeviz/plugins/spectral_extraction/spectral_extraction.py b/jdaviz/configs/cubeviz/plugins/spectral_extraction/spectral_extraction.py
@@ -205,7 +205,6 @@ def _update_mark_scale(self, *args):
             self.aperture.scale_factor = 1.0
         else:
             self.aperture.scale_factor = self.slice_wavelength/self.reference_wavelength
-
         if not self.bg_wavelength_dependent:
             self.background.scale_factor = 1.0
         else:
@@ -346,8 +345,7 @@ def collapse_to_spectrum(self, add_data=True, **kwargs):
             self.hub.broadcast(snackbar_message)
 
         # per https://jwst-docs.stsci.edu/jwst-near-infrared-camera/nircam-performance/nircam-absolute-flux-calibration-and-zeropoints # noqa
-        if 'PIXAR_SR' in spectral_cube.meta:
-            pix_scale_factor = self.aperture.scale_factor * spectral_cube.meta['PIXAR_SR']
+            pix_scale_factor = self.aperture.scale_factor * spectral_cube.meta.get('PIXAR_SR', 1.0)
             collapsed_spec.meta['_pixel_scale_factor'] = pix_scale_factor
 
         return collapsed_spec
@@ -528,24 +526,11 @@ def _live_update(self, event={}):
             mark.visible = True
 
     def translate_units(self, collapsed_spec):
-        sb_to_flux = {
-            u.MJy / u.sr: u.MJy / u.pix,
-            u.Jy / u.sr: u.Jy / u.pix,
-            u.erg / (u.cm**2*u.s*u.AA*u.sr): u.erg / (u.cm**2*u.s*u.AA),
-            u.erg / (u.s*u.cm**2*u.sr): u.erg / (u.s*u.cm**2),
-            u.W / (u.m**2*u.Hz*u.sr): u.W / (u.m**2*u.Hz)
-        }
-
-        flux_to_sb = {value: key for key, value in sb_to_flux.items()}
-
         # remove sr
-        if collapsed_spec._unit in sb_to_flux:
+        if u.sr in collapsed_spec._unit.bases:
             collapsed_spec._data *= collapsed_spec.meta['_pixel_scale_factor']
-            collapsed_spec._unit = sb_to_flux[collapsed_spec._unit]
+            collapsed_spec._unit *= u.sr
         # add sr
-        elif collapsed_spec._unit in flux_to_sb:
+        elif u.sr not in collapsed_spec._unit.bases:
             collapsed_spec._data /= collapsed_spec.meta['_pixel_scale_factor']
-            collapsed_spec._unit = flux_to_sb[collapsed_spec._unit]
-        else:
-            raise ValueError(f"Units '{collapsed_spec._unit}'"
-                             f"are unrecognized surface brightness or flux units.")
+            collapsed_spec._unit /= u.sr
diff --git a/jdaviz/configs/cubeviz/plugins/spectral_extraction/tests/test_spectral_extraction.py b/jdaviz/configs/cubeviz/plugins/spectral_extraction/tests/test_spectral_extraction.py
@@ -382,7 +382,7 @@ def test_unit_translation(cubeviz_helper):
     w = WCS(wcs_dict)
     flux = np.zeros((30, 20, 3001), dtype=np.float32)
     flux[5:15, 1:11, :] = 1
-    cube = Spectrum1D(flux=flux * u.MJy / u.pix, wcs=w, meta=wcs_dict)
+    cube = Spectrum1D(flux=flux * u.MJy, wcs=w, meta=wcs_dict)
     cubeviz_helper.load_data(cube, data_label="test")
 
     center = PixCoord(5, 10)
@@ -414,7 +414,7 @@ def test_unit_translation(cubeviz_helper):
     extract_plg._obj.translate_units(collapsed_spec)
 
     # translating again returns the original units
-    assert collapsed_spec._unit == u.MJy / u.pix
+    assert collapsed_spec._unit == u.MJy
     # returns to the original values
     # which is a value in Jy/pix that we know the outcome after translation
     assert np.allclose(collapsed_spec._data[0], mjy_sr_data1)