diff --git a/specutils/fitting/fitmodels.py b/specutils/fitting/fitmodels.py index a7cc8e37c..f94e86fc4 100644 --- a/specutils/fitting/fitmodels.py +++ b/specutils/fitting/fitmodels.py @@ -286,10 +286,12 @@ def fit_lines(spectrum, model, fitter=fitting.LevMarLSQFitter(), Fitter instance to be used when fitting model to spectrum. exclude_regions : list of `~specutils.SpectralRegion` List of regions to exclude in the fitting. - weights : list or 'unc', optional + weights : array-like or 'unc', optional If 'unc', the unceratinties from the spectrum object are used to - to calculate the weights. If list/ndarray, represents the weights to - use in the fitting. + to calculate the weights. If array-like, represents the weights to + use in the fitting. Note that if a mask is present on the spectrum, it + will be applied to the ``weights`` as it would be to the spectrum + itself. window : `~specutils.SpectralRegion` or list of `~specutils.SpectralRegion` Regions of the spectrum to use in the fitting. If None, then the whole spectrum will be used in the fitting. @@ -403,12 +405,15 @@ def _fit_lines(spectrum, model, fitter=fitting.LevMarLSQFitter(), # Assume that the weights argument is list-like weights = np.array(weights) + mask = spectrum.mask + dispersion = spectrum.spectral_axis dispersion_unit = spectrum.spectral_axis.unit flux = spectrum.flux flux_unit = spectrum.flux.unit + # # Determine the window if it is not None. There # are several options here: @@ -422,40 +427,51 @@ def _fit_lines(spectrum, model, fitter=fitting.LevMarLSQFitter(), # # In this case the window defines the area around the center of each model + window_indices = None if window is not None and isinstance(window, (float, int)): center = model.mean - indices = np.nonzero((spectrum.spectral_axis >= center-window) & - (spectrum.spectral_axis < center+window)) - - dispersion = dispersion[indices] - flux = flux[indices] - - if weights is not None: - weights = weights[indices] + window_indices = np.nonzero((spectrum.spectral_axis >= center-window) & + (spectrum.spectral_axis < center+window)) # In this case the window is the start and end points of where we # should fit elif window is not None and isinstance(window, tuple): - indices = np.nonzero((dispersion >= window[0]) & + window_indices = np.nonzero((dispersion >= window[0]) & (dispersion < window[1])) - dispersion = dispersion[indices] - flux = flux[indices] - - if weights is not None: - weights = weights[indices] - + # in this case the window is spectral regions that determine where + # to fit. elif window is not None and isinstance(window, SpectralRegion): - try: - idx1, idx2 = window.bounds - if idx1 == idx2: - raise Exception("Bad selected region.") - extracted_regions = extract_region(spectrum, window) - dispersion, flux = _combined_region_data(extracted_regions) - dispersion = dispersion * dispersion_unit - flux = flux * flux_unit - except ValueError as e: - return + idx1, idx2 = window.bounds + if idx1 == idx2: + raise IndexError("Tried to fit a region containing no pixels.") + + # HACK WARNING! This uses the extract machinery to create a set of + # indices by making an "index spectrum" + # note that any unit will do but Jy is at least flux-y + # TODO: really the spectral region machinery should have the power + # to create a mask, and we'd just use that... + idxarr = np.arange(spectrum.flux.size).reshape(spectrum.flux.shape) + index_spectrum = Spectrum1D(spectral_axis=spectrum.spectral_axis, + flux=u.Quantity(idxarr, u.Jy, dtype=int)) + + extracted_regions = extract_region(index_spectrum, window) + if isinstance(extracted_regions, list): + if len(extracted_regions) == 0: + raise ValueError('The whole spectrum is windowed out!') + window_indices = np.concatenate([s.flux.value.astype(int) for s in extracted_regions]) + else: + if len(extracted_regions.flux) == 0: + raise ValueError('The whole spectrum is windowed out!') + window_indices = extracted_regions.flux.value.astype(int) + + if window_indices is not None: + dispersion = dispersion[window_indices] + flux = flux[window_indices] + if mask is not None: + mask = mask[window_indices] + if weights is not None: + weights = weights[window_indices] if flux is None or len(flux) == 0: raise Exception("Spectrum flux is empty or None.") @@ -482,6 +498,12 @@ def _fit_lines(spectrum, model, fitter=fitting.LevMarLSQFitter(), # # Do the fitting of spectrum to the model. # + if mask is not None: + nmask = ~mask + dispersion_unitless = dispersion_unitless[nmask] + flux_unitless = flux_unitless[nmask] + if weights is not None: + weights = weights[nmask] fit_model_unitless = fitter(model_unitless, dispersion_unitless, flux_unitless, weights=weights, **kwargs) @@ -500,26 +522,6 @@ def _fit_lines(spectrum, model, fitter=fitting.LevMarLSQFitter(), return fit_model -def _combined_region_data(spec): - if isinstance(spec, list): - - # Merge sub-spec spectral_axis and flux values. - x = np.array([sv for subspec in spec if subspec is not None - for sv in subspec.spectral_axis.value]) - y = np.array([sv for subspec in spec if subspec is not None - for sv in subspec.flux.value]) - else: - if spec is None: - return - x = spec.spectral_axis.value - y = spec.flux.value - - if len(x) == 0: - return - - return x, y - - def _convert(quantity, dispersion_unit, dispersion, flux_unit): """ Convert the quantity to the spectrum's units, and then we will use diff --git a/specutils/tests/spectral_examples.py b/specutils/tests/spectral_examples.py index f188173ad..271d397cf 100644 --- a/specutils/tests/spectral_examples.py +++ b/specutils/tests/spectral_examples.py @@ -1,3 +1,5 @@ +from copy import copy + import numpy as np import astropy.units as u from astropy.modeling import models @@ -32,6 +34,9 @@ class SpectraExamples: 4. s1_AA_nJy_e3 - same as 1, but with a fourth instance of noise dispersion: Angstroms, flux: nJy + + 5. s1_um_mJy_e1_masked - same as 1, but with a random set of pixels + masked. """ def __init__(self): @@ -71,7 +76,7 @@ def __init__(self): flux=self._flux_e2 * u.mJy) # - # Create on spectrum with the same flux but in angstrom units + # Create one spectrum with the same flux but in angstrom units # self.wavelengths_AA = self.wavelengths_um * 10000 @@ -79,7 +84,7 @@ def __init__(self): flux=self._flux_e1 * u.mJy) # - # Create on spectrum with the same flux but in angstrom units and nJy + # Create one spectrum with the same flux but in angstrom units and nJy # self._flux_e4 = (self.base_flux + 400 * np.random.random(self.base_flux.shape)) * 1000000 @@ -87,6 +92,13 @@ def __init__(self): flux=self._flux_e4 * u.nJy) + # + # Create one spectrum like 1 but with a mask + # + self._s1_um_mJy_e1_masked = copy(self._s1_um_mJy_e1) # SHALLOW copy - the data are shared with the above non-masked case + self._s1_um_mJy_e1_masked.mask = (np.random.randn(*self.base_flux.shape) + 1) > 0 + + @property def s1_um_mJy_e1(self): return self._s1_um_mJy_e1 @@ -119,6 +131,11 @@ def s1_AA_nJy_e4(self): def s1_AA_nJy_e4_flux(self): return self._flux_e4 + @property + def s1_um_mJy_e1_masked(self): + return self._s1_um_mJy_e1_masked + + @pytest.fixture def simulated_spectra(): diff --git a/specutils/tests/test_arithmetic.py b/specutils/tests/test_arithmetic.py index b8639e156..490148745 100644 --- a/specutils/tests/test_arithmetic.py +++ b/specutils/tests/test_arithmetic.py @@ -93,3 +93,15 @@ def test_add_diff_spectral_axis(simulated_spectra): # Calculate using the spectrum1d/nddata code spec3 = simulated_spectra.s1_um_mJy_e1 + simulated_spectra.s1_AA_mJy_e3 + + +def test_masks(simulated_spectra): + masked_spec = simulated_spectra.s1_um_mJy_e1_masked + + masked_sum = masked_spec + masked_spec + assert np.all(masked_sum.mask == simulated_spectra.s1_um_mJy_e1_masked.mask) + + masked_sum.mask[:50] = True + masked_diff = masked_sum - masked_spec + assert u.allclose(masked_diff.flux, masked_spec.flux) + assert np.all(masked_diff.mask == masked_sum.mask | masked_spec.mask) diff --git a/specutils/tests/test_fitting.py b/specutils/tests/test_fitting.py index d1fc21c30..1341e4a17 100644 --- a/specutils/tests/test_fitting.py +++ b/specutils/tests/test_fitting.py @@ -617,3 +617,55 @@ def test_spectrum_from_model(): 1.5319218, 1.06886794, 0.71101768, 0.45092638, 0.27264641]) assert np.allclose(spectrum_gaussian.flux.value[::10], flux_expected, atol=1e-5) + + +def test_masking(): + """ + Test fitting spectra with masks + """ + wl, flux = double_peak() + s = Spectrum1D(flux=flux*u.Jy, spectral_axis=wl*u.um) + + # first we fit a single gaussian to the double_peak model, using the + # known-good second peak (but a bit higher in amplitude). It should lock + # in on the *second* peak since it's already close: + g_init = models.Gaussian1D(2.5, 5.5, 0.2) + g_fit1 = fit_lines(s, g_init) + assert u.allclose(g_fit1.mean, 5.5, atol=.1) + + # now create a spectrum where the region around the second peak is masked. + # The fit should now go to the *first* peak + s_msk = Spectrum1D(flux=flux*u.Jy, spectral_axis=wl*u.um, mask=(5.1 < wl)&(wl < 6.1)) + g_fit2 = fit_lines(s_msk, g_init) + assert u.allclose(g_fit2.mean, 4.6, atol=.1) + + # double check that it works with weights as well + g_fit3 = fit_lines(s_msk, g_init, weights=np.ones_like(s_msk.flux.value)) + assert g_fit2.mean == g_fit3.mean + + +def test_window_extras(): + """ + Test that fitting works with masks and weights when a window is present + """ + # similar to the masking test, but add a broad window around the whole thing + wl, flux = double_peak() + g_init = models.Gaussian1D(2.5, 5.5, 0.2) + window_region = SpectralRegion(4*u.um, 8*u.um) + mask = (5.1 < wl) & (wl < 6.1) + + s_msk = Spectrum1D(flux=flux*u.Jy, spectral_axis=wl*u.um, mask=mask) + + g_fit1 = fit_lines(s_msk, g_init, window=window_region) + assert u.allclose(g_fit1.mean, 4.6, atol=.1) + + # check that if we weight instead of masking, we get the same result + s = Spectrum1D(flux=flux*u.Jy, spectral_axis=wl*u.um) + weights = (~mask).astype(float) + g_fit2 = fit_lines(s, g_init, weights=weights, window=window_region) + assert u.allclose(g_fit2.mean, 4.6, atol=.1) + + # and the same with both together + weights = (~mask).astype(float) + g_fit3 = fit_lines(s_msk, g_init, weights=weights, window=window_region) + assert u.allclose(g_fit3.mean, 4.6, atol=.1) diff --git a/specutils/tests/test_region_extract.py b/specutils/tests/test_region_extract.py index 7f3a82f57..d7012a1d6 100644 --- a/specutils/tests/test_region_extract.py +++ b/specutils/tests/test_region_extract.py @@ -174,3 +174,17 @@ def test_linear_excise_invert_from_spectrum(): np.diff(excised_spec[51:61].flux)) assert quantity_allclose(np.diff(excised_spec[80:90].flux), np.diff(excised_spec[81:91].flux)) + + +def test_extract_masked(): + wl = [1, 2, 3, 4]*u.nm + flux = np.arange(4)*u.Jy + mask = [False, False, True, True] + + masked_spec = Spectrum1D(spectral_axis=wl, flux=flux, mask=mask) + region = SpectralRegion(1.5 * u.nm, 3.5 * u.nm) + + extracted = extract_region(masked_spec, region) + + assert np.all(extracted.mask == [False, True]) + assert np.all(extracted.flux.value == [1, 2])