Merge pull request #519 from eteq/mask-fitting

Make line/continuum fitting work in the mask case (and other mask-related tests)

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

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,22 +76,29 @@ 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
         self._s1_AA_mJy_e3 = Spectrum1D(spectral_axis=self.wavelengths_AA * u.AA,
                                         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
         self._s1_AA_nJy_e4 = Spectrum1D(spectral_axis=self.wavelengths_AA * u.AA,
                                         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():

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)

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)

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])