model fitting: ensure specutils has access to equivs & set targ solid…

… angle (spacetelescope#3343)

* ensure specutils has access to custom eqvs and update targ solid angle bug

* fix spacing syntax error

* add spec axis/pixar_sr eqvs to model fit, add test coverage

* add change log

* add indirect conversion handling for PIX2 and test coverage

CHANGES.rst

@@ -76,6 +76,8 @@ Cubeviz
 ^^^^^^^
 - Removed the deprecated ``save as fits`` option from the Collapse, Moment Maps, and Spectral Extraction plugins; use the Export plugin instead. [#3256]
 
+- Fixed bugs where cube model fitting could fail if Jdaviz custom equivalencies were required. [#3343]
+
 Imviz
 ^^^^^
 

jdaviz/configs/default/plugins/model_fitting/model_fitting.py

@@ -27,6 +27,7 @@
 from jdaviz.core.user_api import PluginUserApi
 from jdaviz.core.unit_conversion_utils import (all_flux_unit_conversion_equivs,
                                                flux_conversion_general)
+from jdaviz.core.custom_units_and_equivs import PIX2
 
 __all__ = ['ModelFitting']
 
@@ -982,7 +983,23 @@ def _fit_model_to_cube(self, add_data):
 
         sb_unit = self.app._get_display_unit('sb')
         if spec.flux.unit != sb_unit:
-            spec = spec.with_flux_unit(sb_unit)
+            # ensure specutils has access to jdaviz custom unit equivalencies
+            pixar_sr = spec.meta.get('_pixel_scale_factor', None)
+            equivalencies = all_flux_unit_conversion_equivs(pixar_sr=pixar_sr,
+                                                            cube_wave=spec.spectral_axis)
+
+            pix2_in_flux = 'pix2' in spec.flux.unit.to_string()
+            pix2_in_sb = 'pix2' in sb_unit
+            # Handle various cases when PIX2 angle unit is present in the conversion
+            if pix2_in_flux and pix2_in_sb:
+                spec = spec.with_flux_unit(u.Unit(spec.flux.unit)*PIX2, equivalencies=equivalencies)  # noqa
+                spec = spec.with_flux_unit(u.Unit(sb_unit)*PIX2, equivalencies=equivalencies)
+            elif pix2_in_flux:
+                spec = spec.with_flux_unit(u.Unit(spec.flux.unit) * PIX2, equivalencies=equivalencies)  # noqa
+            elif pix2_in_sb:
+                spec = spec.with_flux_unit(u.Unit(spec.flux.unit) / PIX2, equivalencies=equivalencies)  # noqa
+
+            spec = spec.with_flux_unit(sb_unit, equivalencies=equivalencies)
 
         snackbar_message = SnackbarMessage(
             "Fitting model to cube...",

jdaviz/configs/default/plugins/model_fitting/tests/test_fitting.py

@@ -15,6 +15,7 @@
 from jdaviz.configs.default.plugins.model_fitting import fitting_backend as fb
 from jdaviz.configs.default.plugins.model_fitting import initializers
 from jdaviz.core.custom_units_and_equivs import PIX2
+from jdaviz.conftest import _create_spectrum1d_cube_with_fluxunit
 
 SPECTRUM_SIZE = 200  # length of spectrum
 
@@ -443,8 +444,35 @@ def test_cube_fit_with_subset_and_nans(cubeviz_helper):
     assert np.all(result.get_component("flux").data == 1)
 
 
-def test_cube_fit_after_unit_change(cubeviz_helper, spectrum1d_cube_fluxunit_jy_per_steradian):
-    cubeviz_helper.load_data(spectrum1d_cube_fluxunit_jy_per_steradian, data_label="test")
+def test_fit_with_count_units(cubeviz_helper):
+    flux = np.random.random((7, 8, 9)) * u.count
+    spectral_axis = np.linspace(4000, 5000, flux.shape[-1]) * u.AA
+
+    spec = Spectrum1D(flux=flux, spectral_axis=spectral_axis)
+    cubeviz_helper.load_data(spec, data_label="test")
+
+    mf = cubeviz_helper.plugins["Model Fitting"]
+    mf.cube_fit = True
+    mf.create_model_component("Const1D")
+
+    # ensures specutils.Spectrum1D.with_flux_unit has access to Jdaviz custom equivalencies for
+    # PIX^2 unit
+    with warnings.catch_warnings():
+        warnings.filterwarnings('ignore', message='Model is linear in parameters.*')
+        mf.calculate_fit()
+
+    assert mf._obj.component_models[0]['parameters'][0]['unit'] == 'ct / pix2'
+
+    model_flux = cubeviz_helper.app.data_collection[-1].get_component('flux')
+    assert model_flux.units == 'ct / pix2'
+
+
+@pytest.mark.parametrize("solid_angle_unit", [u.sr, PIX2])
+def test_cube_fit_after_unit_change(cubeviz_helper, solid_angle_unit):
+    cube = _create_spectrum1d_cube_with_fluxunit(fluxunit=u.Jy / solid_angle_unit, shape=(10, 4, 5),
+                                                 with_uncerts=True)
+    cubeviz_helper.load_data(cube, data_label="test")
+    solid_angle_string = str(solid_angle_unit)
 
     uc = cubeviz_helper.plugins['Unit Conversion']
     mf = cubeviz_helper.plugins['Model Fitting']
@@ -453,7 +481,7 @@ def test_cube_fit_after_unit_change(cubeviz_helper, spectrum1d_cube_fluxunit_jy_
 
     mf.create_model_component("Const1D")
     # Check that the parameter is using the current units when initialized
-    assert mf._obj.component_models[0]['parameters'][0]['unit'] == 'MJy / sr'
+    assert mf._obj.component_models[0]['parameters'][0]['unit'] == f'MJy / {solid_angle_string}'
 
     with warnings.catch_warnings():
         warnings.filterwarnings('ignore', message='Model is linear in parameters.*')
@@ -466,18 +494,38 @@ def test_cube_fit_after_unit_change(cubeviz_helper, spectrum1d_cube_fluxunit_jy_
                                       [9.40e-05, 9.90e-05, 1.04e-04, 1.09e-04]])
 
     model_flux = cubeviz_helper.app.data_collection[-1].get_component('flux')
-    assert model_flux.units == 'MJy / sr'
+    assert model_flux.units == f'MJy / {solid_angle_string}'
     assert np.allclose(model_flux.data[:, :, 1], expected_result_slice)
 
     # Switch back to Jy, see that the component didn't change but the output does
     uc.flux_unit = 'Jy'
-    assert mf._obj.component_models[0]['parameters'][0]['unit'] == 'MJy / sr'
+    assert mf._obj.component_models[0]['parameters'][0]['unit'] == f'MJy / {solid_angle_string}'
     with warnings.catch_warnings():
         warnings.filterwarnings('ignore', message='Model is linear in parameters.*')
         mf.calculate_fit()
 
     model_flux = cubeviz_helper.app.data_collection[-1].get_component('flux')
-    assert model_flux.units == 'Jy / sr'
+    assert model_flux.units == f'Jy / {solid_angle_string}'
     assert np.allclose(model_flux.data[:, :, 1], expected_result_slice * 1e6)
 
-    # ToDo: Add a test for a unit change that needs an equivalency
+    # ensure conversions that require the spectral axis/translations are handled by the plugin
+    uc.spectral_y_type = 'Surface Brightness'
+    uc.flux_unit = 'erg / (Angstrom s cm2)'
+
+    mf.reestimate_model_parameters()
+
+    if solid_angle_string == 'sr':
+        expected_unit_string = f'erg / (Angstrom s {solid_angle_string} cm2)'
+    else:
+        expected_unit_string = f'erg / (Angstrom s cm2 {solid_angle_string})'
+
+    assert mf._obj.component_models[0]['parameters'][0]['unit'] == expected_unit_string
+
+    # running this ensures specutils.Spectrum1D.with_flux_unit has Jdaviz custom equivalencies
+    # for spectral axis conversions and scale factor translations
+    with warnings.catch_warnings():
+        warnings.filterwarnings('ignore', message='Model is linear in parameters.*')
+        mf.calculate_fit()
+
+    model_flux = cubeviz_helper.app.data_collection[-1].get_component('flux')
+    assert model_flux.units == expected_unit_string

jdaviz/utils.py

@@ -505,6 +505,7 @@ def _indirect_conversion(values, orig_units, targ_units, eqv,
     elif image_data or (spec_unit and solid_angle_in_spec):
         if not solid_angle_in_targ:
             targ_units /= solid_angle_in_spec
+            solid_angle_in_targ = solid_angle_in_spec
         if ((u.Unit(targ_units) in indirect_units()) or
            (u.Unit(orig_units) in indirect_units())):
             # SB -> Flux -> Flux -> SB