Skip to content

Commit

Permalink
Add support for masked quantities
Browse files Browse the repository at this point in the history
Fixes #202.
  • Loading branch information
lpsinger committed Jan 28, 2025
1 parent 6826f49 commit 24f6e11
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 2 deletions.
12 changes: 10 additions & 2 deletions asdf_astropy/converters/unit/quantity.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,10 @@
import numpy as np
from asdf.extension import Converter
from asdf.tags.core.ndarray import NDArrayType
from astropy.units import Quantity
from astropy.utils.masked import Masked

MaskedQuantity = Masked(Quantity)


class QuantityConverter(Converter):
Expand All @@ -9,11 +14,12 @@ class QuantityConverter(Converter):
# The Distance class has no tag of its own, so we
# just serialize it as a quantity.
"astropy.coordinates.distances.Distance",
MaskedQuantity,
)

def to_yaml_tree(self, obj, tag, ctx):
node = {
"value": obj.value,
"value": np.ma.asarray(obj.value) if isinstance(obj, MaskedQuantity) else obj.value,
"unit": obj.unit,
}

Expand All @@ -39,4 +45,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)
15 changes: 15 additions & 0 deletions asdf_astropy/converters/unit/tests/test_masked.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
import asdf
from astropy import units as u
from astropy.utils.masked import Masked

MaskedQuantity = Masked(u.Quantity)


def test_masked_quantity(tmp_path):
file_path = tmp_path / "test.asdf"
with asdf.AsdfFile() as af:
af["quantity"] = Masked([1, 2, 3], [False, False, True]) * u.yottamole
af.write_to(file_path)

with asdf.open(file_path) as af:
assert isinstance(af["quantity"], MaskedQuantity)

0 comments on commit 24f6e11

Please sign in to comment.