Fix with_spectral_unit (#1119)

* Fix with_spectral_axis when Spectrum1D instantiated with a spectral_axis

* Deprecate old method name, new one is more specific

* Works for WCS-initialized spectra as well

* Remove unneeded imports

* Add kwarg default

* Add copy and deepcopy methods

Codestyle

* Update specutils/spectra/spectrum_mixin.py

Co-authored-by: Derek Homeier <709020+dhomeier@users.noreply.github.com>

* Add changelog entry

* Test to make sure we kept the old WCS in meta

* Better name for meta attribute, add test for it

* Apply suggestions from code review

Co-authored-by: Derek Homeier <709020+dhomeier@users.noreply.github.com>

* Fix codestyle

* Update docstring

Co-authored-by: Derek Homeier <709020+dhomeier@users.noreply.github.com>

* Change docstring to fix docs build failure. Also fixed some incorrect references to SpectralRegion

* Fix one more SpectralRegion reference

---------

Co-authored-by: Derek Homeier <709020+dhomeier@users.noreply.github.com>

CHANGES.rst

@@ -10,6 +10,9 @@ Bug Fixes
 - ``template_correlate`` no longer errors when used on a ``Spectrum1D`` that lacks an
   ``uncertainty`` array. [#1118]
 
+- ``with_spectral_unit`` has been changed to ``with_spectral_axis_unit`` and actually works
+  now. [#1119]
+
 Other Changes and Additions
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
@@ -64,7 +67,7 @@ New Features
 
 - ``wcs1d-fits`` loader now reads and writes celestial components of
   of multi-dimensional WCS, and handles ``mask`` and ``uncertainty``
-  attributes. [#1009] 
+  attributes. [#1009]
 
 - Added support for reading from files with flux in counts. [#1018]
 
@@ -183,7 +186,7 @@ Bug Fixes
   ``Spectrum1D`` without WCS nor spectral axis and the spatial-spatial dimension
   is smaller than spectral dimension. [#926]
 
-- Fixed WCS not accurately reflecting the updated spectral axis after slicing a 
+- Fixed WCS not accurately reflecting the updated spectral axis after slicing a
   ``Spectrum1D``. [#918]
 
 Other Changes and Additions
@@ -212,8 +215,8 @@ New Features
 - Convolution-based smoothing will now apply a 1D kernel to multi-dimensional fluxes
   by convolving along the spectral axis only, rather than raising an error. [#885]
 
-- ``template_comparison`` now handles ``astropy.nddata.Variance`` and 
-  ``astropy.nddata.InverseVariance`` uncertainties instead of assuming 
+- ``template_comparison`` now handles ``astropy.nddata.Variance`` and
+  ``astropy.nddata.InverseVariance`` uncertainties instead of assuming
   the uncertainty is standard deviation. [#899]
 
 Bug Fixes
@@ -253,7 +256,7 @@ New Features
 
 - Allow overriding existing astropy registry elements. [#861]
 
-- ``Spectrum1D`` will now swap the spectral axis with the last axis on initialization 
+- ``Spectrum1D`` will now swap the spectral axis with the last axis on initialization
   if it can be identified from the WCS and is not last, rather than erroring. [#654, #822]
 
 Bug Fixes

specutils/analysis/moment.py

@@ -21,7 +21,7 @@ def moment(spectrum, regions=None, order=0, axis=-1):
     spectrum : `~specutils.spectra.spectrum1d.Spectrum1D`
         The spectrum object over which the width will be calculated.
 
-    regions: `~specutils.utils.SpectralRegion` or list of `~specutils.utils.SpectralRegion`
+    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.
 

specutils/analysis/uncertainty.py

@@ -22,7 +22,7 @@ def snr(spectrum, region=None):
     spectrum : `~specutils.spectra.spectrum1d.Spectrum1D`
         The spectrum object overwhich the equivalent width will be calculated.
 
-    region: `~specutils.utils.SpectralRegion` or list of `~specutils.utils.SpectralRegion`
+    region: `~specutils.SpectralRegion` or list of `~specutils.SpectralRegion`
         Region within the spectrum to calculate the SNR.
 
     Returns
@@ -67,7 +67,7 @@ def _snr_single_region(spectrum, region=None):
     spectrum : `~specutils.spectra.spectrum1d.Spectrum1D`
         The spectrum object overwhich the equivalent width will be calculated.
 
-    region: `~specutils.utils.SpectralRegion`
+    region: `~specutils.SpectralRegion`
         Region within the spectrum to calculate the SNR.
 
     Returns
@@ -109,7 +109,7 @@ def snr_derived(spectrum, region=None):
     spectrum : `~specutils.spectra.spectrum1d.Spectrum1D`
         The spectrum object overwhich the equivalent width will be calculated.
 
-    region: `~specutils.utils.SpectralRegion`
+    region: `~specutils.SpectralRegion`
         Region within the spectrum to calculate the SNR.
 
     Returns
@@ -155,7 +155,7 @@ def _snr_derived(spectrum, region=None):
     spectrum : `~specutils.spectra.spectrum1d.Spectrum1D`
         The spectrum object overwhich the equivalent width will be calculated.
 
-    region: `~specutils.utils.SpectralRegion`
+    region: `~specutils.SpectralRegion`
         Region within the spectrum to calculate the SNR.
 
     Returns

specutils/spectra/spectrum_mixin.py

@@ -5,8 +5,6 @@
 from astropy import units as u
 from astropy.utils.decorators import deprecated
 
-from specutils.utils.wcs_utils import gwcs_from_array
-
 DOPPLER_CONVENTIONS = {}
 DOPPLER_CONVENTIONS['radio'] = u.doppler_radio
 DOPPLER_CONVENTIONS['optical'] = u.doppler_optical
@@ -184,10 +182,20 @@ def velocity(self):
 
         return new_data
 
+    @deprecated('v1.13', alternative="with_spectral_axis_unit")
     def with_spectral_unit(self, unit, velocity_convention=None,
                            rest_value=None):
+        self.with_spectral_axis_unit(unit, velocity_convention=velocity_convention,
+                                     rest_value=rest_value)
+
+    def with_spectral_axis_unit(self, unit, velocity_convention=None,
+                           rest_value=None):
         """
-        Returns a new spectrum with a different spectral axis unit.
+        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
         ----------
@@ -208,14 +216,27 @@ def with_spectral_unit(self, unit, velocity_convention=None,
                      even if your spectrum has air wavelength units
 
         """
-        new_wcs, new_meta = self._new_spectral_wcs(
-            unit=unit,
-            velocity_convention=velocity_convention or self.velocity_convention,
-            rest_value=rest_value or self.rest_value)
 
-        spectrum = self.__class__(flux=self.flux, wcs=new_wcs, meta=new_meta)
+        velocity_convention = velocity_convention if velocity_convention is not None else self.velocity_convention # noqa
+        rest_value = rest_value if rest_value is not None else self.rest_value
+        unit = self._new_wcs_argument_validation(unit, velocity_convention,
+                                                 rest_value)
+
+        # Store the original unit information and WCS for posterity
+        meta = deepcopy(self._meta)
+
+        if 'original_spectral_axis_unit' not in self._meta:
+            orig_unit = self.wcs.unit[0] if hasattr(self.wcs, 'unit') else self.spectral_axis.unit
+            meta['original_spectral_axis_unit'] = orig_unit
+
+        if 'original_wcs' not in self.meta:
+            meta['original_wcs'] = self.wcs.deepcopy()
 
-        return spectrum
+        new_spectral_axis = self.spectral_axis.to(unit, doppler_convention=velocity_convention,
+                                                  doppler_rest=rest_value)
+
+        return self.__class__(flux=self.flux, spectral_axis=new_spectral_axis, meta=meta,
+                              uncertainty=self.uncertainty, mask=self.mask)
 
     def _new_wcs_argument_validation(self, unit, velocity_convention,
                                      rest_value):
@@ -242,52 +263,6 @@ def _new_wcs_argument_validation(self, unit, velocity_convention,
 
         return unit
 
-    def _new_spectral_wcs(self, unit, velocity_convention=None,
-                          rest_value=None):
-        """
-        Returns a new WCS with a different Spectral Axis unit.
-
-        Parameters
-        ----------
-        unit : :class:`~astropy.units.Unit`
-            Any valid spectral unit: velocity, (wave)length, or frequency.
-            Only vacuum units are supported.
-        velocity_convention : 'relativistic', 'radio', or 'optical'
-            The velocity 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 cube'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 cube has air wavelength units
-        """
-
-        unit = self._new_wcs_argument_validation(unit, velocity_convention,
-                                                 rest_value)
-
-        if velocity_convention is not None:
-            equiv = getattr(u, 'doppler_{0}'.format(velocity_convention))
-            rest_value.to(unit, equivalencies=equiv)
-
-        # Store the original unit information for posterity
-        meta = self._meta.copy()
-
-        orig_unit = self.wcs.unit[0] if hasattr(self.wcs, 'unit') else self.spectral_axis.unit
-
-        if 'original_unit' not in self._meta:
-            meta['original_unit'] = orig_unit
-
-        # Create the new wcs object
-        if isinstance(unit, u.UnitBase) and unit.is_equivalent(
-                orig_unit, equivalencies=u.spectral()):
-            return gwcs_from_array(self.spectral_axis), meta
-
-        raise u.UnitConversionError(f"WCS units incompatible: {unit} and {orig_unit}.")
-
     def _check_strictly_increasing_decreasing(self):
         """
         Check that the self._spectral_axis is strictly increasing or decreasing

specutils/tests/test_spectrum1d.py

@@ -150,10 +150,29 @@ def test_spectral_axis_conversions():
     with pytest.raises(ValueError):
         spec.velocity
 
-    spec = Spectrum1D(spectral_axis=np.arange(1, 50) * u.nm,
+    spec = Spectrum1D(spectral_axis=np.arange(100, 150) * u.nm,
                       flux=np.random.randn(49) * u.Jy)
 
-    new_spec = spec.with_spectral_unit(u.GHz)  # noqa
+    new_spec = spec.with_spectral_axis_unit(u.km/u.s, rest_value=125*u.um,
+                                       velocity_convention="relativistic")
+
+    assert new_spec.spectral_axis.unit == u.km/u.s
+    assert new_spec.wcs.world_axis_units[0] == "km.s**-1"
+    # Make sure meta stored the old WCS correctly
+    assert new_spec.meta["original_wcs"].world_axis_units[0] == "nm"
+    assert new_spec.meta["original_spectral_axis_unit"] == "nm"
+
+    wcs_dict = {"CTYPE1": "WAVE", "CRVAL1": 3.622e3, "CDELT1": 8e-2,
+                "CRPIX1": 0, "CUNIT1": "Angstrom"}
+    wcs_spec = Spectrum1D(flux=np.random.randn(49) * u.Jy, wcs=WCS(wcs_dict),
+                          meta={'header': wcs_dict.copy()})
+    new_spec = wcs_spec.with_spectral_axis_unit(u.km/u.s, rest_value=125*u.um,
+                                       velocity_convention="relativistic")
+    new_spec.meta['original_wcs'].wcs.crval = [3.777e-7]
+    new_spec.meta['header']['CRVAL1'] = 3777.0
+
+    assert wcs_spec.wcs.wcs.crval[0] == 3.622e-7
+    assert wcs_spec.meta['header']['CRVAL1'] == 3622.
 
 
 def test_spectral_slice():

specutils/utils/wcs_utils.py

@@ -8,6 +8,44 @@
 from gwcs import coordinate_frames as cf
 
 
+class SpectralGWCS(GWCS):
+    """
+    This is a placeholder lookup-table GWCS created when a :class:`~specutils.Spectrum1D` is
+    instantiated with a ``spectral_axis`` and no WCS.
+    """
+    def __init__(self, *args, **kwargs):
+        self.original_unit = kwargs.pop("original_unit", "")
+        super().__init__(*args, **kwargs)
+
+    def copy(self):
+        """
+        Return a shallow copy of the object.
+
+        Convenience method so user doesn't have to import the
+        :mod:`copy` stdlib module.
+
+        .. warning::
+            Use `deepcopy` instead of `copy` unless you know why you need a
+            shallow copy.
+        """
+        return copy.copy(self)
+
+    def deepcopy(self):
+        """
+        Return a deep copy of the object.
+
+        Convenience method so user doesn't have to import the
+        :mod:`copy` stdlib module.
+        """
+        return copy.deepcopy(self)
+
+    def pixel_to_world(self, *args, **kwargs):
+        if self.original_unit == '':
+            return u.Quantity(super().pixel_to_world_values(*args, **kwargs))
+        return super().pixel_to_world(*args, **kwargs).to(
+            self.original_unit, equivalencies=u.spectral())
+
+
 def refraction_index(wavelength, method='Griesen2006', co2=None):
     """
     Calculates the index of refraction of dry air at standard temperature
@@ -210,14 +248,8 @@ def gwcs_from_array(array):
         forward_transform.inverse = SpectralTabular1D(
                 array[::-1], lookup_table=np.arange(len(array))[::-1])
 
-    class SpectralGWCS(GWCS):
-        def pixel_to_world(self, *args, **kwargs):
-            if orig_array.unit == '':
-                return u.Quantity(super().pixel_to_world_values(*args, **kwargs))
-            return super().pixel_to_world(*args, **kwargs).to(
-                orig_array.unit, equivalencies=u.spectral())
-
-    tabular_gwcs = SpectralGWCS(forward_transform=forward_transform,
+    tabular_gwcs = SpectralGWCS(original_unit = orig_array.unit,
+                                forward_transform=forward_transform,
                                 input_frame=coord_frame,
                                 output_frame=spec_frame)