Add unit support to asdf

sunpy · Nov 6, 2024 · 6d81ee9 · 6d81ee9
1 parent b2fd5e1
commit 6d81ee9
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 8 deletions.
diff --git a/ndcube/asdf/converters/ndcube_converter.py b/ndcube/asdf/converters/ndcube_converter.py
@@ -13,7 +13,8 @@ def from_yaml_tree(self, node, tag, ctx):
         ndcube = NDCube(node["data"],
                         node["wcs"],
                         meta = node.get("meta"),
-                        mask = node.get("mask"))
+                        mask = node.get("mask"),
+                        unit = node.get("unit"))
         if "extra_coords" in node:
             ndcube._extra_coords = node["extra_coords"]
         if "global_coords" in node:
@@ -46,8 +47,10 @@ def to_yaml_tree(self, ndcube, tag, ctx):
         node["meta"] = ndcube.meta
         if ndcube.mask is not None:
             node["mask"] = ndcube.mask
+        if ndcube.unit is not None:
+            node["unit"] = ndcube.unit
 
-        attributes = ['uncertainty', 'unit']
+        attributes = ['uncertainty', 'psf']
         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)

diff --git a/ndcube/asdf/converters/tests/test_ndcube_converter.py b/ndcube/asdf/converters/tests/test_ndcube_converter.py
@@ -8,12 +8,13 @@
 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")
+@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",
+                                "ndcube_gwcs_4d_ln_lt_l_t_unit",
                                 ], indirect=("ndc",))
 @pytest.mark.skipif(Version(gwcs_version) < Version("0.20"), reason="Requires gwcs>=0.20")
 def test_serialization(ndc, tmp_path):

diff --git a/ndcube/asdf/resources/schemas/ndcube-0.1.0.yaml b/ndcube/asdf/resources/schemas/ndcube-0.1.0.yaml
@@ -23,6 +23,10 @@ properties:
     type: object
   mask:
     type: object
+  unit:
+    anyOf:
+     - tag: "tag:stsci.edu:asdf/unit/unit-1.*"
+     - tag: "tag:astropy.org:astropy/units/unit-1.*"
 
 required: [data, wcs]
 additionalProperties: true

diff --git a/ndcube/conftest.py b/ndcube/conftest.py
@@ -463,6 +463,15 @@ def ndcube_gwcs_4d_ln_lt_l_t(gwcs_4d_t_l_lt_ln):
     data_cube = data_nd(shape)
     return NDCube(data_cube, wcs=gwcs_4d_t_l_lt_ln)
 
+
+@pytest.fixture
+def ndcube_gwcs_4d_ln_lt_l_t_unit(gwcs_4d_t_l_lt_ln):
+    shape = (5, 8, 10, 12)
+    gwcs_4d_t_l_lt_ln.array_shape = shape
+    data_cube = data_nd(shape)
+    return NDCube(data_cube, wcs=gwcs_4d_t_l_lt_ln, unit=u.DN)
+
+
 @pytest.fixture
 def ndcube_gwcs_3d_ln_lt_l(gwcs_3d_lt_ln_l):
     shape = (2, 3, 4)