diff options
| -rw-r--r-- | numpy/ma/core.py | 61 | ||||
| -rw-r--r-- | numpy/ma/tests/test_extras.py | 29 |
2 files changed, 65 insertions, 25 deletions
diff --git a/numpy/ma/core.py b/numpy/ma/core.py index 9584e56d2..4245aed49 100644 --- a/numpy/ma/core.py +++ b/numpy/ma/core.py @@ -7110,7 +7110,7 @@ def diag(v, k=0): Examples -------- - + Create an array with negative values masked: >>> import numpy as np @@ -7521,7 +7521,7 @@ def diff(a, /, n=1, axis=-1, prepend=np._NoValue, append=np._NoValue): if len(combined) > 1: a = np.ma.concatenate(combined, axis) - # GH 22465 np.diff without prepend/append preserves the mask + # GH 22465 np.diff without prepend/append preserves the mask return np.diff(a, n, axis) @@ -7752,6 +7752,21 @@ def round_(a, decimals=0, out=None): round = round_ +def _mask_propagate(a, axis): + """ + Mask whole 1-d vectors of an array that contain masked values. + """ + a = array(a, subok=False) + m = getmask(a) + if m is nomask or not m.any() or axis is None: + return a + a._mask = a._mask.copy() + axes = normalize_axis_tuple(axis, a.ndim) + for ax in axes: + a._mask |= m.any(axis=ax, keepdims=True) + return a + + # Needed by dot, so move here from extras.py. It will still be exported # from extras.py for compatibility. def mask_rowcols(a, axis=None): @@ -7770,7 +7785,7 @@ def mask_rowcols(a, axis=None): ---------- a : array_like, MaskedArray The array to mask. If not a MaskedArray instance (or if no array - elements are masked). The result is a MaskedArray with `mask` set + elements are masked), the result is a MaskedArray with `mask` set to `nomask` (False). Must be a 2D array. axis : int, optional Axis along which to perform the operation. If None, applies to a @@ -7827,20 +7842,16 @@ def mask_rowcols(a, axis=None): fill_value=1) """ - a = array(a, subok=False) if a.ndim != 2: raise NotImplementedError("mask_rowcols works for 2D arrays only.") - m = getmask(a) - # Nothing is masked: return a - if m is nomask or not m.any(): - return a - maskedval = m.nonzero() - a._mask = a._mask.copy() - if not axis: - a[np.unique(maskedval[0])] = masked - if axis in [None, 1, -1]: - a[:, np.unique(maskedval[1])] = masked - return a + + if axis is None: + return _mask_propagate(a, axis=(0, 1)) + elif axis == 0: + return _mask_propagate(a, axis=1) + else: + return _mask_propagate(a, axis=0) + # Include masked dot here to avoid import problems in getting it from @@ -7856,10 +7867,6 @@ def dot(a, b, strict=False, out=None): corresponding method, it is recommended that the optional arguments be treated as keyword only. At some point that may be mandatory. - .. note:: - Works only with 2-D arrays at the moment. - - Parameters ---------- a, b : masked_array_like @@ -7903,18 +7910,22 @@ def dot(a, b, strict=False, out=None): fill_value=999999) """ - # !!!: Works only with 2D arrays. There should be a way to get it to run - # with higher dimension - if strict and (a.ndim == 2) and (b.ndim == 2): - a = mask_rowcols(a, 0) - b = mask_rowcols(b, 1) + if strict is True: + if np.ndim(a) == 0 or np.ndim(b) == 0: + pass + elif b.ndim == 1: + a = _mask_propagate(a, a.ndim - 1) + b = _mask_propagate(b, b.ndim - 1) + else: + a = _mask_propagate(a, a.ndim - 1) + b = _mask_propagate(b, b.ndim - 2) am = ~getmaskarray(a) bm = ~getmaskarray(b) if out is None: d = np.dot(filled(a, 0), filled(b, 0)) m = ~np.dot(am, bm) - if d.ndim == 0: + if np.ndim(d) == 0: d = np.asarray(d) r = d.view(get_masked_subclass(a, b)) r.__setmask__(m) diff --git a/numpy/ma/tests/test_extras.py b/numpy/ma/tests/test_extras.py index e59ba3656..c3161ba1f 100644 --- a/numpy/ma/tests/test_extras.py +++ b/numpy/ma/tests/test_extras.py @@ -730,6 +730,35 @@ class TestCompressFunctions: assert_equal(c.mask, [[0, 0, 1], [1, 1, 1], [0, 0, 1]]) c = dot(b, a, strict=False) assert_equal(c, np.dot(b.filled(0), a.filled(0))) + # + a = masked_array(np.arange(8), mask=[1, 0, 0, 0, 0, 0, 0, 0]).reshape(2, 2, 2) + b = masked_array(np.arange(8), mask=[0, 0, 0, 0, 0, 0, 0, 1]).reshape(2, 2, 2) + c = dot(a, b, strict=True) + assert_equal(c.mask, [[[[1, 1], [1, 1]],[[0, 0], [0, 1]]],[[[0, 0], [0, 1]],[[0, 0], [0, 1]]]]) + c = dot(a, b, strict=False) + assert_equal(c.mask, [[[[0, 0], [0, 1]],[[0, 0], [0, 0]]],[[[0, 0], [0, 0]],[[0, 0], [0, 0]]]]) + c = dot(b, a, strict=True) + assert_equal(c.mask, [[[[1, 0], [0, 0]],[[1, 0], [0, 0]]],[[[1, 0], [0, 0]],[[1, 1], [1, 1]]]]) + c = dot(b, a, strict=False) + assert_equal(c.mask, [[[[0, 0], [0, 0]],[[0, 0], [0, 0]]],[[[0, 0], [0, 0]],[[1, 0], [0, 0]]]]) + # + a = masked_array(np.arange(8), mask=[1, 0, 0, 0, 0, 0, 0, 0]).reshape(2, 2, 2) + b = 5. + c = dot(a, b, strict=True) + assert_equal(c.mask, [[[1, 0], [0, 0]],[[0, 0], [0, 0]]]) + c = dot(a, b, strict=False) + assert_equal(c.mask, [[[1, 0], [0, 0]],[[0, 0], [0, 0]]]) + c = dot(b, a, strict=True) + assert_equal(c.mask, [[[1, 0], [0, 0]],[[0, 0], [0, 0]]]) + c = dot(b, a, strict=False) + assert_equal(c.mask, [[[1, 0], [0, 0]],[[0, 0], [0, 0]]]) + # + a = masked_array(np.arange(8), mask=[1, 0, 0, 0, 0, 0, 0, 0]).reshape(2, 2, 2) + b = masked_array(np.arange(2), mask=[0, 1]) + c = dot(a, b, strict=True) + assert_equal(c.mask, [[1, 1], [1, 1]]) + c = dot(a, b, strict=False) + assert_equal(c.mask, [[1, 0], [0, 0]]) def test_dot_returns_maskedarray(self): # See gh-6611 |