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

CHANGES.rst

@@ -66,6 +66,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

@@ -982,7 +982,11 @@ 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)
+            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

@@ -443,6 +443,29 @@ def test_cube_fit_with_subset_and_nans(cubeviz_helper):
     assert np.all(result.get_component("flux").data == 1)
 
 
+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'
+
+
 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")
 
@@ -480,4 +503,18 @@ def test_cube_fit_after_unit_change(cubeviz_helper, spectrum1d_cube_fluxunit_jy_
     assert model_flux.units == 'Jy / sr'
     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()
+    assert mf._obj.component_models[0]['parameters'][0]['unit'] == 'erg / (Angstrom s sr cm2)'
+
+    # 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 == 'erg / (Angstrom s sr cm2)'

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