Merge 3c6f70a into 87dd9ff

specutils/spectra/spectrum1d.py

@@ -3,7 +3,8 @@
 
 import numpy as np
 from astropy import units as u
-from astropy.nddata import NDDataRef, NDUncertainty
+from astropy import constants as cnst
+from astropy.nddata import NDDataRef
 from astropy.utils.decorators import lazyproperty
 
 from ..wcs import WCSAdapter, WCSWrapper
@@ -33,6 +34,10 @@ class Spectrum1D(OneDSpectrumMixin, NDDataRef):
         Any quantity supported by the standard spectral equivalencies
         (wavelength, energy, frequency, wave number). Describes the rest value
         of the spectral axis for use with velocity conversions.
+    redshift
+        See `redshift` for more information.
+    radial_velocity
+        See `radial_velocity` for more information.
     uncertainty : `~astropy.nddata.NDUncertainty`
         Contains uncertainty information along with propagation rules for
         spectrum arithmetic. Can take a unit, but if none is given, will use
@@ -42,7 +47,8 @@ class Spectrum1D(OneDSpectrumMixin, NDDataRef):
         around with the spectrum container object.
     """
     def __init__(self, flux=None, spectral_axis=None, wcs=None,
-                 velocity_convention=None, rest_value=None, **kwargs):
+                 velocity_convention=None, rest_value=None, redshift=None,
+                 radial_velocity=None, **kwargs):
         # Check for pre-defined entries in the kwargs dictionary.
         unknown_kwargs = set(kwargs).difference(
             {'data', 'unit', 'uncertainty', 'meta', 'mask', 'copy'})
@@ -65,7 +71,7 @@ def __init__(self, flux=None, spectral_axis=None, wcs=None,
             raise ValueError("Flux must be a `Quantity` object.")
 
         # In cases of slicing, new objects will be initialized with `data`
-        # instead of `flux`. Ensure we grab the `data` argument.
+        # instead of ``flux``. Ensure we grab the `data` argument.
         if flux is None and 'data' in kwargs:
             flux = kwargs.pop('data')
 
@@ -127,6 +133,17 @@ def __init__(self, flux=None, spectral_axis=None, wcs=None,
             data=flux.value if isinstance(flux, u.Quantity) else flux,
             wcs=wcs, **kwargs)
 
+        # set redshift after super() - necessary because the shape-checking
+        # requires that the flux be initialized
+
+        if redshift is None:
+            self.radial_velocity = radial_velocity
+        elif radial_velocity is None:
+            self.redshift = redshift
+        else:
+            raise ValueError('cannot set both radial_velocity and redshift at '
+                             'the same time.')
+
         if hasattr(self, 'uncertainty') and self.uncertainty is not None:
             if not flux.shape == self.uncertainty.array.shape:
                 raise ValueError('Flux axis ({}) and uncertainty ({}) shapes must be the same.'.format(
@@ -215,6 +232,53 @@ def bin_edges(self):
     def shape(self):
         return self.flux.shape
 
+    @property
+    def redshift(self):
+        """
+        The redshift(s) of the objects represented by this spectrum.  May be
+        scalar (if this spectrum's ``flux`` is 1D) or vector.  Note that
+        the concept of "redshift of a spectrum" can be ambiguous, so the
+        interpretation is set to some extent by either the user, or operations
+        (like template fitting) that set this attribute when they are run on
+        a spectrum.
+        """
+        return self._radial_velocity/cnst.c
+    @redshift.setter
+    def redshift(self, val):
+        if val is None:
+            self._radial_velocity = None
+        else:
+            self.radial_velocity = val * cnst.c
+
+    @property
+    def radial_velocity(self):
+        """
+        The radial velocity(s) of the objects represented by this spectrum.  May
+        be scalar (if this spectrum's ``flux`` is 1D) or vector.  Note that
+        the concept of "RV of a spectrum" can be ambiguous, so the
+        interpretation is set to some extent by either the user, or operations
+        (like template fitting) that set this attribute when they are run on
+        a spectrum.
+        """
+        return self._radial_velocity
+    @radial_velocity.setter
+    def radial_velocity(self, val):
+        if val is None:
+            self._radial_velocity = None
+        else:
+            if not val.unit.is_equivalent(u.km/u.s):
+                raise u.UnitsError('radial_velocity must be a velocity')
+
+            # the trick below checks if the two shapes given are broadcastable onto
+            # each other. See https://stackoverflow.com/questions/47243451/checking-if-two-arrays-are-broadcastable-in-python
+            input_shape = val.shape
+            flux_shape = self.flux.shape[:-1]
+            if not all((m == n) or (m == 1) or (n == 1)
+                   for m, n in zip(input_shape[::-1], flux_shape)):
+                raise ValueError("radial_velocity or redshift must have shape that "
+                                 "is compatible with this spectrum's flux array")
+            self._radial_velocity = val
+
     @staticmethod
     def _compare_wcs(this_operand, other_operand):
         """

specutils/spectra/spectrum_mixin.py

@@ -2,8 +2,6 @@
 from copy import deepcopy
 
 import numpy as np
-from astropy.wcs import WCSSUB_SPECTRAL
-from astropy.units import Unit
 from astropy import units as u
 import astropy.units.equivalencies as eq
 from astropy.utils.decorators import lazyproperty
@@ -167,6 +165,9 @@ def velocity(self):
             (wavelength, energy, frequency, wave number).
         type : {"doppler_relativistic", "doppler_optical", "doppler_radio"}
             The type of doppler spectral equivalency.
+        redshift or radial_velocity
+            If present, this shift is applied to the final output velocity to
+            get into the rest frame of the object.
 
         Returns
         -------
@@ -185,6 +186,10 @@ def velocity(self):
 
         new_data = self.spectral_axis.to(u.km/u.s, equivalencies=equiv)
 
+        # if redshift/rv is present, apply it:
+        if self._radial_velocity is not None:
+            new_data += self.radial_velocity
+
         return new_data
 
     def with_spectral_unit(self, unit, velocity_convention=None,

specutils/tests/test_spectrum1d.py

@@ -121,6 +121,63 @@ def test_spectral_axis_conversions():
     new_spec = spec.with_spectral_unit(u.GHz)
 
 
+def test_redshift():
+    spec = Spectrum1D(flux=np.array([26.0, 30., 44.5]) * u.Jy,
+                      spectral_axis=np.array([4000, 6000, 8000]) * u.AA,
+                      velocity_convention='optical',
+                      rest_value=6000 * u.AA)
+
+    assert u.allclose(spec.velocity, [-99930.8, 0, 99930.8]*u.km/u.s,
+                     atol=0.5*u.km/u.s)
+
+    spec = Spectrum1D(flux=np.array([26.0, 30., 44.5]) * u.Jy,
+                      spectral_axis=np.array([4000, 6000, 8000]) * u.AA,
+                      velocity_convention='optical',
+                      rest_value=6000 * u.AA,
+                      redshift= 0.1)
+
+    assert u.allclose(spec.velocity, [-69951.3, 29979.2, 129910.1]*u.km/u.s,
+                     atol=0.5*u.km/u.s)
+
+    #-------------------------
+
+    spec = Spectrum1D(flux=np.array([26.0, 30.0, 44.5]) * u.Jy,
+                      spectral_axis=np.array([10.5, 11.0, 11.5]) * u.GHz,
+                      velocity_convention='radio',
+                      rest_value=11.0 * u.GHz)
+
+    assert u.allclose(spec.velocity, [13626., 0, -13626]*u.km/u.s,
+                     atol=1*u.km/u.s)
+
+    spec = Spectrum1D(flux=np.array([26.0, 30.0, 44.5]) * u.Jy,
+                      spectral_axis=np.array([10.5, 11.0, 11.5]) * u.GHz,
+                      velocity_convention='radio',
+                      rest_value=11.0 * u.GHz,
+                      redshift= 0.1)
+
+    assert u.allclose(spec.velocity, [43606., 29979., 16352.]*u.km/u.s,
+                      atol=1*u.km/u.s)
+
+    #------------------------- radial velocity mode
+
+    spec = Spectrum1D(flux=np.array([26.0, 30., 44.5]) * u.Jy,
+                      spectral_axis=np.array([4000, 6000, 8000]) * u.AA,
+                      velocity_convention='optical',
+                      rest_value=6000 * u.AA)
+
+    assert u.allclose(spec.velocity, [-99930.8, 0.0, 99930.8]*u.km/u.s,
+                      atol=0.5*u.km/u.s)
+
+    spec = Spectrum1D(flux=np.array([26.0, 30., 44.5]) * u.Jy,
+                      spectral_axis=np.array([4000, 6000, 8000]) * u.AA,
+                      velocity_convention='optical',
+                      rest_value=6000 * u.AA,
+                      radial_velocity=1000.*u.km/u.s)
+
+    assert u.allclose(spec.velocity, [-98930.8, 1000.0, 100930.8]*u.km/u.s,
+                      atol=0.5*u.km/u.s)
+
+
 def test_flux_unit_conversion():
     # By default the flux units should be set to Jy
     s = Spectrum1D(flux=np.array([26.0, 44.5]) * u.Jy,