diff options
| -rw-r--r-- | numpy/core/src/multiarray/ctors.c | 10 | ||||
| -rw-r--r-- | numpy/core/src/multiarray/item_selection.c | 219 | ||||
| -rw-r--r-- | numpy/core/src/multiarray/methods.c | 41 | ||||
| -rw-r--r-- | numpy/core/src/multiarray/multiarraymodule.c | 3 | ||||
| -rw-r--r-- | numpy/core/tests/test_maskna.py | 61 |
5 files changed, 286 insertions, 48 deletions
diff --git a/numpy/core/src/multiarray/ctors.c b/numpy/core/src/multiarray/ctors.c index 16e912c95..ac3a6648b 100644 --- a/numpy/core/src/multiarray/ctors.c +++ b/numpy/core/src/multiarray/ctors.c @@ -236,7 +236,7 @@ _update_descr_and_dimensions(PyArray_Descr **des, npy_intp *newdims, newnd = oldnd + numnew; - if (newnd > MAX_DIMS) { + if (newnd > NPY_MAXDIMS) { goto finish; } if (tuple) { @@ -2453,7 +2453,7 @@ PyArray_FromInterface(PyObject *input) char *data; Py_ssize_t buffer_len; int res, i, n; - intp dims[MAX_DIMS], strides[MAX_DIMS]; + intp dims[NPY_MAXDIMS], strides[NPY_MAXDIMS]; int dataflags = NPY_ARRAY_BEHAVED; /* Get the memory from __array_data__ and __array_offset__ */ @@ -2732,7 +2732,7 @@ PyArray_FromDimsAndDataAndDescr(int nd, int *d, { PyObject *ret; int i; - npy_intp newd[MAX_DIMS]; + npy_intp newd[NPY_MAXDIMS]; char msg[] = "PyArray_FromDimsAndDataAndDescr: use PyArray_NewFromDescr."; if (DEPRECATE(msg) < 0) { @@ -3822,14 +3822,14 @@ PyArray_CheckAxis(PyArrayObject *arr, int *axis, int flags) PyObject *temp1, *temp2; int n = PyArray_NDIM(arr); - if (*axis == MAX_DIMS || n == 0) { + if (*axis == NPY_MAXDIMS || n == 0) { if (n != 1) { temp1 = PyArray_Ravel(arr,0); if (temp1 == NULL) { *axis = 0; return NULL; } - if (*axis == MAX_DIMS) { + if (*axis == NPY_MAXDIMS) { *axis = PyArray_NDIM((PyArrayObject *)temp1)-1; } } diff --git a/numpy/core/src/multiarray/item_selection.c b/numpy/core/src/multiarray/item_selection.c index 5be85aaef..8e861d460 100644 --- a/numpy/core/src/multiarray/item_selection.c +++ b/numpy/core/src/multiarray/item_selection.c @@ -32,10 +32,11 @@ PyArray_TakeFrom(PyArrayObject *self0, PyObject *indices0, int axis, intp nd, i, j, n, m, max_item, tmp, chunk, nelem; intp shape[MAX_DIMS]; char *src, *dest; - int err; + int err, use_maskna = 0; indices = NULL; - self = (PyArrayObject *)PyArray_CheckAxis(self0, &axis, NPY_ARRAY_CARRAY); + self = (PyArrayObject *)PyArray_CheckAxis(self0, &axis, + NPY_ARRAY_CARRAY | NPY_ARRAY_ALLOWNA); if (self == NULL) { return NULL; } @@ -45,6 +46,9 @@ PyArray_TakeFrom(PyArrayObject *self0, PyObject *indices0, int axis, if (indices == NULL) { goto fail; } + + + n = m = chunk = 1; nd = PyArray_NDIM(self) + PyArray_NDIM(indices) - 1; for (i = 0; i < nd; i++) { @@ -75,14 +79,24 @@ PyArray_TakeFrom(PyArrayObject *self0, PyObject *indices0, int axis, if (obj == NULL) { goto fail; } + + /* Allocate an NA mask if necessary */ + if (PyArray_HASMASKNA(self)) { + if (PyArray_AllocateMaskNA(obj, 1, 0, 1) < 0) { + goto fail; + } + use_maskna = 1; + } } else { - int flags = NPY_ARRAY_CARRAY | NPY_ARRAY_UPDATEIFCOPY; + int flags = NPY_ARRAY_CARRAY | + NPY_ARRAY_UPDATEIFCOPY | + NPY_ARRAY_ALLOWNA; if ((PyArray_NDIM(out) != nd) || !PyArray_CompareLists(PyArray_DIMS(out), shape, nd)) { PyErr_SetString(PyExc_ValueError, - "bad shape in output array"); + "output array does not match result of ndarray.take"); goto fail; } @@ -100,6 +114,18 @@ PyArray_TakeFrom(PyArrayObject *self0, PyObject *indices0, int axis, if (obj == NULL) { goto fail; } + + if (PyArray_HASMASKNA(self)) { + if (PyArray_HASMASKNA(obj)) { + use_maskna = 1; + } + else if (PyArray_ContainsNA(self)) { + PyErr_SetString(PyExc_ValueError, + "Cannot assign NA value to an array which " + "does not support NAs"); + goto fail; + } + } } max_item = PyArray_DIMS(self)[axis]; @@ -109,7 +135,121 @@ PyArray_TakeFrom(PyArrayObject *self0, PyObject *indices0, int axis, dest = PyArray_DATA(obj); func = PyArray_DESCR(self)->f->fasttake; - if (func == NULL) { + if (use_maskna) { + char *dst_maskna = NULL, *src_maskna = NULL; + npy_intp itemsize = PyArray_DESCR(obj)->elsize; + PyArray_MaskedStridedTransferFn *maskedstransfer = NULL; + NpyAuxData *transferdata = NULL; + int needs_api = 0; + + if (PyArray_GetMaskedDTypeTransferFunction( + 1, + itemsize, + itemsize, + 1, + PyArray_DESCR(obj), + PyArray_DESCR(obj), + PyArray_MASKNA_DTYPE(obj), + 0, + &maskedstransfer, &transferdata, + &needs_api) != NPY_SUCCEED) { + goto fail; + } + + + src_maskna = PyArray_MASKNA_DATA(self); + dst_maskna = PyArray_MASKNA_DATA(obj); + if (PyDataType_REFCHK(PyArray_DESCR(self))) { + /* + * TODO: Should use PyArray_GetDTypeTransferFunction + * instead of raw memmove to remedy this. + */ + PyErr_SetString(PyExc_RuntimeError, + "ndarray.take doesn't support object arrays with " + "masks yet"); + NPY_AUXDATA_FREE(transferdata); + goto fail; + } + + switch(clipmode) { + case NPY_RAISE: + for (i = 0; i < n; i++) { + for (j = 0; j < m; j++) { + tmp = ((intp *)(PyArray_DATA(indices)))[j]; + if (tmp < 0) { + tmp = tmp + max_item; + } + if ((tmp < 0) || (tmp >= max_item)) { + PyErr_SetString(PyExc_IndexError, + "index out of range "\ + "for array"); + NPY_AUXDATA_FREE(transferdata); + goto fail; + } + maskedstransfer(dest, itemsize, + src + tmp*chunk, itemsize, + (npy_mask *)src_maskna, 1, + nelem, itemsize, transferdata); + dest += chunk; + memmove(dst_maskna, src_maskna + tmp*nelem, nelem); + dst_maskna += nelem; + } + src += chunk*max_item; + src_maskna += nelem*max_item; + } + break; + case NPY_WRAP: + for (i = 0; i < n; i++) { + for (j = 0; j < m; j++) { + tmp = ((intp *)(PyArray_DATA(indices)))[j]; + if (tmp < 0) { + while (tmp < 0) { + tmp += max_item; + } + } + else if (tmp >= max_item) { + while (tmp >= max_item) { + tmp -= max_item; + } + } + maskedstransfer(dest, itemsize, + src + tmp*chunk, itemsize, + (npy_mask *)src_maskna, 1, + nelem, itemsize, transferdata); + dest += chunk; + memmove(dst_maskna, src_maskna + tmp*nelem, nelem); + dst_maskna += nelem; + } + src += chunk*max_item; + src_maskna += nelem*max_item; + } + break; + case NPY_CLIP: + for (i = 0; i < n; i++) { + for (j = 0; j < m; j++) { + tmp = ((intp *)(PyArray_DATA(indices)))[j]; + if (tmp < 0) { + tmp = 0; + } + else if (tmp >= max_item) { + tmp = max_item - 1; + } + maskedstransfer(dest, itemsize, + src + tmp*chunk, itemsize, + (npy_mask *)src_maskna, 1, + nelem, itemsize, transferdata); + dest += chunk; + memmove(dst_maskna, src_maskna + tmp*nelem, nelem); + dst_maskna += nelem; + } + src += chunk*max_item; + src_maskna += nelem*max_item; + } + break; + } + NPY_AUXDATA_FREE(transferdata); + } + else if (func == NULL) { switch(clipmode) { case NPY_RAISE: for (i = 0; i < n; i++) { @@ -1667,14 +1807,26 @@ PyArray_Compress(PyArrayObject *self, PyObject *condition, int axis, PyArrayObject *cond; PyObject *res, *ret; - cond = (PyArrayObject *)PyArray_FROM_O(condition); - if (cond == NULL) { - return NULL; + if (PyArray_Check(condition)) { + cond = (PyArrayObject *)condition; + Py_INCREF(cond); } + else { + PyArray_Descr *dtype = PyArray_DescrFromType(NPY_BOOL); + if (dtype == NULL) { + return NULL; + } + cond = (PyArrayObject *)PyArray_FromAny(condition, dtype, + 0, 0, NPY_ARRAY_ALLOWNA, NULL); + if (cond == NULL) { + return NULL; + } + } + if (PyArray_NDIM(cond) != 1) { Py_DECREF(cond); PyErr_SetString(PyExc_ValueError, - "condition must be 1-d array"); + "condition must be a 1-d array"); return NULL; } @@ -1742,7 +1894,8 @@ count_boolean_trues(int ndim, char *data, npy_intp *ashape, npy_intp *astrides) } /*NUMPY_API - * Counts the number of non-zero elements in the array + * Counts the number of non-zero elements in the array. Raises + * an error if the array contains an NA. * * Returns -1 on error. */ @@ -1759,6 +1912,16 @@ PyArray_CountNonzero(PyArrayObject *self) char **dataptr; npy_intp *strideptr, *innersizeptr; + /* If 'self' has an NA mask, make sure it has no NA values */ + if (PyArray_HASMASKNA(self)) { + if (PyArray_ContainsNA(self)) { + PyErr_SetString(PyExc_ValueError, + "Cannot count the number of nonzeros in an array " + "which contains an NA"); + return -1; + } + } + /* Special low-overhead version specific to the boolean type */ if (PyArray_DESCR(self)->type_num == NPY_BOOL) { return count_boolean_trues(PyArray_NDIM(self), PyArray_DATA(self), @@ -1789,9 +1952,14 @@ PyArray_CountNonzero(PyArrayObject *self) return 0; } - /* Otherwise create and use an iterator to count the nonzeros */ - iter = NpyIter_New(self, NPY_ITER_READONLY| - NPY_ITER_EXTERNAL_LOOP| + /* + * Otherwise create and use an iterator to count the nonzeros. + * Can ignore any NAs because we already checked PyArray_ContainsNA + * earlier. + */ + iter = NpyIter_New(self, NPY_ITER_READONLY | + NPY_ITER_IGNORE_MASKNA | + NPY_ITER_EXTERNAL_LOOP | NPY_ITER_REFS_OK, NPY_KEEPORDER, NPY_NO_CASTING, NULL); @@ -1844,7 +2012,7 @@ PyArray_Nonzero(PyArrayObject *self) PyArray_NonzeroFunc *nonzero = PyArray_DESCR(self)->f->nonzero; char *data; npy_intp stride, count; - npy_intp nonzero_count = PyArray_CountNonzero(self); + npy_intp nonzero_count; npy_intp *multi_index; NpyIter *iter; @@ -1852,6 +2020,16 @@ PyArray_Nonzero(PyArrayObject *self) NpyIter_GetMultiIndexFunc *get_multi_index; char **dataptr; + /* + * First count the number of non-zeros in 'self'. If 'self' contains + * an NA value, this will raise an error, so after this call + * we can assume 'self' contains no NAs. + */ + nonzero_count = PyArray_CountNonzero(self); + if (nonzero_count < 0) { + return NULL; + } + /* Allocate the result as a 2D array */ ret_dims[0] = nonzero_count; ret_dims[1] = (ndim == 0) ? 1 : ndim; @@ -1881,10 +2059,15 @@ PyArray_Nonzero(PyArrayObject *self) goto finish; } - /* Build an iterator tracking a multi-index, in C order */ - iter = NpyIter_New(self, NPY_ITER_READONLY| - NPY_ITER_MULTI_INDEX| - NPY_ITER_ZEROSIZE_OK| + /* + * Build an iterator tracking a multi-index, in C order. We + * can ignore NAs because the PyArray_CountNonzero call checked + * that there were no NAs already. + */ + iter = NpyIter_New(self, NPY_ITER_READONLY | + NPY_ITER_IGNORE_MASKNA | + NPY_ITER_MULTI_INDEX | + NPY_ITER_ZEROSIZE_OK | NPY_ITER_REFS_OK, NPY_CORDER, NPY_NO_CASTING, NULL); diff --git a/numpy/core/src/multiarray/methods.c b/numpy/core/src/multiarray/methods.c index 7f850f4dd..c728f17e0 100644 --- a/numpy/core/src/multiarray/methods.c +++ b/numpy/core/src/multiarray/methods.c @@ -46,9 +46,6 @@ NpyArg_ParseKeywords(PyObject *keys, const char *format, char **kwlist, ...) return ret; } -/* Should only be used if x is known to be an nd-array */ -#define _ARET(x) PyArray_Return((PyArrayObject *)(x)) - static PyObject * array_take(PyArrayObject *self, PyObject *args, PyObject *kwds) { @@ -65,7 +62,8 @@ array_take(PyArrayObject *self, PyObject *args, PyObject *kwds) PyArray_ClipmodeConverter, &mode)) return NULL; - return _ARET(PyArray_TakeFrom(self, indices, dimension, out, mode)); + return PyArray_Return((PyArrayObject *) + PyArray_TakeFrom(self, indices, dimension, out, mode)); } static PyObject * @@ -103,7 +101,7 @@ array_reshape(PyArrayObject *self, PyObject *args, PyObject *kwds) static char *keywords[] = {"order", NULL}; PyArray_Dims newshape; PyObject *ret; - PyArray_ORDER order = PyArray_CORDER; + PyArray_ORDER order = NPY_CORDER; Py_ssize_t n = PyTuple_Size(args); if (!NpyArg_ParseKeywords(kwds, "|O&", keywords, @@ -222,7 +220,7 @@ array_argmax(PyArrayObject *self, PyObject *args, PyObject *kwds) PyArray_OutputConverter, &out)) return NULL; - return _ARET(PyArray_ArgMax(self, axis, out)); + return PyArray_Return((PyArrayObject *)PyArray_ArgMax(self, axis, out)); } static PyObject * @@ -237,7 +235,7 @@ array_argmin(PyArrayObject *self, PyObject *args, PyObject *kwds) PyArray_OutputConverter, &out)) return NULL; - return _ARET(PyArray_ArgMin(self, axis, out)); + return PyArray_Return((PyArrayObject *)PyArray_ArgMin(self, axis, out)); } static PyObject * @@ -1042,7 +1040,7 @@ array_getarray(PyArrayObject *self, PyObject *args) static PyObject * array_copy(PyArrayObject *self, PyObject *args, PyObject *kwds) { - PyArray_ORDER order = PyArray_CORDER; + PyArray_ORDER order = NPY_CORDER; static char *kwlist[] = {"order", NULL}; if (!PyArg_ParseTupleAndKeywords(args, kwds, "|O&", kwlist, @@ -1087,7 +1085,7 @@ array_resize(PyArrayObject *self, PyObject *args, PyObject *kwds) return NULL; } - ret = PyArray_Resize(self, &newshape, refcheck, PyArray_CORDER); + ret = PyArray_Resize(self, &newshape, refcheck, NPY_CORDER); PyDimMem_FREE(newshape.ptr); if (ret == NULL) { return NULL; @@ -1108,7 +1106,7 @@ array_repeat(PyArrayObject *self, PyObject *args, PyObject *kwds) { PyArray_AxisConverter, &axis)) { return NULL; } - return _ARET(PyArray_Repeat(self, repeats, axis)); + return PyArray_Return((PyArrayObject *)PyArray_Repeat(self, repeats, axis)); } static PyObject * @@ -1135,7 +1133,7 @@ array_choose(PyArrayObject *self, PyObject *args, PyObject *kwds) return NULL; } - return _ARET(PyArray_Choose(self, choices, out, clipmode)); + return PyArray_Return((PyArrayObject *)PyArray_Choose(self, choices, out, clipmode)); } static PyObject * @@ -1241,7 +1239,7 @@ array_argsort(PyArrayObject *self, PyObject *args, PyObject *kwds) Py_XDECREF(PyArray_DESCR(self)); ((PyArrayObject_fieldaccess *)self)->descr = saved; } - return _ARET(res); + return PyArray_Return((PyArrayObject *)res); } static PyObject * @@ -1256,7 +1254,7 @@ array_searchsorted(PyArrayObject *self, PyObject *args, PyObject *kwds) PyArray_SearchsideConverter, &side)) { return NULL; } - return _ARET(PyArray_SearchSorted(self, keys, side)); + return PyArray_Return((PyArrayObject *)PyArray_SearchSorted(self, keys, side)); } static void @@ -2047,7 +2045,8 @@ array_compress(PyArrayObject *self, PyObject *args, PyObject *kwds) PyArray_OutputConverter, &out)) { return NULL; } - return _ARET(PyArray_Compress(self, condition, axis, out)); + return PyArray_Return( + (PyArrayObject *)PyArray_Compress(self, condition, axis, out)); } @@ -2082,7 +2081,7 @@ array_trace(PyArrayObject *self, PyObject *args, PyObject *kwds) rtype = _CHKTYPENUM(dtype); Py_XDECREF(dtype); - return _ARET(PyArray_Trace(self, offset, axis1, axis2, rtype, out)); + return PyArray_Return((PyArrayObject *)PyArray_Trace(self, offset, axis1, axis2, rtype, out)); } #undef _CHKTYPENUM @@ -2105,7 +2104,7 @@ array_clip(PyArrayObject *self, PyObject *args, PyObject *kwds) PyErr_SetString(PyExc_ValueError, "One of max or min must be given."); return NULL; } - return _ARET(PyArray_Clip(self, min, max, out)); + return PyArray_Return((PyArrayObject *)PyArray_Clip(self, min, max, out)); } @@ -2135,14 +2134,14 @@ array_diagonal(PyArrayObject *self, PyObject *args, PyObject *kwds) &axis2)) { return NULL; } - return _ARET(PyArray_Diagonal(self, offset, axis1, axis2)); + return PyArray_Return((PyArrayObject *)PyArray_Diagonal(self, offset, axis1, axis2)); } static PyObject * array_flatten(PyArrayObject *self, PyObject *args, PyObject *kwds) { - PyArray_ORDER order = PyArray_CORDER; + PyArray_ORDER order = NPY_CORDER; static char *kwlist[] = {"order", NULL}; if (!PyArg_ParseTupleAndKeywords(args, kwds, "|O&", kwlist, @@ -2156,7 +2155,7 @@ array_flatten(PyArrayObject *self, PyObject *args, PyObject *kwds) static PyObject * array_ravel(PyArrayObject *self, PyObject *args, PyObject *kwds) { - PyArray_ORDER order = PyArray_CORDER; + PyArray_ORDER order = NPY_CORDER; static char *kwlist[] = {"order", NULL}; if (!PyArg_ParseTupleAndKeywords(args, kwds, "|O&", kwlist, @@ -2179,7 +2178,7 @@ array_round(PyArrayObject *self, PyObject *args, PyObject *kwds) PyArray_OutputConverter, &out)) { return NULL; } - return _ARET(PyArray_Round(self, decimals, out)); + return PyArray_Return((PyArrayObject *)PyArray_Round(self, decimals, out)); } @@ -2467,5 +2466,3 @@ NPY_NO_EXPORT PyMethodDef array_methods[] = { METH_VARARGS | METH_KEYWORDS, NULL}, {NULL, NULL, 0, NULL} /* sentinel */ }; - -#undef _ARET diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c index c7e0cfa9d..9f73139b9 100644 --- a/numpy/core/src/multiarray/multiarraymodule.c +++ b/numpy/core/src/multiarray/multiarraymodule.c @@ -1995,7 +1995,8 @@ array_count_nonzero(PyObject *NPY_UNUSED(self), PyObject *args) return NULL; } - array = (PyArrayObject *)PyArray_FromAny(array_in, NULL, 0, 0, 0, NULL); + array = (PyArrayObject *)PyArray_FromAny(array_in, NULL, + 0, 0, NPY_ARRAY_ALLOWNA, NULL); if (array == NULL) { return NULL; } diff --git a/numpy/core/tests/test_maskna.py b/numpy/core/tests/test_maskna.py index 393edf74d..a7127868a 100644 --- a/numpy/core/tests/test_maskna.py +++ b/numpy/core/tests/test_maskna.py @@ -517,7 +517,64 @@ def test_array_maskna_view_array_assignment_1D(): # TODO: fancy indexing is next... -def test_ufunc_1D(): +def test_maskna_nonzero_1D(): + a = np.zeros((5,), maskna=True) + + # The nonzeros without any NAs + assert_equal(np.count_nonzero(a), 0) + assert_equal(np.nonzero(a)[0], []) + a[2] = 3 + assert_equal(np.count_nonzero(a), 1) + assert_equal(np.nonzero(a)[0], [2]) + a[3:] = 2 + assert_equal(np.count_nonzero(a), 3) + assert_equal(np.nonzero(a)[0], [2,3,4]) + + # The nonzeros with an NA + a[2] = np.NA + assert_raises(ValueError, np.count_nonzero, a) + assert_raises(ValueError, np.nonzero, a) + +def test_maskna_take_1D(): + a = np.arange(5, maskna=True) + b = np.arange(3) + c = b.view(maskna=True) + + # Take without any NAs + assert_equal(a.take([0,2,4]), [0,2,4]) + + # Take without any NAs, into non-NA output parameter + a.take([0,2,4], out=b) + assert_equal(b, [0,2,4]) + + # Take without any NAs, into NA output parameter + b[...] = 1 + c[...] = np.NA + a.take([0,2,4], out=c) + assert_equal(c, [0,2,4]) + + # Take with some NAs + a[2] = np.NA + a[3] = np.NA + ret = a.take([0,2,4]) + assert_equal([ret[0], ret[2]], [0,4]) + assert_equal(np.isna(ret), [0,1,0]) + + # Take with some NAs, into NA output parameter + b[...] = 1 + c[...] = np.NA + a.take([0,2,4], out=c) + assert_equal(b, [0,1,4]) + assert_equal([c[0], c[2]], [0,4]) + assert_equal(np.isna(c), [0,1,0]) + + c[...] = 1 + a.take([0,2,4], out=c) + assert_equal(b, [0,1,4]) + assert_equal([c[0], c[2]], [0,4]) + assert_equal(np.isna(c), [0,1,0]) + +def test_maskna_ufunc_1D(): a = np.arange(3, maskna=True) b = np.arange(3) @@ -526,7 +583,7 @@ def test_ufunc_1D(): assert_(c.flags.maskna) #assert_equal(c, [0,2,4]) -def test_ufunc_add_reduce_1D(): +def test_maskna_ufunc_add_reduce_1D(): a = np.arange(3.0, maskna=True) b = np.array(0.5) c_orig = np.array(0.5) |