Support serialization of astropy.wcs.WCS objects to ASDF.

CHANGES.rst

@@ -8,6 +8,8 @@
 
 - replace usages of ``copy_arrays`` with ``memmap`` [#230]
 
+- Add support for astropy.wcs.WCS and astropy.wcs.wcsapi.SlicedLowLevelWCS [#235]
+
 0.6.1 (2024-04-05)
 ------------------
 

asdf_astropy/converters/__init__.py

@@ -15,6 +15,7 @@
     "FitsConverter",
     "AsdfFitsConverter",
     "AstropyFitsConverter",
+    "FitsWCSConverter",
     "ColumnConverter",
     "AstropyTableConverter",
     "AsdfTableConverter",
@@ -51,7 +52,7 @@
     SkyCoordConverter,
     SpectralCoordConverter,
 )
-from .fits import AsdfFitsConverter, AstropyFitsConverter, FitsConverter
+from .fits import AsdfFitsConverter, AstropyFitsConverter, FitsConverter, FitsWCSConverter
 from .table import AsdfTableConverter, AstropyTableConverter, ColumnConverter, NdarrayMixinConverter
 from .time import TimeConverter, TimeDeltaConverter
 from .transform import (

asdf_astropy/converters/fits/__init__.py

@@ -2,6 +2,8 @@
     "FitsConverter",
     "AsdfFitsConverter",
     "AstropyFitsConverter",
+    "FitsWCSConverter",
 ]
 
 from .fits import AsdfFitsConverter, AstropyFitsConverter, FitsConverter
+from .fitswcs import FitsWCSConverter

asdf_astropy/converters/fits/fitswcs.py

@@ -0,0 +1,17 @@
+from asdf.extension import Converter
+
+
+class FitsWCSConverter(Converter):
+    tags = ("tag:astropy.org:astropy/fits/fitswcs-*",)
+    types = ("astropy.wcs.wcs.WCS",)
+
+    def from_yaml_tree(self, node, tag, ctx):
+        from astropy.wcs import WCS
+
+        primary_hdu = node["hdu"][0]
+        return WCS(primary_hdu.header)
+
+    def to_yaml_tree(self, wcs, tag, ctx):
+        node = {}
+        node["hdu"] = wcs.to_fits()
+        return node

asdf_astropy/converters/fits/tests/test_fitswcs.py

@@ -0,0 +1,31 @@
+import warnings
+
+import asdf
+import pytest
+from astropy.io import fits
+from astropy.utils.data import download_file
+from astropy.wcs import WCS, FITSFixedWarning
+
+
+def create_wcs():
+    urls = [
+        "http://data.astropy.org/tutorials/FITS-cubes/reduced_TAN_C14.fits",
+        "http://data.astropy.org/tutorials/FITS-images/HorseHead.fits",
+    ]
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", FITSFixedWarning)
+        return [WCS(fits.open(download_file(url, cache=True))[0].header) for url in urls]
+
+
+@pytest.mark.filterwarnings("ignore::astropy.wcs.wcs.FITSFixedWarning")
+@pytest.mark.parametrize("wcs", create_wcs())
+def test_wcs_serialization(wcs, tmp_path):
+    file_path = tmp_path / "test_wcs.asdf"
+    with asdf.AsdfFile() as af:
+        af["wcs"] = wcs
+        af.write_to(file_path)
+
+    with asdf.open(file_path) as af:
+        loaded_wcs = af["wcs"]
+        assert wcs.to_header() == loaded_wcs.to_header()

asdf_astropy/converters/slicedwcs/__init__.py

@@ -0,0 +1,5 @@
+__all__ = [
+    "SlicedWCSConverter",
+]
+
+from .slicedwcs import SlicedWCSConverter

asdf_astropy/converters/slicedwcs/slicedwcs.py

@@ -0,0 +1,34 @@
+from asdf.extension import Converter
+
+
+class SlicedWCSConverter(Converter):
+    tags = ("tag:astropy.org:astropy/slicedwcs/slicedwcs-*",)
+    types = ("astropy.wcs.wcsapi.wrappers.sliced_wcs.SlicedLowLevelWCS",)
+
+    def from_yaml_tree(self, node, tag, ctx):
+        from astropy.wcs.wcsapi.wrappers.sliced_wcs import SlicedLowLevelWCS
+
+        wcs = node["wcs"]
+        slice_array = []
+        slice_array = [
+            s if isinstance(s, int) else slice(s["start"], s["stop"], s["step"]) for s in node["slices_array"]
+        ]
+        return SlicedLowLevelWCS(wcs, slice_array)
+
+    def to_yaml_tree(self, sl, tag, ctx):
+        node = {}
+        node["wcs"] = sl._wcs
+        node["slices_array"] = []
+
+        for s in sl._slices_array:
+            if isinstance(s, slice):
+                node["slices_array"].append(
+                    {
+                        "start": s.start if s.start is not None else None,
+                        "stop": s.stop if s.stop is not None else None,
+                        "step": s.step if s.step is not None else None,
+                    },
+                )
+            else:
+                node["slices_array"].append(s)
+        return node

asdf_astropy/converters/slicedwcs/test/test_slicedwcs.py

@@ -0,0 +1,35 @@
+import asdf
+import numpy as np
+import pytest
+from astropy.wcs import WCS
+from astropy.wcs.wcsapi.wrappers.sliced_wcs import SlicedLowLevelWCS
+
+
+def create_wcs():
+    wcs = WCS(naxis=4)
+    wcs.wcs.ctype = "RA---CAR", "DEC--CAR", "FREQ", "TIME"
+    wcs.wcs.cunit = "deg", "deg", "Hz", "s"
+    wcs.wcs.cdelt = -2.0, 2.0, 3.0e9, 1
+    wcs.wcs.crval = 4.0, 0.0, 4.0e9, 3
+    wcs.wcs.crpix = 6.0, 7.0, 11.0, 11.0
+    wcs.wcs.cname = "Right Ascension", "Declination", "Frequency", "Time"
+    wcs0 = SlicedLowLevelWCS(wcs, 1)
+    wcs1 = SlicedLowLevelWCS(wcs, [slice(None), slice(None), slice(None), 10])
+    wcs3 = SlicedLowLevelWCS(SlicedLowLevelWCS(wcs, slice(None)), [slice(3), slice(None), slice(None), 10])
+    wcs_ellipsis = SlicedLowLevelWCS(wcs, [Ellipsis, slice(5, 10)])
+    wcs2 = SlicedLowLevelWCS(wcs, np.s_[:, 2, 3, :])
+    return [wcs0, wcs1, wcs2, wcs_ellipsis, wcs3]
+
+
+@pytest.mark.filterwarnings("ignore::astropy.wcs.wcs.FITSFixedWarning")
+@pytest.mark.parametrize("sl_wcs", create_wcs())
+def test_sliced_wcs_serialization(sl_wcs, tmp_path):
+    file_path = tmp_path / "test_slicedwcs.asdf"
+    with asdf.AsdfFile() as af:
+        af["sl_wcs"] = sl_wcs
+        af.write_to(file_path)
+
+    with asdf.open(file_path) as af:
+        loaded_sl_wcs = af["sl_wcs"]
+        assert sl_wcs._wcs.to_header() == loaded_sl_wcs._wcs.to_header()
+        assert sl_wcs._slices_array == loaded_sl_wcs._slices_array

asdf_astropy/extensions.py

@@ -13,6 +13,8 @@
 from .converters.coordinates.sky_coord import SkyCoordConverter
 from .converters.coordinates.spectral_coord import SpectralCoordConverter
 from .converters.fits.fits import AsdfFitsConverter, AstropyFitsConverter
+from .converters.fits.fitswcs import FitsWCSConverter
+from .converters.slicedwcs.slicedwcs import SlicedWCSConverter
 from .converters.table.table import AsdfTableConverter, AstropyTableConverter, ColumnConverter, NdarrayMixinConverter
 from .converters.time.time import TimeConverter
 from .converters.time.time_delta import TimeDeltaConverter
@@ -482,6 +484,8 @@
     AstropyTableConverter(),
     AstropyFitsConverter(),
     NdarrayMixinConverter(),
+    FitsWCSConverter(),
+    SlicedWCSConverter(),
 ]
 
 _COORDINATES_MANIFEST_URIS = [

asdf_astropy/resources/manifests/astropy-1.1.0.yaml

@@ -54,3 +54,20 @@ tags:
   title: NdarrayMixin column.
   description: |-
     Represents an astropy.table.NdarrayMixin instance.
+- tag_uri: tag:astropy.org:astropy/slicedwcs/slicedwcs-1.0.0
+  schema_uri: http://astropy.org/schemas/astropy/slicedwcs/slicedwcs-1.0.0
+  title: Represents an instance of SlicedLowLevelWCS
+  description: |-
+    The SlicedLowLevelWCS class is a wrapper class for WCS that applies slices
+    to the WCS, allowing certain pixel and world dimensions to be retained or
+    dropped.
+
+    It manages the slicing and coordinate transformations while preserving
+    the underlying WCS object.
+- tag_uri: tag:astropy.org:astropy/fits/fitswcs-1.0.0
+  schema_uri: http://astropy.org/schemas/astropy/fits/fitswcs-1.0.0
+  title: FITS WCS (World Coordinate System) Converter
+  description: |-
+    Represents the FITS WCS object, the HDUlist of the FITS header is preserved
+    during serialization and during deserialization the WCS object is recreated
+    from the HDUlist.

asdf_astropy/resources/schemas/fits/fitswcs-1.0.0.yaml

@@ -0,0 +1,18 @@
+%YAML 1.1
+---
+$schema: "http://stsci.edu/schemas/yaml-schema/draft-01"
+id: "http://astropy.org/schemas/astropy/fits/fitswcs-1.0.0"
+
+title:
+  Represents the fits object
+
+description:
+  Represents the fits object
+
+allOf:
+  - type: object
+    properties:
+      hdu:
+        tag: "tag:astropy.org:astropy/fits/fits-*"
+
+    required: ["hdu"]

asdf_astropy/resources/schemas/slicedwcs/slicedwcs-1.0.0.yaml

@@ -0,0 +1,24 @@
+%YAML 1.1
+---
+$schema: "http://stsci.edu/schemas/yaml-schema/draft-01"
+id: "http://astropy.org/schemas/astropy/slicedwcs/slicedwcs-1.0.0"
+
+title:
+  Represents the SlicedLowLevelWCS object
+
+description:
+  Represents the SlicedLowLevelWCS object
+
+allOf:
+  - type: object
+    properties:
+      wcs:
+        tag: "tag:astropy.org:astropy/fits/fitswcs-1*"
+      slices_array:
+        type: array
+        items:
+          - oneOf:
+            - type: integer
+            - type: object
+
+    required: ["wcs", "slices_array"]