Add serialization logic for the NDCubeSequence and NDCollection

Author: ViciousEagle03

ndcube/asdf/converters/ndcollection_converter.py

@@ -0,0 +1,30 @@
+from asdf.extension import Converter
+
+
+class NDCollectionConverter(Converter):
+    tags = ["tag:sunpy.org:ndcube/ndcube/ndcollection-*"]
+    types = ["ndcube.ndcollection.NDCollection"]
+
+    def from_yaml_tree(self, node, tag, ctx):
+        from ndcube.ndcollection import NDCollection
+
+        key_value_pairs = list(zip(node["keys"], node["value"]))
+        aligned_axes = list(node.get("aligned_axes").values())
+        aligned_axes = tuple(tuple(lst) for lst in aligned_axes)
+        ndcollection = NDCollection(key_value_pairs,
+                            meta=node.get("meta"),
+                            aligned_axes = aligned_axes)
+        return ndcollection
+
+
+    def to_yaml_tree(self, ndcollection, tag, ctx):
+
+        node = {}
+        node["keys"] = tuple(ndcollection.keys())
+        node["value"] = tuple(ndcollection.values())
+        if ndcollection.meta is not None:
+            node["meta"] = ndcollection.meta
+        if ndcollection._aligned_axes is not None:
+            node["aligned_axes"] = ndcollection._aligned_axes
+
+        return node

ndcube/asdf/converters/ndcubesequence_converter.py

@@ -0,0 +1,25 @@
+from asdf.extension import Converter
+
+
+class NDCubeSequenceConverter(Converter):
+    tags = ["tag:sunpy.org:ndcube/ndcube/ndcube_sequence-*"]
+    types = ["ndcube.ndcube_sequence.NDCubeSequence"]
+
+    def from_yaml_tree(self, node, tag, ctx):
+        from ndcube.ndcube_sequence import NDCubeSequence
+
+        return NDCubeSequence(node["data"],
+                              meta=node.get("meta"),
+                              common_axis=node.get("common_axis"))
+
+
+    def to_yaml_tree(self, ndcseq, tag, ctx):
+
+        node = {}
+        node["data"] = ndcseq.data
+        if ndcseq.meta is not None:
+            node["meta"] = ndcseq.meta
+        if ndcseq._common_axis is not None:
+            node["common_axis"] = ndcseq._common_axis
+
+        return node

ndcube/asdf/converters/tests/test_ndcollection_converter.py

@@ -0,0 +1,51 @@
+import pytest
+from gwcs import __version__ as gwcs_version
+from packaging.version import Version
+
+import asdf
+
+from ndcube.ndcollection import NDCollection
+from ndcube.ndcube_sequence import NDCubeSequence
+from ndcube.tests.helpers import assert_collections_equal
+
+
+@pytest.fixture
+def create_ndcollection_cube(ndcube_gwcs_3d_ln_lt_l, ndcube_gwcs_3d_ln_lt_l_ec_q_t_gc, ndcube_gwcs_3d_ln_lt_l_ec_dropped_dim):
+    aligned_axes = ((1, 2), (1, 2), (1, 2))
+    cube_collection = NDCollection([("cube0", ndcube_gwcs_3d_ln_lt_l),
+                                    ("cube1", ndcube_gwcs_3d_ln_lt_l_ec_q_t_gc),
+                                    ("cube2", ndcube_gwcs_3d_ln_lt_l_ec_dropped_dim)],
+                                   aligned_axes=aligned_axes)
+
+    return cube_collection
+
+
+@pytest.mark.skipif(Version(gwcs_version) < Version("0.20"), reason="Requires gwcs>=0.20")
+def test_serialization_cube(create_ndcollection_cube, tmp_path):
+    ndcollection = create_ndcollection_cube
+    file_path = tmp_path / "test.asdf"
+    with asdf.AsdfFile() as af:
+        af["ndcube_gwcs"] = ndcollection
+        af.write_to(file_path)
+
+    with asdf.open(file_path) as af:
+        assert_collections_equal(af["ndcube_gwcs"], ndcollection)
+@pytest.fixture
+def create_ndcollection_sequence(ndcube_gwcs_3d_ln_lt_l, ndcube_gwcs_3d_ln_lt_l_ec_dropped_dim):
+
+    sequence02 = NDCubeSequence([ndcube_gwcs_3d_ln_lt_l, ndcube_gwcs_3d_ln_lt_l_ec_dropped_dim])
+    sequence20 = NDCubeSequence([ndcube_gwcs_3d_ln_lt_l_ec_dropped_dim, ndcube_gwcs_3d_ln_lt_l])
+    seq_collection = NDCollection([("seq0", sequence02), ("seq1", sequence20)], aligned_axes="all")
+    return seq_collection
+
+
+@pytest.mark.skipif(Version(gwcs_version) < Version("0.20"), reason="Requires gwcs>=0.20")
+def test_serialization_sequence(create_ndcollection_sequence, tmp_path):
+    ndcollection = create_ndcollection_sequence
+    file_path = tmp_path / "test.asdf"
+    with asdf.AsdfFile() as af:
+        af["ndcube_gwcs"] = ndcollection
+        af.write_to(file_path)
+
+    with asdf.open(file_path) as af:
+        assert_collections_equal(af["ndcube_gwcs"], ndcollection)

ndcube/asdf/converters/tests/test_ndcubesequence_converter.py

@@ -0,0 +1,20 @@
+import pytest
+from gwcs import __version__ as gwcs_version
+from packaging.version import Version
+
+import asdf
+
+from ndcube.ndcube_sequence import NDCubeSequence
+from ndcube.tests.helpers import assert_cubesequences_equal
+
+
+@pytest.mark.skipif(Version(gwcs_version) < Version("0.20"), reason="Requires gwcs>=0.20")
+def test_serialization(ndcube_gwcs_3d_ln_lt_l, ndcube_gwcs_3d_ln_lt_l_ec_q_t_gc, tmp_path):
+    file_path = tmp_path / "test.asdf"
+    ndcseq = NDCubeSequence([ndcube_gwcs_3d_ln_lt_l, ndcube_gwcs_3d_ln_lt_l_ec_q_t_gc], common_axis=1)
+    with asdf.AsdfFile() as af:
+        af["ndcube_gwcs"] = ndcseq
+        af.write_to(file_path)
+
+    with asdf.open(file_path) as af:
+        assert_cubesequences_equal(af["ndcube_gwcs"], ndcseq)

ndcube/asdf/entry_points.py

@@ -33,20 +33,23 @@ def get_extensions():
     """
     from ndcube.asdf.converters.extracoords_converter import ExtraCoordsConverter
     from ndcube.asdf.converters.globalcoords_converter import GlobalCoordsConverter
+    from ndcube.asdf.converters.ndcollection_converter import NDCollectionConverter
     from ndcube.asdf.converters.ndcube_converter import NDCubeConverter
+    from ndcube.asdf.converters.ndcubesequence_converter import NDCubeSequenceConverter
     from ndcube.asdf.converters.tablecoord_converter import (
         QuantityTableCoordinateConverter,
         SkyCoordTableCoordinateConverter,
         TimeTableCoordConverter,
     )
-
     ndcube_converters = [
         NDCubeConverter(),
         ExtraCoordsConverter(),
         TimeTableCoordConverter(),
         QuantityTableCoordinateConverter(),
         SkyCoordTableCoordinateConverter(),
         GlobalCoordsConverter(),
+        NDCubeSequenceConverter(),
+        NDCollectionConverter(),
         ]
     _manifest_uri = "asdf://sunpy.org/ndcube/manifests/ndcube-0.1.0"
 

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

@@ -23,3 +23,9 @@ 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/ndcube/ndcube_sequence-0.1.0"
+    schema_uri: "asdf://sunpy.org/ndcube/schemas/ndcube_sequence-0.1.0"
+
+  - tag_uri:  "tag:sunpy.org:ndcube/ndcube/ndcollection-0.1.0"
+    schema_uri: "asdf://sunpy.org/ndcube/schemas/ndcollection-0.1.0"

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

@@ -0,0 +1,30 @@
+%YAML 1.1
+---
+$schema: "http://stsci.edu/schemas/yaml-schema/draft-01"
+id: "asdf://sunpy.org/ndcube/schemas/ndcollection-0.1.0"
+
+title:
+  Represents the ndcube.ndcollection.NDCollection object
+
+description:
+  Represents the ndcube ndcube.ndcollection.NDCollection object
+
+type: object
+properties:
+  keys:
+    type: array
+  value:
+    type: array
+    items:
+        - type: object
+          oneOf:
+            - tag: "tag:sunpy.org:ndcube/ndcube/ndcube-0.*"
+            - tag: "tag:sunpy.org:ndcube/ndcube/ndcube_sequence-0.*"
+  aligned_axes:
+    anyOf:
+    - type: object
+    - type: string
+
+required: [keys, value]
+additionalProperties: true
+...

ndcube/asdf/resources/schemas/ndcube_sequence-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/ndcube_sequence-0.1.0"
+
+title:
+  Represents the ndcube.ndcube_sequence.NDCubeSequence object
+
+description:
+  Represents the ndcube.ndcube_sequence.NDCubeSequence object
+
+type: object
+properties:
+  data:
+    type: array
+    items:
+      tag: "tag:sunpy.org:ndcube/ndcube/ndcube-0.*"
+  meta:
+    type: object
+  common_axis:
+    type: integer
+
+required: [data]
+additionalProperties: true
+...