diff --git a/docs/basic_analysis.rst b/docs/basic_analysis.rst index 3d0a30a12..e88d7a47b 100644 --- a/docs/basic_analysis.rst +++ b/docs/basic_analysis.rst @@ -42,6 +42,39 @@ Currently, specutils supports basic signal-to-noise ratio calculations. >>> snr(spec) #doctest:+SKIP +Centroid +-------- + +Currently, specutils supports basic centroid calculations. + +.. code-block:: python + + >>> import numpy as np + >>> import astropy.units as u + >>> from specutils.spectra import Spectrum1D + >>> from astropy.nddata import StdDevUncertainty + >>> from specutils.analysis import centroid + + >>> spec = Spectrum1D(spectral_axis=np.arange(50), flux=(3+np.random.randn(50))*u.Jy) + >>> centroid(spec, region=None) #doctest:+SKIP + + +And if the spectrum contains a continuum, then it should be subtracted first: +.. code-block:: python + + >>> import numpy as np + >>> import astropy.units as u + >>> from specutils.spectra import Spectrum1D + >>> from astropy.nddata import StdDevUncertainty + >>> from specutils.fitting import fit_generic_continuum + >>> from specutils.analysis import centroid + + >>> spec = Spectrum1D(spectral_axis=np.arange(50), flux=(10+np.random.randn(50))*u.Jy) + >>> continuum_baseline = fit_generic_continuum(spec) #doctest:+SKIP + >>> continuum_flux = continuum_baseline(spec.spectral_axis.value) #doctest:+SKIP + >>> continuum = Spectrum1D(spectral_axis=spec.spectral_axis, flux=continuum_flux) #doctest:+SKIP + >>> c = centroid(spec-continuum, region=None) #doctest:+SKIP + Reference/API ------------- .. automodapi:: specutils.analysis diff --git a/specutils/analysis/__init__.py b/specutils/analysis/__init__.py index 47dd37e9d..2514ef9df 100644 --- a/specutils/analysis/__init__.py +++ b/specutils/analysis/__init__.py @@ -2,3 +2,4 @@ from .equivalent_width import * # noqa from .snr import snr # noqa +from .centroid import centroid # noqa diff --git a/specutils/analysis/centroid.py b/specutils/analysis/centroid.py new file mode 100644 index 000000000..736269394 --- /dev/null +++ b/specutils/analysis/centroid.py @@ -0,0 +1,83 @@ +import numpy as np +from ..spectra import SpectralRegion + +__all__ = ['centroid'] + + +def centroid(spectrum, region): + """ + Calculate the centroid of a region, or regions, of the spectrum. + + Parameters + ---------- + spectrum : `~specutils.spectra.spectrum1d.Spectrum1D` + The spectrum object overwhich the centroid will be calculated. + + region: `~specutils.utils.SpectralRegion` or list of `~specutils.utils.SpectralRegion` + Region within the spectrum to calculate the centroid. + + Returns + ------- + centroid : float or list (based on region input) + Centroid of the spectrum or within the regions + + Notes + ----- + The spectrum will need to be continuum subtracted before calling + this method. See the + `analysis documentation `_ for more information. + + """ + + # No region, therefore whole spectrum. + if region is None: + return _centroid_single_region(spectrum) + + # Single region + elif isinstance(region, SpectralRegion): + return _centroid_single_region(spectrum, region=region) + + # List of regions + elif isinstance(region, list): + return [_centroid_single_region(spectrum, region=reg) + for reg in region] + + +def _centroid_single_region(spectrum, region=None): + """ + Calculate the centroid of the spectrum based on the flux and uncertainty + in the spectrum. + + Parameters + ---------- + spectrum : `~specutils.spectra.spectrum1d.Spectrum1D` + The spectrum object overwhich the centroid will be calculated. + + region: `~specutils.utils.SpectralRegion` + Region within the spectrum to calculate the centroid. + + Returns + ------- + centroid : float or list (based on region input) + Centroid of the spectrum or within the regions + + Notes + ----- + This is a helper function for the above `centroid()` method. + + """ + + if region is not None: + calc_spectrum = region.extract(spectrum) + else: + calc_spectrum = spectrum + + flux = calc_spectrum.flux + dispersion = calc_spectrum.spectral_axis + + if len(flux.shape) > 1: + dispersion = np.tile(dispersion, [flux.shape[0], 1]) + + # the axis=-1 will enable this to run on single-dispersion, single-flux + # and single-dispersion, multiple-flux + return np.sum(flux * dispersion, axis=-1) / np.sum(flux, axis=-1) diff --git a/specutils/tests/test_analysis.py b/specutils/tests/test_analysis.py index 9aaa771ce..4b84d70db 100644 --- a/specutils/tests/test_analysis.py +++ b/specutils/tests/test_analysis.py @@ -7,7 +7,7 @@ from ..tests.spectral_examples import simulated_spectra from ..spectra.spectrum1d import Spectrum1D import astropy.units as u -from ..analysis import equivalent_width, snr +from ..analysis import equivalent_width, snr, centroid from ..manipulation import noise_region_uncertainty from ..spectra import SpectralRegion from astropy.nddata import StdDevUncertainty @@ -43,7 +43,7 @@ def test_snr(simulated_spectra): """ np.random.seed(42) - + # # Set up the data and add the uncertainty and calculate the expected SNR # @@ -178,7 +178,7 @@ def test_snr_single_region_with_noise_region(simulated_spectra): region = SpectralRegion(0.52*u.um, 0.59*u.um) noise_region = SpectralRegion(0.40*u.um, 0.45*u.um) - + # # Set up the data # @@ -206,3 +206,66 @@ def test_snr_single_region_with_noise_region(simulated_spectra): assert np.allclose(spec_snr.value, spec_snr_expected.value) +def test_centroid(simulated_spectra): + """ + Test the simple version of the spectral centroid. + """ + + np.random.seed(42) + + # + # Set up the data and add the uncertainty and calculate the expected SNR + # + + spectrum = simulated_spectra.s1_um_mJy_e1 + uncertainty = StdDevUncertainty(0.1*np.random.random(len(spectrum.flux))*u.mJy) + spectrum.uncertainty = uncertainty + + wavelengths = spectrum.spectral_axis + flux = spectrum.flux + + spec_centroid_expected = np.sum(flux * wavelengths) / np.sum(flux) + + # + # SNR of the whole spectrum + # + + spec_centroid = centroid(spectrum, None) + + assert isinstance(spec_centroid, u.Quantity) + assert np.allclose(spec_centroid.value, spec_centroid_expected.value) + + +def test_inverted_centroid(simulated_spectra): + """ + Ensures the centroid calculation also works for *inverted* spectra - i.e. + continuum-subtracted absorption lines. + """ + spectrum = simulated_spectra.s1_um_mJy_e1 + spec_centroid_expected = (np.sum(spectrum.flux * spectrum.spectral_axis) / + np.sum(spectrum.flux)) + + spectrum_inverted = Spectrum1D(spectral_axis=spectrum.spectral_axis, + flux=-spectrum.flux) + spec_centroid_inverted = centroid(spectrum_inverted, None) + assert np.allclose(spec_centroid_inverted.value, spec_centroid_expected.value) + + +def test_centroid_multiple_flux(simulated_spectra): + """ + Test the simple version of the spectral SNR, with multiple flux per single dispersion. + """ + + # + # Set up the data and add the uncertainty and calculate the expected SNR + # + + np.random.seed(42) + + spec = Spectrum1D(spectral_axis=np.arange(10) * u.um, + flux=np.random.sample((5, 10)) * u.Jy) + + centroid_spec = centroid(spec, None) + + assert np.allclose(centroid_spec.value, np.array([4.46190995, 4.17223565, 4.37778249, 4.51595259, 4.7429066])) + assert centroid_spec.unit == u.um