Skip to content

Commit 3c6b6ba

Browse files
mhvkmhvk
committed
BUG Allow subclasses in MaskedArray ufuncs -- for non-ndarray _data
1 parent 8c86a0a commit 3c6b6ba

3 files changed

Lines changed: 147 additions & 104 deletions

File tree

‎numpy/ma/core.py‎

Lines changed: 94 additions & 94 deletions
Original file line numberDiff line numberDiff line change
@@ -880,14 +880,10 @@ def __call__ (self, a, *args, **kwargs):
880880
except TypeError:
881881
pass
882882
# Transform to
883-
if isinstance(a, MaskedArray):
884-
subtype = type(a)
885-
else:
886-
subtype = MaskedArray
887-
result = result.view(subtype)
888-
result._mask = m
889-
result._update_from(a)
890-
return result
883+
masked_result = result.view(get_masked_subclass(a))
884+
masked_result._mask = m
885+
masked_result._update_from(result)
886+
return masked_result
891887
#
892888
def __str__ (self):
893889
return "Masked version of %s. [Invalid values are masked]" % str(self.f)
@@ -928,8 +924,12 @@ def __init__ (self, mbfunc, fillx=0, filly=0):
928924
def __call__ (self, a, b, *args, **kwargs):
929925
"Execute the call behavior."
930926
# Get the data, as ndarray
931-
(da, db) = (getdata(a, subok=False), getdata(b, subok=False))
932-
# Get the mask
927+
(da, db) = (getdata(a), getdata(b))
928+
# Get the result
929+
with np.errstate():
930+
np.seterr(divide='ignore', invalid='ignore')
931+
result = self.f(da, db, *args, **kwargs)
932+
# Get the mask for the result
933933
(ma, mb) = (getmask(a), getmask(b))
934934
if ma is nomask:
935935
if mb is nomask:
@@ -940,66 +940,64 @@ def __call__ (self, a, b, *args, **kwargs):
940940
m = umath.logical_or(ma, getmaskarray(b))
941941
else:
942942
m = umath.logical_or(ma, mb)
943-
# Get the result
944-
with np.errstate(divide='ignore', invalid='ignore'):
945-
result = self.f(da, db, *args, **kwargs)
946-
# check it worked
947-
if result is NotImplemented:
948-
return NotImplemented
949943
# Case 1. : scalar
950944
if not result.ndim:
951945
if m:
952946
return masked
953947
return result
954948
# Case 2. : array
955949
# Revert result to da where masked
956-
if m is not nomask:
957-
np.copyto(result, da, casting='unsafe', where=m)
950+
if m is not nomask and m.any():
951+
# any errors, just abort; impossible to guarantee masked values
952+
try:
953+
np.copyto(result, 0, casting='unsafe', where=m)
954+
# avoid using "*" since this may be overlaid
955+
masked_da = umath.multiply(m, da)
956+
# only add back if it can be cast safely
957+
if np.can_cast(masked_da.dtype, result.dtype, casting='safe'):
958+
result += masked_da
959+
except:
960+
pass
958961
# Transforms to a (subclass of) MaskedArray
959-
result = result.view(get_masked_subclass(a, b))
960-
result._mask = m
961-
# Update the optional info from the inputs
962-
if isinstance(b, MaskedArray):
963-
if isinstance(a, MaskedArray):
964-
result._update_from(a)
965-
else:
966-
result._update_from(b)
967-
elif isinstance(a, MaskedArray):
968-
result._update_from(a)
969-
return result
970-
962+
masked_result = result.view(get_masked_subclass(a, b))
963+
masked_result._mask = m
964+
masked_result._update_from(result)
965+
return masked_result
971966

972967
def reduce(self, target, axis=0, dtype=None):
973968
"""Reduce `target` along the given `axis`."""
974-
if isinstance(target, MaskedArray):
975-
tclass = type(target)
976-
else:
977-
tclass = MaskedArray
969+
tclass = get_masked_subclass(target)
978970
m = getmask(target)
979971
t = filled(target, self.filly)
980972
if t.shape == ():
981973
t = t.reshape(1)
982974
if m is not nomask:
983975
m = make_mask(m, copy=1)
984976
m.shape = (1,)
977+
985978
if m is nomask:
986-
return self.f.reduce(t, axis).view(tclass)
987-
t = t.view(tclass)
988-
t._mask = m
989-
tr = self.f.reduce(getdata(t), axis, dtype=dtype or t.dtype)
990-
mr = umath.logical_and.reduce(m, axis)
991-
tr = tr.view(tclass)
992-
if mr.ndim > 0:
993-
tr._mask = mr
994-
return tr
995-
elif mr:
996-
return masked
997-
return tr
979+
tr = self.f.reduce(t, axis)
980+
mr = nomask
981+
else:
982+
tr = self.f.reduce(t, axis, dtype=dtype or t.dtype)
983+
mr = umath.logical_and.reduce(m, axis)
998984

999-
def outer (self, a, b):
985+
if not tr.shape:
986+
if mr:
987+
return masked
988+
else:
989+
return tr
990+
masked_tr = tr.view(tclass)
991+
masked_tr._mask = mr
992+
masked_tr._update_from(tr)
993+
return masked_tr
994+
995+
def outer(self, a, b):
1000996
"""Return the function applied to the outer product of a and b.
1001997
1002998
"""
999+
(da, db) = (getdata(a), getdata(b))
1000+
d = self.f.outer(da, db)
10031001
ma = getmask(a)
10041002
mb = getmask(b)
10051003
if ma is nomask and mb is nomask:
@@ -1010,31 +1008,28 @@ def outer (self, a, b):
10101008
m = umath.logical_or.outer(ma, mb)
10111009
if (not m.ndim) and m:
10121010
return masked
1013-
(da, db) = (getdata(a), getdata(b))
1014-
d = self.f.outer(da, db)
1015-
# check it worked
1016-
if d is NotImplemented:
1017-
return NotImplemented
10181011
if m is not nomask:
10191012
np.copyto(d, da, where=m)
1020-
if d.shape:
1021-
d = d.view(get_masked_subclass(a, b))
1022-
d._mask = m
1023-
return d
1013+
if not d.shape:
1014+
return d
1015+
masked_d = d.view(get_masked_subclass(a, b))
1016+
masked_d._mask = m
1017+
masked_d._update_from(d)
1018+
return masked_d
10241019

1025-
def accumulate (self, target, axis=0):
1020+
def accumulate(self, target, axis=0):
10261021
"""Accumulate `target` along `axis` after filling with y fill
10271022
value.
10281023
10291024
"""
1030-
if isinstance(target, MaskedArray):
1031-
tclass = type(target)
1032-
else:
1033-
tclass = MaskedArray
1025+
tclass = get_masked_subclass(target)
10341026
t = filled(target, self.filly)
1035-
return self.f.accumulate(t, axis).view(tclass)
1027+
result = self.f.accumulate(t, axis)
1028+
masked_result = result.view(tclass)
1029+
masked_result._update_from(result)
1030+
return masked_result
10361031

1037-
def __str__ (self):
1032+
def __str__(self):
10381033
return "Masked version of " + str(self.f)
10391034

10401035

@@ -1074,19 +1069,15 @@ def __init__ (self, dbfunc, domain, fillx=0, filly=0):
10741069

10751070
def __call__(self, a, b, *args, **kwargs):
10761071
"Execute the call behavior."
1077-
# Get the data and the mask
1078-
(da, db) = (getdata(a, subok=False), getdata(b, subok=False))
1079-
(ma, mb) = (getmask(a), getmask(b))
1072+
# Get the data
1073+
(da, db) = (getdata(a), getdata(b))
10801074
# Get the result
10811075
with np.errstate(divide='ignore', invalid='ignore'):
10821076
result = self.f(da, db, *args, **kwargs)
1083-
# check it worked
1084-
if result is NotImplemented:
1085-
return NotImplemented
1086-
# Get the mask as a combination of ma, mb and invalid
1077+
# Get the mask as a combination of the source masks and invalid
10871078
m = ~umath.isfinite(result)
1088-
m |= ma
1089-
m |= mb
1079+
m |= getmask(a)
1080+
m |= getmask(b)
10901081
# Apply the domain
10911082
domain = ufunc_domain.get(self.f, None)
10921083
if domain is not None:
@@ -1097,18 +1088,23 @@ def __call__(self, a, b, *args, **kwargs):
10971088
return masked
10981089
else:
10991090
return result
1100-
# When the mask is True, put back da
1101-
np.copyto(result, da, casting='unsafe', where=m)
1102-
result = result.view(get_masked_subclass(a, b))
1103-
result._mask = m
1104-
if isinstance(b, MaskedArray):
1105-
if isinstance(a, MaskedArray):
1106-
result._update_from(a)
1107-
else:
1108-
result._update_from(b)
1109-
elif isinstance(a, MaskedArray):
1110-
result._update_from(a)
1111-
return result
1091+
# When the mask is True, put back da if possible
1092+
# any errors, just abort; impossible to guarantee masked values
1093+
try:
1094+
np.copyto(result, 0, casting='unsafe', where=m)
1095+
# avoid using "*" since this may be overlaid
1096+
masked_da = umath.multiply(m, da)
1097+
# only add back if it can be cast safely
1098+
if np.can_cast(masked_da.dtype, result.dtype, casting='safe'):
1099+
result += masked_da
1100+
except:
1101+
pass
1102+
1103+
# Transforms to a (subclass of) MaskedArray
1104+
masked_result = result.view(get_masked_subclass(a, b))
1105+
masked_result._mask = m
1106+
masked_result._update_from(result)
1107+
return masked_result
11121108

11131109
def __str__ (self):
11141110
return "Masked version of " + str(self.f)
@@ -1361,7 +1357,7 @@ def getmaskarray(arr):
13611357
"""
13621358
mask = getmask(arr)
13631359
if mask is nomask:
1364-
mask = make_mask_none(np.shape(arr), getdata(arr).dtype)
1360+
mask = make_mask_none(np.shape(arr), getattr(arr, 'dtype', None))
13651361
return mask
13661362

13671363
def is_mask(m):
@@ -3756,34 +3752,38 @@ def __ne__(self, other):
37563752
return check
37573753
#
37583754
def __add__(self, other):
3759-
"Add other to self, and return a new masked array."
3755+
"Add self to other, and return a new masked array."
37603756
if self._delegate_binop(other):
37613757
return NotImplemented
37623758
return add(self, other)
37633759
#
37643760
def __radd__(self, other):
37653761
"Add other to self, and return a new masked array."
3766-
return add(self, other)
3762+
# In analogy with __rsub__ and __rdiv__, use original order:
3763+
# we get here from `other + self`.
3764+
return add(other, self)
37673765
#
37683766
def __sub__(self, other):
3769-
"Subtract other to self, and return a new masked array."
3767+
"Subtract other from self, and return a new masked array."
37703768
if self._delegate_binop(other):
37713769
return NotImplemented
37723770
return subtract(self, other)
37733771
#
37743772
def __rsub__(self, other):
3775-
"Subtract other to self, and return a new masked array."
3773+
"Subtract self from other, and return a new masked array."
37763774
return subtract(other, self)
37773775
#
37783776
def __mul__(self, other):
3779-
"Multiply other by self, and return a new masked array."
3777+
"Multiply self by other, and return a new masked array."
37803778
if self._delegate_binop(other):
37813779
return NotImplemented
37823780
return multiply(self, other)
37833781
#
37843782
def __rmul__(self, other):
37853783
"Multiply other by self, and return a new masked array."
3786-
return multiply(self, other)
3784+
# In analogy with __rsub__ and __rdiv__, use original order:
3785+
# we get here from `other * self`.
3786+
return multiply(other, self)
37873787
#
37883788
def __div__(self, other):
37893789
"Divide other into self, and return a new masked array."
@@ -3798,7 +3798,7 @@ def __truediv__(self, other):
37983798
return true_divide(self, other)
37993799
#
38003800
def __rtruediv__(self, other):
3801-
"Divide other into self, and return a new masked array."
3801+
"Divide self into other, and return a new masked array."
38023802
return true_divide(other, self)
38033803
#
38043804
def __floordiv__(self, other):
@@ -3808,7 +3808,7 @@ def __floordiv__(self, other):
38083808
return floor_divide(self, other)
38093809
#
38103810
def __rfloordiv__(self, other):
3811-
"Divide other into self, and return a new masked array."
3811+
"Divide self into other, and return a new masked array."
38123812
return floor_divide(other, self)
38133813
#
38143814
def __pow__(self, other):
@@ -3818,7 +3818,7 @@ def __pow__(self, other):
38183818
return power(self, other)
38193819
#
38203820
def __rpow__(self, other):
3821-
"Raise self to the power other, masking the potential NaNs/Infs"
3821+
"Raise other to the power self, masking the potential NaNs/Infs"
38223822
return power(other, self)
38233823
#............................................
38243824
def __iadd__(self, other):

‎numpy/ma/tests/test_core.py‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1754,6 +1754,29 @@ def __rmul__(self, other):
17541754
assert_(me * a == "My mul")
17551755
assert_(a * me == "My rmul")
17561756

1757+
# and that __array_priority__ is respected
1758+
class MyClass2(object):
1759+
__array_priority__ = 100
1760+
1761+
def __mul__(self, other):
1762+
return "Me2mul"
1763+
1764+
def __rmul__(self, other):
1765+
return "Me2rmul"
1766+
1767+
def __rdiv__(self, other):
1768+
return "Me2rdiv"
1769+
1770+
__rtruediv__ = __rdiv__
1771+
1772+
me_too = MyClass2()
1773+
assert_(a.__mul__(me_too) is NotImplemented)
1774+
assert_(all(multiply.outer(a, me_too) == "Me2rmul"))
1775+
assert_(a.__truediv__(me_too) is NotImplemented)
1776+
assert_(me_too * a == "Me2mul")
1777+
assert_(a * me_too == "Me2rmul")
1778+
assert_(a / me_too == "Me2rdiv")
1779+
17571780

17581781
#------------------------------------------------------------------------------
17591782
class TestMaskedArrayInPlaceArithmetics(TestCase):

0 commit comments

Comments
 (0)