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,8 @@
     "FitsConverter",
     "AsdfFitsConverter",
     "AstropyFitsConverter",
+    "FitsWCSConverter",
+    "SlicedWCSConverter",
     "ColumnConverter",
     "AstropyTableConverter",
     "AsdfTableConverter",
@@ -74,3 +76,4 @@
     UnitsMappingConverter,
 )
 from .unit import EquivalencyConverter, MagUnitConverter, QuantityConverter, UnitConverter
+from .wcs import FitsWCSConverter, SlicedWCSConverter

asdf_astropy/converters/wcs/__init__.py

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

asdf_astropy/converters/wcs/fitswcs.py

@@ -0,0 +1,19 @@
+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
+
+        return WCS(node["hdu"][0].header, fobj=node["hdu"])
+
+    def to_yaml_tree(self, wcs, tag, ctx):
+        node = {}
+        if wcs.sip is not None:
+            node["hdu"] = wcs.to_fits(relax=True)
+        else:
+            node["hdu"] = wcs.to_fits()
+        return node

asdf_astropy/converters/wcs/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,
+                        "stop": s.stop,
+                        "step": s.step,
+                    },
+                )
+            else:
+                node["slices_array"].append(s)
+        return node

asdf_astropy/converters/wcs/tests/test_fitswcs.py

@@ -0,0 +1,89 @@
+import asdf
+import numpy as np
+import pytest
+from astropy import wcs
+
+from asdf_astropy.testing.helpers import assert_hdu_list_equal
+
+
+def create_sip_distortion_wcs():
+    rng = np.random.default_rng(42)
+    twcs = wcs.WCS(naxis=2)
+    twcs.wcs.crval = [251.29, 57.58]
+    twcs.wcs.cdelt = [1, 1]
+    twcs.wcs.crpix = [507, 507]
+    twcs.wcs.pc = np.array([[7.7e-6, 3.3e-5], [3.7e-5, -6.8e-6]])
+    twcs._naxis = [1014, 1014]
+    twcs.wcs.ctype = ["RA---TAN-SIP", "DEC--TAN-SIP"]
+
+    # Generate random SIP coefficients
+    a = rng.uniform(low=-1e-5, high=1e-5, size=(5, 5))
+    b = rng.uniform(low=-1e-5, high=1e-5, size=(5, 5))
+
+    # Assign SIP coefficients
+    twcs.sip = wcs.Sip(a, b, None, None, twcs.wcs.crpix)
+    twcs.wcs.set()
+
+    return (twcs,)
+
+
+@pytest.mark.xfail(reason="Fails due to normalization differences when using wcs.to_fits().")
+@pytest.mark.parametrize("wcs", create_sip_distortion_wcs())
+@pytest.mark.filterwarnings("ignore::astropy.wcs.wcs.FITSFixedWarning")
+@pytest.mark.filterwarnings(
+    "ignore:Some non-standard WCS keywords were excluded:astropy.utils.exceptions.AstropyWarning",
+)
+def test_sip_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_hdu_list_equal(wcs.to_fits(relax=True), loaded_wcs.to_fits(relax=True))
+
+
+def create_tabular_wcs():
+    # Creates a WCS object with distortion lookup tables
+    img_world_wcs = wcs.WCS(naxis=2)
+    img_world_wcs.wcs.crpix = 1, 1
+    img_world_wcs.wcs.crval = 0, 0
+    img_world_wcs.wcs.cdelt = 1, 1
+
+    # Create maps with zero distortion except at one particular pixel
+    x_dist_array = np.zeros((25, 25))
+    x_dist_array[10, 20] = 0.5
+    map_x = wcs.DistortionLookupTable(
+        x_dist_array.astype(np.float32),
+        (5, 10),
+        (10, 20),
+        (2, 2),
+    )
+    y_dist_array = np.zeros((25, 25))
+    y_dist_array[10, 5] = 0.7
+    map_y = wcs.DistortionLookupTable(
+        y_dist_array.astype(np.float32),
+        (5, 10),
+        (10, 20),
+        (3, 3),
+    )
+
+    img_world_wcs.cpdis1 = map_x
+    img_world_wcs.cpdis2 = map_y
+
+    return (img_world_wcs,)
+
+
+@pytest.mark.parametrize("wcs", create_tabular_wcs())
+@pytest.mark.filterwarnings("ignore::astropy.wcs.wcs.FITSFixedWarning")
+def test_twcs_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()
+        assert_hdu_list_equal(wcs.to_fits(), loaded_wcs.to_fits())

asdf_astropy/converters/wcs/tests/test_slicedwcs.py

@@ -0,0 +1,38 @@
+import asdf
+import numpy as np
+import pytest
+from astropy.wcs import WCS
+from astropy.wcs.wcsapi.wrappers.sliced_wcs import SlicedLowLevelWCS
+
+from asdf_astropy.testing.helpers import assert_hdu_list_equal
+
+
+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_hdu_list_equal(sl_wcs._wcs.to_fits(), loaded_sl_wcs._wcs.to_fits())
+        assert sl_wcs._slices_array == loaded_sl_wcs._slices_array

asdf_astropy/extensions.py

@@ -31,6 +31,8 @@
 from .converters.unit.magunit import MagUnitConverter
 from .converters.unit.quantity import QuantityConverter
 from .converters.unit.unit import UnitConverter
+from .converters.wcs.fitswcs import FitsWCSConverter
+from .converters.wcs.slicedwcs import SlicedWCSConverter
 
 __all__ = [
     "TRANSFORM_CONVERTERS",
@@ -482,6 +484,8 @@
     AstropyTableConverter(),
     AstropyFitsConverter(),
     NdarrayMixinConverter(),
+    FitsWCSConverter(),
+    SlicedWCSConverter(),
 ]
 
 _COORDINATES_MANIFEST_URIS = [
@@ -499,6 +503,7 @@
 
 
 _ASTROPY_EXTENSION_MANIFEST_URIS = [
+    "asdf://astropy.org/astropy/manifests/astropy-1.3.0",
     "asdf://astropy.org/astropy/manifests/astropy-1.2.0",
     "asdf://astropy.org/astropy/manifests/astropy-1.1.0",
     "asdf://astropy.org/astropy/manifests/astropy-1.0.0",

asdf_astropy/resources/manifests/astropy-1.3.0.yaml

@@ -0,0 +1,73 @@
+id: asdf://astropy.org/astropy/manifests/astropy-1.3.0
+extension_uri: asdf://astropy.org/astropy/extensions/astropy-1.3.0
+title: Astropy extension 1.3.0
+description: |-
+  A set of tags for serializing astropy objects.  This does not include most
+  model classes, which are handled by an implementation of the ASDF
+  transform extension.
+asdf_standard_requirement:
+  gte: 1.5.0
+tags:
+- tag_uri: tag:astropy.org:astropy/time/timedelta-1.0.0
+  schema_uri: http://astropy.org/schemas/astropy/time/timedelta-1.0.0
+  title: Represents an instance of TimeDelta from astropy
+  description: |-
+    Represents the time difference between two times.
+- tag_uri: tag:astropy.org:astropy/fits/fits-1.0.0
+  schema_uri: http://astropy.org/schemas/astropy/fits/fits-1.0.0
+  title: A FITS file inside of an ASDF file.
+  description: |-
+    This schema is useful for distributing ASDF files that can
+    automatically be converted to FITS files by specifying the exact
+    content of the resulting FITS file.
+
+    Not all kinds of data in FITS are directly representable in ASDF.
+    For example, applying an offset and scale to the data using the
+    `BZERO` and `BSCALE` keywords.  In these cases, it will not be
+    possible to store the data in the native format from FITS and also
+    be accessible in its proper form in the ASDF file.
+
+    Only image and binary table extensions are supported.
+- tag_uri: tag:astropy.org:astropy/table/table-1.1.0
+  schema_uri: http://astropy.org/schemas/astropy/table/table-1.1.0
+  title: A table.
+  description: |-
+    A table is represented as a list of columns, where each entry is a
+    [column](ref:http://stsci.edu/schemas/asdf/core/column-1.0.0)
+    object, containing the data and some additional information.
+
+    The data itself may be stored inline as text, or in binary in either
+    row- or column-major order by use of the `strides` property on the
+    individual column arrays.
+
+    Each column in the table must have the same first (slowest moving)
+    dimension.
+- tag_uri: tag:astropy.org:astropy/transform/units_mapping-1.0.0
+  schema_uri: http://astropy.org/schemas/astropy/transform/units_mapping-1.0.0
+  title: Mapper that operates on the units of the input.
+  description: |-
+    This transform operates on the units of the input, first converting to
+    the expected input units, then assigning replacement output units without
+    further conversion.
+- tag_uri: tag:astropy.org:astropy/table/ndarraymixin-1.0.0
+  schema_uri: http://astropy.org/schemas/astropy/table/ndarraymixin-1.0.0
+  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 WCS object, the HDUlist of the FITS header is preserved
+  during serialization and during deserialization the WCS object is recreated
+  from the HDUlist.
+
+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,40 @@
+%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: 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.
+
+allOf:
+  - type: object
+    properties:
+      wcs:
+        tag: "tag:astropy.org:astropy/fits/fitswcs-1*"
+      slices_array:
+        type: array
+        items:
+          - oneOf:
+            - type: integer
+            - type: object
+              properties:
+                start:
+                  anyOf:
+                    - type: integer
+                    - type: "null"
+                stop:
+                  anyOf:
+                    - type: integer
+                    - type: "null"
+                step:
+                  anyOf:
+                    - type: integer
+                    - type: "null"
+
+
+    required: ["wcs", "slices_array"]