Only reshape here in 3D case

spacetelescope · Sep 16, 2024 · 4c824f6 · 4c824f6
1 parent 2601799
commit 4c824f6
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/jdaviz/configs/default/plugins/model_fitting/model_fitting.py b/jdaviz/configs/default/plugins/model_fitting/model_fitting.py
@@ -516,9 +516,11 @@ def _initialize_model_component(self, model_comp, comp_label, poly_order=None):
                 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))
+            init_y = masked_spectrum.flux[~mask]
+            if mask.ndim == 3:
+                init_y = init_y.reshape(orig_flux_shape[0],
+                                        orig_flux_shape[1],
+                                        len(init_x))
         else:
             init_x = masked_spectrum.spectral_axis
             init_y = masked_spectrum.flux