summaryrefslogtreecommitdiff
path: root/numpy/core
diff options
context:
space:
mode:
Diffstat (limited to 'numpy/core')
-rw-r--r--numpy/core/src/multiarray/nditer_constr.c6
-rw-r--r--numpy/core/src/umath/ufunc_object.c60
-rw-r--r--numpy/core/tests/test_maskna.py29
3 files changed, 69 insertions, 26 deletions
diff --git a/numpy/core/src/multiarray/nditer_constr.c b/numpy/core/src/multiarray/nditer_constr.c
index 10a9d091a..e8385ab85 100644
--- a/numpy/core/src/multiarray/nditer_constr.c
+++ b/numpy/core/src/multiarray/nditer_constr.c
@@ -1093,6 +1093,8 @@ npyiter_prepare_one_operand(PyArrayObject **op,
if (PyArray_Check(*op)) {
+ npy_uint32 tmp;
+
if (((*op_itflags) & NPY_OP_ITFLAG_WRITE) &&
(!PyArray_CHKFLAGS(*op, NPY_ARRAY_WRITEABLE))) {
PyErr_SetString(PyExc_ValueError,
@@ -1118,7 +1120,9 @@ npyiter_prepare_one_operand(PyArrayObject **op,
return 0;
}
/* Arrays with NA masks must have USE_MASKNA specified */
- if ((op_flags & NPY_ITER_USE_MASKNA) == 0 && PyArray_HASMASKNA(*op)) {
+ tmp = op_flags & (NPY_ITER_USE_MASKNA | NPY_ITER_IGNORE_MASKNA);
+ if (tmp != NPY_ITER_USE_MASKNA && tmp != NPY_ITER_IGNORE_MASKNA &&
+ PyArray_HASMASKNA(*op)) {
PyErr_SetString(PyExc_ValueError,
"Operand has an NA mask but the operation does "
"not support NA via the flag USE_MASKNA");
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index 6cf5fbbba..cdcc4ade6 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -2645,7 +2645,7 @@ conform_reduce_result(int ndim, npy_bool *axis_flags, PyArrayObject *out)
static PyArrayObject *
allocate_or_conform_reduce_result(PyArrayObject *arr, PyArrayObject *out,
npy_bool *axis_flags, PyArray_Descr * otype_dtype,
- int keepmask)
+ int addmask)
{
if (out == NULL) {
PyArrayObject *result;
@@ -2654,7 +2654,7 @@ allocate_or_conform_reduce_result(PyArrayObject *arr, PyArrayObject *out,
result = allocate_reduce_result(arr, axis_flags, otype_dtype);
/* Allocate an NA mask if necessary */
- if (keepmask && result != NULL && PyArray_HASMASKNA(arr)) {
+ if (addmask && result != NULL) {
if (PyArray_AllocateMaskNA(result, 1, 0, 1) < 0) {
Py_DECREF(result);
return NULL;
@@ -2864,16 +2864,27 @@ PyUFunc_Reduce(PyUFuncObject *self, PyArrayObject *arr, PyArrayObject *out,
use_maskna = PyArray_HASMASKNA(arr);
+ /* Detect whether to ignore the MASKNA */
+ if (use_maskna) {
+ if (!skipna && out != NULL && !PyArray_HASMASKNA(out)) {
+ if (PyArray_ContainsNA(arr)) {
+ PyErr_SetString(PyExc_ValueError,
+ "Cannot assign NA value to an array which "
+ "does not support NAs");
+ goto fail;
+ }
+ else {
+ use_maskna = 0;
+ }
+ }
+ }
+
/* Get the appropriate ufunc inner loop */
if (use_maskna) {
retcode = get_masked_binary_op_function(self, arr, otype,
&otype_dtype, &maskedinnerloop, &maskedinnerloopdata);
}
else {
- /*
- * TODO: Switch to using the type resolution function like
- * in the masked case.
- */
int otype_final = otype;
retcode = get_binary_op_function(self, &otype_final,
&innerloop, &innerloopdata);
@@ -2917,42 +2928,41 @@ PyUFunc_Reduce(PyUFuncObject *self, PyArrayObject *arr, PyArrayObject *out,
/* Allocate an output or conform 'out' to 'self' */
result = allocate_or_conform_reduce_result(arr, out,
- axis_flags, otype_dtype, !skipna);
+ axis_flags, otype_dtype, !skipna && use_maskna);
if (result == NULL) {
return NULL;
}
/* Prepare the NA mask if there is one */
if (use_maskna) {
- //printf("doing masked %s.reduce\n", ufunc_name); fflush(stdout);
/*
* Do the reduction on the NA mask before the data. This way
* we can avoid modifying the outputs which end up masked, obeying
* the required NA masking semantics.
*/
if (!skipna) {
- if (PyArray_HASMASKNA(result)) {
- if (PyArray_ReduceMaskNAArray(ndim, PyArray_DIMS(arr),
- PyArray_MASKNA_DTYPE(arr),
- PyArray_MASKNA_DATA(arr),
- PyArray_MASKNA_STRIDES(arr),
- PyArray_MASKNA_DTYPE(result),
- PyArray_MASKNA_DATA(result),
- PyArray_MASKNA_STRIDES(result)) < 0) {
- goto fail;
+ int idim;
+ npy_intp result_strides[NPY_MAXDIMS];
+ /* Need to make sure the appropriate strides are 0 in 'result' */
+ for (idim = 0; idim < PyArray_NDIM(arr); ++idim) {
+ if (PyArray_DIMS(result)[idim] == 1) {
+ result_strides[idim] = 0;
+ }
+ else {
+ result_strides[idim] = PyArray_MASKNA_STRIDES(result)[idim];
}
}
- else if (PyArray_ContainsNA(arr)) {
- PyErr_SetString(PyExc_ValueError,
- "Cannot assign NA value to an array which "
- "does not support NAs");
+ if (PyArray_ReduceMaskNAArray(ndim, PyArray_DIMS(arr),
+ PyArray_MASKNA_DTYPE(arr),
+ PyArray_MASKNA_DATA(arr),
+ PyArray_MASKNA_STRIDES(arr),
+ PyArray_MASKNA_DTYPE(result),
+ PyArray_MASKNA_DATA(result),
+ result_strides) < 0) {
goto fail;
}
- else {
- use_maskna = 0;
- }
- /* Short circuit any calculation if the result is a single NA */
+ /* Short circuit any calculation if the result 0-dim NA */
if (PyArray_SIZE(result) == 1 &&
!NpyMaskValue_IsExposed(
(npy_mask)*PyArray_MASKNA_DATA(result))) {
diff --git a/numpy/core/tests/test_maskna.py b/numpy/core/tests/test_maskna.py
index 1207760d1..ec8fda00f 100644
--- a/numpy/core/tests/test_maskna.py
+++ b/numpy/core/tests/test_maskna.py
@@ -520,5 +520,34 @@ def test_ufunc_1D():
assert_(c.flags.maskna)
#assert_equal(c, [0,2,4])
+def test_ufunc_reduce_1D():
+ a = np.arange(3.0, maskna=True)
+ b = np.array(0.5)
+ c = np.array(0.5, maskna=True)
+
+ # Since 'a' has no NA values, this should work
+ np.add.reduce(a, out=b)
+ assert_equal(b, 3.0)
+
+ # With an NA value, the reduce should throw with the non-NA output param
+ a[1] = np.NA
+ assert_raises(ValueError, np.add.reduce, a, out=b)
+
+ # With an NA value, the output parameter can still be an NA-array
+ np.add.reduce(a, out=c)
+ assert_(np.isna(c))
+
+ # Should not touch the out= element when assigning NA
+ b[...] = 1.0
+ d = b.view(maskna=True)
+ np.add.reduce(a, out=d)
+ assert_(np.isna(d))
+ assert_equal(b, 1.0)
+
+ # Without an output parameter, return NA
+ ret = np.add.reduce(a)
+ assert_(np.isna(ret))
+
+
if __name__ == "__main__":
run_module_suite()