Merge pull request #796 from pllim/fix-cubeviz-model-fitting-error

Cubeviz: Allow model fitting to operate on spectral Subset

jdaviz/configs/default/plugins/collapse/collapse.py

@@ -49,6 +49,7 @@ def __init__(self, *args, **kwargs):
                            handler=self._on_data_updated)
 
         self._selected_data = None
+        self._spectral_subsets = {}
         self._label_counter = 0
 
     def _on_data_updated(self, msg):
@@ -68,6 +69,7 @@ def _on_data_item_selected(self, event):
 
         # Also set the spectral min and max to default to the full range
         cube = self._selected_data.get_object(cls=SpectralCube)
+        self.selected_subset = "None"
         self.spectral_min = cube.spectral_axis[0].value
         self.spectral_max = cube.spectral_axis[-1].value
         self.spectral_unit = str(cube.spectral_axis.unit)
@@ -77,13 +79,12 @@ def _on_data_item_selected(self, event):
     @observe("selected_subset")
     def _on_subset_selected(self, event):
         # If "None" selected, reset based on bounds of selected data
-        self._selected_subset = self.selected_subset
-        if self._selected_subset == "None":
+        if self.selected_subset == "None":
             cube = self._selected_data.get_object(cls=SpectralCube)
             self.spectral_min = cube.spectral_axis[0].value
             self.spectral_max = cube.spectral_axis[-1].value
         else:
-            spec_sub = self._spectral_subsets[self._selected_subset]
+            spec_sub = self._spectral_subsets[self.selected_subset]
             unit = u.Unit(self.spectral_unit)
             spec_reg = SpectralRegion.from_center(spec_sub.center.x * unit,
                                                   spec_sub.width * unit)
@@ -94,15 +95,13 @@ def vue_list_subsets(self, event):
         """Populate the spectral subset selection dropdown"""
         temp_subsets = self.app.get_subsets_from_viewer("spectrum-viewer",
                                                         subset_type="spectral")
-        temp_list = ["None"]
         temp_dict = {}
         # Attempt to filter out spatial subsets
         for key, region in temp_subsets.items():
             if type(region) == RectanglePixelRegion:
                 temp_dict[key] = region
-                temp_list.append(key)
         self._spectral_subsets = temp_dict
-        self.spectral_subset_items = temp_list
+        self.spectral_subset_items = ["None"] + sorted(temp_dict.keys())
 
     def vue_collapse(self, *args, **kwargs):
         try:

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

@@ -14,7 +14,7 @@
 
 
 def fit_model_to_spectrum(spectrum, component_list, expression,
-                          run_fitter=False):
+                          run_fitter=False, window=None, n_cpu=None):
     """
     Fits an astropy CompoundModel to a Spectrum1D instance.
 
@@ -24,7 +24,7 @@ def fit_model_to_spectrum(spectrum, component_list, expression,
 
     Parameters
     ----------
-    spectrum : :class:`specutils.spectrum.Spectrum1D`
+    spectrum : :class:`specutils.spectra.Spectrum1D`
         The spectrum to be fitted.
     component_list : list
         Spectral model subcomponents stored in a list.
@@ -38,6 +38,14 @@ def fit_model_to_spectrum(spectrum, component_list, expression,
         When False (the default), the function composes the compound
         model and returns it without fitting - This is currently being
         ignored for 3D fits.
+    window : `None` or :class:`specutils.spectra.SpectralRegion`
+        See :func:`specutils.fitting.fitmodels.fit_lines`.
+    n_cpu : `None` or int
+        **This is only used for spectral cube fitting.**
+        Number of cores to use for multiprocessing.
+        Using all the cores at once is not recommended.
+        If `None`, it will use max cores minus one.
+        Set this to 1 for debugging.
 
     Returns
     -------
@@ -47,43 +55,45 @@ def fit_model_to_spectrum(spectrum, component_list, expression,
         3D spectral cube input, instead o model instances for every
         spaxel, a list with 2D arrays, each one storing fitted parameter
         values for all spaxels, is returned.
-    :class:`specutils.spectrum.Spectrum1D`
+    :class:`specutils.spectra.Spectrum1D`
         The realization of the fitted model as a spectrum. The spectrum
         will be 1D or 3D depending on the input spectrum's shape.
     """
     # Initial guess for the fit.
     initial_model = _build_model(component_list, expression)
 
     if len(spectrum.shape) > 1:
-        return _fit_3D(initial_model, spectrum)
+        return _fit_3D(initial_model, spectrum, window=window, n_cpu=n_cpu)
     else:
-        return _fit_1D(initial_model, spectrum, run_fitter)
+        return _fit_1D(initial_model, spectrum, run_fitter, window=window)
 
 
-def _fit_1D(initial_model, spectrum, run_fitter):
+def _fit_1D(initial_model, spectrum, run_fitter, window=None):
     """
     Fits an astropy CompoundModel to a Spectrum1D instance.
 
     Parameters
     ----------
-    spectrum : :class:`specutils.spectrum.Spectrum1D`
-        The spectrum to be fitted.
     initial_model : :class: `astropy.modeling.CompoundModel`
         Initial guess for the model to be fitted.
+    spectrum : :class:`specutils.spectra.Spectrum1D`
+        The spectrum to be fitted.
     run_fitter : bool
         When False (the default), the function composes the compound
         model and returns it without fitting.
+    window : `None` or :class:`specutils.spectra.SpectralRegion`
+        See :func:`specutils.fitting.fitmodels.fit_lines`.
 
     Returns
     -------
     :class: `astropy.modeling.CompoundModel`
         The model resulting from the fit.
-    :class:`specutils.spectrum.Spectrum1D`
+    :class:`specutils.spectra.Spectrum1D`
         The realization of the fitted model as a spectrum.
 
     """
     if run_fitter:
-        output_model = fit_lines(spectrum, initial_model)
+        output_model = fit_lines(spectrum, initial_model, window=window)
         output_values = output_model(spectrum.spectral_axis)
     else:
         # Return without fitting.
@@ -97,26 +107,37 @@ def _fit_1D(initial_model, spectrum, run_fitter):
     return output_model, output_spectrum
 
 
-def _fit_3D(initial_model, spectrum):
+def _fit_3D(initial_model, spectrum, window=None, n_cpu=None):
     """
     Fits an astropy CompoundModel to every spaxel in a cube
     using a multiprocessor pool running in parallel. Computes
     realizations of the models over each spaxel.
 
     Parameters
     ----------
-    spectrum : :class:`specutils.spectrum.Spectrum1D`
+    initial_model : :class: `astropy.modeling.CompoundModel`
+        Initial guess for the model to be fitted.
+    spectrum : :class:`specutils.spectra.Spectrum1D`
         The spectrum that stores the cube in its 'flux' attribute.
+    window : `None` or :class:`specutils.spectra.SpectralRegion`
+        See :func:`specutils.fitting.fitmodels.fit_lines`.
+    n_cpu : `None` or int
+        Number of cores to use for multiprocessing.
+        Using all the cores at once is not recommended.
+        If `None`, it will use max cores minus one.
+        Set this to 1 for debugging.
 
     Returns
     -------
     :list: a list that stores 2D arrays. Each array contains one
         parameter from `astropy.modeling.CompoundModel` instances
         fitted to every spaxel in the input cube.
-    :class:`specutils.spectrum.Spectrum1D`
+    :class:`specutils.spectra.Spectrum1D`
         The spectrum that stores the fitted model values in its 'flux'
         attribute.
     """
+    if n_cpu is None:
+        n_cpu = mp.cpu_count() - 1
 
     # Generate list of all spaxels to be fitted
     spaxels = _generate_spaxel_list(spectrum)
@@ -145,19 +166,20 @@ def collect_result(results):
     # Run multiprocessor pool to fit each spaxel and
     # compute model values on that same spaxel.
     results = []
-    pool = Pool(mp.cpu_count() - 1)
+    pool = Pool(n_cpu)
 
     # The communicate overhead of spawning a process for each *individual*
     # parameter set is prohibitively high (it's actually faster to run things
     # sequentially). Instead, chunk the spaxel list based on the number of
     # available processors, and have each processor do the model fitting
     # on the entire subset of spaxel tuples, then return the set of results.
-    for spx in np.array_split(spaxels, mp.cpu_count() - 1):
+    for spx in np.array_split(spaxels, n_cpu):
         # Worker for the multiprocess pool.
         worker = SpaxelWorker(spectrum.flux,
                               spectrum.spectral_axis,
                               initial_model,
-                              param_set=spx)
+                              param_set=spx,
+                              window=window)
         r = pool.apply_async(worker, callback=collect_result)
         results.append(r)
     for r in results:
@@ -187,11 +209,12 @@ class SpaxelWorker:
     instance. We need to use the current model instance while
     it still exists.
     """
-    def __init__(self, flux_cube, wave_array, initial_model, param_set):
+    def __init__(self, flux_cube, wave_array, initial_model, param_set, window=None):
         self.cube = flux_cube
         self.wave = wave_array
         self.model = initial_model
         self.param_set = param_set
+        self.window = window
 
     def __call__(self):
         results = {'x': [], 'y': [], 'fitted_model': [], 'fitted_values': []}
@@ -211,7 +234,7 @@ def __call__(self):
             flux = self.cube[x, y, :]  # transposed!
             sp = Spectrum1D(spectral_axis=self.wave, flux=flux)
 
-            fitted_model = fit_lines(sp, self.model)
+            fitted_model = fit_lines(sp, self.model, window=self.window)
 
             fitted_values = fitted_model(self.wave)
 
@@ -303,7 +326,7 @@ def _generate_spaxel_list(spectrum):
 
     Parameters
     ----------
-    spectrum : :class:`specutils.spectrum.Spectrum1D`
+    spectrum : :class:`specutils.spectra.Spectrum1D`
         The spectrum that stores the cube in its 'flux' attribute.
 
     Returns

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

@@ -7,10 +7,12 @@
 from glue.core.message import (SubsetCreateMessage,
                                SubsetDeleteMessage,
                                SubsetUpdateMessage)
-from specutils import Spectrum1D
+from regions import RectanglePixelRegion
+from specutils import Spectrum1D, SpectralRegion
 from specutils.utils import QuantityModel
-from traitlets import Bool, Int, List, Unicode
+from traitlets import Any, Bool, Int, List, Unicode, observe
 from glue.core.data import Data
+from glue.core.subset import Subset, RangeSubsetState
 
 from jdaviz.core.events import AddDataMessage, RemoveDataMessage, SnackbarMessage
 from jdaviz.core.registries import tray_registry
@@ -37,6 +39,12 @@ class ModelFitting(TemplateMixin):
     template = load_template("model_fitting.vue", __file__).tag(sync=True)
     dc_items = List([]).tag(sync=True)
 
+    spectral_min = Any().tag(sync=True)
+    spectral_max = Any().tag(sync=True)
+    spectral_unit = Unicode().tag(sync=True)
+    spectral_subset_items = List(["None"]).tag(sync=True)
+    selected_subset = Unicode("None").tag(sync=True)
+
     model_label = Unicode().tag(sync=True)
     cube_fit = Bool(False).tag(sync=True)
     temp_name = Unicode().tag(sync=True)
@@ -63,6 +71,8 @@ def __init__(self, *args, **kwargs):
         self._display_order = False
         self.model_label = "Model"
         self._selected_data_label = None
+        self._spectral_subsets = {}
+        self._window = None
         if self.app.state.settings.get("configuration") == "cubeviz":
             self.cube_fit = True
 
@@ -112,7 +122,9 @@ def _on_viewer_data_changed(self, msg=None):
         viewer = self.app.get_viewer('spectrum-viewer')
 
         self.dc_items = [layer_state.layer.label
-                         for layer_state in viewer.state.layers]
+                         for layer_state in viewer.state.layers
+                         if (not isinstance(layer_state.layer, Subset)
+                             or not isinstance(layer_state.layer.subset_state, RangeSubsetState))]
 
     def _param_units(self, param, order=0):
         """Helper function to handle units that depend on x and y"""
@@ -240,6 +252,41 @@ def vue_data_selected(self, event):
 
         self._spectrum1d = selected_spec
 
+        # Also set the spectral min and max to default to the full range
+        self.selected_subset = "None"
+        self._window = None
+        self.spectral_min = selected_spec.spectral_axis[0].value
+        self.spectral_max = selected_spec.spectral_axis[-1].value
+        self.spectral_unit = str(selected_spec.spectral_axis.unit)
+
+    def vue_list_subsets(self, event):
+        """Populate the spectral subset selection dropdown"""
+        temp_subsets = self.app.get_subsets_from_viewer("spectrum-viewer",
+                                                        subset_type="spectral")
+        temp_dict = {}
+        # Attempt to filter out spatial subsets
+        for key, region in temp_subsets.items():
+            if type(region) == RectanglePixelRegion:
+                temp_dict[key] = region
+        self._spectral_subsets = temp_dict
+        self.spectral_subset_items = ["None"] + sorted(temp_dict.keys())
+
+    @observe("selected_subset")
+    def _on_subset_selected(self, event):
+        # If "None" selected, reset based on bounds of selected data
+        if self.selected_subset == "None":
+            self._window = None
+            self.spectral_min = self._spectrum1d.spectral_axis[0].value
+            self.spectral_max = self._spectrum1d.spectral_axis[-1].value
+        else:
+            spec_sub = self._spectral_subsets[self.selected_subset]
+            unit = u.Unit(self.spectral_unit)
+            spreg = SpectralRegion.from_center(spec_sub.center.x * unit,
+                                               spec_sub.width * unit)
+            self._window = (spreg.lower, spreg.upper)
+            self.spectral_min = spreg.lower.value
+            self.spectral_max = spreg.upper.value
+
     def vue_model_selected(self, event):
         # Add the model selected to the list of models
         self.temp_model = event
@@ -349,7 +396,8 @@ def vue_model_fitting(self, *args, **kwargs):
                 self._spectrum1d,
                 models_to_fit,
                 self.model_equation,
-                run_fitter=True)
+                run_fitter=True,
+                window=self._window)
         except AttributeError:
             msg = SnackbarMessage("Unable to fit: model equation may be invalid",
                                   color="error", sender=self)
@@ -375,7 +423,10 @@ def vue_fit_model_to_cube(self, *args, **kwargs):
 
         if self._warn_if_no_equation():
             return
-        data = self.app.data_collection[self._selected_data_label]
+        if self._selected_data_label in self.app.data_collection.labels:
+            data = self.app.data_collection[self._selected_data_label]
+        else:  # User selected some subset from spectrum viewer, just use original cube
+            data = self.app.data_collection[0]
 
         # First, ensure that the selected data is cube-like. It is possible
         # that the user has selected a pre-existing 1d data object.
@@ -435,7 +486,8 @@ def vue_fit_model_to_cube(self, *args, **kwargs):
                 spec,
                 models_to_fit,
                 self.model_equation,
-                run_fitter=True)
+                run_fitter=True,
+                window=self._window)
         except ValueError:
             snackbar_message = SnackbarMessage(
                 "Cube fitting failed",
@@ -489,7 +541,8 @@ def vue_register_spectrum(self, event):
         else:
             model, spectrum = fit_model_to_spectrum(self._spectrum1d,
                                                     self._initialized_models.values(),
-                                                    self.model_equation)
+                                                    self.model_equation,
+                                                    window=self._window)
 
         self.n_models += 1
         label = self.model_label