Patch 01

CHANGES.rst

@@ -1,3 +1,8 @@
+0.8.0 (unreleased)
+------------------
+
+- Add support for masked quantities, angles, and coordinates [#749]
+
 0.7.0 (2024-11-13)
 ------------------
 

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):

asdf_astropy/converters/coordinates/tests/test_masked.py

@@ -0,0 +1,38 @@
+import asdf
+import numpy as np
+import pytest
+from astropy import units as u
+from astropy.coordinates import Angle, Latitude, Longitude, SkyCoord
+from astropy.utils.masked import Masked, get_data_and_mask
+
+
+@pytest.mark.parametrize("angle_class", [Angle, Latitude, Longitude])
+def test_masked_angle(tmp_path, angle_class):
+    masked_class = Masked(angle_class)
+    file_path = tmp_path / "test.asdf"
+    with asdf.AsdfFile() as af:
+        af["angle"] = Masked(angle_class([1, 2, 3] * u.deg), [False, False, True])
+        af.write_to(file_path)
+
+    with asdf.open(file_path) as af:
+        assert isinstance(af["angle"], masked_class)
+
+
+def test_masked_skycoord(tmp_path):
+    ra_deg = [0, 1, 2]
+    dec_deg = [2, 1, 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)

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,
         }
 
@@ -28,6 +32,7 @@ def from_yaml_tree(self, node, tag, ctx):
         # astropy 6.1 changed Quantity in a similar way
         import numpy as np
         from astropy.units import Quantity
+        from astropy.utils.masked.core import MaskedQuantity
 
         copy = None if np.lib.NumpyVersion(np.__version__) >= "2.0.0b1" else False
 
@@ -39,4 +44,6 @@ def from_yaml_tree(self, node, tag, ctx):
             value = value._make_array()
             dtype = value.dtype
 
-        return Quantity(value, unit=node["unit"], copy=copy, dtype=dtype)
+        class_ = MaskedQuantity if isinstance(value, np.ma.MaskedArray) else Quantity
+
+        return class_(value, unit=node["unit"], copy=copy, dtype=dtype)

asdf_astropy/converters/unit/tests/test_masked.py

@@ -0,0 +1,22 @@
+import asdf
+import numpy as np
+from astropy import units as u
+from astropy.utils.masked import Masked, get_data_and_mask
+
+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)

pyproject.toml

@@ -21,7 +21,8 @@ dependencies = [
   "asdf-coordinates-schemas>=0.3",
   "asdf-transform-schemas>=0.5",
   "asdf-standard>=1.1.0",
-  "astropy>=5.2.0",
+  # "astropy>=5.2.0",
+  "astropy @ git+https://github.com/lpsinger/astropy@masked-classes-importable",
   "numpy>=1.24",
   "packaging>=19",
 ]