Merge pull request #246 from braingram/fits_wcs

Fits wcs

CHANGES.rst

@@ -12,6 +12,8 @@
 
 - Add support for astropy.nddata.uncertainty classes [#239]
 
+- Add support for astropy.wcs.WCS and astropy.wcs.wcsapi.SlicedLowLevelWCS [#246]
+
 0.6.1 (2024-04-05)
 ------------------
 

asdf_astropy/converters/__init__.py

@@ -15,6 +15,8 @@
     "FitsConverter",
     "AsdfFitsConverter",
     "AstropyFitsConverter",
+    "WCSConverter",
+    "SlicedWCSConverter",
     "ColumnConverter",
     "AstropyTableConverter",
     "AsdfTableConverter",
@@ -76,3 +78,4 @@
     UnitsMappingConverter,
 )
 from .unit import EquivalencyConverter, MagUnitConverter, QuantityConverter, UnitConverter
+from .wcs import SlicedWCSConverter, WCSConverter

asdf_astropy/converters/wcs/__init__.py

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

asdf_astropy/converters/wcs/slicedwcs.py

@@ -0,0 +1,34 @@
+from asdf.extension import Converter
+
+
+class SlicedWCSConverter(Converter):
+    tags = ("tag:astropy.org:astropy/wcs/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 = [
+            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):
+        slices_array = []
+
+        for s in sl._slices_array:
+            if isinstance(s, slice):
+                slices_array.append(
+                    {
+                        "start": s.start,
+                        "stop": s.stop,
+                        "step": s.step,
+                    },
+                )
+            else:
+                slices_array.append(s)
+        return {
+            "wcs": sl._wcs,
+            "slices_array": slices_array,
+        }

asdf_astropy/converters/wcs/tests/test_slicedwcs.py

@@ -0,0 +1,37 @@
+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_wcs_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.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_wcs_equal(sl_wcs._wcs, loaded_sl_wcs._wcs)
+        assert sl_wcs._slices_array == loaded_sl_wcs._slices_array

asdf_astropy/converters/wcs/tests/test_wcs.py

@@ -0,0 +1,120 @@
+import warnings
+
+import numpy as np
+import pytest
+from astropy.io import fits
+from astropy.utils.data import get_pkg_data_filename, get_pkg_data_filenames
+from astropy.wcs import WCS, DistortionLookupTable, FITSFixedWarning, Sip
+
+from asdf_astropy.testing.helpers import assert_wcs_roundtrip
+
+_astropy_test_header_filenames = list(get_pkg_data_filenames("tests/data/maps", "astropy.wcs", "*.hdr")) + list(
+    get_pkg_data_filenames("tests/data/spectra", "astropy.wcs", "*.hdr"),
+)
+
+_astropy_test_fits_filenames = [
+    get_pkg_data_filename(f"tests/data/{fn}", "astropy.wcs")
+    for fn in [
+        "ie6d07ujq_wcs.fits",
+        "j94f05bgq_flt.fits",
+        "sip.fits",
+        "sip2.fits",
+    ]
+]
+
+
+def create_empty_wcs():
+    return WCS()
+
+
+def create_wcs_with_attrs():
+    wcs = WCS(naxis=3)
+    wcs.pixel_shape = [100, 200, 300]
+    wcs.pixel_bounds = [[11, 22], [33, 45], [55, 67]]
+    return wcs
+
+
+def create_sip_distortion_wcs():
+    rng = np.random.default_rng(42)
+    wcs = WCS(naxis=2)
+    wcs.wcs.crval = [251.29, 57.58]
+    wcs.wcs.cdelt = [1, 1]
+    wcs.wcs.crpix = [507, 507]
+    wcs.wcs.pc = np.array([[7.7e-6, 3.3e-5], [3.7e-5, -6.8e-6]])
+    wcs._naxis = [1014, 1014]
+    wcs.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
+    wcs.sip = Sip(a, b, None, None, wcs.wcs.crpix)
+    wcs.wcs.set()
+
+    return wcs
+
+
+def create_tabular_wcs():
+    # Creates a WCS object with distortion lookup tables
+    img_world_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 = 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 = 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("version", ["1.5.0", "1.6.0"])
+@pytest.mark.parametrize(
+    "wcs_gen",
+    [create_empty_wcs, create_wcs_with_attrs, create_tabular_wcs, create_sip_distortion_wcs],
+)
+def test_roundtrip(wcs_gen, tmp_path, version):
+    wcs = wcs_gen()
+    assert_wcs_roundtrip(wcs, tmp_path, version)
+
+
+@pytest.mark.parametrize("fn", _astropy_test_header_filenames)
+@pytest.mark.parametrize("version", ["1.5.0", "1.6.0"])
+def test_astropy_data_header_roundtrip(fn, tmp_path, version):
+    with open(fn) as f:
+        header = f.read()
+
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", category=FITSFixedWarning)
+        wcs = WCS(header)
+
+    assert_wcs_roundtrip(wcs, tmp_path, version)
+
+
+@pytest.mark.parametrize("fn", _astropy_test_fits_filenames)
+@pytest.mark.parametrize("version", ["1.5.0", "1.6.0"])
+def test_astropy_data_fits_roundtrip(fn, tmp_path, version):
+    with fits.open(fn) as ff:
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=FITSFixedWarning)
+            wcs = WCS(ff[0].header, ff)
+
+        assert_wcs_roundtrip(wcs, tmp_path, version)

asdf_astropy/converters/wcs/wcs.py

@@ -0,0 +1,37 @@
+from asdf.extension import Converter
+
+# These attributes don't end up in the hdulist and
+# instead will be stored in "attrs"
+_WCS_ATTRS = ("naxis", "colsel", "keysel", "key", "pixel_bounds")
+
+
+class WCSConverter(Converter):
+    tags = ("tag:astropy.org:astropy/wcs/wcs-*",)
+    types = ("astropy.wcs.wcs.WCS",)
+
+    def from_yaml_tree(self, node, tag, ctx):
+        from astropy.wcs import WCS
+
+        hdulist = node["hdulist"]
+        attrs = node["attrs"]
+
+        if naxis := attrs.pop("naxis"):
+            hdulist[0].header["naxis"] = naxis
+
+        pixel_bounds = attrs.pop("pixel_bounds")
+
+        wcs = WCS(hdulist[0].header, fobj=hdulist, **attrs)
+
+        if wcs.sip is not None:
+            # work around a bug in astropy where sip headers lose precision
+            # see https://github.com/astropy/astropy/issues/17334
+            wcs.sip = wcs._read_sip_kw(hdulist[0].header, attrs.get("key", " "))
+            wcs.wcs.set()
+
+        wcs.pixel_bounds = pixel_bounds
+        return wcs
+
+    def to_yaml_tree(self, wcs, tag, ctx):
+        hdulist = wcs.to_fits(relax=True)
+        attrs = {a: getattr(wcs, a) for a in _WCS_ATTRS if hasattr(wcs, a)}
+        return {"hdulist": hdulist, "attrs": attrs}

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.slicedwcs import SlicedWCSConverter
+from .converters.wcs.wcs import WCSConverter
 
 __all__ = [
     "TRANSFORM_CONVERTERS",
@@ -478,6 +480,8 @@
     AstropyFitsConverter(),
     NdarrayMixinConverter(),
     UncertaintyConverter(),
+    WCSConverter(),
+    SlicedWCSConverter(),
 ]
 
 _COORDINATES_MANIFEST_URIS = [

asdf_astropy/resources/manifests/astropy-1.2.0.yaml

@@ -59,3 +59,20 @@ tags:
   title: Represent an astropy.nddata uncertainty
   description: |-
     Represents an instance of an astropy.nddata uncertainty
+- tag_uri: tag:astropy.org:astropy/wcs/slicedwcs-1.0.0
+  schema_uri: http://astropy.org/schemas/astropy/wcs/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/wcs/wcs-1.0.0
+  schema_uri: http://astropy.org/schemas/astropy/wcs/wcs-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/manifests/astropy-1.3.0.yaml

@@ -54,8 +54,20 @@ tags:
   title: NdarrayMixin column.
   description: |-
     Represents an astropy.table.NdarrayMixin instance.
-- tag_uri: tag:astropy.org:astropy/nddata/uncertainty-1.0.0
-  schema_uri: http://astropy.org/schemas/astropy/nddata/uncertainty-1.0.0
-  title: Represent an astropy.nddata uncertainty
+- tag_uri: tag:astropy.org:astropy/wcs/slicedwcs-1.0.0
+  schema_uri: http://astropy.org/schemas/astropy/wcs/slicedwcs-1.0.0
+  title: Represents an instance of SlicedLowLevelWCS
   description: |-
-    Represents an instance of an astropy.nddata uncertainty
+    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/wcs/wcs-1.0.0
+  schema_uri: http://astropy.org/schemas/astropy/wcs/wcs-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/wcs/slicedwcs-1.0.0.yaml

@@ -0,0 +1,41 @@
+%YAML 1.1
+---
+$schema: "http://stsci.edu/schemas/yaml-schema/draft-01"
+id: "http://astropy.org/schemas/astropy/wcs/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/wcs/wcs-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"]

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

@@ -0,0 +1,23 @@
+%YAML 1.1
+---
+$schema: "http://stsci.edu/schemas/yaml-schema/draft-01"
+id: "http://astropy.org/schemas/astropy/wcs/wcs-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:
+      hdulist:
+        title: "HDUList produced by WCS.to_fits"
+        tag: "tag:astropy.org:astropy/fits/fits-*"
+      attrs:
+        title: "extra WCS attributes not contained in hdulist"
+        type: object
+
+    required: ["hdulist", "attrs"]