Issue 1119 changing spectral axis

src/dysh/spectra/spectrum.py

@@ -1097,6 +1097,73 @@ def rest_value(self, value: Quantity):
         self.meta["RESTFREQ"] = value.to("Hz").value
         self.meta["RESTFRQ"] = value.to("Hz").value
 
+    def with_spectral_axis_unit(self, *args, **kwargs):
+        """
+        Returns a new spectrum with a different spectral axis unit. Note that this creates a new
+        object using the converted spectral axis and thus drops the original WCS, if it existed,
+        replacing it with a lookup-table :class:`~gwcs.wcs.WCS` based on the new spectral axis. The
+        original WCS will be stored in the ``original_wcs`` entry of the new object's ``meta``
+        dictionary.
+
+        Parameters
+        ----------
+        unit : :class:`~astropy.units.Unit`
+            Any valid spectral unit: velocity, (wave)length, or frequency.
+            Only vacuum units are supported.
+        doppler_convention : {None, 'relativistic', 'radio', 'optical'}
+            The Doppler convention to use for the output velocity axis.
+            Required if the output type is velocity. This can be either one
+            of the above strings, or an `astropy.units` equivalency.
+        rest_value : :class:`~astropy.units.Quantity`
+            A rest wavelength or frequency with appropriate units.  Required if
+            output type is velocity.  The spectrum's WCS should include this
+            already if the *input* type is velocity, but the WCS's rest
+            wavelength/frequency can be overridden with this parameter.
+
+            .. note: This must be the rest frequency/wavelength *in vacuum*,
+                     even if your spectrum has air wavelength units
+
+        """
+        # The docstring should be kept up to date.
+        # Keep docstring empty to re-use the one from specutils.
+        # Keeping the specutils doc introduces yet another name
+        # for what we call `doppler_convention`, `velocity_convention`.
+        # This is a source of frustration...
+        # Rename kwargs to match specutils signature.
+        rename_dict = {"doppler_convention": "velocity_convention"}
+        new_kwargs = {rename_dict[key]: value for key, value in kwargs.items()}
+        spec = super().with_spectral_axis_unit(*args, **new_kwargs)
+        # Copy observer and target to support further frame conversions.
+        # See issue #1119.
+        spec._observer = self.observer.copy()
+        spec._spectral_axis.observer = self.spectral_axis.observer.copy()
+        spec._target = self.target.copy()
+        spec._spectral_axis.target = self.spectral_axis.target.copy()
+        return spec
+
+    def set_spectral_axis(self, unit=None, toframe=None, doppler_convention=None) -> None:
+        """
+        Modifies the spectral axis to have units of `unit` in reference frame `toframe`
+        using Doppler convention `doppler_convention`. This changes the `Spectrum`.
+
+        Parameters
+        ----------
+        unit : `~astropy.units.quantity.Quantity` or str that can be converted to Quantity
+            The unit to which the axis is to be converted.
+        toframe : str
+            The coordinate frame to convert to, e.g. 'hcrs', 'icrs'.
+        doppler_convention : {None, 'relativistic', 'radio', 'optical'}
+            The Doppler convention to use when converting to velocity.
+            One of 'optical', 'radio', or 'relativistic'.
+        """
+        if unit is None and toframe is not None:
+            self.set_frame(toframe)
+        if unit is None and doppler_convention is not None:
+            self.set_convention(doppler_convention)
+        if unit is not None:
+            sa = self.velocity_axis_to(unit=unit, toframe=toframe, doppler_convention=doppler_convention)
+            self._spectral_axis = sa
+
     def axis_velocity(self, unit=KMS):
         """Get the spectral axis in velocity units.
         *Note*: This is not the same as `Spectrum.velocity`, which includes the source radial velocity.
@@ -1123,7 +1190,7 @@ def velocity_axis_to(self, unit=KMS, toframe=None, doppler_convention=None):
             The unit to which the spectral axis is to be converted.
         toframe : str
             The coordinate frame to convert to, e.g. 'hcrs', 'icrs'.
-        doppler_convention : None or {'optical', 'radio', 'relativistic'}
+        doppler_convention : {None, 'relativistic', 'radio', 'optical'}
             The Doppler convention to use when converting to velocity.
             One of 'optical', 'radio', or 'relativistic'.
 

src/dysh/spectra/tests/test_spectrum.py

@@ -1300,3 +1300,23 @@ def test_header(self, capsys):
         s.header()
         captured = capsys.readouterr()
         assert captured.out == expected
+
+    def test_set_spectral_axis(self):
+        """
+        Test that the spectral_axis is updated.
+        """
+
+        sp = Spectrum.fake_spectrum()
+        sp.set_spectral_axis(toframe="galactocentric")
+        assert sp.observer.name == "galactocentric"
+        # Try all three options in case the defaults change.
+        if sp.doppler_convention != "radio":
+            expected = "radio"
+        elif sp.doppler_convention != "optical":
+            expected = "optical"
+        else:
+            expected = "relativistic"
+        sp.set_spectral_axis(doppler_convention=expected)
+        assert sp.doppler_convention == expected
+        sp.set_spectral_axis(unit="km/s")
+        assert sp.spectral_axis.unit.to_string() == "km / s"