diff --git a/CHANGES.rst b/CHANGES.rst index 4027b550..a380fbb0 100644 --- a/CHANGES.rst +++ b/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) ------------------ diff --git a/asdf_astropy/converters/__init__.py b/asdf_astropy/converters/__init__.py index 92d7730a..3334fa0b 100644 --- a/asdf_astropy/converters/__init__.py +++ b/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 diff --git a/asdf_astropy/converters/wcs/__init__.py b/asdf_astropy/converters/wcs/__init__.py new file mode 100644 index 00000000..73e65a8f --- /dev/null +++ b/asdf_astropy/converters/wcs/__init__.py @@ -0,0 +1,6 @@ +__all__ = [ + "WCSConverter", + "SlicedWCSConverter", +] +from .slicedwcs import SlicedWCSConverter +from .wcs import WCSConverter diff --git a/asdf_astropy/converters/wcs/slicedwcs.py b/asdf_astropy/converters/wcs/slicedwcs.py new file mode 100644 index 00000000..34eea6d6 --- /dev/null +++ b/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, + } diff --git a/asdf_astropy/converters/wcs/tests/__init__.py b/asdf_astropy/converters/wcs/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/asdf_astropy/converters/wcs/tests/test_slicedwcs.py b/asdf_astropy/converters/wcs/tests/test_slicedwcs.py new file mode 100644 index 00000000..f845c2b2 --- /dev/null +++ b/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 diff --git a/asdf_astropy/converters/wcs/tests/test_wcs.py b/asdf_astropy/converters/wcs/tests/test_wcs.py new file mode 100644 index 00000000..70206fd5 --- /dev/null +++ b/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) diff --git a/asdf_astropy/converters/wcs/wcs.py b/asdf_astropy/converters/wcs/wcs.py new file mode 100644 index 00000000..3ed9f7ad --- /dev/null +++ b/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} diff --git a/asdf_astropy/extensions.py b/asdf_astropy/extensions.py index ffaf0904..0bda0584 100644 --- a/asdf_astropy/extensions.py +++ b/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 = [ diff --git a/asdf_astropy/resources/manifests/astropy-1.2.0.yaml b/asdf_astropy/resources/manifests/astropy-1.2.0.yaml index 80f63bb1..e6d8839d 100644 --- a/asdf_astropy/resources/manifests/astropy-1.2.0.yaml +++ b/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. diff --git a/asdf_astropy/resources/manifests/astropy-1.3.0.yaml b/asdf_astropy/resources/manifests/astropy-1.3.0.yaml index 53c3720a..964412fe 100644 --- a/asdf_astropy/resources/manifests/astropy-1.3.0.yaml +++ b/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. diff --git a/asdf_astropy/resources/schemas/wcs/slicedwcs-1.0.0.yaml b/asdf_astropy/resources/schemas/wcs/slicedwcs-1.0.0.yaml new file mode 100644 index 00000000..b3464da2 --- /dev/null +++ b/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"] diff --git a/asdf_astropy/resources/schemas/wcs/wcs-1.0.0.yaml b/asdf_astropy/resources/schemas/wcs/wcs-1.0.0.yaml new file mode 100644 index 00000000..3ab83217 --- /dev/null +++ b/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"] diff --git a/asdf_astropy/testing/helpers.py b/asdf_astropy/testing/helpers.py index c0cd064c..f4d18921 100644 --- a/asdf_astropy/testing/helpers.py +++ b/asdf_astropy/testing/helpers.py @@ -152,6 +152,27 @@ def assert_table_equal(a, b): ) +def assert_wcs_equal(a, b): + from asdf.tags.core.ndarray import NDArrayType + + from asdf_astropy.converters.wcs.wcs import _WCS_ATTRS + + assert type(a) == type(b) + assert_hdu_list_equal(a.to_fits(relax=True), b.to_fits(relax=True)) + for attr in _WCS_ATTRS: + in_a = hasattr(a, attr) + in_b = hasattr(b, attr) + assert in_a == in_b + if not in_a: + continue + a_val = getattr(a, attr) + b_val = getattr(b, attr) + if isinstance(a_val, (np.ndarray, NDArrayType)) or isinstance(b_val, (np.ndarray, NDArrayType)): + np.testing.assert_array_equal(a_val, b_val) + else: + assert a_val == b_val + + def assert_table_roundtrip(table, tmp_path): """ Assert that a table can be written to an ASDF file and read back @@ -237,3 +258,15 @@ def assert_model_roundtrip(model, tmp_path, version=None): with asdf.open(path) as af: assert_model_equal(model, af["model"]) return af["model"] + + +def assert_wcs_roundtrip(wcs, tmp_path, version=None): + import asdf + + path = str(tmp_path / "test.asdf") + + with asdf.AsdfFile({"wcs": wcs}, version=version) as af: + af.write_to(path) + + with asdf.open(path) as af: + assert_wcs_equal(wcs, af["wcs"]) diff --git a/docs/asdf-astropy/schemas.rst b/docs/asdf-astropy/schemas.rst index 30d5ea4c..2f670a9d 100644 --- a/docs/asdf-astropy/schemas.rst +++ b/docs/asdf-astropy/schemas.rst @@ -81,3 +81,17 @@ units/equivalency-1.0.0 .. asdf-autoschemas:: units/equivalency-1.0.0 + +WCS +--- + +The following schemas are associated with **astropy** objects from the +:ref:`astropy-wcs` submodule: + +wcs/wcs-1.0.0 +^^^^^^^^^^^^^^^ + +.. asdf-autoschemas:: + + wcs/wcs-1.0.0 + wcs/slicedwcs-1.0.0 diff --git a/pyproject.toml b/pyproject.toml index 5cfadc13..d817bebe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -97,6 +97,7 @@ extend-ignore = [ "SLF001", # private-member-access "FBT001", # boolean positional arguments in function definition "RUF012", # mutable class attributes should be annotated with typing.ClassVar + "PTH123", # replace open with Path.open ] extend-exclude = ["docs/*"]