diff options
Diffstat (limited to 'numpy')
| -rw-r--r-- | numpy/core/src/umath/ufunc_object.c | 22 | ||||
| -rw-r--r-- | numpy/core/tests/test_maskna.py | 20 |
2 files changed, 35 insertions, 7 deletions
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c index cdcc4ade6..d1b76d070 100644 --- a/numpy/core/src/umath/ufunc_object.c +++ b/numpy/core/src/umath/ufunc_object.c @@ -2980,6 +2980,16 @@ PyUFunc_Reduce(PyUFuncObject *self, PyArrayObject *arr, PyArrayObject *out, "isn't implemented yet"); goto fail; } + + /* + * If the result has a mask (i.e. from the out= parameter), + * Set it to all exposed. + */ + if (PyArray_HASMASKNA(result)) { + if (PyArray_AssignMaskNA(result, 1) < 0) { + goto fail; + } + } } } @@ -3017,18 +3027,18 @@ PyUFunc_Reduce(PyUFuncObject *self, PyArrayObject *arr, PyArrayObject *out, /* Add mask-related flags */ if (use_maskna) { if (skipna) { - /* Need the input's mask to determine what to skip */ - op_flags[0] |= NPY_ITER_USE_MASKNA; /* The output's mask has been set to all exposed already */ - op_flags[1] |= NPY_ITER_IGNORE_MASKNA; + op_flags[0] |= NPY_ITER_IGNORE_MASKNA; + /* Need the input's mask to determine what to skip */ + op_flags[1] |= NPY_ITER_USE_MASKNA; /* TODO: allocate a temporary buffer for inverting the mask */ } else { - /* The input's mask is already incorporated in the output's mask */ - op_flags[0] |= NPY_ITER_IGNORE_MASKNA; /* Iterate over the output's mask */ - op_flags[1] |= NPY_ITER_USE_MASKNA; + op_flags[0] |= NPY_ITER_USE_MASKNA; + /* The input's mask is already incorporated in the output's mask */ + op_flags[1] |= NPY_ITER_IGNORE_MASKNA; } } else { diff --git a/numpy/core/tests/test_maskna.py b/numpy/core/tests/test_maskna.py index ec8fda00f..5ad6f0990 100644 --- a/numpy/core/tests/test_maskna.py +++ b/numpy/core/tests/test_maskna.py @@ -523,7 +523,8 @@ def test_ufunc_1D(): def test_ufunc_reduce_1D(): a = np.arange(3.0, maskna=True) b = np.array(0.5) - c = np.array(0.5, maskna=True) + c_orig = np.array(0.5) + c = c_orig.view(maskna=True) # Since 'a' has no NA values, this should work np.add.reduce(a, out=b) @@ -534,7 +535,9 @@ def test_ufunc_reduce_1D(): assert_raises(ValueError, np.add.reduce, a, out=b) # With an NA value, the output parameter can still be an NA-array + c_orig[...] = 0.5 np.add.reduce(a, out=c) + assert_equal(c_orig, 0.5) assert_(np.isna(c)) # Should not touch the out= element when assigning NA @@ -548,6 +551,21 @@ def test_ufunc_reduce_1D(): ret = np.add.reduce(a) assert_(np.isna(ret)) + # With 'skipna=True' + ret = np.add.reduce(a, skipna=True) + assert_equal(ret, 2.0) + + # With 'skipna=True', and out= parameter + b[...] = 0.5 + np.add.reduce(a, skipna=True, out=b) + assert_equal(b, 2.0) + + # With 'skipna=True', and out= parameter with a mask + c[...] = 0.5 + c[...] = np.NA + np.add.reduce(a, skipna=True, out=c) + assert_(not np.isna(c)) + assert_equal(c, 2.0) if __name__ == "__main__": run_module_suite() |