From 2601799f3fa24ec1065ad7bf178c68c2bd44944e Mon Sep 17 00:00:00 2001 From: Ricky O'Steen Date: Mon, 16 Sep 2024 12:43:50 -0400 Subject: [PATCH] Handle linear component estimation for cube case --- .../plugins/model_fitting/initializers.py | 2 +- .../plugins/model_fitting/model_fitting.py | 17 +++++++++++++++-- 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/jdaviz/configs/default/plugins/model_fitting/initializers.py b/jdaviz/configs/default/plugins/model_fitting/initializers.py index be6b55ec4b..0ec0d1c2b4 100644 --- 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 index 5288b0933e..213976cab5 100644 --- 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