-
-
Notifications
You must be signed in to change notification settings - Fork 124
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
5 changed files
with
242 additions
and
1 deletion.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,80 @@ | ||
""" | ||
A module for analysis tools dealing with uncertainties or error analysis in | ||
spectra. | ||
""" | ||
|
||
import copy | ||
import numpy as np | ||
import operator | ||
|
||
__all__ = ['snr_threshold'] | ||
|
||
|
||
def snr_threshold(spectrum, value, op=operator.gt): | ||
""" | ||
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. | ||
op: One of operator.gt, operator.ge, operator.lt, operator.le or | ||
the str equivalent '>', '>=', '<', '<=' | ||
The mathematical operator to apply for thresholding. | ||
Returns | ||
------- | ||
spectrum: `~specutils.Spectrum1D` | ||
Output 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. | ||
""" | ||
|
||
# Setup the mapping | ||
operator_mapping = { | ||
'>': operator.gt, | ||
'<': operator.lt, | ||
'>=': operator.ge, | ||
'<=': operator.le | ||
} | ||
|
||
if not hasattr(spectrum, 'uncertainty') or spectrum.uncertainty is None: | ||
raise Exception("S/N thresholding requires the uncertainty be defined.") | ||
|
||
if (op not in [operator.gt, operator.ge, operator.lt, operator.le] and | ||
op not in operator_mapping.keys()): | ||
raise ValueError('Threshold operator must be a string or operator that represents ' + | ||
'greater-than, less-than, greater-than-or-equal or ' + | ||
'less-than-or-equal') | ||
|
||
# If the operator passed in is a string, then map to the | ||
# operator method. | ||
if isinstance(op, str): | ||
op = operator_mapping[op] | ||
|
||
# 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 = op(data / (spectrum.uncertainty.quantity), value) | ||
|
||
spectrum_out = copy.copy(spectrum) | ||
spectrum_out._mask = mask | ||
|
||
return spectrum_out |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,123 @@ | ||
import operator | ||
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])]) | ||
|
||
spectrum_masked = snr_threshold(spectrum, 50, operator.gt) | ||
assert all([x==y for x,y in zip(spectrum_masked.mask, [True, False, True, True, False, False, True, True, True, False])]) | ||
|
||
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])]) | ||
|
||
spectrum_masked = snr_threshold(spectrum, 50, operator.ge) | ||
assert all([x==y for x,y in zip(spectrum_masked.mask, [True, False, True, True, False, False, True, True, True, False])]) | ||
|
||
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])]) | ||
|
||
spectrum_masked = snr_threshold(spectrum, 50, operator.lt) | ||
assert all([not x==y for x,y in zip(spectrum_masked.mask, [True, False, True, True, False, False, True, True, True, False])]) | ||
|
||
spectrum_masked = snr_threshold(spectrum, 50, '<') | ||
assert all([not x==y for x,y in zip(spectrum_masked.mask, [True, False, True, True, False, False, True, True, True, False])]) | ||
|
||
spectrum_masked = snr_threshold(spectrum, 50, operator.le) | ||
assert all([not x==y for x,y in zip(spectrum_masked.mask, [True, False, True, True, False, False, True, True, True, False])]) | ||
|
||
spectrum_masked = snr_threshold(spectrum, 50, '<=') | ||
assert all([not 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())]) |