Merge 28ea2c9 into 7eaeb57

specutils/analysis/__init__.py

@@ -2,3 +2,4 @@
 from .uncertainty import *  # noqa
 from .location import *  # noqa
 from .width import *  # noqa
+from .higherlevel import *  # noqa

specutils/analysis/higherlevel.py

@@ -0,0 +1,45 @@
+"""
+A module for analysis tools dealing with uncertainties or error analysis in
+spectra.
+"""
+
+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.spectra.spectrum1d.Spectrum1D`
+        The spectrum object overwhich the equivalent width will be calculated.
+
+    value: ``float``
+        Threshold value to be applied to flux / uncertainty.
+
+    Returns
+    -------
+    spectrum_masked : `~astropy.units.Quantity` or list (based on region input)
+        Signal to noise ratio of the spectrum or within the regions
+
+    Notes
+    -----
+    The spectrum will need to have the uncertainty defined in order for the SNR
+    to be calculated. If the goal is instead signal to noise *per pixel*, this
+    should be computed directly as ``spectrum.flux / spectrum.uncertainty``.
+
+    """
+
+    if not hasattr(spectrum, 'uncertainty') or spectrum.uncertainty is None:
+        raise Exception("S/N thresholding requires the uncertainty be defined.")
+
+    mask = (spectrum.flux / (spectrum.uncertainty.array*spectrum.uncertainty.unit)) > value
+
+    spectrum.mask = mask
+
+    return spectrum

specutils/tests/test_higherlevel.py

@@ -0,0 +1,50 @@
+import pytest
+import numpy as np
+
+import astropy.units as u
+from astropy.modeling import models
+from astropy.nddata import StdDevUncertainty
+from astropy.tests.helper import quantity_allclose
+
+from ..spectra import Spectrum1D, SpectralRegion
+from ..analysis import (snr_threshold)
+
+
+def test_snr_threshold():
+
+    np.random.seed(42)
+
+    # Setup 1D spectrum
+    wavelengths = np.linspace(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())])