Centroid calculation

docs/basic_analysis.rst

@@ -42,6 +42,39 @@ Currently, specutils supports basic signal-to-noise ratio calculations.
     >>> snr(spec) #doctest:+SKIP
     <Quantity 149.97247134>
 
+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
+    <Quantity 24.39045495 Angstrom>
+
+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

specutils/analysis/__init__.py

@@ -2,3 +2,4 @@
 
 from .equivalent_width import *  # noqa
 from .snr import snr  # noqa
+from .centroid import centroid  # noqa

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 <https://specutils.readthedocs.io/en/latest/basic_analysis.html>`_ 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)

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