Repository: astropy/asdf-astropy

Issue:
Support MaskedQuantity
Astropy (v5.0+) has the `Masked` class (https://docs.astropy.org/en/stable/utils/masked/index.html), which can produce `MaskedQuantity`. It would be great if ASDF could serialize Masked objects.

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/astropy/asdf-astropy/pull/253?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astropy) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.41%. Comparing base [(`b36f1ab`)](https://app.codecov.io/gh/astropy/asdf-astropy/commit/b36f1ab8c75bcadfb7beb5bcf67b89cd2fa3468a?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astropy) to head [(`3397b36`)](https://app.codecov.io/gh/astropy/asdf-astropy/commit/3397b362a967ae77abbea1f9f476befa69cbf97a?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astropy).
> Report is 5 commits behind head on main.

<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##             main     #253   +/-   ##
=======================================
  Coverage   99.40%   99.41%           
=======================================
  Files          62       63    +1     
  Lines        2200     2216   +16     
=======================================
+ Hits         2187     2203   +16     
  Misses         13       13           
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/astropy/asdf-astropy/pull/253?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astropy).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astropy).

lpsinger: > I read this as masked stating that the `MaskedQuantity` class exists at `astropy.utils.masked.core.MaskedQuantity` but it doesn't. For asdf to check for instances to handle with the Converter it needs to have access to the class.

That's right. `Masked` is apparently a class factory, and types are dynamically created here: https://github.com/astropy/astropy/blob/v7.0.0/astropy/utils/masked/core.py#L266-L271

The fully qualified name `astropy.utils.masked.core.MaskedQuantity` seems to be totally fictional because it is not added to the module dictionary. Even if it wasn't, I don't see how that fully qualified name would distinguish between instances of the class `Masked(astropy.units.Quantity)` vs `Masked(foo.bar.Quantity)` (a class with a name that is coincidentally the same but comes from a completely different module.

Is there a way to configure a converter to decide whether it is capable of handling instances of a class by calling a function?
braingram: > Is there a way to configure a converter to decide whether it is capable of handling instances of a class by calling a function?

Not that I'm aware of. The matching of converter to object uses the object type:
https://github.com/asdf-format/asdf/blob/a0898c2d8d5afad0f8cf9a0d60847263d786f9e5/asdf/yamlutil.py#L297
and the assumption is that either:
- asdf has indexed the class (it was provided via `Converter.tags` as is done in this PR at the expense of importing the required modules to make the class at "integration" time)
- asdf has the "path" via which the class can be imported (so asdf can check if the contained module is imported and index the converter as needed)

(see this [comment](https://github.com/asdf-format/asdf/blob/a0898c2d8d5afad0f8cf9a0d60847263d786f9e5/asdf/extension/_manager.py#L64) for more details).

I suspect that more extensive changes to asdf will be needed to support `Masked` since the classes are not publicly importable via astropy (we can't provide a "path" to the class in `Converter.tags`).
lpsinger: > I suspect that more extensive changes to asdf will be needed to support `Masked` since the classes are not publicly importable via astropy (we can't provide a "path" to the class in `Converter.tags`).

What do you have in mind?

It is turning out to be surprisingly simple to implement `Masked` classes; the significant hurdle is how to register the classes without importing them.
lpsinger: One option could be to have Astropy add those dynamically-created subclasses to the module dictionary, so that (for example) `Masked(astropy.units.Quantity)` is available as `astropy.utils.masked.Masked_astropy_units_Quantity` (would need some way to transform the fully qualified name of the data class so that it is a valid Python identifier).
braingram: > > I suspect that more extensive changes to asdf will be needed to support `Masked` since the classes are not publicly importable via astropy (we can't provide a "path" to the class in `Converter.tags`).
> 
> What do you have in mind?

I will try to carve out some time to look into this soon. At the moment I don't have a clear solution to propose.


lpsinger: > > > I suspect that more extensive changes to asdf will be needed to support `Masked` since the classes are not publicly importable via astropy (we can't provide a "path" to the class in `Converter.tags`).
> > 
> > 
> > What do you have in mind?
> 
> I will try to carve out some time to look into this soon. At the moment I don't have a clear solution to propose.

For a possible workaround, see https://github.com/astropy/astropy/pull/17685.
pllim: Are you able to run a job here using your branch at https://github.com/astropy/astropy/pull/17685 and see if it is actually green? Thanks! (devdeps is usually the easiest to hack for such things)
lpsinger: > Are you able to run a job here using your branch at [astropy/astropy#17685](https://github.com/astropy/astropy/pull/17685) and see if it is actually green? Thanks! (devdeps is usually the easiest to hack for such things)

... maybe? See latest commit 🤞 
lpsinger: > Oh right, your fork won't have the right tags for setuptools_scm. You can make a fake tag on your fork, or otherwise relax any astropy pins here. The former might be easier.

I'm not really sure how to to make this work in the CI against a branch of astropy. But the tests I added are passing on my computer.
pllim: `astropy==0.2.dev37528+g4b8a62975`

Try grab v7.1.dev tag from astropy proper and push that tag to your fork.  Then rerun the jobs.
pllim: https://github.com/astropy/astropy/releases/tag/v7.1.dev
braingram: https://github.com/astropy/asdf-astropy/pull/254/commits/dfee9194c90c7e6da17106f4487a77be11a3addb
is one way to appease (parts of) the CI. With that change the tests (added in this PR) are passing with the astropy PR:
https://github.com/astropy/asdf-astropy/actions/runs/13018230750/job/36312572977?pr=254
@lpsinger would you integrate that change with this PR (and drop the tox changes). Not everything will pass (the py310 jobs will fail with astropy dev which doesn't support python 3.10) but hopefully it will illustrate that the astropy changes make this PR possible.
braingram: Overall this looks good to me. Once the astropy PR is merged we can:
- restore the pyproject.toml astropy version
- update this PR with the appropriate version checks to get things working with all supported (and dev) astropy versions

Thanks again for the PR and adding support for `Masked` classes!

Patch:
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index c6e0288..f8730e5 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -68,8 +68,9 @@ jobs:
       # Any env name which does not start with `pyXY` will use this Python version.
       default_python: '3.13'
       envs: |
-        - macos: py311-test-parallel
-        - windows: py311-test-parallel
+        - macos: py311-test-parallel-cov
+        - windows: py311-test-parallel-cov
+      coverage: codecov
 
   dev:
     needs: [core, asdf-schemas]
@@ -82,8 +83,9 @@ jobs:
       # Any env name which does not start with `pyXY` will use this Python version.
       default_python: '3.13'
       envs: |
-        - linux: py313-test-predeps-parallel
-        - linux: py313-test-devdeps-parallel
+        - linux: py313-test-predeps-parallel-cov
+        - linux: py313-test-devdeps-parallel-cov
+      coverage: codecov
 
   oldest:
     needs: [core, asdf-schemas]
diff --git a/CHANGES.rst b/CHANGES.rst
index c9faa72..3db138f 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -3,6 +3,7 @@
 
 - drop support for python 3.10 and numpy 1.24. [#255]
 - remove soon to be deprecated and non-working use of astropy TestRunner. [#263]
+- Add support for masked quantities, angles, and coordinates for astropy >= 7.1. [#253]
 
 0.7.1 (2025-02-12)
 ------------------
diff --git a/asdf_astropy/converters/coordinates/angle.py b/asdf_astropy/converters/coordinates/angle.py
index bcd70b0..3a50486 100644
--- a/asdf_astropy/converters/coordinates/angle.py
+++ b/asdf_astropy/converters/coordinates/angle.py
@@ -6,6 +6,7 @@ class AngleConverter(QuantityConverter):
     types = (
         "astropy.coordinates.angles.Angle",
         "astropy.coordinates.angles.core.Angle",
+        "astropy.utils.masked.core.MaskedAngle",
     )
 
     def from_yaml_tree(self, node, tag, ctx):
@@ -19,6 +20,7 @@ class LatitudeConverter(QuantityConverter):
     types = (
         "astropy.coordinates.angles.Latitude",
         "astropy.coordinates.angles.core.Latitude",
+        "astropy.utils.masked.core.MaskedLatitude",
     )
 
     def from_yaml_tree(self, node, tag, ctx):
@@ -32,6 +34,7 @@ class LongitudeConverter(QuantityConverter):
     types = (
         "astropy.coordinates.angles.Longitude",
         "astropy.coordinates.angles.core.Longitude",
+        "astropy.utils.masked.core.MaskedLongitude",
     )
 
     def to_yaml_tree(self, obj, tag, ctx):
diff --git a/asdf_astropy/converters/coordinates/earth_location.py b/asdf_astropy/converters/coordinates/earth_location.py
index 7420043..453b7cb 100644
--- a/asdf_astropy/converters/coordinates/earth_location.py
+++ b/asdf_astropy/converters/coordinates/earth_location.py
@@ -3,10 +3,15 @@ from asdf.extension import Converter
 
 class EarthLocationConverter(Converter):
     tags = ("tag:astropy.org:astropy/coordinates/earthlocation-*",)
-    types = ("astropy.coordinates.earth.EarthLocation",)
+    types = (
+        "astropy.coordinates.earth.EarthLocation",
+        "astropy.utils.masked.core.MaskedEarthLocation",
+    )
 
     def to_yaml_tree(self, obj, tag, ctx):
-        return obj.info._represent_as_dict()
+        result = obj.info._represent_as_dict()
+        result.pop("__class__", None)
+        return result
 
     def from_yaml_tree(self, node, tag, ctx):
         from astropy.coordinates.earth import EarthLocation
diff --git a/asdf_astropy/converters/unit/quantity.py b/asdf_astropy/converters/unit/quantity.py
index ea6065f..e77dfa4 100644
--- a/asdf_astropy/converters/unit/quantity.py
+++ b/asdf_astropy/converters/unit/quantity.py
@@ -9,11 +9,15 @@ class QuantityConverter(Converter):
         # The Distance class has no tag of its own, so we
         # just serialize it as a quantity.
         "astropy.coordinates.distances.Distance",
+        "astropy.utils.masked.core.MaskedQuantity",
     )
 
     def to_yaml_tree(self, obj, tag, ctx):
+        import numpy as np
+        from astropy.utils.masked import Masked
+
         node = {
-            "value": obj.value,
+            "value": np.ma.asarray(obj.value) if isinstance(obj, Masked) else obj.value,
             "unit": obj.unit,
         }
 
@@ -39,4 +43,15 @@ class QuantityConverter(Converter):
             value = value._make_array()
             dtype = value.dtype
 
-        return Quantity(value, unit=node["unit"], copy=copy, dtype=dtype)
+        if isinstance(value, np.ma.MaskedArray):
+            from astropy.utils import minversion
+            from astropy.utils.masked import Masked
+
+            if not minversion("astropy", "7.1.dev"):  # ASTROPY_GE_7_1
+                msg = "MaskedQuantity support requires astropy 7.1 or later"
+                raise NotImplementedError(msg)
+            class_ = Masked(Quantity)
+        else:
+            class_ = Quantity
+
+        return class_(value, unit=node["unit"], copy=copy, dtype=dtype)


Test patch:
diff --git a/asdf_astropy/converters/coordinates/tests/test_masked.py b/asdf_astropy/converters/coordinates/tests/test_masked.py
new file mode 100644
index 0000000..927bb8b
--- /dev/null
+++ b/asdf_astropy/converters/coordinates/tests/test_masked.py
@@ -0,0 +1,77 @@
+import asdf
+import numpy as np
+import pytest
+from astropy import units as u
+from astropy.coordinates import Angle, EarthLocation, Latitude, Longitude, SkyCoord
+from astropy.utils.masked import Masked
+
+from asdf_astropy.tests.versions import ASTROPY_GE_7_1
+
+if ASTROPY_GE_7_1:
+    from astropy.utils.masked import get_data_and_mask
+else:
+    pytest.skip(reason="MaskedQuantity support was added in astropy 7.1", allow_module_level=True)
+
+
+@pytest.mark.parametrize("angle_class", [Angle, Latitude, Longitude])
+def test_masked_angle(tmp_path, angle_class):
+    angle = [1, 2, 3]
+    mask = [False, False, True]
+    masked_class = Masked(angle_class)
+    file_path = tmp_path / "test.asdf"
+    with asdf.AsdfFile() as af:
+        af["angle"] = Masked(angle_class(angle * u.deg), mask)
+        af.write_to(file_path)
+
+    with asdf.open(file_path) as af:
+        assert isinstance(af["angle"], masked_class)
+        out_data, out_mask = get_data_and_mask(af["angle"].to_value(u.deg))
+        np.testing.assert_array_equal(out_data, angle)
+        np.testing.assert_array_equal(out_mask, mask)
+
+
+def test_masked_skycoord(tmp_path):
+    ra_deg = [0.0, 1.0, 2.0]
+    dec_deg = [2.0, 1.0, 0.0]
+    mask = [False, False, True]
+    file_path = tmp_path / "test.asdf"
+    with asdf.AsdfFile() as af:
+        af["coord"] = SkyCoord(Masked(ra_deg, mask), Masked(dec_deg, mask), unit=u.deg)
+        af.write_to(file_path)
+
+    with asdf.open(file_path) as af:
+        assert af["coord"].ra.unit == u.deg
+        assert af["coord"].dec.unit == u.deg
+
+        out_data, out_mask = get_data_and_mask(af["coord"].ra.deg)
+        np.testing.assert_array_equal(out_data, ra_deg)
+        np.testing.assert_array_equal(out_mask, mask)
+        out_data, out_mask = get_data_and_mask(af["coord"].dec.deg)
+        np.testing.assert_array_equal(out_data, dec_deg)
+        np.testing.assert_array_equal(out_mask, mask)
+
+
+def test_masked_earth_location(tmp_path):
+    x = [0.0, 1.0, 2.0]
+    y = [2.0, 1.0, 0.0]
+    z = [1.0, 2.0, 0.0]
+    mask = [False, False, True]
+    file_path = tmp_path / "test.asdf"
+    with asdf.AsdfFile() as af:
+        af["loc"] = EarthLocation.from_geocentric(Masked(x * u.m, mask), Masked(y * u.m, mask), Masked(z * u.m, mask))
+        af.write_to(file_path)
+
+    with asdf.open(file_path) as af:
+        assert af["loc"].x.unit == u.m
+        assert af["loc"].y.unit == u.m
+        assert af["loc"].z.unit == u.m
+
+        out_data, out_mask = get_data_and_mask(af["loc"].x.to_value(u.m))
+        np.testing.assert_array_equal(out_data, x)
+        np.testing.assert_array_equal(out_mask, mask)
+        out_data, out_mask = get_data_and_mask(af["loc"].y.to_value(u.m))
+        np.testing.assert_array_equal(out_data, y)
+        np.testing.assert_array_equal(out_mask, mask)
+        out_data, out_mask = get_data_and_mask(af["loc"].z.to_value(u.m))
+        np.testing.assert_array_equal(out_data, z)
+        np.testing.assert_array_equal(out_mask, mask)
diff --git a/asdf_astropy/converters/unit/tests/test_masked.py b/asdf_astropy/converters/unit/tests/test_masked.py
new file mode 100644
index 0000000..3169c4f
--- /dev/null
+++ b/asdf_astropy/converters/unit/tests/test_masked.py
@@ -0,0 +1,31 @@
+import asdf
+import numpy as np
+import pytest
+from astropy import units as u
+from astropy.utils.masked import Masked
+
+from asdf_astropy.tests.versions import ASTROPY_GE_7_1
+
+if ASTROPY_GE_7_1:
+    from astropy.utils.masked import get_data_and_mask
+else:
+    pytest.skip(reason="MaskedQuantity support was added in astropy 7.1", allow_module_level=True)
+
+MaskedQuantity = Masked(u.Quantity)
+
+
+def test_masked_quantity(tmp_path):
+    data = [1, 2, 3]
+    mask = [False, False, True]
+    file_path = tmp_path / "test.asdf"
+    with asdf.AsdfFile() as af:
+        af["quantity"] = Masked(data, mask) * u.yottamole
+        af.write_to(file_path)
+
+    with asdf.open(file_path) as af:
+        assert isinstance(af["quantity"], MaskedQuantity)
+        assert af["quantity"].unit == u.yottamole
+
+        result_data, result_mask = get_data_and_mask(af["quantity"].value)
+        np.testing.assert_array_equal(result_data, data)
+        np.testing.assert_array_equal(result_mask, mask)
diff --git a/asdf_astropy/converters/unit/tests/test_quantity.py b/asdf_astropy/converters/unit/tests/test_quantity.py
index 71c07ae..1d3bb66 100644
--- a/asdf_astropy/converters/unit/tests/test_quantity.py
+++ b/asdf_astropy/converters/unit/tests/test_quantity.py
@@ -7,6 +7,8 @@ from astropy.units import Quantity
 from astropy.utils.introspection import minversion
 from numpy.testing import assert_array_equal
 
+from asdf_astropy.tests.versions import ASTROPY_GE_7_1
+
 
 def asdf_open_memory_mapping_kwarg(memmap: bool) -> dict:
     if minversion("asdf", "3.1.0"):
@@ -149,3 +151,22 @@ def test_no_memmap(tmp_path):
         assert af.tree["quantity"][-1, -1] != new_value
         assert (af.tree["quantity"] != new_quantity).any()
         assert (af.tree["quantity"] == quantity).all()
+
+
+@pytest.mark.skipif(ASTROPY_GE_7_1, reason="MaskedQuantity support was added in astropy 7.1")
+def test_masked_quantity_raises():
+    yaml = """
+quantity: !unit/quantity-1.1.0
+  unit: !unit/unit-1.0.0 Ymol
+  value: !core/ndarray-1.0.0
+    data: [1.0, 2.0, null]
+    mask: !core/ndarray-1.0.0
+      data: [false, false, true]
+      datatype: bool8
+      shape: [3]
+    datatype: float64
+    shape: [3]
+"""
+    buff = helpers.yaml_to_asdf(yaml, version="1.5.0")
+    with pytest.raises(NotImplementedError, match="MaskedQuantity support requires astropy 7.1 or later"):
+        asdf.open(buff)
diff --git a/asdf_astropy/tests/versions.py b/asdf_astropy/tests/versions.py
new file mode 100644
index 0000000..3f03e69
--- /dev/null
+++ b/asdf_astropy/tests/versions.py
@@ -0,0 +1,3 @@
+from astropy.utils import minversion
+
+ASTROPY_GE_7_1 = minversion("astropy", "7.1.dev")