diff --git a/ndcube/asdf/converters/ndcube_converter.py b/ndcube/asdf/converters/ndcube_converter.py index c88b8bdff..32e572cfd 100644 --- 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 index ad5df5ee7..649b08b2e 100644 --- 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 index 5529ccbfa..5ef9a976e 100644 --- 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 index 6d0dcf9c5..2d3cd359f 100644 --- 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)