Skip to content

Commit bca63dc

Browse files
committed
Ensure mask calculation is correct also for multi-op ufunc with scalars
1 parent 2b47bcf commit bca63dc

2 files changed

Lines changed: 24 additions & 1 deletion

File tree

‎astropy/utils/masked/core.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -697,7 +697,8 @@ def _combine_masks(self, masks, out=None, where=True, copy=True):
697697
np.copyto(out, masks[0], where=where)
698698
return out
699699

700-
out = np.logical_or(masks[0], masks[1], out=out, where=where)
700+
# [...] at the end to ensure we have an array, not a scalar.
701+
out = np.logical_or(masks[0], masks[1], out=out, where=where)[...]
701702
for mask in masks[2:]:
702703
np.logical_or(out, mask, out=out, where=where)
703704
return out

‎astropy/utils/masked/tests/test_functions.py‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
coverage. Complete coverage of all numpy functions is done
66
with less detailed tests in test_function_helpers.
77
"""
8+
import erfa.ufunc as erfa_ufunc
89
import numpy as np
910
import pytest
1011
from numpy.testing import assert_array_equal
@@ -151,6 +152,27 @@ def test_3op_ufunc(self):
151152
assert_array_equal(ma_mb.unmasked, expected_data)
152153
assert_array_equal(ma_mb.mask, expected_mask)
153154

155+
def test_multi_op_ufunc(self):
156+
mask = [True, False, False]
157+
iy = Masked([2000, 2001, 2002], mask=mask)
158+
im = Masked([1, 2, 3], mask=mask)
159+
idy = Masked([10, 20, 25], mask=mask)
160+
ihr = Masked([11, 12, 13], mask=[False, False, True])
161+
imn = np.array([50, 51, 52])
162+
isc = np.array([12.5, 13.6, 14.7])
163+
result = erfa_ufunc.dtf2d("utc", iy, im, idy, ihr, imn, isc)
164+
# Also test scalar
165+
result0 = erfa_ufunc.dtf2d("utc", iy[0], im[0], idy[0], ihr[0], imn[0], isc[0])
166+
expected = erfa_ufunc.dtf2d(
167+
"utc", iy.unmasked, im.unmasked, idy.unmasked, ihr.unmasked, imn, isc
168+
)
169+
expected_mask = np.array([True, False, True])
170+
for res, res0, exp in zip(result, result0, expected):
171+
assert_array_equal(res.unmasked, exp)
172+
assert_array_equal(res.mask, expected_mask)
173+
assert res0.unmasked == exp[0]
174+
assert res0.mask == expected_mask[0]
175+
154176
@pytest.mark.parametrize("axis", (0, 1, None))
155177
def test_add_reduce(self, axis):
156178
ma_reduce = np.add.reduce(self.ma, axis=axis)

0 commit comments

Comments
 (0)