Merge 78691cb into 769e5ad

docs/manipulation.rst

@@ -246,6 +246,40 @@ Or similarly, expressed in pixels:
     >>> spec_w_unc.uncertainty # doctest: +ELLIPSIS
     StdDevUncertainty([0.18714535, ..., 0.18714535])
 
+S/N Threshold Mask
+------------------
+
+It is useful to be able to find all the spaxels in an ND spectrum
+in which the signal to noise ratio is greater than some threshold.
+This method implements this functionality so that a `~specutils.Spectrum1D`
+object, `~specutils.SpectrumCollection` or an :class:`~astropy.nddata.NDData` derived
+object may be passed in as the first parameter. The second parameter
+is a floating point threshold. 
+
+For example, first a spectrum with flux and uncertainty is created, and 
+then call the ``snr_threshold`` method:
+
+.. code-block:: python
+
+    >>> import numpy as np
+    >>> from astropy.nddata import StdDevUncertainty
+    >>> import astropy.units as u
+    >>> from specutils import Spectrum1D
+    >>> from specutils.manipulation import snr_threshold
+    >>> np.random.seed(42)
+    >>> wavelengths = np.arange(0, 10)*u.um
+    >>> flux = 100*np.abs(np.random.randn(10))*u.Jy
+    >>> uncertainty = StdDevUncertainty(np.abs(np.random.randn(10))*u.Jy)
+    >>> spectrum = Spectrum1D(spectral_axis=wavelengths, flux=flux, uncertainty=uncertainty)
+    >>> spectrum_masked = snr_threshold(spectrum, 50) #doctest:+SKIP
+
+The output ``spectrum_masked`` is a shallow copy of the input ``spectrum``
+with the ``mask`` attribute set to True where the S/N is greater than 50
+and False elsewhere.
+
+.. note:: The mask attribute is the only attribute modified by ``snr_threshold()``. To
+             retrieve the masked flux data use ``spectrum.masked.flux_masked``.
+
 Reference/API
 -------------
 

specutils/manipulation/__init__.py

@@ -2,4 +2,5 @@
 from .estimate_uncertainty import *  # noqa
 from .extract_spectral_region import *  # noqa
 from .utils import *  # noqa
+from .manipulation import *  # noqa
 from .resample import  *

specutils/manipulation/manipulation.py

@@ -0,0 +1,56 @@
+"""
+A module for analysis tools dealing with uncertainties or error analysis in
+spectra.
+"""
+
+import copy
+import numpy as np
+
+__all__ = ['snr_threshold']
+
+
+def snr_threshold(spectrum, value):
+    """
+    Calculate the mean S/N of the spectrum based on the flux and uncertainty
+    in the spectrum. This will be calculated over the regions, if they
+    are specified.
+
+    Parameters
+    ----------
+    spectrum : `~specutils.Spectrum1D`, `~specutils.SpectrumCollection` or `~astropy.nddata.NDData`
+        The spectrum object overwhich the S/N threshold will be calculated.
+
+    value: ``float``
+        Threshold value to be applied to flux / uncertainty.
+
+    Returns
+    -------
+    spectrum: `~specutils.Spectrum1D`
+        Ouput object with ``spectrum.mask`` set based on threshold.
+
+    Notes
+    -----
+    The input object will need to have the uncertainty defined in order for the SNR
+    to be calculated.
+
+    """
+
+    if not hasattr(spectrum, 'uncertainty') or spectrum.uncertainty is None:
+        raise Exception("S/N thresholding requires the uncertainty be defined.")
+
+    # Spectrum1D
+    if hasattr(spectrum, 'flux'):
+        data = spectrum.flux
+
+    # NDData
+    elif hasattr(spectrum, 'data'):
+        data = spectrum.data * (spectrum.unit if spectrum.unit is not None else 1)
+    else:
+        raise ValueError('Could not find data attribute.')
+
+    mask = (data / (spectrum.uncertainty.array*spectrum.uncertainty.unit)) > value
+
+    spectrum_out = copy.copy(spectrum)
+    spectrum_out._mask = mask
+
+    return spectrum_out

specutils/spectra/spectrum1d.py

@@ -4,7 +4,7 @@
 import numpy as np
 from astropy import units as u
 from astropy import constants as cnst
-from astropy.nddata import NDDataRef
+from astropy.nddata import NDDataRef, NDUncertainty
 from astropy.utils.decorators import lazyproperty
 
 from ..wcs import WCSAdapter, WCSWrapper

specutils/spectra/spectrum_mixin.py

@@ -90,6 +90,45 @@ def flux(self):
         """
         return u.Quantity(self.data, unit=self.unit, copy=False)
 
+    @property
+    def flux_masked(self):
+        """
+        Converts the stored data and unit information into a quantity masked
+        by the ``mask`` attribute.
+
+        Returns
+        -------
+        `~astropy.units.Quantity`
+            Spectral data as a quantity.
+        """
+
+        data = self.data
+        if self.mask is not None:
+            data[self.mask] = np.nan
+
+        return u.Quantity(data, unit=self.flux.unit, copy=False)
+
+    @property
+    def uncertainty_masked(self):
+        """
+        Converts the stored data and unit information into a quantity masked
+        by the ``mask`` attribute.
+
+        Returns
+        -------
+        `~astropy.units.Quantity`
+            Spectral data as a quantity.
+        """
+
+        if not hasattr(self, 'uncertainty') or self.uncertainty is None:
+            return None
+
+        uncertainty = self.uncertainty
+        if self.mask is not None:
+            uncertainty.array[self.mask] = np.nan
+
+        return uncertainty
+
     def new_flux_unit(self, unit, equivalencies=None, suppress_conversion=False):
         """
         Converts the flux data to the specified unit.  This is an in-place

specutils/tests/test_manipulation.py

@@ -0,0 +1,100 @@
+import pytest
+import numpy as np
+
+import astropy.units as u
+from astropy.modeling import models
+from astropy.nddata import StdDevUncertainty, NDData
+from astropy.tests.helper import quantity_allclose
+from specutils.wcs.wcs_wrapper import WCSWrapper
+
+from ..spectra import Spectrum1D, SpectralRegion, SpectrumCollection
+from ..manipulation import snr_threshold
+
+
+def test_snr_threshold():
+
+    np.random.seed(42)
+
+    # Setup 1D spectrum
+    wavelengths = np.arange(0, 10)*u.um
+    flux = 100*np.abs(np.random.randn(10))*u.Jy
+    uncertainty = StdDevUncertainty(np.abs(np.random.randn(10))*u.Jy)
+    spectrum = Spectrum1D(spectral_axis=wavelengths, flux=flux, uncertainty=uncertainty)
+
+    spectrum_masked = snr_threshold(spectrum, 50)
+
+    assert all([x==y for x,y in zip(spectrum_masked.mask, [True, False, True, True, False, False, True, True, True, False])])
+
+
+    # Setup 3D spectrum
+    np.random.seed(42)
+    wavelengths = np.arange(0, 10)*u.um
+    flux = 100*np.abs(np.random.randn(3, 4, 10))*u.Jy
+    uncertainty = StdDevUncertainty(np.abs(np.random.randn(3, 4, 10))*u.Jy)
+    spectrum = Spectrum1D(spectral_axis=wavelengths, flux=flux, uncertainty=uncertainty)
+
+    spectrum_masked = snr_threshold(spectrum, 50)
+
+    masked_true = np.array([[[ True, False, False,  True, False, False,  True,  True,  True, True],
+        [False,  True, False,  True,  True, False,  True,  True,  True, True],
+        [ True, False, False,  True,  True, False,  True, False,  True, True],
+        [ True,  True, False,  True,  True,  True, False,  True,  True, False]],
+        [[ True, False, False, False,  True,  True,  True,  True,  True, True],
+        [False, False,  True,  True,  True,  True,  True, False,  True, False],
+        [ True, False,  True,  True,  True,  True, False,  True, False, False],
+        [ True,  True, False,  True,  True,  True, False,  True,  True, True]],
+        [[ True,  True,  True, False,  True,  True,  True,  True,  True, False],
+        [False,  True,  True,  True,  True,  True, False,  True, False, True],
+        [ True, False, False, False, False, False,  True, False, False, False],
+        [ True, False,  True,  True, False, False, False,  True,  True, True]]])
+
+    assert all([x==y for x,y in zip(spectrum_masked.mask.ravel(), masked_true.ravel())])
+
+
+    # Setup 3D NDData
+    np.random.seed(42)
+    flux = 100*np.abs(np.random.randn(3, 4, 10))*u.Jy
+    uncertainty = StdDevUncertainty(np.abs(np.random.randn(3, 4, 10))*u.Jy)
+    spectrum = NDData(data=flux, uncertainty=uncertainty)
+
+    spectrum_masked = snr_threshold(spectrum, 50)
+
+    masked_true = np.array([[[ True, False, False,  True, False, False,  True,  True,  True, True],
+        [False,  True, False,  True,  True, False,  True,  True,  True, True],
+        [ True, False, False,  True,  True, False,  True, False,  True, True],
+        [ True,  True, False,  True,  True,  True, False,  True,  True, False]],
+        [[ True, False, False, False,  True,  True,  True,  True,  True, True],
+        [False, False,  True,  True,  True,  True,  True, False,  True, False],
+        [ True, False,  True,  True,  True,  True, False,  True, False, False],
+        [ True,  True, False,  True,  True,  True, False,  True,  True, True]],
+        [[ True,  True,  True, False,  True,  True,  True,  True,  True, False],
+        [False,  True,  True,  True,  True,  True, False,  True, False, True],
+        [ True, False, False, False, False, False,  True, False, False, False],
+        [ True, False,  True,  True, False, False, False,  True,  True, True]]])
+
+    assert all([x==y for x,y in zip(spectrum_masked.mask.ravel(), masked_true.ravel())])
+
+
+    # Test SpectralCollection
+    np.random.seed(42)
+    flux = u.Quantity(np.random.sample((5, 10)), unit='Jy')
+    spectral_axis = u.Quantity(np.arange(50).reshape((5, 10)), unit='AA')
+    wcs = np.array([WCSWrapper.from_array(x).wcs for x in spectral_axis])
+    uncertainty = StdDevUncertainty(np.random.sample((5, 10)), unit='Jy')
+    mask = np.ones((5, 10)).astype(bool)
+    meta = [{'test': 5, 'info': [1, 2, 3]} for i in range(5)]
+
+    spec_coll = SpectrumCollection(
+        flux=flux, spectral_axis=spectral_axis, wcs=wcs,
+        uncertainty=uncertainty, mask=mask, meta=meta)
+
+    spec_coll_masked = snr_threshold(spec_coll, 3)
+    print(spec_coll_masked.mask)
+
+    ma = np.array([[False, False, False, False, False, False, False,  True,  True, False],
+                  [False,  True, False, False, False, False, False, False,  True, False],
+                  [False, False,  True, False, False, False, False,  True, False, False],
+                  [False, False, False,  True,  True, False, False, False, False, False],
+                  [False, False, False, False, False, False, False, False,  True, False]])
+
+    assert all([x==y for x,y in zip(spec_coll_masked.mask.ravel(), ma.ravel())])

specutils/tests/test_spectrum1d.py

@@ -368,3 +368,51 @@ def test_str():
     spec_single_flux = Spectrum1D(1 * u.Jy)
     result = str(spec_single_flux)
     assert result == 'Spectrum1D (length=1)\nflux:   {}'.format(spec_single_flux.flux)
+
+
+def test_flux_mask():
+
+    np.random.seed(42)
+    data = np.random.randn(50)
+    mask = np.random.randn(50) > 0.5
+    spec = Spectrum1D(spectral_axis=np.arange(50) * u.AA,
+                      flux=data * u.Jy,
+                      mask=mask)
+
+    flux_masked = np.array([0.49671415, -0.1382643 , 0.64768854, np.nan, np.nan,
+                            np.nan,  1.57921282,  0.76743473, -0.46947439, np.nan,
+                            -0.46341769, -0.46572975, 0.24196227, -1.91328024, np.nan,
+                            np.nan, -1.01283112, np.nan, -0.90802408, -1.4123037 ,
+                            1.46564877, np.nan,  0.0675282 , np.nan, -0.54438272,
+                            np.nan, -1.15099358,  0.37569802, -0.60063869, -0.29169375,
+                            -0.60170661,  1.85227818, np.nan, -1.05771093,  0.82254491,
+                            -1.22084365, np.nan, -1.95967012, -1.32818605, np.nan,
+                            0.73846658, np.nan, -0.11564828, -0.3011037 , -1.47852199,
+                            -0.71984421, -0.46063877,  1.05712223,  0.34361829, -1.76304016])
+
+    assert np.allclose(flux_masked, spec.flux_masked.value, equal_nan=True)
+    assert spec.flux_masked.unit == u.Jy
+
+
+    np.random.seed(42)
+    uncertainty = np.random.randn(50)
+    mask = np.random.randn(50) > 0.5
+    data = np.random.randn(50)
+    spec = Spectrum1D(spectral_axis=np.arange(50) * u.AA,
+                      flux=data * u.Jy,
+                      uncertainty=StdDevUncertainty(uncertainty * u.Jy),
+                      mask=mask)
+
+    uncertainty_masked = np.array([0.49671415, -0.1382643 , 0.64768854, np.nan, np.nan,
+                            np.nan,  1.57921282,  0.76743473, -0.46947439, np.nan,
+                            -0.46341769, -0.46572975, 0.24196227, -1.91328024, np.nan,
+                            np.nan, -1.01283112, np.nan, -0.90802408, -1.4123037 ,
+                            1.46564877, np.nan,  0.0675282 , np.nan, -0.54438272,
+                            np.nan, -1.15099358,  0.37569802, -0.60063869, -0.29169375,
+                            -0.60170661,  1.85227818, np.nan, -1.05771093,  0.82254491,
+                            -1.22084365, np.nan, -1.95967012, -1.32818605, np.nan,
+                            0.73846658, np.nan, -0.11564828, -0.3011037 , -1.47852199,
+                            -0.71984421, -0.46063877,  1.05712223,  0.34361829, -1.76304016])
+
+    assert np.allclose(uncertainty_masked, spec.uncertainty_masked.array, equal_nan=True)
+    assert spec.uncertainty_masked.unit == u.Jy