Asdf-Support

.github/workflows/ci.yml

@@ -55,6 +55,7 @@ jobs:
         - windows: py311
         - macos: py310
         - linux: py310-oldestdeps
+        - linux: asdf_schemas
     secrets:
       CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
 

ndcube/asdf/converters/extracoords_converter.py

@@ -0,0 +1,28 @@
+from asdf.extension import Converter
+
+
+class ExtraCoordsConverter(Converter):
+    tags = ["tag:sunpy.org:ndcube/extra_coords/extra_coords/extracoords-*"]
+    types = ["ndcube.extra_coords.extra_coords.ExtraCoords"]
+
+    def from_yaml_tree(self, node, tag, ctx):
+        from ndcube.extra_coords.extra_coords import ExtraCoords
+        extra_coords = ExtraCoords()
+        extra_coords._wcs = node.get("wcs")
+        extra_coords._mapping = node.get("mapping")
+        extra_coords._lookup_tables = node.get("lookup_tables", [])
+        extra_coords._dropped_tables = node.get("dropped_tables")
+        extra_coords._ndcube = node.get("ndcube")
+        return extra_coords
+
+    def to_yaml_tree(self, extracoords, tag, ctx):
+        node = {}
+        if extracoords._wcs is not None:
+            node["wcs"] = extracoords._wcs
+        if extracoords._mapping is not None:
+            node["mapping"] = extracoords._mapping
+        if extracoords._lookup_tables:
+            node["lookup_tables"] = extracoords._lookup_tables
+        node["dropped_tables"] = extracoords._dropped_tables
+        node["ndcube"] = extracoords._ndcube
+        return node

ndcube/asdf/converters/globalcoords_converter.py

@@ -0,0 +1,24 @@
+from asdf.extension import Converter
+
+
+class GlobalCoordsConverter(Converter):
+    tags = ["tag:sunpy.org:ndcube/global_coords/globalcoords-*"]
+    types = ["ndcube.global_coords.GlobalCoords"]
+
+    def from_yaml_tree(self, node, tag, ctx):
+        from ndcube.global_coords import GlobalCoords
+
+        globalcoords = GlobalCoords()
+        if "internal_coords" in node:
+            globalcoords._internal_coords = node["internal_coords"]
+        globalcoords._ndcube = node["ndcube"]
+
+        return globalcoords
+
+    def to_yaml_tree(self, globalcoords, tag, ctx):
+        node = {}
+        node["ndcube"] = globalcoords._ndcube
+        if globalcoords._internal_coords:
+            node["internal_coords"] = globalcoords._internal_coords
+
+        return node

ndcube/asdf/converters/ndcube_converter.py

@@ -0,0 +1,55 @@
+import warnings
+
+from asdf.extension import Converter
+
+
+class NDCubeConverter(Converter):
+    tags = ["tag:sunpy.org:ndcube/ndcube/ndcube-*"]
+    types = ["ndcube.ndcube.NDCube"]
+
+    def from_yaml_tree(self, node, tag, ctx):
+        from ndcube.ndcube import NDCube
+
+        ndcube = NDCube(node["data"],
+                        node["wcs"],
+                        meta = node.get("meta"),
+                        mask = node.get("mask"))
+        if "extra_coords" in node:
+            ndcube._extra_coords = node["extra_coords"]
+        if "global_coords" in node:
+            ndcube._global_coords = node["global_coords"]
+
+        return ndcube
+
+    def to_yaml_tree(self, ndcube, tag, ctx):
+        """
+        Notes
+        -----
+        This methods serializes the primary components of the NDCube object,
+        including the `data`, `wcs`, `extra_coords`, and `global_coords` attributes.
+        Issues a warning if unsupported attributes (uncertainty, mask, meta, unit) are present,
+        which are not currently serialized to ASDF.
+
+        Warnings
+        --------
+        UserWarning
+            Warns if the NDCube object has attributes 'uncertainty', 'mask',
+            or 'unit' that are present but not being saved in the ASDF serialization.
+            This ensures that users are aware of potentially important information
+            that is not included in the serialized output.
+        """
+        node = {}
+        node["data"] = ndcube.data
+        node["wcs"] = ndcube.wcs
+        node["extra_coords"] = ndcube.extra_coords
+        node["global_coords"] = ndcube.global_coords
+        node["meta"] = ndcube.meta
+        if ndcube.mask is not None:
+            node["mask"] = ndcube.mask
+
+        attributes = ['uncertainty', 'unit']
+        for attr in attributes:
+            if getattr(ndcube, attr) is not None:
+                warnings.warn(f"Attribute '{attr}' is present but not being saved in ASDF serialization.", UserWarning)
+
+        return node

ndcube/asdf/converters/tablecoord_converter.py

@@ -0,0 +1,80 @@
+from asdf.extension import Converter
+
+
+class TimeTableCoordConverter(Converter):
+    tags = ["tag:sunpy.org:ndcube/extra_coords/table_coord/timetablecoordinate-*"]
+    types = ["ndcube.extra_coords.table_coord.TimeTableCoordinate"]
+
+    def from_yaml_tree(self, node, tag, ctx):
+        from ndcube.extra_coords.table_coord import TimeTableCoordinate
+
+        names = node.get("names")
+        physical_types = node.get("physical_types")
+        reference_time = node.get("reference_time")
+        timetablecoordinate = TimeTableCoordinate(
+            node["table"], names=names, physical_types=physical_types, reference_time=reference_time)
+
+        return timetablecoordinate
+
+    def to_yaml_tree(self, timetablecoordinate, tag, ctx):
+        node = {}
+        node["table"] = timetablecoordinate.table
+        node["names"] = timetablecoordinate.names
+        if timetablecoordinate.physical_types is not None:
+            node["physical_types"] = timetablecoordinate.physical_types
+        node["reference_time"] = timetablecoordinate.reference_time
+
+        return node
+
+
+class QuantityTableCoordinateConverter(Converter):
+    tags = ["tag:sunpy.org:ndcube/extra_coords/table_coord/quantitytablecoordinate-*"]
+    types = ["ndcube.extra_coords.table_coord.QuantityTableCoordinate"]
+
+    def from_yaml_tree(self, node, tag, ctx):
+        from ndcube.extra_coords.table_coord import QuantityTableCoordinate
+
+        names = node.get("names")
+        mesh = node.get("mesh")
+        physical_types = node.get("physical_types")
+        quantitytablecoordinate = QuantityTableCoordinate(*node["table"],
+                                                          names=names, physical_types=physical_types)
+        quantitytablecoordinate.unit = node["unit"]
+        quantitytablecoordinate.mesh = mesh
+        return quantitytablecoordinate
+
+    def to_yaml_tree(self, quantitytablecoordinate, tag, ctx):
+        node = {}
+        node["unit"] = quantitytablecoordinate.unit
+        node["table"] = quantitytablecoordinate.table
+        node["names"] = quantitytablecoordinate.names
+        node["mesh"] = quantitytablecoordinate.mesh
+        if quantitytablecoordinate.physical_types is not None:
+            node["physical_types"] = quantitytablecoordinate.physical_types
+
+        return node
+
+
+class SkyCoordTableCoordinateConverter(Converter):
+    tags = ["tag:sunpy.org:ndcube/extra_coords/table_coord/skycoordtablecoordinate-*"]
+    types = ["ndcube.extra_coords.table_coord.SkyCoordTableCoordinate"]
+
+    def from_yaml_tree(self, node, tag, ctx):
+        from ndcube.extra_coords.table_coord import SkyCoordTableCoordinate
+
+        names = node.get("names")
+        mesh = node.get("mesh")
+        physical_types = node.get("physical_types")
+        skycoordinatetablecoordinate = SkyCoordTableCoordinate(node["table"], mesh=mesh,
+                                                               names=names, physical_types=physical_types)
+        return skycoordinatetablecoordinate
+
+    def to_yaml_tree(self, skycoordinatetablecoordinate, tag, ctx):
+        node = {}
+        node["table"] = skycoordinatetablecoordinate.table
+        node["names"] = skycoordinatetablecoordinate.names
+        node["mesh"] = skycoordinatetablecoordinate.mesh
+        if skycoordinatetablecoordinate.physical_types is not None:
+            node["physical_types"] = skycoordinatetablecoordinate.physical_types
+
+        return node

ndcube/asdf/converters/tests/test_ndcube_converter.py

@@ -0,0 +1,49 @@
+import numpy as np
+import pytest
+from gwcs import __version__ as gwcs_version
+from packaging.version import Version
+
+import asdf
+
+from ndcube.tests.helpers import assert_cubes_equal
+
+
+@pytest.mark.parametrize("ndc",[("ndcube_gwcs_2d_ln_lt_mask"),
+                                ("ndcube_gwcs_3d_ln_lt_l"),
+                                ("ndcube_gwcs_3d_ln_lt_l_ec_dropped_dim"),
+                                ("ndcube_gwcs_3d_ln_lt_l_ec_q_t_gc"),
+                                ("ndcube_gwcs_3d_rotated"),
+                                ("ndcube_gwcs_4d_ln_lt_l_t")
+                                ], indirect=("ndc",))
+@pytest.mark.skipif(Version(gwcs_version) < Version("0.20"), reason="Requires gwcs>=0.20")
+def test_serialization(ndc, tmp_path):
+    file_path = tmp_path / "test.asdf"
+    with asdf.AsdfFile() as af:
+        af["ndcube_gwcs"] = ndc
+        af.write_to(file_path)
+
+    with asdf.open(file_path) as af:
+        assert_cubes_equal(af["ndcube_gwcs"], ndc)
+
+
+@pytest.mark.xfail(reason="Serialization of sliced ndcube not supported")
+def test_serialization_sliced_ndcube(ndcube_gwcs_3d_ln_lt_l, tmp_path):
+    sndc = ndcube_gwcs_3d_ln_lt_l[np.s_[0, :, :]]
+    file_path = tmp_path / "test.asdf"
+    with asdf.AsdfFile() as af:
+        af["ndcube_gwcs"] = sndc
+        af.write_to(file_path)
+
+    with asdf.open(file_path) as af:
+        assert_cubes_equal(af["ndcube_gwcs"], sndc)
+
+
+@pytest.mark.xfail(reason="Serialization of ndcube with .wcs attribute as astropy.wcs.wcs.WCS not supported")
+def test_serialization_ndcube_wcs(ndcube_3d_ln_lt_l, tmp_path):
+    file_path = tmp_path / "test.asdf"
+    with asdf.AsdfFile() as af:
+        af["ndcube"] = ndcube_3d_ln_lt_l
+        af.write_to(file_path)
+
+    with asdf.open(file_path) as af:
+        assert_cubes_equal(af["ndcube"], ndcube_3d_ln_lt_l)

ndcube/asdf/entry_points.py

@@ -0,0 +1,55 @@
+"""
+This file contains the entry points for asdf.
+"""
+import importlib.resources as importlib_resources
+
+from asdf.extension import ManifestExtension
+from asdf.resource import DirectoryResourceMapping
+
+
+def get_resource_mappings():
+    """
+    Get the resource mapping instances for myschemas
+    and manifests.  This method is registered with the
+    asdf.resource_mappings entry point.
+
+    Returns
+    -------
+    list of collections.abc.Mapping
+    """
+    from ndcube.asdf import resources
+    resources_root = importlib_resources.files(resources)
+    return [
+        DirectoryResourceMapping(
+            resources_root / "schemas", "asdf://sunpy.org/ndcube/schemas/"),
+        DirectoryResourceMapping(
+            resources_root / "manifests", "asdf://sunpy.org/ndcube/manifests/"),
+    ]
+
+
+def get_extensions():
+    """
+    Get the list of extensions.
+    """
+    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.tablecoord_converter import (
+        QuantityTableCoordinateConverter,
+        SkyCoordTableCoordinateConverter,
+        TimeTableCoordConverter,
+    )
+
+    ndcube_converters = [
+        NDCubeConverter(),
+        ExtraCoordsConverter(),
+        TimeTableCoordConverter(),
+        QuantityTableCoordinateConverter(),
+        SkyCoordTableCoordinateConverter(),
+        GlobalCoordsConverter(),
+        ]
+    _manifest_uri = "asdf://sunpy.org/ndcube/manifests/ndcube-0.1.0"
+
+    return [
+        ManifestExtension.from_uri(_manifest_uri, converters=ndcube_converters)
+    ]

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

@@ -0,0 +1,25 @@
+%YAML 1.1
+---
+id: asdf://sunpy.org/ndcube/manifests/ndcube-0.1.0
+extension_uri: asdf://sunpy.org/extensions/ndcube-0.1.0
+title: NDCube ASDF Manifest
+description: ASDF schemas and tags for NDCube classes.
+
+tags:
+  - tag_uri: "tag:sunpy.org:ndcube/ndcube/ndcube-0.1.0"
+    schema_uri: "asdf://sunpy.org/ndcube/schemas/ndcube-0.1.0"
+
+  - tag_uri:  "tag:sunpy.org:ndcube/extra_coords/extra_coords/extracoords-0.1.0"
+    schema_uri: "asdf://sunpy.org/ndcube/schemas/extra_coords-0.1.0"
+
+  - tag_uri:  "tag:sunpy.org:ndcube/extra_coords/table_coord/timetablecoordinate-0.1.0"
+    schema_uri: "asdf://sunpy.org/ndcube/schemas/timetablecoordinate-0.1.0"
+
+  - tag_uri:  "tag:sunpy.org:ndcube/extra_coords/table_coord/quantitytablecoordinate-0.1.0"
+    schema_uri: "asdf://sunpy.org/ndcube/schemas/quantitytablecoordinate-0.1.0"
+
+  - tag_uri:  "tag:sunpy.org:ndcube/extra_coords/table_coord/skycoordtablecoordinate-0.1.0"
+    schema_uri: "asdf://sunpy.org/ndcube/schemas/skycoordtablecoordinate-0.1.0"
+
+  - tag_uri:  "tag:sunpy.org:ndcube/global_coords/globalcoords-0.1.0"
+    schema_uri: "asdf://sunpy.org/ndcube/schemas/global_coords-0.1.0"

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

@@ -0,0 +1,37 @@
+%YAML 1.1
+---
+$schema: "http://stsci.edu/schemas/yaml-schema/draft-01"
+id: "asdf://sunpy.org/ndcube/schemas/extra_coords-0.1.0"
+
+title:
+  Represents the ndcube ExtraCoords object
+
+description:
+  Represents the ndcube ExtraCoords object
+
+type: object
+properties:
+  wcs:
+    tag: "tag:stsci.edu:gwcs/wcs-1.*"
+  mapping:
+    type: array
+  lookup_tables:
+    type: array
+    items:
+        type: array
+        items:
+            - oneOf:
+              - type: number
+              - type: array
+            - oneOf:
+              - tag: "tag:sunpy.org:ndcube/extra_coords/table_coord/quantitytablecoordinate-0.*"
+              - tag: "tag:sunpy.org:ndcube/extra_coords/table_coord/skycoordtablecoordinate-0.*"
+              - tag: "tag:sunpy.org:ndcube/extra_coords/table_coord/timetablecoordinate-0.*"
+  dropped_tables:
+    type: array
+  ndcube:
+    tag: "tag:sunpy.org:ndcube/ndcube/ndcube-0.*"
+
+required: [ndcube]
+additionalProperties: false
+...

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

@@ -0,0 +1,29 @@
+%YAML 1.1
+---
+$schema: "http://stsci.edu/schemas/yaml-schema/draft-01"
+id: "asdf://sunpy.org/ndcube/schemas/global_coords-0.1.0"
+
+title:
+  Represents the ndcube GlobalCoords object
+
+description:
+  Represents the ndcube GlobalCoords object
+
+type: object
+properties:
+  internal_coords:
+    type: object
+    additionalProperties:
+      type: array
+      items:
+          - type: string
+          - type: object
+            oneOf:
+              - tag: "tag:stsci.edu:asdf/unit/quantity-*"
+              - tag: "tag:astropy.org:astropy/coordinates/skycoord-*"
+  ndcube:
+    tag: "tag:sunpy.org:ndcube/ndcube/ndcube-0.*"
+
+required: [ndcube]
+additionalProperties: false
+...