diff options
Diffstat (limited to 'numpy')
| -rw-r--r-- | numpy/core/arrayprint.py | 6 | ||||
| -rw-r--r-- | numpy/core/code_generators/numpy_api.py | 4 | ||||
| -rw-r--r-- | numpy/core/src/multiarray/convert_datatype.c | 2 | ||||
| -rw-r--r-- | numpy/core/src/multiarray/na_mask.c | 8 | ||||
| -rw-r--r-- | numpy/core/src/multiarray/shape.c | 196 | ||||
| -rw-r--r-- | numpy/core/src/umath/ufunc_object.c | 298 | ||||
| -rw-r--r-- | numpy/core/tests/test_maskna.py | 19 |
7 files changed, 416 insertions, 117 deletions
diff --git a/numpy/core/arrayprint.py b/numpy/core/arrayprint.py index fcda825c7..525fcab8a 100644 --- a/numpy/core/arrayprint.py +++ b/numpy/core/arrayprint.py @@ -218,13 +218,16 @@ def _boolFormatter(x): def _array2string(a, max_line_width, precision, suppress_small, separator=' ', prefix="", formatter=None): + print "DEBUG: in array2string!" if max_line_width is None: max_line_width = _line_width + print "DEBUG: A" if precision is None: precision = _float_output_precision + print "DEBUG: B" if suppress_small is None: suppress_small = _float_output_suppress_small @@ -238,6 +241,7 @@ def _array2string(a, max_line_width, precision, suppress_small, separator=' ', summary_insert = "" data = ravel(a) + print "DEBUG: making formatdict" formatdict = {'bool' : _boolFormatter, 'int' : IntegerFormat(data), 'float' : FloatFormat(data, precision, suppress_small), @@ -249,6 +253,8 @@ def _array2string(a, max_line_width, precision, suppress_small, separator=' ', 'timedelta' : TimedeltaFormat(data), 'numpystr' : repr, 'str' : str} + print "DEBUG: made formatdict" + if formatter is not None: fkeys = [k for k in formatter.keys() if formatter[k] is not None] if 'all' in fkeys: diff --git a/numpy/core/code_generators/numpy_api.py b/numpy/core/code_generators/numpy_api.py index 0dadd0574..a066e6968 100644 --- a/numpy/core/code_generators/numpy_api.py +++ b/numpy/core/code_generators/numpy_api.py @@ -326,10 +326,12 @@ multiarray_funcs_api = { 'PyArray_AllocateMaskNA': 286, 'NpyIter_GetFirstMaskNAOp': 288, 'NpyIter_GetMaskNAIndexArray': 289, - 'PyArray_ReduceMaskArray': 290, + 'PyArray_ReduceMaskNAArray': 290, 'PyArray_CreateSortedStridePerm': 291, 'PyArray_FillWithZero': 292, 'PyArray_FillWithOne': 293, + 'PyArray_AssignNA': 294, + 'PyArray_AssignMaskNA': 295, } ufunc_types_api = { diff --git a/numpy/core/src/multiarray/convert_datatype.c b/numpy/core/src/multiarray/convert_datatype.c index 1eb9931a4..e5cd965aa 100644 --- a/numpy/core/src/multiarray/convert_datatype.c +++ b/numpy/core/src/multiarray/convert_datatype.c @@ -1836,7 +1836,7 @@ PyArray_ConvertToCommonType(PyObject *op, int *retn) /* Make sure all arrays are actual array objects. */ for (i = 0; i < n; i++) { - int flags = NPY_ARRAY_CARRAY; + int flags = NPY_ARRAY_CARRAY | NPY_ARRAY_ALLOWNA; if ((otmp = PySequence_GetItem(op, i)) == NULL) { goto fail; diff --git a/numpy/core/src/multiarray/na_mask.c b/numpy/core/src/multiarray/na_mask.c index e649e9e1b..70b6e7838 100644 --- a/numpy/core/src/multiarray/na_mask.c +++ b/numpy/core/src/multiarray/na_mask.c @@ -123,7 +123,8 @@ fill_raw_byte_array(int ndim, npy_intp *shape, return 0; } -/* +/*NUMPY_API + * * Assigns the mask value to all the NA mask elements of * the array. * @@ -290,7 +291,8 @@ PyArray_AllocateMaskNA(PyArrayObject *arr, return 0; } -/* +/*NUMPY_API + * * Assigns the given NA value to all the elements in the array. If * 'arr' has a mask, masks all the elements of the array. * @@ -463,7 +465,7 @@ PyArray_IsNA(PyObject *obj) * Returns 0 on success, -1 on failure. */ NPY_NO_EXPORT int -PyArray_ReduceMaskArray(int ndim, npy_intp *shape, +PyArray_ReduceMaskNAArray(int ndim, npy_intp *shape, PyArray_Descr *src_dtype, char *src_data, npy_intp *src_strides, PyArray_Descr *dst_dtype, char *dst_data, npy_intp *dst_strides) { diff --git a/numpy/core/src/multiarray/shape.c b/numpy/core/src/multiarray/shape.c index da44f032c..56c08304c 100644 --- a/numpy/core/src/multiarray/shape.c +++ b/numpy/core/src/multiarray/shape.c @@ -18,17 +18,19 @@ #include "shape.h" -#define PyAO PyArrayObject +#define PyArrayObject PyArrayObject static int -_check_ones(PyArrayObject *self, int newnd, intp* newdims, intp *strides); +_check_ones(PyArrayObject *self, int newnd, + npy_intp* newdims, npy_intp *strides, npy_intp *masknastrides); static int -_fix_unknown_dimension(PyArray_Dims *newshape, intp s_original); +_fix_unknown_dimension(PyArray_Dims *newshape, npy_intp s_original); static int -_attempt_nocopy_reshape(PyArrayObject *self, int newnd, intp* newdims, - intp *newstrides, int is_f_order); +_attempt_nocopy_reshape(PyArrayObject *self, int newnd, npy_intp* newdims, + npy_intp *newstrides, npy_intp *newmasknastrides, + int is_f_order); static void _putzero(char *optr, PyObject *zero, PyArray_Descr *dtype); @@ -43,15 +45,15 @@ NPY_NO_EXPORT PyObject * PyArray_Resize(PyArrayObject *self, PyArray_Dims *newshape, int refcheck, NPY_ORDER order) { - intp oldsize, newsize; + npy_intp oldsize, newsize; int new_nd=newshape->len, k, n, elsize; int refcnt; - intp* new_dimensions=newshape->ptr; - intp new_strides[MAX_DIMS]; + npy_intp* new_dimensions=newshape->ptr; + npy_intp new_strides[NPY_MAXDIMS]; size_t sd; - intp *dimptr; + npy_intp *dimptr; char *new_data; - intp largest; + npy_intp largest; if (!PyArray_ISONESEGMENT(self)) { PyErr_SetString(PyExc_ValueError, @@ -158,8 +160,8 @@ PyArray_Resize(PyArrayObject *self, PyArray_Dims *newshape, int refcheck, sd = (size_t) PyArray_DESCR(self)->elsize; sd = (size_t) _array_fill_strides(new_strides, new_dimensions, new_nd, sd, PyArray_FLAGS(self), &(((PyArrayObject_fieldaccess *)self)->flags)); - memmove(PyArray_DIMS(self), new_dimensions, new_nd*sizeof(intp)); - memmove(PyArray_STRIDES(self), new_strides, new_nd*sizeof(intp)); + memmove(PyArray_DIMS(self), new_dimensions, new_nd*sizeof(npy_intp)); + memmove(PyArray_STRIDES(self), new_strides, new_nd*sizeof(npy_intp)); Py_INCREF(Py_None); return Py_None; } @@ -178,29 +180,33 @@ NPY_NO_EXPORT PyObject * PyArray_Newshape(PyArrayObject *self, PyArray_Dims *newdims, NPY_ORDER order) { - intp i; - intp *dimensions = newdims->ptr; + npy_intp i; + npy_intp *dimensions = newdims->ptr; PyArrayObject *ret; - int n = newdims->len; - Bool same, incref = TRUE; - intp *strides = NULL; - intp newstrides[MAX_DIMS]; - int flags; + int ndim = newdims->len; + npy_bool same, incref = TRUE; + npy_intp *strides = NULL; + npy_intp newstrides[NPY_MAXDIMS]; + npy_intp newmasknastrides[NPY_MAXDIMS]; + int flags, build_maskna_strides = 0; + + printf("in newshape\n"); fflush(stdout); if (order == NPY_ANYORDER) { order = PyArray_ISFORTRAN(self); } /* Quick check to make sure anything actually needs to be done */ - if (n == PyArray_NDIM(self)) { + if (ndim == PyArray_NDIM(self)) { same = TRUE; i = 0; - while (same && i < n) { + while (same && i < ndim) { if (PyArray_DIM(self,i) != dimensions[i]) { same=FALSE; } i++; } if (same) { + printf("returning view\n"); fflush(stdout); return PyArray_View(self, NULL, NULL); } } @@ -213,13 +219,16 @@ PyArray_Newshape(PyArrayObject *self, PyArray_Dims *newdims, * In this case we don't need to do anything but update strides and * dimensions. So, we can handle non single-segment cases. */ - i = _check_ones(self, n, dimensions, newstrides); + i = _check_ones(self, ndim, dimensions, newstrides, newmasknastrides); if (i == 0) { + printf("setting strides to newstrides\n"); fflush(stdout); strides = newstrides; } - flags = PyArray_FLAGS(self); + flags = PyArray_FLAGS(self) & ~(NPY_ARRAY_OWNMASKNA | + NPY_ARRAY_MASKNA); if (strides == NULL) { + printf("strides are null\n"); fflush(stdout); /* * we are really re-shaping not just adding ones to the shape somewhere * fix any -1 dimensions and check new-dimensions against old size @@ -239,29 +248,35 @@ PyArray_Newshape(PyArrayObject *self, PyArray_Dims *newdims, (PyArray_CHKFLAGS(self, NPY_ARRAY_F_CONTIGUOUS) && order == NPY_CORDER)) && (PyArray_NDIM(self) > 1))) { int success = 0; - success = _attempt_nocopy_reshape(self,n,dimensions, - newstrides,order); + success = _attempt_nocopy_reshape(self, ndim, dimensions, + newstrides, newmasknastrides, order); if (success) { + printf("nocopy reshape succeeded\n"); fflush(stdout); /* no need to copy the array after all */ strides = newstrides; - flags = PyArray_FLAGS(self); } else { - PyObject *new; - new = PyArray_NewCopy(self, order); - if (new == NULL) { + printf("nocopy reshape failed\n"); fflush(stdout); + PyObject *newcopy; + newcopy = PyArray_NewCopy(self, order); + if (newcopy == NULL) { return NULL; } incref = FALSE; - self = (PyArrayObject *)new; - flags = PyArray_FLAGS(self); + self = (PyArrayObject *)newcopy; + build_maskna_strides = 1; } + flags = PyArray_FLAGS(self) & ~(NPY_ARRAY_OWNMASKNA | + NPY_ARRAY_MASKNA); + } + else { + build_maskna_strides = 1; } /* We always have to interpret the contiguous buffer correctly */ /* Make sure the flags argument is set. */ - if (n > 1) { + if (ndim > 1) { if (order == NPY_FORTRANORDER) { flags &= ~NPY_ARRAY_C_CONTIGUOUS; flags |= NPY_ARRAY_F_CONTIGUOUS; @@ -272,7 +287,7 @@ PyArray_Newshape(PyArrayObject *self, PyArray_Dims *newdims, } } } - else if (n > 0) { + else if (ndim > 0) { /* * replace any 0-valued strides with * appropriate value to preserve contiguousness @@ -281,17 +296,17 @@ PyArray_Newshape(PyArrayObject *self, PyArray_Dims *newdims, if (strides[0] == 0) { strides[0] = PyArray_DESCR(self)->elsize; } - for (i = 1; i < n; i++) { + for (i = 1; i < ndim; i++) { if (strides[i] == 0) { strides[i] = strides[i-1] * dimensions[i-1]; } } } else { - if (strides[n-1] == 0) { - strides[n-1] = PyArray_DESCR(self)->elsize; + if (strides[ndim-1] == 0) { + strides[ndim-1] = PyArray_DESCR(self)->elsize; } - for (i = n - 2; i > -1; i--) { + for (i = ndim - 2; i > -1; i--) { if (strides[i] == 0) { strides[i] = strides[i+1] * dimensions[i+1]; } @@ -300,9 +315,9 @@ PyArray_Newshape(PyArrayObject *self, PyArray_Dims *newdims, } Py_INCREF(PyArray_DESCR(self)); - ret = (PyAO *)PyArray_NewFromDescr(Py_TYPE(self), + ret = (PyArrayObject *)PyArray_NewFromDescr(Py_TYPE(self), PyArray_DESCR(self), - n, dimensions, + ndim, dimensions, strides, PyArray_DATA(self), flags, (PyObject *)self); @@ -310,6 +325,7 @@ PyArray_Newshape(PyArrayObject *self, PyArray_Dims *newdims, if (ret == NULL) { goto fail; } + if (incref) { Py_INCREF(self); } @@ -317,6 +333,37 @@ PyArray_Newshape(PyArrayObject *self, PyArray_Dims *newdims, Py_DECREF(ret); return NULL; } + + /* If there's an NA mask, make sure to view it too */ + if (PyArray_HASMASKNA(self)) { + PyArrayObject_fieldaccess *fa = (PyArrayObject_fieldaccess *)ret; + fa->maskna_dtype = PyArray_MASKNA_DTYPE(self); + Py_INCREF(fa->maskna_dtype); + fa->maskna_data = PyArray_MASKNA_DATA(self); + if (build_maskna_strides) { + npy_intp stride = 1; + if (order == NPY_FORTRANORDER) { + printf("building fortran strides\n"); fflush(stdout); + for (i = 0; i < ndim; ++i) { + fa->maskna_strides[i] = stride; + printf("stride %d\n", stride); + stride *= fa->dimensions[i]; + } + } + else { + for (i = ndim; i >= 0; --i) { + fa->maskna_strides[i] = stride; + stride *= fa->dimensions[i]; + } + } + } + else { + memcpy(fa->maskna_strides, newmasknastrides, + fa->nd * sizeof(npy_intp)); + } + fa->flags |= NPY_ARRAY_MASKNA; + } + PyArray_UpdateFlags(ret, NPY_ARRAY_C_CONTIGUOUS | NPY_ARRAY_F_CONTIGUOUS); return (PyObject *)ret; @@ -350,12 +397,14 @@ PyArray_Reshape(PyArrayObject *self, PyObject *shape) /* inserts 0 for strides where dimension will be 1 */ static int -_check_ones(PyArrayObject *self, int newnd, intp* newdims, intp *strides) +_check_ones(PyArrayObject *self, int newnd, + npy_intp* newdims, npy_intp *strides, npy_intp *masknastrides) { int nd; - intp *dims; - Bool done=FALSE; + npy_intp *dims; + npy_bool done=FALSE; int j, k; + int has_maskna = PyArray_HASMASKNA(self); nd = PyArray_NDIM(self); dims = PyArray_DIMS(self); @@ -363,11 +412,17 @@ _check_ones(PyArrayObject *self, int newnd, intp* newdims, intp *strides) for (k = 0, j = 0; !done && (j < nd || k < newnd);) { if ((j<nd) && (k<newnd) && (newdims[k] == dims[j])) { strides[k] = PyArray_STRIDES(self)[j]; + if (has_maskna) { + masknastrides[k] = PyArray_MASKNA_STRIDES(self)[j]; + } j++; k++; } else if ((k < newnd) && (newdims[k] == 1)) { strides[k] = 0; + if (has_maskna) { + masknastrides[k] = 0; + } k++; } else if ((j<nd) && (dims[j] == 1)) { @@ -431,20 +486,25 @@ _putzero(char *optr, PyObject *zero, PyArray_Descr *dtype) * stride of the next-fastest index. */ static int -_attempt_nocopy_reshape(PyArrayObject *self, int newnd, intp* newdims, - intp *newstrides, int is_f_order) +_attempt_nocopy_reshape(PyArrayObject *self, int newnd, npy_intp* newdims, + npy_intp *newstrides, npy_intp *newmasknastrides, + int is_f_order) { int oldnd; - intp olddims[MAX_DIMS]; - intp oldstrides[MAX_DIMS]; + npy_intp olddims[NPY_MAXDIMS]; + npy_intp oldstrides[NPY_MAXDIMS], oldmasknastrides[NPY_MAXDIMS]; int oi, oj, ok, ni, nj, nk; int np, op; + int has_maskna = PyArray_HASMASKNA(self); oldnd = 0; for (oi = 0; oi < PyArray_NDIM(self); oi++) { if (PyArray_DIMS(self)[oi]!= 1) { olddims[oldnd] = PyArray_DIMS(self)[oi]; oldstrides[oldnd] = PyArray_STRIDES(self)[oi]; + if (has_maskna) { + oldmasknastrides[oldnd] = PyArray_MASKNA_STRIDES(self)[oi]; + } oldnd++; } } @@ -495,14 +555,18 @@ _attempt_nocopy_reshape(PyArrayObject *self, int newnd, intp* newdims, for (ok = oi; ok < oj - 1; ok++) { if (is_f_order) { - if (oldstrides[ok+1] != olddims[ok]*oldstrides[ok]) { + if (oldstrides[ok+1] != olddims[ok]*oldstrides[ok] || + (has_maskna && oldmasknastrides[ok+1] != + olddims[ok]*oldmasknastrides[ok])) { /* not contiguous enough */ return 0; } } else { /* C order */ - if (oldstrides[ok] != olddims[ok+1]*oldstrides[ok+1]) { + if (oldstrides[ok] != olddims[ok+1]*oldstrides[ok+1] || + (has_maskna && oldmasknastrides[ok] != + olddims[ok+1]*oldmasknastrides[ok+1])) { /* not contiguous enough */ return 0; } @@ -514,6 +578,13 @@ _attempt_nocopy_reshape(PyArrayObject *self, int newnd, intp* newdims, for (nk = ni + 1; nk < nj; nk++) { newstrides[nk] = newstrides[nk - 1]*newdims[nk - 1]; } + if (has_maskna) { + newmasknastrides[ni] = oldmasknastrides[oi]; + for (nk = ni + 1; nk < nj; nk++) { + newmasknastrides[nk] = + newmasknastrides[nk - 1]*newdims[nk - 1]; + } + } } else { /* C order */ @@ -521,6 +592,13 @@ _attempt_nocopy_reshape(PyArrayObject *self, int newnd, intp* newdims, for (nk = nj - 1; nk > ni; nk--) { newstrides[nk - 1] = newstrides[nk]*newdims[nk]; } + if (has_maskna) { + newmasknastrides[nj - 1] = oldmasknastrides[oj - 1]; + for (nk = nj - 1; nk > ni; nk--) { + newmasknastrides[nk - 1] = + newmasknastrides[nk]*newdims[nk]; + } + } } ni = nj++; oi = oj++; @@ -540,10 +618,10 @@ _attempt_nocopy_reshape(PyArrayObject *self, int newnd, intp* newdims, } static int -_fix_unknown_dimension(PyArray_Dims *newshape, intp s_original) +_fix_unknown_dimension(PyArray_Dims *newshape, npy_intp s_original) { - intp *dimensions; - intp i_unknown, s_known; + npy_intp *dimensions; + npy_intp i_unknown, s_known; int i, n; static char msg[] = "total size of new array must be unchanged"; @@ -596,8 +674,8 @@ PyArray_Squeeze(PyArrayObject *self) { int nd = PyArray_NDIM(self); int newnd = nd; - intp dimensions[MAX_DIMS]; - intp strides[MAX_DIMS]; + npy_intp dimensions[NPY_MAXDIMS]; + npy_intp strides[NPY_MAXDIMS]; int i, j; PyArrayObject *ret; PyArray_Descr *dtype; @@ -643,7 +721,7 @@ NPY_NO_EXPORT PyObject * PyArray_SwapAxes(PyArrayObject *ap, int a1, int a2) { PyArray_Dims new_axes; - intp dims[MAX_DIMS]; + npy_intp dims[NPY_MAXDIMS]; int n, i, val; PyObject *ret; @@ -699,9 +777,9 @@ PyArray_SwapAxes(PyArrayObject *ap, int a1, int a2) NPY_NO_EXPORT PyObject * PyArray_Transpose(PyArrayObject *ap, PyArray_Dims *permute) { - intp *axes, axis; - intp i, n; - intp permutation[MAX_DIMS], reverse_permutation[MAX_DIMS]; + npy_intp *axes, axis; + npy_intp i, n; + npy_intp permutation[NPY_MAXDIMS], reverse_permutation[NPY_MAXDIMS]; PyArrayObject *ret = NULL; if (permute == NULL) { @@ -839,7 +917,7 @@ NPY_NO_EXPORT PyObject * PyArray_Ravel(PyArrayObject *a, NPY_ORDER order) { PyArray_Dims newdim = {NULL,1}; - intp val[1] = {-1}; + npy_intp val[1] = {-1}; newdim.ptr = val; @@ -918,7 +996,7 @@ NPY_NO_EXPORT PyObject * PyArray_Flatten(PyArrayObject *a, NPY_ORDER order) { PyArrayObject *ret; - intp size; + npy_intp size; if (order == NPY_ANYORDER) { order = PyArray_ISFORTRAN(a) ? NPY_FORTRANORDER : NPY_CORDER; diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c index 1709795f1..1dd0ef730 100644 --- a/numpy/core/src/umath/ufunc_object.c +++ b/numpy/core/src/umath/ufunc_object.c @@ -2433,6 +2433,76 @@ get_binary_op_function(PyUFuncObject *self, int *otype, } /* + * Given the output type, finds the specified binary op, and + * returns a masked inner loop. The ufunc must have nin==2 + * and nout==1. The function may modify otype if the given + * type isn't found. + * + * Returns 0 on success, -1 on failure. + */ +static int +get_masked_binary_op_function(PyUFuncObject *self, PyArrayObject *arr, + int otype, + PyArray_Descr **out_dtype, + PyUFuncGenericMaskedFunction *out_innerloop, + NpyAuxData **out_innerloopdata) +{ + int i, retcode; + PyArrayObject *op[3] = {arr, arr, arr}; + PyArray_Descr *dtype[3] = {NULL, NULL, NULL}; + PyObject *type_tup = NULL; + + NPY_UF_DBG_PRINT1("Getting masked binary op function for type number %d\n", + *otype); + + *out_dtype = NULL; + + /* Build a type tuple if otype is specified */ + if (otype != NPY_NOTYPE) { + PyArray_Descr *otype_dtype = PyArray_DescrFromType(otype); + if (otype_dtype == NULL) { + return -1; + } + type_tup = Py_BuildValue("(N)", otype_dtype); + if (type_tup == NULL) { + return -1; + } + } + + /* Use the type resolution function to find our loop */ + retcode = self->type_resolution_masked_function(self, NPY_SAFE_CASTING, + op, type_tup, dtype, + out_innerloop, out_innerloopdata); + Py_XDECREF(type_tup); + if (retcode == -1) { + return -1; + } + else if (retcode == -2) { + PyErr_SetString(PyExc_RuntimeError, + "type resolution returned NotImplemented"); + return -1; + } + + /* The selected dtypes should all be equivalent */ + if (!PyArray_EquivTypes(dtype[0], dtype[1]) || + !PyArray_EquivTypes(dtype[1], dtype[2])) { + for (i = 0; i < 3; ++i) { + Py_DECREF(dtype[i]); + } + PyErr_SetString(PyExc_RuntimeError, + "could not find a masked binary loop appropriate for " + "reduce ufunc"); + return -1; + } + + *out_dtype = dtype[0]; + Py_DECREF(dtype[1]); + Py_DECREF(dtype[2]); + + return 0; +} + +/* * Allocates a result array for a reduction operation, with * dimensions matching 'arr' except set to 1 with 0 stride * whereever axis_flags is True. Dropping the reduction axes @@ -2574,34 +2644,15 @@ conform_reduce_result(int ndim, npy_bool *axis_flags, PyArrayObject *out) /* Allocate 'result' or conform 'out' to 'self' (in 'result') */ static PyArrayObject * allocate_or_conform_reduce_result(PyArrayObject *arr, PyArrayObject *out, - npy_bool *axis_flags, int otype_final, int keepmask) + npy_bool *axis_flags, PyArray_Descr * otype_dtype, + int keepmask) { if (out == NULL) { PyArrayObject *result; - PyArray_Descr *dtype = NULL; - /* - * Set up the output data type, using the input's exact - * data type if the type number didn't change to preserve - * metadata - */ - if (PyArray_DESCR(arr)->type_num == otype_final) { - if (PyArray_ISNBO(PyArray_DESCR(arr)->byteorder)) { - dtype = PyArray_DESCR(arr); - Py_INCREF(dtype); - } - else { - dtype = PyArray_DescrNewByteorder(PyArray_DESCR(arr), - NPY_NATIVE); - } - } - else { - dtype = PyArray_DescrFromType(otype_final); - } - if (dtype == NULL) { - return NULL; - } - result = allocate_reduce_result(arr, axis_flags, dtype); + printf("allocating result\n"); fflush(stdout); + Py_INCREF(otype_dtype); + result = allocate_reduce_result(arr, axis_flags, otype_dtype); /* Allocate an NA mask if necessary */ if (keepmask && result != NULL && PyArray_HASMASKNA(arr)) { @@ -2614,6 +2665,7 @@ allocate_or_conform_reduce_result(PyArrayObject *arr, PyArrayObject *out, return result; } else { + printf("conforming result\n"); fflush(stdout); return conform_reduce_result(PyArray_NDIM(arr), axis_flags, out); } } @@ -2767,16 +2819,20 @@ static PyObject * PyUFunc_Reduce(PyUFuncObject *self, PyArrayObject *arr, PyArrayObject *out, int naxes, int *axes, int otype, int skipna) { - int ndim; - int iaxes; + int iaxes, ndim, retcode; + PyArray_Descr *otype_dtype = NULL; npy_bool axis_flags[NPY_MAXDIMS]; PyArrayObject *arr_view = NULL, *result = NULL; - int otype_final; - /* The selected inner loop */ + /* The normal selected inner loop */ PyUFuncGenericFunction innerloop = NULL; void *innerloopdata = NULL; + /* The masked selected inner loop */ + int use_maskna = 0; + PyUFuncGenericMaskedFunction maskedinnerloop = NULL; + NpyAuxData *maskedinnerloopdata = NULL; + char *ufunc_name = self->name ? self->name : "(unknown)"; NPY_BEGIN_THREADS_DEF; @@ -2808,10 +2864,45 @@ PyUFunc_Reduce(PyUFuncObject *self, PyArrayObject *arr, PyArrayObject *out, axis_flags[axis] = 1; } + use_maskna = PyArray_HASMASKNA(arr); + /* Get the appropriate ufunc inner loop */ - otype_final = otype; - if (get_binary_op_function(self, &otype_final, - &innerloop, &innerloopdata) < 0) { + 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); + + /* + * Set up the output data type, using the input's exact + * data type if the type number didn't change to preserve + * metadata + */ + if (PyArray_DESCR(arr)->type_num == otype_final) { + if (PyArray_ISNBO(PyArray_DESCR(arr)->byteorder)) { + otype_dtype = PyArray_DESCR(arr); + Py_INCREF(otype_dtype); + } + else { + otype_dtype = PyArray_DescrNewByteorder(PyArray_DESCR(arr), + NPY_NATIVE); + } + } + else { + otype_dtype = PyArray_DescrFromType(otype_final); + } + if (otype_dtype == NULL) { + return NULL; + } + } + if (retcode < 0) { PyArray_Descr *dtype = PyArray_DescrFromType(otype); PyErr_Format(PyExc_ValueError, "could not find a matching type for %s.reduce, " @@ -2828,11 +2919,62 @@ 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_final, !skipna); + axis_flags, otype_dtype, !skipna); 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; + } + } + else 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; + } + + /* Short circuit any calculation if the result is a single NA */ + if (PyArray_SIZE(result) == 1 && + !NpyMaskValue_IsExposed( + (npy_mask)*PyArray_MASKNA_DATA(result))) { + goto finish; + } + } + else { + /* + * If there's no identity, validate that there is no + * reduction which is all NA values. + */ + if (self->identity == PyUFunc_None) { + PyErr_SetString(PyExc_ValueError, + "skipna=True together with a non-identity reduction " + "isn't implemented yet"); + goto fail; + } + } + } + /* * Initialize 'result' to the identity or initial elements * copied from 'arr', and create a view of 'arr' containing @@ -2849,11 +2991,8 @@ PyUFunc_Reduce(PyUFuncObject *self, PyArrayObject *arr, PyArrayObject *out, op[1] = arr_view; /* op is length 3 in case the inner loop wanted these as its data */ op[2] = result; - op_dtypes[0] = PyArray_DescrFromType(otype_final); - if (op_dtypes[0] == NULL) { - goto fail; - } - op_dtypes[1] = op_dtypes[0]; + op_dtypes[0] = otype_dtype; + op_dtypes[1] = otype_dtype; flags = NPY_ITER_BUFFERED | NPY_ITER_EXTERNAL_LOOP | @@ -2867,6 +3006,32 @@ PyUFunc_Reduce(PyUFuncObject *self, PyArrayObject *arr, PyArrayObject *out, op_flags[1] = NPY_ITER_READONLY | NPY_ITER_ALIGNED; + /* 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; + + /* 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; + } + } + else { + /* + * If 'out' had no mask, and 'arr' did, we checked that 'arr' + * contains no NA values and can ignore the masks. + */ + op_flags[0] |= NPY_ITER_IGNORE_MASKNA; + op_flags[1] |= NPY_ITER_IGNORE_MASKNA; + } + iter = NpyIter_MultiNew(2, op, flags, NPY_KEEPORDER, NPY_UNSAFE_CASTING, op_flags, @@ -2894,23 +3059,44 @@ PyUFunc_Reduce(PyUFuncObject *self, PyArrayObject *arr, PyArrayObject *out, count_ptr = NpyIter_GetInnerLoopSizePtr(iter); needs_api = NpyIter_IterationNeedsAPI(iter) || - otype_final == NPY_OBJECT; + PyDataType_REFCHK(otype_dtype); if (!needs_api) { NPY_BEGIN_THREADS; } - do { - /* Turn the two items into three for the inner loop */ - dataptr_copy[0] = dataptr[0]; - dataptr_copy[1] = dataptr[1]; - dataptr_copy[2] = dataptr[0]; - stride_copy[0] = stride[0]; - stride_copy[1] = stride[1]; - stride_copy[2] = stride[0]; - innerloop(dataptr_copy, count_ptr, - stride_copy, innerloopdata); - } while (iternext(iter)); + /* Straightforward reduction */ + if (!use_maskna) { + do { + /* Turn the two items into three for the inner loop */ + dataptr_copy[0] = dataptr[0]; + dataptr_copy[1] = dataptr[1]; + dataptr_copy[2] = dataptr[0]; + stride_copy[0] = stride[0]; + stride_copy[1] = stride[1]; + stride_copy[2] = stride[0]; + innerloop(dataptr_copy, count_ptr, + stride_copy, innerloopdata); + } while (iternext(iter)); + } + /* Masked reduction */ + else { + do { + /* Turn the two items into three for the inner loop */ + dataptr_copy[0] = dataptr[0]; + dataptr_copy[1] = dataptr[1]; + dataptr_copy[2] = dataptr[0]; + stride_copy[0] = stride[0]; + stride_copy[1] = stride[1]; + stride_copy[2] = stride[0]; + /* + * If skipna=True, this masks based on the mask in 'arr', + * otherwise it masks based on the mask in 'result' + */ + maskedinnerloop(dataptr_copy, dataptr[2], count_ptr, + stride_copy, stride[2], maskedinnerloopdata); + } while (iternext(iter)); + } if (!needs_api) { NPY_END_THREADS; @@ -2921,6 +3107,8 @@ PyUFunc_Reduce(PyUFuncObject *self, PyArrayObject *arr, PyArrayObject *out, } } +finish: + /* Strip out the extra 'one' dimensions in the result */ if (out == NULL) { strip_flagged_dimensions(result, axis_flags); @@ -2931,9 +3119,12 @@ PyUFunc_Reduce(PyUFuncObject *self, PyArrayObject *arr, PyArrayObject *out, Py_INCREF(result); } - NpyIter_Deallocate(iter); - Py_DECREF(arr_view); - Py_DECREF(op_dtypes[0]); + if (iter != NULL) { + NpyIter_Deallocate(iter); + } + Py_XDECREF(arr_view); + Py_XDECREF(otype_dtype); + NPY_AUXDATA_FREE(maskedinnerloopdata); return (PyObject *)result; fail: @@ -2942,7 +3133,8 @@ fail: } Py_XDECREF(result); Py_XDECREF(arr_view); - Py_XDECREF(op_dtypes[0]); + Py_XDECREF(otype_dtype); + NPY_AUXDATA_FREE(maskedinnerloopdata); return NULL; } diff --git a/numpy/core/tests/test_maskna.py b/numpy/core/tests/test_maskna.py index c2ae14d86..dbd435c8c 100644 --- a/numpy/core/tests/test_maskna.py +++ b/numpy/core/tests/test_maskna.py @@ -225,6 +225,25 @@ def test_array_maskna_array_function_1D(): assert_(c.flags.ownmaskna) assert_(not (c is b_view)) +def test_array_maskna_reshape(): + # Simple reshape 1D -> 2D + a = np.arange(6, maskna=True) + a[1] = np.NA + a[5] = np.NA + + b = a.reshape(2,3) + assert_(b.base is a) + assert_equal(b.shape, (2,3)) + assert_(b.flags.maskna) + assert_(not b.flags.ownmaskna) + assert_equal(np.isna(b), [[0,1,0],[0,0,1]]) + + b = a.reshape(2,3,order='F') + assert_equal(b.shape, (2,3)) + assert_(b.flags.maskna) + assert_(not b.flags.ownmaskna) + assert_equal(np.isna(b), [[0,0,0],[1,0,1]]) + def test_array_maskna_view_NA_assignment_1D(): a = np.arange(10) a_ref = a.copy() |