Adding redshift and radial velocity properties to SpectrumCollection

CHANGES.rst

@@ -1,3 +1,18 @@
+2.5.0 (unreleased)
+------------------
+
+New Features
+^^^^^^^^^^^^
+
+- ``SpectrumCollection`` now allows setting scalar redshift or radial velocity,
+  analagously to ``Spectrum``. [#1332]
+
+Bug Fixes
+^^^^^^^^^
+
+Other Changes and Additions
+^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
 2.4.0 (2026-06-01)
 ------------------
 

docs/spectrum_collection.rst

@@ -96,6 +96,14 @@ and ``M`` is the length of the output dispersion grid.
     >>> print(spec_coll.flux.shape)
     (2, 50)
 
+It is important to understand that `~specutils.SpectrumCollection` does *not* store
+the data as a collection of individual `~specutils.Spectrum`
+objects, but rather contains single multi-dimensional flux and spectral axis arrays, and
+constructs a `~specutils.Spectrum` on the fly from the corresponding slices of these
+arrays when the collection is indexed. Thus, doing operations on the returned
+`~specutils.Spectrum` returned by indexing a `~specutils.SpectrumCollection`
+in this case does not affect the `~specutils.SpectrumCollection` object itself.
+
 
 Reference/API
 -------------

docs/types_of_spectra.rst

@@ -30,18 +30,19 @@ with their corresponding ``specutils`` representations:
 2. A set of fluxes that can be represented in an array-like form of shape
    ``n x m (x ...)``,  with a spectral axis strictly of length ``n`` (and a
    matched WCS). In ``specutils`` this is represented by the
-   `~specutils.Spectrum` object where ``len(flux.shape) > 1`` . In this sense
-   the "1D" refers to the spectral axis, *not* the flux. Note that 
+   `~specutils.Spectrum` object where ``len(flux.shape) > 1`` . Note that
    `~specutils.Spectrum` subclasses `NDCube <https://github.com/sunpy/ndcube>`_,
    which provideds utilities useful for these sorts of multidimensional fluxes.
 3. A set of fluxes  of shape ``n x m (x ...)``, and a set of spectral axes that
    are the same shape. This is distinguished from the above cases because there
-   are as many spectral axes as there are spectra.  In this sense it is a
-   collection of spectra, so can be thought of as a collection of
-   `~specutils.Spectrum` objects.  But because it is often more performant to
+   are as many spectral axes as there are spectra.  Because it is often more performant to
    store the collection together as just one set of flux and spectral axis
    arrays, this case is represented by a separate object in ``specutils``:
-   `~specutils.SpectrumCollection`.
+   `~specutils.SpectrumCollection`. Note that `~specutils.SpectrumCollection`
+   does *not* store the data as a collection of individual `~specutils.Spectrum`
+   objects, but rather constructs the returned `~specutils.Spectrum` on the fly when
+   the collection is indexed. Thus, doing operations on the returned `~specutils.Spectrum`
+   in this case does not affect the `~specutils.SpectrumCollection` object itself.
 4. An arbitrary collection of fluxes that are not all the same spectral length
    even in the spectral axis direction.  That is, case 3, but "ragged" in the
    sense that not all the spectra are length ``n``.  Because there is no

specutils/spectra/spectrum.py

@@ -7,11 +7,10 @@
 from astropy.utils.decorators import lazyproperty
 from astropy.utils.decorators import deprecated
 from astropy.nddata import NDUncertainty, NDIOMixin, NDArithmeticMixin, NDDataArray
-from astropy.wcs import WCS
 from gwcs.wcs import WCS as GWCS
 
 from .spectral_axis import SpectralAxis
-from .spectrum_mixin import OneDSpectrumMixin
+from .spectrum_mixin import OneDSpectrumMixin, RedshiftMixin
 from .spectral_region import SpectralRegion
 from ..utils.wcs_utils import gwcs_from_array
 
@@ -20,7 +19,7 @@
 __all__ = ['Spectrum1D', 'Spectrum']
 
 
-class Spectrum(OneDSpectrumMixin, NDCube, NDIOMixin, NDArithmeticMixin):
+class Spectrum(OneDSpectrumMixin, RedshiftMixin, NDCube, NDIOMixin, NDArithmeticMixin):
     """
     Spectrum container for N-dimensional data with one spectral axis.
 
@@ -743,117 +742,6 @@ def shape(self):
     def spectral_axis_direction(self):
         return self._spectral_axis_direction
 
-    @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.spectral_axis.redshift
-
-    @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.spectral_axis.radial_velocity
-
-    def set_redshift_to(self, redshift):
-        """
-        This sets the redshift of the spectrum to be `redshift` *without*
-        changing the values of the `spectral_axis`.
-
-        If you want to shift the `spectral_axis` based on this value, use
-        `shift_spectrum_to`.
-        """
-        new_spec_coord = self.spectral_axis.replicate(redshift=redshift)
-        self._spectral_axis = new_spec_coord
-
-    def set_radial_velocity_to(self, radial_velocity):
-        """
-        This sets the radial velocity of the spectrum to be `radial_velocity`
-        *without* changing the values of the `spectral_axis`.
-
-        If you want to shift the `spectral_axis` based on this value, use
-        `shift_spectrum_to`.
-        """
-        new_spec_coord = self.spectral_axis.replicate(
-            radial_velocity=radial_velocity
-        )
-        self._spectral_axis = new_spec_coord
-
-    def shift_spectrum_to(self, *, redshift=None, radial_velocity=None):
-        """
-        This shifts in-place the values of the `spectral_axis`, given either a
-        redshift or radial velocity.
-
-        If you do *not* want to change the `spectral_axis`, use
-        `set_redshift_to` or `set_radial_velocity_to`.
-        """
-        if redshift is not None and radial_velocity is not None:
-            raise ValueError(
-                "Only one of redshift or radial_velocity can be used."
-            )
-
-        old_redshift = self.redshift
-
-        if redshift is not None:
-            # with_radial_velocity_shift(redshift) looks wrong but astropy SpectralCoord handles
-            # redshift input to that method
-            new_spectral_axis = self.spectral_axis.with_radial_velocity_shift(
-                -self.spectral_axis.radial_velocity
-            ).with_radial_velocity_shift(redshift)
-            self._spectral_axis = new_spectral_axis
-        elif radial_velocity is not None:
-            if not radial_velocity.unit.is_equivalent(u.km/u.s):
-                raise u.UnitsError("Radial velocity must be a velocity.")
-
-            new_spectral_axis = self.spectral_axis.with_radial_velocity_shift(
-                -self.spectral_axis.radial_velocity
-            ).with_radial_velocity_shift(radial_velocity)
-            self._spectral_axis = new_spectral_axis
-            redshift = radial_velocity.to(u.Unit(''), u.doppler_redshift())
-        else:
-            raise ValueError("One of redshift or radial_velocity must be set.")
-
-        # Also store an updated WCS if we can update it.
-        if isinstance(self.wcs, WCS):
-            wcs_spectral_index = self.wcs.wcs.spec + 1
-            h = self.wcs.to_header()
-            spec_ctype = h[f'CTYPE{wcs_spectral_index}']
-            z_factor = (1 + redshift) / (1 + old_redshift)
-            if spec_ctype[0:4] != 'WAVE':
-                # Frequency, wavenumber and energy all invert this factor. Note that the FITS
-                # keyword for wavenumber is WAVN, which won't match here.
-                z_factor = 1 / z_factor
-            new_crval = h[f'CRVAL{wcs_spectral_index}'] * z_factor
-            h[f'CRVAL{wcs_spectral_index}'] = new_crval.value
-            pc_key = f'PC{wcs_spectral_index}_{wcs_spectral_index}'
-            if pc_key in h:
-                h[pc_key] *= z_factor
-            if f'CDELT{wcs_spectral_index}' in h:
-                new_cdelt = h[f'CDELT{wcs_spectral_index}'] * z_factor
-                h[f'CDELT{wcs_spectral_index}'] = new_cdelt.value
-            # WCS doesn't allow updating, but you can set it to None and then assign a new value
-            self.wcs = None
-            self.wcs = WCS(h, preserve_units=True)
-        else:
-            # I don't know how to update a GWCS cleanly so for now, we replace it and store the
-            # old one to retain any spatial information in the original
-            self._original_wcs = self.wcs
-            self.wcs = None
-            self.wcs = gwcs_from_array(new_spectral_axis, self.flux.shape,
-                                    spectral_axis_index=self.spectral_axis_index)
-
     def with_spectral_axis_last(self):
         """
         Convenience method to return a new copy of the Spectrum with the spectral axis last.

specutils/spectra/spectrum_collection.py

@@ -6,12 +6,13 @@
 from astropy.coordinates import SpectralCoord
 
 from .spectrum import Spectrum
+from .spectrum_mixin import RedshiftMixin
 from astropy.nddata import NDIOMixin
 
 __all__ = ['SpectrumCollection']
 
 
-class SpectrumCollection(NDIOMixin):
+class SpectrumCollection(NDIOMixin, RedshiftMixin):
     """
     A class to represent a heterogeneous set of spectra that are the same length
     but have different spectral axes. Spectra that meet this requirement can be
@@ -48,20 +49,58 @@ class SpectrumCollection(NDIOMixin):
     mask : ndarray or None
         The n-dimensional mask information associated with each spectrum. If
         present, must match the dimensionality of ``flux``.
+    redshift : float or None
+        A single value to apply to every spectrum in the collection.
+        Cannot be specified if radial_velocity is specified.
+    radial_velocity : `~astropy.units.Quantity`
+        A single value to apply to all spectra in the collection.
+        Cannot be specified if redshift is specified.
+    velocity_convention : {"relativistic", "optical", "radio"}
+        Convention used for velocity conversions.
+    rest_value : `~astropy.units.Quantity`
+        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.
+
     meta : list
         The list of dictionaries containing meta data to be associated with
         each spectrum in the collection.
     """
     def __init__(self, flux, spectral_axis=None, wcs=None, uncertainty=None,
-                 mask=None, meta=None, spectral_axis_index=None):
+                 mask=None, meta=None, spectral_axis_index=None, redshift=None,
+                 radial_velocity=None, rest_value=None, velocity_convention=None):
         # Check for quantity
         if not isinstance(flux, u.Quantity):
             raise u.UnitsError("Flux must be a `Quantity`.")
 
+        # Ensure that only one or neither of these parameters is set
+        if redshift is not None and radial_velocity is not None:
+            raise ValueError("Cannot set both radial_velocity and redshift at "
+                             "the same time.")
+
+        if redshift is not None and not (isinstance(redshift, (int, float)) or
+                (isinstance(redshift, u.Quantity) and redshift.ndim == 0)):
+            raise ValueError("Only single-value redshifts are supported at this time.")
+
+        if radial_velocity is not None and (not isinstance(radial_velocity, u.Quantity)
+                                            or radial_velocity.ndim > 0):
+            raise ValueError("radial_velocity must be a scalar `Quantity`.")
+
         if spectral_axis is not None:
             if not isinstance(spectral_axis, u.Quantity):
                 raise u.UnitsError("Spectral axis must be a `Quantity`.")
-            spectral_axis = SpectralCoord(spectral_axis)
+            if isinstance(spectral_axis, SpectralCoord):
+                if redshift is None:
+                    redshift = spectral_axis.redshift
+                else:
+                    if redshift != spectral_axis.redshift:
+                        raise ValueError("Cannot set a different redshift than defined on the"
+                                         " spectral_axis.")
+            else:
+                spectral_axis = SpectralCoord(spectral_axis, redshift=redshift,
+                                            radial_velocity=radial_velocity,
+                                            doppler_rest=rest_value,
+                                            doppler_convention=velocity_convention)
 
             # Ensure that the input values are the same shape
             if not (flux.shape == spectral_axis.shape):
@@ -104,6 +143,8 @@ def __getitem__(self, key):
         uncertainty = None if self.uncertainty is None else self.uncertainty[key]
         wcs = None if self.wcs is None else self.wcs[key]
         mask = None if self.mask is None else self.mask[key]
+        # Currently only allow scalar redshift
+        redshift = self.redshift
         if self.meta is None:
             meta = None
         else:
@@ -114,7 +155,7 @@ def __getitem__(self, key):
 
         return Spectrum(flux=flux, spectral_axis=spectral_axis,
                           uncertainty=uncertainty, wcs=wcs, mask=mask,
-                          meta=meta)
+                          meta=meta, redshift=redshift)
 
     @classmethod
     def from_spectra(cls, spectra):
@@ -181,13 +222,18 @@ def from_spectra(cls, spectra):
             warnings.warn("Not all spectra have associated masks, "
                           "skipping masks.")
 
+        # Check that all redshifts are equal
+        redshift = spectra[0].redshift
+        if (not all((x.redshift == redshift for x in spectra))):
+            warnings.warn("Not all redshifts are equal, using redshift of first spectrum.")
+
         # Store the wcs and meta as lists
         wcs = [spec.wcs for spec in spectra]
         meta = [spec.meta for spec in spectra]
 
         return cls(flux=flux, spectral_axis=spectral_axis,
                    uncertainty=uncertainty, wcs=wcs, mask=mask, meta=meta,
-                   spectral_axis_index=spectral_axis_index)
+                   spectral_axis_index=spectral_axis_index, redshift=redshift)
 
     @property
     def flux(self):
@@ -245,6 +291,30 @@ def meta(self):
         """A dictionary of metadata for theis spectrum collection, or `None`."""
         return self._meta
 
+    @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.spectral_axis.redshift
+
+    @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.spectral_axis.radial_velocity
+
     @property
     def shape(self):
         """