From 4c824f6c755968cfc07d5f0fa46048f677995dce Mon Sep 17 00:00:00 2001 From: Ricky O'Steen Date: Mon, 16 Sep 2024 12:58:12 -0400 Subject: [PATCH] Only reshape here in 3D case --- .../default/plugins/model_fitting/model_fitting.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/jdaviz/configs/default/plugins/model_fitting/model_fitting.py b/jdaviz/configs/default/plugins/model_fitting/model_fitting.py index 213976cab5..02f4a23e7b 100644 --- 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