use SpectralCoord instead of Quantity in world_axis_object_components

Author: nden

astropy/wcs/wcsapi/fitswcs.py

@@ -12,6 +12,7 @@
 from astropy.coordinates import ICRS, Galactic, SpectralCoord
 from astropy.coordinates.spectral_coordinate import (
     attach_zero_velocities,
+    SpectralCoord,
     update_differentials_to_match,
 )
 from astropy.units import allclose as quantity_allclose
@@ -605,7 +606,7 @@ def redshift_from_spectralcoord(spectralcoord):
                             - 1.0
                         )
 
-                classes["spectral"] = (u.Quantity, (), {}, spectralcoord_from_redshift)
+                classes["spectral"] = (SpectralCoord, (), {}, spectralcoord_from_redshift)
                 components[self.wcs.spec] = ("spectral", 0, redshift_from_spectralcoord)
 
             elif ctype == "BETA":
@@ -652,7 +653,7 @@ def beta_from_spectralcoord(spectralcoord):
                             / C_SI
                         )
 
-                classes["spectral"] = (u.Quantity, (), {}, spectralcoord_from_beta)
+                classes["spectral"] = (SpectralCoord, (), {}, spectralcoord_from_beta)
                 components[self.wcs.spec] = ("spectral", 0, beta_from_spectralcoord)
 
             else:
@@ -732,7 +733,7 @@ def value_from_spectralcoord(spectralcoord):
                             observer
                         ).to_value(**kwargs)
 
-                classes["spectral"] = (u.Quantity, (), {}, spectralcoord_from_value)
+                classes["spectral"] = (SpectralCoord, (), {}, spectralcoord_from_value)
                 components[self.wcs.spec] = ("spectral", 0, value_from_spectralcoord)
 
         # We can then make sure we correctly return Time objects where appropriate

astropy/wcs/wcsapi/tests/test_fitswcs.py

@@ -277,7 +277,7 @@ def test_spectral_cube():
     assert isinstance(wcs.world_axis_object_classes["celestial"][2]["frame"], Galactic)
     assert wcs.world_axis_object_classes["celestial"][2]["unit"] == (u.deg, u.deg)
 
-    assert wcs.world_axis_object_classes["spectral"][0] is Quantity
+    assert wcs.world_axis_object_classes["spectral"][0] is SpectralCoord
     assert wcs.world_axis_object_classes["spectral"][1] == ()
     assert wcs.world_axis_object_classes["spectral"][2] == {}
 
@@ -397,7 +397,7 @@ def test_spectral_cube_nonaligned():
     assert isinstance(wcs.world_axis_object_classes["celestial"][2]["frame"], Galactic)
     assert wcs.world_axis_object_classes["celestial"][2]["unit"] == (u.deg, u.deg)
 
-    assert wcs.world_axis_object_classes["spectral"][0] is Quantity
+    assert wcs.world_axis_object_classes["spectral"][0] is SpectralCoord
     assert wcs.world_axis_object_classes["spectral"][1] == ()
     assert wcs.world_axis_object_classes["spectral"][2] == {}
 

astropy/wcs/wcsapi/wrappers/tests/test_sliced_wcs.py

@@ -5,7 +5,7 @@
 from numpy.testing import assert_allclose, assert_equal
 
 import astropy.units as u
-from astropy.coordinates import ICRS, Galactic, SkyCoord
+from astropy.coordinates import ICRS, Galactic, SkyCoord, SpectralCoord
 from astropy.io.fits import Header
 from astropy.io.fits.verify import VerifyWarning
 from astropy.time import Time
@@ -148,7 +148,7 @@ def test_ellipsis():
     assert isinstance(wcs.world_axis_object_classes["celestial"][2]["frame"], Galactic)
     assert wcs.world_axis_object_classes["celestial"][2]["unit"] == (u.deg, u.deg)
 
-    assert wcs.world_axis_object_classes["spectral"][0] is Quantity
+    assert wcs.world_axis_object_classes["spectral"][0] is SpectralCoord
     assert wcs.world_axis_object_classes["spectral"][1] == ()
     assert wcs.world_axis_object_classes["spectral"][2] == {}
 
@@ -307,7 +307,7 @@ def test_spectral_range():
     assert isinstance(wcs.world_axis_object_classes["celestial"][2]["frame"], Galactic)
     assert wcs.world_axis_object_classes["celestial"][2]["unit"] == (u.deg, u.deg)
 
-    assert wcs.world_axis_object_classes["spectral"][0] is Quantity
+    assert wcs.world_axis_object_classes["spectral"][0] is SpectralCoord
     assert wcs.world_axis_object_classes["spectral"][1] == ()
     assert wcs.world_axis_object_classes["spectral"][2] == {}
 
@@ -387,7 +387,7 @@ def test_celestial_slice():
     assert isinstance(wcs.world_axis_object_classes["celestial"][2]["frame"], Galactic)
     assert wcs.world_axis_object_classes["celestial"][2]["unit"] == (u.deg, u.deg)
 
-    assert wcs.world_axis_object_classes["spectral"][0] is Quantity
+    assert wcs.world_axis_object_classes["spectral"][0] is SpectralCoord
     assert wcs.world_axis_object_classes["spectral"][1] == ()
     assert wcs.world_axis_object_classes["spectral"][2] == {}
 
@@ -469,7 +469,7 @@ def test_celestial_range():
     assert isinstance(wcs.world_axis_object_classes["celestial"][2]["frame"], Galactic)
     assert wcs.world_axis_object_classes["celestial"][2]["unit"] == (u.deg, u.deg)
 
-    assert wcs.world_axis_object_classes["spectral"][0] is Quantity
+    assert wcs.world_axis_object_classes["spectral"][0] is SpectralCoord
     assert wcs.world_axis_object_classes["spectral"][1] == ()
     assert wcs.world_axis_object_classes["spectral"][2] == {}
 
@@ -560,7 +560,7 @@ def test_celestial_range_rot():
     assert isinstance(wcs.world_axis_object_classes["celestial"][2]["frame"], Galactic)
     assert wcs.world_axis_object_classes["celestial"][2]["unit"] == (u.deg, u.deg)
 
-    assert wcs.world_axis_object_classes["spectral"][0] is Quantity
+    assert wcs.world_axis_object_classes["spectral"][0] is SpectralCoord
     assert wcs.world_axis_object_classes["spectral"][1] == ()
     assert wcs.world_axis_object_classes["spectral"][2] == {}
 
@@ -1015,7 +1015,7 @@ def test_dropped_dimensions_4d(cube_4d_fitswcs):
     assert wao_classes["celestial"][1] == ()
     assert isinstance(wao_classes["celestial"][2]["frame"], ICRS)
     assert wao_classes["celestial"][2]["unit"] == (u.deg, u.deg)
-    assert wao_classes["spectral"][0:3] == (u.Quantity, (), {})
+    assert wao_classes["spectral"][0:3] == (SpectralCoord, (), {})
 
     assert wao_components[0] == ("celestial", 0, "spherical.lon.degree")
     assert wao_components[1] == ("celestial", 1, "spherical.lat.degree")
@@ -1040,7 +1040,7 @@ def test_dropped_dimensions_4d(cube_4d_fitswcs):
     assert wao_components[1][0] == "time"
     assert wao_components[1][1] == 0
 
-    assert wao_classes["spectral"][0:3] == (u.Quantity, (), {})
+    assert wao_classes["spectral"][0:3] == (SpectralCoord, (), {})
     assert wao_classes["time"][0:3] == (Time, (), {})