diff options
| -rw-r--r-- | doc/neps/missing-data.rst | 18 | ||||
| -rw-r--r-- | numpy/core/src/umath/ufunc_object.c | 81 |
2 files changed, 76 insertions, 23 deletions
diff --git a/doc/neps/missing-data.rst b/doc/neps/missing-data.rst index bc61ebe1b..051cd9351 100644 --- a/doc/neps/missing-data.rst +++ b/doc/neps/missing-data.rst @@ -301,18 +301,20 @@ As part of the implementation, ufuncs and other operations will have to be extended to support masked computation. Because this is a useful feature in general, even outside the context of a masked array, in addition to working with masked arrays ufuncs -will take an optional 'mask=' parameter which allows the use -of boolean arrays to choose where a computation should be done. -This functions similar to a "where" clause on the ufunc.:: +will take an optional 'where=' parameter which allows the use +of boolean arrays to choose where a computation should be done.:: - >>> np.add(a, b, out=b, mask=(a > threshold)) + >>> np.add(a, b, out=b, where=(a > threshold)) -A benefit of having this 'mask=' parameter is that it provides a way +A benefit of having this 'where=' parameter is that it provides a way to temporarily treat an object with a mask without ever creating a -masked array object. +masked array object. In the example above, this would only do the +add for the array elements with True in the 'where' clause, and neither +'a' nor 'b' need to be masked arrays. -If the 'out' parameter isn't specified, use of the 'mask=' parameter -will produce an array with a mask as the result. +If the 'out' parameter isn't specified, use of the 'where=' parameter +will produce an array with a mask as the result, with missing values +for everywhere the 'where' clause had the value False. For boolean operations, the R project special cases logical_and and logical_or so that logical_and(NA, False) is False, and diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c index 8300c5d5e..3c605e1f3 100644 --- a/numpy/core/src/umath/ufunc_object.c +++ b/numpy/core/src/umath/ufunc_object.c @@ -721,7 +721,8 @@ static int get_ufunc_arguments(PyUFuncObject *self, NPY_CASTING *out_casting, PyObject **out_extobj, PyObject **out_typetup, - int *out_subok) + int *out_subok, + PyArrayObject **out_wheremask) { npy_intp i, nargs, nin = self->nin; PyObject *obj, *context; @@ -732,6 +733,12 @@ static int get_ufunc_arguments(PyUFuncObject *self, ufunc_name = self->name ? self->name : "<unnamed ufunc>"; + *out_extobj = NULL; + *out_typetup = NULL; + if (out_wheremask != NULL) { + *out_wheremask = NULL; + } + /* Check number of arguments */ nargs = PyTuple_Size(args); if ((nargs < nin) || (nargs > self->nargs)) { @@ -843,6 +850,25 @@ static int get_ufunc_arguments(PyUFuncObject *self, bad_arg = 0; } break; + case 'd': + /* Another way to specify 'sig' */ + if (strncmp(str,"dtype",5) == 0) { + /* Allow this parameter to be None */ + PyArray_Descr *dtype; + if (!PyArray_DescrConverter2(value, &dtype)) { + goto fail; + } + if (dtype != NULL) { + if (*out_typetup != NULL) { + PyErr_SetString(PyExc_RuntimeError, + "cannot specify both 'sig' and 'dtype'"); + goto fail; + } + *out_typetup = Py_BuildValue("(N)", dtype); + } + bad_arg = 0; + } + break; case 'e': /* * Overrides the global parameters buffer size, @@ -910,24 +936,27 @@ static int get_ufunc_arguments(PyUFuncObject *self, bad_arg = 0; } break; - case 'd': - /* Another way to specify 'sig' */ - if (strncmp(str,"dtype",5) == 0) { - /* Allow this parameter to be None */ + case 'w': + /* + * Provides a boolean array 'where=' mask if + * out_wheremask is supplied. + */ + if (out_wheremask != NULL && + strncmp(str,"where",5) == 0) { PyArray_Descr *dtype; - if (!PyArray_DescrConverter2(value, &dtype)) { + dtype = PyArray_DescrFromType(NPY_BOOL); + if (dtype == NULL) { goto fail; } - if (dtype != NULL) { - if (*out_typetup != NULL) { - PyErr_SetString(PyExc_RuntimeError, - "cannot specify both 'sig' and 'dtype'"); - goto fail; - } - *out_typetup = Py_BuildValue("(N)", dtype); + *out_wheremask = (PyArrayObject *)PyArray_FromAny( + value, dtype, + 0, 0, 0, NULL); + if (*out_wheremask == NULL) { + goto fail; } bad_arg = 0; } + break; } if (bad_arg) { @@ -943,6 +972,14 @@ static int get_ufunc_arguments(PyUFuncObject *self, fail: Py_XDECREF(str_key_obj); + Py_XDECREF(*out_extobj); + *out_extobj = NULL; + Py_XDECREF(*out_typetup); + *out_typetup = NULL; + if (out_wheremask != NULL) { + Py_XDECREF(*out_wheremask); + *out_wheremask = NULL; + } return -1; } @@ -1499,7 +1536,8 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *self, /* Get all the arguments */ retval = get_ufunc_arguments(self, args, kwds, - op, &order, &casting, &extobj, &type_tup, &subok); + op, &order, &casting, &extobj, + &type_tup, &subok, NULL); if (retval < 0) { goto fail; } @@ -1839,6 +1877,16 @@ PyUFunc_GenericFunction(PyUFuncObject *self, PyUFuncGenericFunction innerloop = NULL; void *innerloopdata = NULL; + /* + * The selected masked inner loop, when the 'where=' + * parameter or arrays with missing values are in op. + */ + PyUFuncGenericMaskedFunction masked_innerloop = NULL; + NpyAuxData *masked_innerloopdata = NULL; + + /* The mask provided in the 'where=' parameter */ + PyArrayObject *wheremask = NULL; + /* The __array_prepare__ function to call for each output */ PyObject *arr_prep[NPY_MAXARGS]; /* @@ -1885,7 +1933,8 @@ PyUFunc_GenericFunction(PyUFuncObject *self, /* Get all the arguments */ retval = get_ufunc_arguments(self, args, kwds, - op, &order, &casting, &extobj, &type_tup, &subok); + op, &order, &casting, &extobj, + &type_tup, &subok, &wheremask); if (retval < 0) { goto fail; } @@ -2006,6 +2055,7 @@ PyUFunc_GenericFunction(PyUFuncObject *self, Py_XDECREF(errobj); Py_XDECREF(type_tup); Py_XDECREF(arr_prep_args); + Py_XDECREF(wheremask); NPY_UF_DBG_PRINT("Returning Success\n"); @@ -2022,6 +2072,7 @@ fail: Py_XDECREF(errobj); Py_XDECREF(type_tup); Py_XDECREF(arr_prep_args); + Py_XDECREF(wheremask); return retval; } |