Merge pull request spacetelescope#2685 from pllim/exorcise-cube-fit-s…

…pinner-of-death

BUG: Cube fit should work even without valid 3D WCS

CHANGES.rst

@@ -121,6 +121,9 @@ Cubeviz
 - Fixes Spectral Extraction's assumptions of one data per viewer, and flux data only in
   flux-viewer/uncertainty data only in uncert-viewer. [#2646]
 
+- Fixed a bug where cube model fitting could fail (endless spinner) if input cube
+  has invalid 3D WCS. [#2685]
+
 Imviz
 ^^^^^
 

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

@@ -6,7 +6,6 @@
 from specutils import Spectrum1D
 from specutils.utils import QuantityModel
 from traitlets import Bool, List, Unicode, observe
-from glue.core.data import Data
 
 from jdaviz.core.events import SnackbarMessage, GlobalDisplayUnitChanged
 from jdaviz.core.registries import tray_registry
@@ -947,17 +946,9 @@ def _fit_model_to_cube(self, add_data):
                 temp_label = "{} ({}, {})".format(self.results_label, m["x"], m["y"])
                 self.app.fitted_models[temp_label] = m["model"]
 
-        count = max(map(lambda s: int(next(iter(re.findall(r"\d$", s)), 0)),
-                        self.data_collection.labels)) + 1
-
-        label = self.app.return_data_label(f"{self.results_label}[Cube]", ext=count)
-
-        # Create new glue data object
-        output_cube = Data(label=label,
-                           coords=fitted_spectrum.wcs,
-                           flux=fitted_spectrum.flux.value)
-        output_cube.get_component('flux').units = fitted_spectrum.flux.unit.to_string()
+        output_cube = Spectrum1D(flux=fitted_spectrum.flux, wcs=fitted_spectrum.wcs)
 
+        # Create new data entry for glue
         if add_data:
             self.add_results.add_results_from_plugin(output_cube)
             self._set_default_results_label()

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

@@ -150,15 +150,34 @@ def test_register_cube_model(cubeviz_helper, spectrum1d_cube):
     modelfit_plugin._obj.results_label = test_label
     # changing the lable should set auto to False, but the event may not have triggered yet
     modelfit_plugin._obj.results_label_auto = False
-    modelfit_plugin._obj.cube_fit = True
+    modelfit_plugin.cube_fit = True
     assert modelfit_plugin._obj.results_label_default == 'cube-fit model'
     assert modelfit_plugin._obj.results_label == test_label
     with warnings.catch_warnings():
-        warnings.filterwarnings('ignore', message=r'.*Model is linear in parameters.*')
+        warnings.filterwarnings('ignore', message='.*Model is linear in parameters.*')
         modelfit_plugin.calculate_fit()
     assert test_label in cubeviz_helper.app.data_collection
 
 
+def test_fit_cube_no_wcs(cubeviz_helper):
+    # This is like when user do something to a cube outside of Jdaviz
+    # and then load it back into a new instance of Cubeviz for further analysis.
+    sp = Spectrum1D(flux=np.ones((7, 8, 9)) * u.nJy)  # ny, nx, nz
+    cubeviz_helper.load_data(sp, data_label="test_cube")
+    mf = cubeviz_helper.plugins['Model Fitting']
+    mf.create_model_component('Linear1D')
+    mf.cube_fit = True
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", message="Model is linear in parameters.*")
+        fitted_model, output_cube = mf.calculate_fit(add_data=True)
+    assert len(fitted_model) == 56  # ny * nx
+    # Make sure shapes are all self-consistent within Cubeviz instance.
+    fitted_data = cubeviz_helper.app.data_collection["cube-fit model"]
+    assert fitted_data.shape == (8, 7, 9)  # nx, ny, nz
+    assert fitted_data.shape == cubeviz_helper.app.data_collection[0].shape
+    assert fitted_data.shape == output_cube.shape
+
+
 def test_refit_plot_options(specviz_helper, spectrum1d):
     specviz_helper.load_data(spectrum1d)
     modelfit_plugin = specviz_helper.plugins['Model Fitting']