Handle linear component estimation for cube case

spacetelescope · Sep 16, 2024 · 2601799 · 2601799
1 parent 37ac6c7
commit 2601799
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 3 deletions.
diff --git a/jdaviz/configs/default/plugins/model_fitting/initializers.py b/jdaviz/configs/default/plugins/model_fitting/initializers.py
@@ -79,7 +79,7 @@ def initialize(self, instance, x, y):
         """
         if y.ndim == 3:
             # For cube fitting, need to collapse before this calculation
-            y = np.nanmean(y, axis=(0,1))
+            y = np.nanmean(y, axis=(0, 1))
         slope, intercept = np.polynomial.Polynomial.fit(x.value.flatten(), y.value.flatten(), 1)
 
         instance.slope.value = slope

diff --git a/jdaviz/configs/default/plugins/model_fitting/model_fitting.py b/jdaviz/configs/default/plugins/model_fitting/model_fitting.py
@@ -509,12 +509,25 @@ def _initialize_model_component(self, model_comp, comp_label, poly_order=None):
             masked_spectrum = self._apply_subset_masks(self.dataset.selected_spectrum,
                                                        self.spectral_subset)
         mask = masked_spectrum.mask
+        if mask is not None:
+            if mask.ndim == 3:
+                spectral_mask = mask.all(axis=(0, 1))
+            else:
+                spectral_mask = mask
+            init_x = masked_spectrum.spectral_axis[~spectral_mask]
+            orig_flux_shape = masked_spectrum.flux.shape
+            init_y = masked_spectrum.flux[~mask].reshape(orig_flux_shape[0],
+                                                         orig_flux_shape[1],
+                                                         len(init_x))
+        else:
+            init_x = masked_spectrum.spectral_axis
+            init_y = masked_spectrum.flux
+
         initialized_model = initialize(
             MODELS[model_comp](name=comp_label,
                                **initial_values,
                                **new_model.get("model_kwargs", {})),
-            masked_spectrum.spectral_axis[~mask] if mask is not None else masked_spectrum.spectral_axis,  # noqa
-            masked_spectrum.flux[~mask] if mask is not None else masked_spectrum.flux)
+            init_x, init_y)
 
         # need to loop over parameters again as the initializer may have overridden
         # the original default value. However, if we toggled cube_fit, we may need to override