Merge pull request #1176 from rosteen/fix-model-error

Fix model fitting plugin bugs

CHANGES.rst

@@ -43,6 +43,9 @@ Cubeviz
 ^^^^^^^
 
 - Fixed linking of data to allow contour over-plotting. [#1154]
+- Fixed an error trace when fitting a model to a spatial subset. [#1176]
+- Fixed the model fitting plugin data dropdown not populating with spatial
+  subsets properly. [#1176]
 
 Imviz
 ^^^^^

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

@@ -9,6 +9,7 @@
 from glue.core.data import Data
 from glue.core.subset import Subset, RangeSubsetState, OrState, AndState
 from glue.core.link_helpers import LinkSame
+from glue.core.message import SubsetDeleteMessage, SubsetUpdateMessage
 
 from jdaviz.core.events import AddDataMessage, RemoveDataMessage, SnackbarMessage
 from jdaviz.core.registries import tray_registry
@@ -86,6 +87,10 @@ def __init__(self, *args, **kwargs):
                            handler=self._on_viewer_data_changed)
         self.hub.subscribe(self, RemoveDataMessage,
                            handler=self._on_viewer_data_changed)
+        self.hub.subscribe(self, SubsetUpdateMessage,
+                           handler=self._on_viewer_data_changed)
+        self.hub.subscribe(self, SubsetDeleteMessage,
+                           handler=self._on_viewer_data_changed)
 
     def _on_viewer_data_changed(self, msg=None):
         """
@@ -108,20 +113,14 @@ def _on_viewer_data_changed(self, msg=None):
         self._viewer_id = self.app._viewer_item_by_reference(
             'spectrum-viewer').get('id')
 
-        # Subsets are global and are not linked to specific viewer instances,
-        # so it's not required that we match any specific ids for that case.
-        # However, if the msg is not none, check to make sure that it's the
-        # viewer we care about.
-        if msg is not None and msg.viewer_id != self._viewer_id:
-            return
-
         viewer = self.app.get_viewer('spectrum-viewer')
 
         self.dc_items = [layer_state.layer.label
                          for layer_state in viewer.state.layers
-                         if (not isinstance(layer_state.layer, Subset)
-                             or not isinstance(layer_state.layer.subset_state,
-                                               (RangeSubsetState, OrState, AndState)))]
+                         if ((not isinstance(layer_state.layer, Subset)
+                              or not isinstance(layer_state.layer.subset_state,
+                                                (RangeSubsetState, OrState, AndState)))
+                             and layer_state.layer.label not in self.app.fitted_models.keys())]
 
     def _param_units(self, param, model_type=None):
         """Helper function to handle units that depend on x and y"""
@@ -428,8 +427,8 @@ def vue_model_fitting(self, *args, **kwargs):
         self._fitted_model = fitted_model
         self._fitted_spectrum = fitted_spectrum
 
-        self.vue_register_spectrum({"spectrum": fitted_spectrum})
         self.app.fitted_models[self.model_label] = fitted_model
+        self.vue_register_spectrum({"spectrum": fitted_spectrum})
 
         # Update component model parameters with fitted values
         if type(self._fitted_model) == QuantityModel:
@@ -563,9 +562,10 @@ def vue_register_spectrum(self, event):
 
         self.app.add_data(spectrum, label)
 
-        # Make sure we link the result spectrum to the data we're fitting
-        data_fitted = self.app.session.data_collection[self.selected_data]
-        data_id = data_fitted.world_component_ids[0]
+        # Link the result spectrum to the reference data of the spectrum viewer
+
+        ref_data = self.app.get_viewer('spectrum-viewer').state.reference_data
+        data_id = ref_data.world_component_ids[0]
         model_id = self.app.session.data_collection[label].world_component_ids[0]
         self.app.session.data_collection.add_link(LinkSame(data_id, model_id))