Add support for ASDF serialisation and deserialisation

.github/workflows/ci.yml

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

.ruff.toml

@@ -82,6 +82,11 @@ extend-ignore = [
 "test_*.py" = [
   "E402", # Module level import not at top of cell
 ]
+# ASDF converters delay imports for performance reasons
+"ndcube/asdf/**.py" = [
+  "E402", # Module level import not at top of cell
+  "PLC0415", # `import` should be at the top-level of a file
+]
 
 [lint.pydocstyle]
 convention = "numpy"

changelog/776.breaking.rst

@@ -0,0 +1,9 @@
+The minimum supported version of some dependencies has increased:
+
+* astropy >= 5.3
+* gwcs >= 0.20
+* numpy >= 1.25
+* scipy >= 1.11
+* matplotlib >= 3.8
+* mpl_animators >= 1.1
+* reproject >= 0.11

changelog/776.feature.rst

@@ -0,0 +1 @@
+Add support for serialization of most `ndcube` objects to ASDF files.

docs/conf.py

@@ -94,7 +94,8 @@
     'sunpy': ('https://docs.sunpy.org/en/stable/', None),
     'mpl_animators': ('https://docs.sunpy.org/projects/mpl-animators/en/stable/', None),
     'gwcs': ('https://gwcs.readthedocs.io/en/stable/', None),
-    'reproject': ("https://reproject.readthedocs.io/en/stable/", None)
+    'reproject': ("https://reproject.readthedocs.io/en/stable/", None),
+    'asdf': ("https://www.asdf-format.org/projects/asdf/en/stable/", None),
     }
 
 # -- Options for HTML output -------------------------------------------------

docs/explaining_ndcube/asdf_serialization.rst

@@ -0,0 +1,67 @@
+.. _asdf_serialization:
+
+*************************
+Saving ND objects to ASDF
+*************************
+
+:ref:`asdf` is an extensible format for validating and saving complex scientific data along with its metadata.
+`ndcube` provides schemas and converters for all the ND objects (`~ndcube.NDCube`, `~ndcube.NDCubeSequence` and `~ndcube.NDCollection`) as well as for various WCS and table objects required by them.
+To make use of these, simply save an ND object to an ASDF file and it will be correctly serialized.
+ASDF files save a "tree" which is a `dict`.
+You can save any number of cubes in your ASDF by adding them to the dictionary.
+
+.. expanding-code-block:: python
+  :summary: Click to reveal/hide instantiation of the NDCube.
+
+  >>> import numpy as np
+  >>> import asdf
+  >>> import astropy.wcs
+  >>> from ndcube import NDCube
+
+  >>> # Define data array.
+  >>> data = np.random.rand(4, 4, 5)
+
+  >>> # Define WCS transformations in an astropy WCS object.
+  >>> wcs = astropy.wcs.WCS(naxis=3)
+  >>> wcs.wcs.ctype = 'WAVE', 'HPLT-TAN', 'HPLN-TAN'
+  >>> wcs.wcs.cunit = 'Angstrom', 'deg', 'deg'
+  >>> wcs.wcs.cdelt = 0.2, 0.5, 0.4
+  >>> wcs.wcs.crpix = 0, 2, 2
+  >>> wcs.wcs.crval = 10, 0.5, 1
+  >>> wcs.wcs.cname = 'wavelength', 'HPC lat', 'HPC lon'
+
+  >>> # Now instantiate the NDCube
+  >>> my_cube = NDCube(data, wcs=wcs)
+
+
+.. code-block:: python
+
+  >>> my_tree = {"mycube": my_cube}
+  >>> with asdf.AsdfFile(tree=my_tree) as f:  # doctest:  +SKIP
+  ...     f.write_to("somefile.asdf")  # doctest:  +SKIP
+
+
+What's Supported and What Isn't
+===============================
+
+We aim to support all features of `ndcube` when saving and loading to ASDF.
+However, because it is possible to create `ndcube` objects with many different components (for example dask arrays) which aren't part of the `ndcube` package these may not be supported.
+Many common components of `ndcube` classes are supported in the `asdf_astropy <https://asdf-astropy.readthedocs.io/en/stable/>`__ package, such as `astropy.wcs.WCS`, `astropy.wcs.wcsapi.SlicedLowLevelWCS` and uncertainty objects.
+
+The only component of the `ndcube.NDCube` class which is never saved is the ``.psf`` attribute.
+
+`ndcube` implements converters and schemas for the following objects:
+
+* `~ndcube.NDCube`
+* `~ndcube.NDCubeSequence`
+* `~ndcube.NDCollection`
+* `~ndcube.NDMeta`
+* `~ndcube.GlobalCoords`
+* `~ndcube.ExtraCoords`
+* `~ndcube.extra_coords.TimeTableCoordinate`
+* `~ndcube.extra_coords.QuantityTableCoordinate`
+* `~ndcube.extra_coords.SkyCoordTableCoordinate`
+* `~ndcube.extra_coords.MultipleTableCoordinate`
+* `~ndcube.wcs.wrappers.ReorderedLowLevelWCS`
+* `~ndcube.wcs.wrappers.ResampledLowLevelWCS`
+* `~ndcube.wcs.wrappers.CompoundLowLevelWCS`

docs/explaining_ndcube/index.rst

@@ -14,3 +14,4 @@ Explaining ``ndcube``
    tabular_coordinates
    reproject
    visualization
+   asdf_serialization

ndcube/asdf/converters/compoundwcs_converter.py

@@ -0,0 +1,18 @@
+from asdf.extension import Converter
+
+
+class CompoundConverter(Converter):
+    tags = ["tag:sunpy.org:ndcube/compoundwcs-*"]
+    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/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/ndcollection_converter.py

@@ -0,0 +1,23 @@
+from asdf.extension import Converter
+
+
+class NDCollectionConverter(Converter):
+    tags = ["tag:sunpy.org:ndcube/ndcollection-*"]
+    types = ["ndcube.ndcollection.NDCollection"]
+
+    def from_yaml_tree(self, node, tag, ctx):
+        from ndcube.ndcollection import NDCollection
+
+        aligned_axes = list(node.get("aligned_axes").values())
+        aligned_axes = tuple(tuple(lst) for lst in aligned_axes)
+        return NDCollection(node["items"], meta=node.get("meta"), aligned_axes=aligned_axes)
+
+    def to_yaml_tree(self, ndcollection, tag, ctx):
+        node = {}
+        node["items"] = dict(ndcollection)
+        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/ndcube_converter.py

@@ -0,0 +1,64 @@
+import warnings
+
+from asdf.extension import Converter
+
+
+class NDCubeConverter(Converter):
+    tags = ["tag:sunpy.org: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"),
+            unit=node.get("unit"),
+            uncertainty=node.get("uncertainty"),
+        )
+        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 are present.
+
+        Warnings
+        --------
+        UserWarning
+            Warns if the NDCube object has a 'psf' attribute that will not be
+            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
+        # NDData always has .wcs as a high level wcs
+        node["wcs"] = ndcube.wcs.low_level_wcs
+        if not ndcube.extra_coords.is_empty:
+            node["extra_coords"] = ndcube.extra_coords
+        if ndcube.global_coords._all_coords:
+            node["global_coords"] = ndcube.global_coords
+        if ndcube.meta:
+            node["meta"] = ndcube.meta
+        if ndcube.mask is not None:
+            node["mask"] = ndcube.mask
+        if ndcube.unit is not None:
+            node["unit"] = ndcube.unit
+        if ndcube.uncertainty is not None:
+            node["uncertainty"] = ndcube.uncertainty
+
+        if getattr(ndcube, 'psf') is not None:
+            warnings.warn("Attribute 'psf' is present but not being saved in ASDF serialization.", UserWarning)
+
+        return node

ndcube/asdf/converters/ndcubesequence_converter.py

@@ -0,0 +1,23 @@
+from asdf.extension import Converter
+
+
+class NDCubeSequenceConverter(Converter):
+    tags = ["tag:sunpy.org:ndcube/ndcubesequence-*"]
+    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/ndmeta_converter.py

@@ -0,0 +1,24 @@
+import numpy as np
+
+from asdf.extension import Converter
+
+
+class NDMetaConverter(Converter):
+    tags = ["tag:sunpy.org:ndcube/meta/ndmeta-*"]
+    types = ["ndcube.meta.NDMeta"]
+
+    def from_yaml_tree(self, node, tag, ctx):
+        from ndcube.meta import NDMeta
+        axes = {k: np.array(v) for k, v in node["axes"].items()}
+        meta = NDMeta(node["meta"], node["key_comments"], axes, node["data_shape"])
+        meta._original_meta = node["original_meta"]
+        return meta
+
+    def to_yaml_tree(self, meta, tag, ctx):
+        node = {}
+        node["meta"] = dict(meta)
+        node["key_comments"] = meta.key_comments
+        node["axes"] = meta.axes
+        node["data_shape"] = meta.data_shape
+        node["original_meta"] = meta._original_meta  # not the MappingProxy object
+        return node

ndcube/asdf/converters/reorderedwcs_converter.py

@@ -0,0 +1,22 @@
+from asdf.extension import Converter
+
+
+class ReorderedConverter(Converter):
+    tags = ["tag:sunpy.org:ndcube/reorderedwcs-*"]
+    types = ["ndcube.wcs.wrappers.reordered_wcs.ReorderedLowLevelWCS"]
+
+    def from_yaml_tree(self, node, tag, ctx):
+        from ndcube.wcs.wrappers import ReorderedLowLevelWCS
+
+        return ReorderedLowLevelWCS(
+            wcs=node["wcs"],
+            pixel_order=node.get("pixel_order"),
+            world_order=node.get("world_order"),
+        )
+
+    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,23 @@
+from asdf.extension import Converter
+
+
+class ResampledConverter(Converter):
+    tags = ["tag:sunpy.org:ndcube/resampledwcs-*"]
+    types = ["ndcube.wcs.wrappers.resampled_wcs.ResampledLowLevelWCS"]
+
+    def from_yaml_tree(self, node, tag, ctx):
+        from ndcube.wcs.wrappers import ResampledLowLevelWCS
+
+        return ResampledLowLevelWCS(
+            wcs=node["wcs"],
+            offset=node.get("offset"),
+            factor=node.get("factor"),
+        )
+
+    def to_yaml_tree(self, resampledwcs, tag, ctx):
+        node = {}
+        node["wcs"] = resampledwcs._wcs
+        node["factor"] = resampledwcs._factor
+        node["offset"] = resampledwcs._offset
+
+        return node