From d26971f50716fe9210756bf251d28720adcbc21f Mon Sep 17 00:00:00 2001 From: Ricky O'Steen <39831871+rosteen@users.noreply.github.com> Date: Fri, 13 Sep 2024 10:41:35 -0400 Subject: [PATCH] Fix NaN handling in cube fitting (#3191) * Fix NaN handling in cube fitting Remove debugging prints, add comment for context Codestyle, changelog Remove unit conversion stuff to just fix NaN handling * Move changelog * Add test * Codestyle * Adding test for spectrum with existing mask * Fix a couple bugs with a loading a masked cube * Linking fix for case with extra components like mask * Test mask addition behavior with a subset instead * Codestyle * Ignore UserWarning for Ubuntu tests * Update jdaviz/configs/default/plugins/model_fitting/tests/test_fitting.py Co-authored-by: P. L. Lim <2090236+pllim@users.noreply.github.com> * Do this as in above test --------- Co-authored-by: P. L. Lim <2090236+pllim@users.noreply.github.com> --- CHANGES.rst | 3 +- jdaviz/app.py | 5 ++- .../spectral_extraction.py | 3 ++ .../plugins/model_fitting/initializers.py | 4 +- .../plugins/model_fitting/model_fitting.py | 8 +++- .../model_fitting/tests/test_fitting.py | 39 +++++++++++++++++++ 6 files changed, 56 insertions(+), 6 deletions(-) diff --git a/CHANGES.rst b/CHANGES.rst index aa9eabf67f..1c86f1c661 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -142,7 +142,6 @@ Cubeviz - No longer incorrectly swap RA and Dec axes when loading Spectrum1D objects. [#3133] - Imviz ^^^^^ @@ -239,6 +238,8 @@ Bug Fixes Cubeviz ^^^^^^^ +- Fixed fitting a model to the entire cube when NaNs are present. [#3191] + Imviz ^^^^^ diff --git a/jdaviz/app.py b/jdaviz/app.py index 160eb9fae0..6df7f76217 100644 --- a/jdaviz/app.py +++ b/jdaviz/app.py @@ -735,8 +735,9 @@ def _link_new_data(self, reference_data=None, data_to_be_linked=None): return elif self.config == 'cubeviz' and linked_data.ndim == 1: - ref_wavelength_component = dc[0].components[-2] - ref_flux_component = dc[0].components[-1] + # Don't want to use negative indices in case there are extra components like a mask + ref_wavelength_component = dc[0].components[5] + ref_flux_component = dc[0].components[6] linked_wavelength_component = dc[-1].components[1] linked_flux_component = dc[-1].components[-1] diff --git a/jdaviz/configs/cubeviz/plugins/spectral_extraction/spectral_extraction.py b/jdaviz/configs/cubeviz/plugins/spectral_extraction/spectral_extraction.py index b82acec4f7..b649e9f06a 100644 --- a/jdaviz/configs/cubeviz/plugins/spectral_extraction/spectral_extraction.py +++ b/jdaviz/configs/cubeviz/plugins/spectral_extraction/spectral_extraction.py @@ -471,6 +471,9 @@ def _extract_from_aperture(self, cube, uncert_cube, aperture, wcs = cube.meta['_orig_spec'].wcs.spectral elif hasattr(cube.coords, 'spectral'): wcs = cube.coords.spectral + elif hasattr(cube.coords, 'spectral_wcs'): + # This is the attribute for a PaddedSpectrumWCS in the 3D case + wcs = cube.coords.spectral_wcs else: wcs = None diff --git a/jdaviz/configs/default/plugins/model_fitting/initializers.py b/jdaviz/configs/default/plugins/model_fitting/initializers.py index c7aa327e71..1548873927 100644 --- a/jdaviz/configs/default/plugins/model_fitting/initializers.py +++ b/jdaviz/configs/default/plugins/model_fitting/initializers.py @@ -119,7 +119,7 @@ def initialize(self, instance, x, y): instance: `~astropy.modeling.Model` The initialized model. """ - y_mean = np.mean(y) + y_mean = np.nanmean(y) x_range = x[-1] - x[0] position = x_range / 2.0 + x[0] @@ -190,7 +190,7 @@ def initialize(self, instance, x, y): # width can be estimated by the weighted # 2nd moment of the X coordinate. - dx = x - np.mean(x) + dx = x - np.nanmean(x) fwhm = 2 * np.sqrt(np.sum((dx * dx) * y) / np.sum(y)) # amplitude is derived from area. diff --git a/jdaviz/configs/default/plugins/model_fitting/model_fitting.py b/jdaviz/configs/default/plugins/model_fitting/model_fitting.py index 64d2f68571..d5f2aa8155 100644 --- a/jdaviz/configs/default/plugins/model_fitting/model_fitting.py +++ b/jdaviz/configs/default/plugins/model_fitting/model_fitting.py @@ -493,7 +493,7 @@ def _initialize_model_component(self, model_comp, comp_label, poly_order=None): masked_spectrum.flux[~mask] if mask is not None else masked_spectrum.flux) # need to loop over parameters again as the initializer may have overridden - # the original default value + # the original default value. for param_name in get_model_parameters(model_cls, new_model["model_kwargs"]): param_quant = getattr(initialized_model, param_name) new_model["parameters"].append({"name": param_name, @@ -917,6 +917,12 @@ def _fit_model_to_cube(self, add_data): # Apply masks from selected spectral subset spec = self._apply_subset_masks(spec, self.spectral_subset) + # Also mask out NaNs for fitting. Simply adding filter_non_finite to the cube fit + # didn't work out of the box, so doing this for now. + if spec.mask is None: + spec.mask = np.isnan(spec.flux) + else: + spec.mask = spec.mask | np.isnan(spec.flux) try: fitted_model, fitted_spectrum = fit_model_to_spectrum( diff --git a/jdaviz/configs/default/plugins/model_fitting/tests/test_fitting.py b/jdaviz/configs/default/plugins/model_fitting/tests/test_fitting.py index cbddd12c59..9044ad5a4c 100644 --- a/jdaviz/configs/default/plugins/model_fitting/tests/test_fitting.py +++ b/jdaviz/configs/default/plugins/model_fitting/tests/test_fitting.py @@ -8,6 +8,7 @@ from astropy.nddata import StdDevUncertainty from astropy.tests.helper import assert_quantity_allclose from astropy.wcs import WCS +from glue.core.roi import XRangeROI from numpy.testing import assert_allclose, assert_array_equal from specutils.spectra import Spectrum1D @@ -378,3 +379,41 @@ def test_incompatible_units(specviz_helper, spectrum1d): with warnings.catch_warnings(): warnings.filterwarnings('ignore', message='Model is linear in parameters.*') mf.calculate_fit(add_data=True) + + +def test_cube_fit_with_nans(cubeviz_helper): + flux = np.ones((7, 8, 9)) * u.nJy + flux[:, :, 0] = np.nan + spec = Spectrum1D(flux=flux) + cubeviz_helper.load_data(spec, data_label="test") + + mf = cubeviz_helper.plugins["Model Fitting"] + mf.cube_fit = True + mf.create_model_component("Const1D") + with warnings.catch_warnings(): + warnings.filterwarnings('ignore', message='Model is linear in parameters.*') + mf.calculate_fit() + result = cubeviz_helper.app.data_collection['model'] + assert np.all(result.get_component("flux").data == 1) + + +def test_cube_fit_with_subset_and_nans(cubeviz_helper): + # Also test with existing mask + flux = np.ones((7, 8, 9)) * u.nJy + flux[:, :, 0] = np.nan + spec = Spectrum1D(flux=flux) + spec.flux[5, 5, 7] = 10 * u.nJy + cubeviz_helper.load_data(spec, data_label="test") + + sv = cubeviz_helper.app.get_viewer('spectrum-viewer') + sv.apply_roi(XRangeROI(0, 5)) + + mf = cubeviz_helper.plugins["Model Fitting"] + mf.cube_fit = True + mf.spectral_subset = 'Subset 1' + mf.create_model_component("Const1D") + with warnings.catch_warnings(): + warnings.filterwarnings('ignore', message='Model is linear in parameters.*') + mf.calculate_fit() + result = cubeviz_helper.app.data_collection['model'] + assert np.all(result.get_component("flux").data == 1)