Add support for the serialization of the ndcube WCS wrappers.

ndcube/asdf/converters/compoundwcs_converter.py

@@ -0,0 +1,18 @@
+from asdf.extension import Converter
+
+
+class CompoundConverter(Converter):
+    tags = ["tag:sunpy.org:ndcube/compoundwcs-0.1.0"]
+    types = ["ndcube.wcs.wrappers.compound_wcs.CompoundLowLevelWCS"]
+
+    def from_yaml_tree(self, node, tag, ctx):
+        from ndcube.wcs.wrappers import CompoundLowLevelWCS
+
+        return(CompoundLowLevelWCS(*node["wcs"], mapping=node.get("mapping"), pixel_atol=node.get("atol")))
+
+    def to_yaml_tree(self, compoundwcs, tag, ctx):
+        node = {}
+        node["wcs"] = compoundwcs._wcs
+        node["mapping"] = compoundwcs.mapping.mapping
+        node["atol"] = compoundwcs.atol
+        return node

ndcube/asdf/converters/ndcube_converter.py

@@ -39,9 +39,14 @@ def to_yaml_tree(self, ndcube, tag, ctx):
             This ensures that users are aware of potentially important information
             that is not included in the serialized output.
         """
+        from astropy.wcs.wcsapi import HighLevelWCSWrapper
+
         node = {}
         node["data"] = ndcube.data
-        node["wcs"] = ndcube.wcs
+        if isinstance(ndcube.wcs, HighLevelWCSWrapper):
+            node["wcs"] = ndcube.wcs._low_level_wcs
+        else:
+            node["wcs"] = ndcube.wcs
         node["extra_coords"] = ndcube.extra_coords
         node["global_coords"] = ndcube.global_coords
         node["meta"] = ndcube.meta

ndcube/asdf/converters/reorderedwcs_converter.py

@@ -0,0 +1,21 @@
+from asdf.extension import Converter
+
+
+class ReorderedConverter(Converter):
+    tags = ["tag:sunpy.org:ndcube/reorderedwcs-0.1.0"]
+    types = ["ndcube.wcs.wrappers.reordered_wcs.ReorderedLowLevelWCS"]
+
+    def from_yaml_tree(self, node, tag, ctx):
+        from ndcube.wcs.wrappers import ReorderedLowLevelWCS
+
+        reorderedwcs = ReorderedLowLevelWCS(wcs=node["wcs"],
+                                            pixel_order=node.get("pixel_order"),
+                                            world_order=node.get("world_order")
+                                            )
+        return reorderedwcs
+    def to_yaml_tree(self, reorderedwcs, tag, ctx):
+        node = {}
+        node["wcs"] = reorderedwcs._wcs
+        node["pixel_order"] = reorderedwcs._pixel_order
+        node["world_order"] = reorderedwcs._world_order
+        return node

ndcube/asdf/converters/resampled_converter.py

@@ -0,0 +1,22 @@
+from asdf.extension import Converter
+
+
+class ResampledConverter(Converter):
+    tags = ["tag:sunpy.org:ndcube/resampledwcs-0.1.0"]
+    types = ["ndcube.wcs.wrappers.resampled_wcs.ResampledLowLevelWCS"]
+
+    def from_yaml_tree(self, node, tag, ctx):
+        from ndcube.wcs.wrappers import ResampledLowLevelWCS
+
+        resampledwcs = ResampledLowLevelWCS(wcs=node["wcs"],
+                                            offset=node.get("offset"),
+                                            factor=node.get("factor"),
+                                            )
+        return resampledwcs
+    def to_yaml_tree(self, resampledwcs, tag, ctx):
+        node = {}
+        node["wcs"] = resampledwcs._wcs
+        node["factor"] = resampledwcs._factor
+        node["offset"] = resampledwcs._offset
+
+        return node

ndcube/asdf/converters/tests/test_ndcube_wcs_wrappers.py

@@ -0,0 +1,92 @@
+"""
+Tests for roundtrip serialization of NDCube with various GWCS types.
+
+TODO: Add tests for the roundtrip serialization of NDCube with ResampledLowLevelWCS, ReorderedLowLevelWCS, and CompoundLowLevelWCS when using astropy.wcs.WCS.
+"""
+
+import pytest
+from gwcs import __version__ as gwcs_version
+from packaging.version import Version
+
+import asdf
+
+from ndcube import NDCube
+from ndcube.conftest import data_nd
+from ndcube.tests.helpers import assert_cubes_equal
+from ndcube.wcs.wrappers import CompoundLowLevelWCS, ReorderedLowLevelWCS, ResampledLowLevelWCS
+
+
+@pytest.fixture
+def create_ndcube_resampledwcs(gwcs_3d_lt_ln_l):
+    shape = (2, 3, 4)
+    new_wcs = ResampledLowLevelWCS(wcs = gwcs_3d_lt_ln_l, factor=2 ,offset = 1)
+    data = data_nd(shape)
+    return NDCube(data = data, wcs =new_wcs)
+
+
+@pytest.mark.skipif(Version(gwcs_version) < Version("0.20"), reason="Requires gwcs>=0.20")
+def test_serialization_resampled(create_ndcube_resampledwcs, tmp_path):
+    ndc = create_ndcube_resampledwcs
+    file_path = tmp_path / "test.asdf"
+    with asdf.AsdfFile() as af:
+        af["ndcube"] = ndc
+        af.write_to(file_path)
+
+    with asdf.open(file_path) as af:
+        loaded_ndcube = af["ndcube"]
+
+        loaded_resampledwcs = loaded_ndcube.wcs.low_level_wcs
+        resampledwcs = ndc.wcs.low_level_wcs
+        assert (loaded_resampledwcs._factor == resampledwcs._factor).all()
+        assert (loaded_resampledwcs._offset == resampledwcs._offset).all()
+
+        assert_cubes_equal(loaded_ndcube, ndc)
+
+@pytest.fixture
+def create_ndcube_reorderedwcs(gwcs_3d_lt_ln_l):
+    shape = (2, 3, 4)
+    new_wcs = ReorderedLowLevelWCS(wcs = gwcs_3d_lt_ln_l, pixel_order=[1, 2, 0] ,world_order=[2, 0, 1])
+    data = data_nd(shape)
+    return NDCube(data = data, wcs =new_wcs)
+
+
+
+@pytest.mark.skipif(Version(gwcs_version) < Version("0.20"), reason="Requires gwcs>=0.20")
+def test_serialization_reordered(create_ndcube_reorderedwcs, tmp_path):
+    ndc = create_ndcube_reorderedwcs
+    file_path = tmp_path / "test.asdf"
+    with asdf.AsdfFile() as af:
+        af["ndcube"] = ndc
+        af.write_to(file_path)
+
+    with asdf.open(file_path) as af:
+        loaded_ndcube = af["ndcube"]
+
+        loaded_reorderedwcs = loaded_ndcube.wcs.low_level_wcs
+        reorderedwcs = ndc.wcs.low_level_wcs
+        assert (loaded_reorderedwcs._pixel_order == reorderedwcs._pixel_order)
+        assert (loaded_reorderedwcs._world_order == reorderedwcs._world_order)
+
+        assert_cubes_equal(loaded_ndcube, ndc)
+
+@pytest.fixture
+def create_ndcube_compoundwcs(gwcs_2d_lt_ln, time_and_simple_extra_coords_2d):
+
+    shape = (1, 2, 3, 4)
+    new_wcs = CompoundLowLevelWCS(gwcs_2d_lt_ln, time_and_simple_extra_coords_2d.wcs, mapping = [0, 1, 2, 3])
+    data = data_nd(shape)
+    return NDCube(data = data, wcs = new_wcs)
+
+@pytest.mark.skipif(Version(gwcs_version) < Version("0.20"), reason="Requires gwcs>=0.20")
+def test_serialization_compoundwcs(create_ndcube_compoundwcs, tmp_path):
+    ndc = create_ndcube_compoundwcs
+    file_path = tmp_path / "test.asdf"
+    with asdf.AsdfFile() as af:
+        af["ndcube"] = ndc
+        af.write_to(file_path)
+
+    with asdf.open(file_path) as af:
+        loaded_ndcube = af["ndcube"]
+        assert_cubes_equal(loaded_ndcube, ndc)
+        assert (loaded_ndcube.wcs.low_level_wcs.mapping.mapping == ndc.wcs.low_level_wcs.mapping.mapping)
+        assert (loaded_ndcube.wcs.low_level_wcs.atol == ndc.wcs.low_level_wcs.atol)

ndcube/asdf/entry_points.py

@@ -31,9 +31,12 @@ def get_extensions():
     """
     Get the list of extensions.
     """
+    from ndcube.asdf.converters.compoundwcs_converter import CompoundConverter
     from ndcube.asdf.converters.extracoords_converter import ExtraCoordsConverter
     from ndcube.asdf.converters.globalcoords_converter import GlobalCoordsConverter
     from ndcube.asdf.converters.ndcube_converter import NDCubeConverter
+    from ndcube.asdf.converters.reorderedwcs_converter import ReorderedConverter
+    from ndcube.asdf.converters.resampled_converter import ResampledConverter
     from ndcube.asdf.converters.tablecoord_converter import (
         QuantityTableCoordinateConverter,
         SkyCoordTableCoordinateConverter,
@@ -47,6 +50,9 @@ def get_extensions():
         QuantityTableCoordinateConverter(),
         SkyCoordTableCoordinateConverter(),
         GlobalCoordsConverter(),
+        ResampledConverter(),
+        ReorderedConverter(),
+        CompoundConverter(),
         ]
     _manifest_uri = "asdf://sunpy.org/ndcube/manifests/ndcube-0.1.0"
 

ndcube/asdf/resources/manifests/ndcube-0.1.0.yaml

@@ -23,3 +23,12 @@ tags:
 
   - tag_uri:  "tag:sunpy.org:ndcube/global_coords/globalcoords-0.1.0"
     schema_uri: "asdf://sunpy.org/ndcube/schemas/global_coords-0.1.0"
+
+  - tag_uri:  "tag:sunpy.org:ndcube/resampledwcs-0.1.0"
+    schema_uri: "asdf://sunpy.org/ndcube/schemas/resampledwcs-0.1.0"
+
+  - tag_uri:  "tag:sunpy.org:ndcube/reorderedwcs-0.1.0"
+    schema_uri: "asdf://sunpy.org/ndcube/schemas/reorderedwcs-0.1.0"
+
+  - tag_uri:  "tag:sunpy.org:ndcube/compoundwcs-0.1.0"
+    schema_uri: "asdf://sunpy.org/ndcube/schemas/compoundwcs-0.1.0"

ndcube/asdf/resources/schemas/compoundwcs-0.1.0.yaml

@@ -0,0 +1,25 @@
+%YAML 1.1
+---
+$schema: "http://stsci.edu/schemas/yaml-schema/draft-01"
+id: "asdf://sunpy.org/ndcube/schemas/Compoundwcs-0.1.0"
+
+title:
+  Represents the ndcube CompoundLowLevelWCS object
+
+description:
+  Represents the ndcube CompoundLowLevelWCS object
+
+type: object
+properties:
+  wcs:
+    type: array
+    items:
+      tag: "tag:stsci.edu:gwcs/wcs-1.*"
+  mapping:
+    type: array
+  atol:
+    type: number
+
+required: [wcs]
+additionalProperties: true
+...

ndcube/asdf/resources/schemas/ndcube-0.1.0.yaml

@@ -14,7 +14,11 @@ properties:
   data:
     description: "Must be compatible with ASDF serialization/deserialization and supported by NDCube."
   wcs:
-    tag: "tag:stsci.edu:gwcs/wcs-1.*"
+    anyOf:
+      - tag: "tag:stsci.edu:gwcs/wcs-1.*"
+      - tag: "tag:sunpy.org:ndcube/resampledwcs-*"
+      - tag: "tag:sunpy.org:ndcube/reorderedwcs-*"
+      - tag: "tag:sunpy.org:ndcube/compoundwcs-*"
   extra_coords:
     tag: "tag:sunpy.org:ndcube/extra_coords/extra_coords/extracoords-0.*"
   global_coords:

ndcube/asdf/resources/schemas/reorderedwcs-0.1.0.yaml

@@ -0,0 +1,23 @@
+%YAML 1.1
+---
+$schema: "http://stsci.edu/schemas/yaml-schema/draft-01"
+id: "asdf://sunpy.org/ndcube/schemas/resampledwcs-0.1.0"
+
+title:
+  Represents the ndcube ReorderedLowLevelWCS object
+
+description:
+  Represents the ndcube ReorderedLowLevelWCS object
+
+type: object
+properties:
+  wcs:
+    tag: "tag:stsci.edu:gwcs/wcs-1.*"
+  pixel_order:
+    type: array
+  world_order:
+    type: array
+
+required: [wcs]
+additionalProperties: true
+...

ndcube/asdf/resources/schemas/resampledwcs-0.1.0.yaml

@@ -0,0 +1,23 @@
+%YAML 1.1
+---
+$schema: "http://stsci.edu/schemas/yaml-schema/draft-01"
+id: "asdf://sunpy.org/ndcube/schemas/resampledwcs-0.1.0"
+
+title:
+  Represents the ndcube ResampledLowLevelWCS object
+
+description:
+  Represents the ndcube ResampledLowLevelWCS object
+
+type: object
+properties:
+  wcs:
+    tag: "tag:stsci.edu:gwcs/wcs-1.*"
+  factor:
+    tag: "tag:stsci.edu:asdf/core/ndarray-1.0.0"
+  offset:
+    tag: "tag:stsci.edu:asdf/core/ndarray-1.0.0"
+
+required: [wcs]
+additionalProperties: true
+...

ndcube/extra_coords/table_coord.py

@@ -11,7 +11,7 @@
 from astropy.coordinates import SkyCoord
 from astropy.modeling import models
 from astropy.modeling.models import tabular_model
-from astropy.modeling.tabular import _Tabular
+from astropy.modeling.tabular import Tabular1D, Tabular2D, _Tabular
 from astropy.time import Time
 from astropy.wcs.wcsapi.wrappers.sliced_wcs import combine_slices, sanitize_slices
 
@@ -136,7 +136,17 @@
         raise TypeError("lookup_table must be a Quantity.")  # pragma: no cover
 
     ndim = lookup_table.ndim
-    TabularND = tabular_model(ndim, name=f"Tabular{ndim}D")
+
+    # Use existing Tabular1D and Tabular2D classes for 1D and 2D models
+    # to ensure compatibility with asdf-astropy converters and avoid
+    # dynamically generated classes that are not recognized by asdf.
+    # See PR #751 for details on the issue this addresses.
+    if ndim == 1:
+        TabularND = Tabular1D
+    elif ndim == 2:
+        TabularND = Tabular2D
+    else:
+        TabularND = tabular_model(ndim, name=f"Tabular{ndim}D")
 
     # The integer location is at the centre of the pixel.
     points = [(np.arange(size) - 0) * points_unit for size in lookup_table.shape]