Source code for specutils.analysis.moment

"""
A module for analysis tools focused on determining the moment of
spectral features.
"""

import numpy as np
from ..manipulation import extract_region
from ..spectra import SpectrumCollection
from .utils import computation_wrapper


__all__ = ['moment']


[docs] def moment(spectrum, regions=None, order=0, axis='spectral'): """ Estimate the moment of the spectrum. Parameters ---------- spectrum : `~specutils.spectra.spectrum.Spectrum` The spectrum object over which the width will be calculated. regions: `~specutils.SpectralRegion` or list of `~specutils.SpectralRegion` Region within the spectrum to calculate the gaussian sigma width. If regions is `None`, computation is performed over entire spectrum. order : int The order of the moment to be calculated. Default=0 axis : int, str Axis along which a moment is calculated. Default='spectral', which computes along the spectral axis. Returns ------- moment: `float` or list (based on region input) Moment of the spectrum. Returns None if (order < 0 or None) """ if isinstance(spectrum, SpectrumCollection): return [computation_wrapper(_compute_moment, spec, regions,order=order, axis=axis) for spec in spectrum] return computation_wrapper(_compute_moment, spectrum, regions, order=order, axis=axis)
def _compute_moment(spectrum, regions=None, order=0, axis='spectral'): """ This is a helper function for the above `moment()` method. """ if int(order) != order or order < 0: raise ValueError("Order must be a positive integer.") if axis == "spectral": if isinstance(spectrum, SpectrumCollection): axes = [spec.spectral_axis_index for spec in spectrum] if not np.all([x==axes[0] for x in axes]): raise ValueError("All spectra in SpectrumCollection must have the same " "spectral_axis_index for simultaneous moment calculation.") # SpectumCollection adds a leading axis when it stacks the spectra. axis = axes[0]+1 else: axis = spectrum.spectral_axis_index if regions is not None: calc_spectrum = extract_region(spectrum, regions) else: calc_spectrum = spectrum # Ignore masks for now. This should be fully addressed when # specutils gets revamped to handle multi-dimensional masks. flux = calc_spectrum.flux spectral_axis = calc_spectrum.spectral_axis # If axis is not the spectral axis, we simply use pixel values for dx and dispersion if (axis != spectrum.spectral_axis_index and not (axis == -1 and spectrum.spectral_axis_index == spectrum.flux.ndim-1)): dx = np.ones(flux.shape[axis]) dispersion = np.arange(flux.shape[axis]) + 1 else: dx = np.abs(np.diff(spectral_axis.bin_edges)) dispersion = spectral_axis # We now have to account for the desired axis being anywhere, not always last if len(flux.shape) > len(spectral_axis.shape): for i in range(flux.ndim): if i != axis: dx = np.expand_dims(dx, i) dx = np.repeat(dx, flux.shape[i], i) dispersion = np.expand_dims(dispersion, i) dispersion = np.repeat(dispersion, flux.shape[i], i) m0 = np.sum(flux * dx, axis=axis) if order == 0: return m0 if order == 1: return np.sum(flux * dispersion * dx, axis=axis) / m0 if order > 1: # By setting keepdims to True, the axes which are reduced are # left in the result as dimensions with size one. This means # that we can broadcast m1 correctly against dispersion. m1 = (np.sum(flux * dispersion * dx, axis=axis, keepdims=True) / np.sum(flux * dx, axis=axis, keepdims=True)) return np.sum(flux * dx * (dispersion - m1) ** order, axis=axis) / m0