summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
Diffstat
-rw-r--r--numpy/core/code_generators/numpy_api.py2
-rw-r--r--numpy/core/src/multiarray/nditer_api.c64
-rw-r--r--numpy/core/src/multiarray/nditer_constr.c113
-rw-r--r--numpy/core/src/multiarray/nditer_pywrap.c105
4 files changed, 246 insertions, 38 deletions
diff --git a/numpy/core/code_generators/numpy_api.py b/numpy/core/code_generators/numpy_api.py
index ec8388311..49d38c32f 100644
--- a/numpy/core/code_generators/numpy_api.py
+++ b/numpy/core/code_generators/numpy_api.py
@@ -324,6 +324,8 @@ multiarray_funcs_api = {
'PyArray_HasNASupport': 284,
'PyArray_ContainsNA': 285,
'PyArray_AllocateMaskNA': 286,
+ 'NpyIter_GetFirstMaskNAOp': 288,
+ 'NpyIter_GetMaskNAIndices': 289,
}
ufunc_types_api = {
diff --git a/numpy/core/src/multiarray/nditer_api.c b/numpy/core/src/multiarray/nditer_api.c
index 96715acd8..ddd227bf6 100644
--- a/numpy/core/src/multiarray/nditer_api.c
+++ b/numpy/core/src/multiarray/nditer_api.c
@@ -272,7 +272,9 @@ NpyIter_Reset(NpyIter *iter, char **errmsg)
}
/*NUMPY_API
- * Resets the iterator to its initial state, with new base data pointers
+ * Resets the iterator to its initial state, with new base data pointers.
+ * This function requires great caution, even more so if any
+ * NPY_ITER_USE_MASKNA operands were specified.
*
* If errmsg is non-NULL, it should point to a variable which will
* receive the error message, and no Python exception will be set.
@@ -748,6 +750,35 @@ NpyIter_GetNOp(NpyIter *iter)
}
/*NUMPY_API
+ * Gets the index of the first operand which is the
+ * mask for an NPY_ITER_USE_MASKNA operand.
+ */
+NPY_NO_EXPORT int
+NpyIter_GetFirstMaskNAOp(NpyIter *iter)
+{
+ return NIT_FIRST_MASKNA_OP(iter);
+}
+
+/*NUMPY_API
+ * Gets the correspondences between the operands with
+ * NPY_ITER_USEMASKNA set and their corresponding masks.
+ *
+ * If i < NpyIter_GetFirstMaskNAOp(iter), then
+ * NpyIter_GetMaskNAIndices(iter)[i] is either -1 or
+ * an index >= NpyIter_GetFirstMaskNAOp(iter) of the corresponding
+ * mask.
+ *
+ * If i >= NpyIter_GetFirstMaskNAOp(iter), then
+ * NpyIter_GetMaskNAIndices(iter)[i] is the index
+ * of the corresponding maskna operand for the mask.
+ */
+NPY_NO_EXPORT npy_int8 *
+NpyIter_GetMaskNAIndices(NpyIter *iter)
+{
+ return NIT_MASKNA_INDICES(iter);
+}
+
+/*NUMPY_API
* Gets the number of elements being iterated
*/
NPY_NO_EXPORT npy_intp
@@ -1004,6 +1035,7 @@ NpyIter_GetIterView(NpyIter *iter, npy_intp i)
npy_uint32 itflags = NIT_ITFLAGS(iter);
int idim, ndim = NIT_NDIM(iter);
int nop = NIT_NOP(iter);
+ int first_maskna_op = NIT_FIRST_MASKNA_OP(iter);
npy_intp shape[NPY_MAXDIMS], strides[NPY_MAXDIMS];
PyArrayObject *obj, *view;
@@ -1012,8 +1044,9 @@ NpyIter_GetIterView(NpyIter *iter, npy_intp i)
NpyIter_AxisData *axisdata;
npy_intp sizeof_axisdata;
int writeable;
+ npy_int8 *maskna_indices = NIT_MASKNA_INDICES(iter);
- if (i < 0 || i >= nop) {
+ if (i < 0 || i >= first_maskna_op) {
PyErr_SetString(PyExc_IndexError,
"index provided for an iterator view was out of bounds");
return NULL;
@@ -1034,9 +1067,11 @@ NpyIter_GetIterView(NpyIter *iter, npy_intp i)
sizeof_axisdata = NIT_AXISDATA_SIZEOF(itflags, ndim, nop);
/* Retrieve the shape and strides from the axisdata */
- for (idim = 0; idim < ndim; ++idim, NIT_ADVANCE_AXISDATA(axisdata, 1)) {
+ for (idim = 0; idim < ndim; ++idim) {
shape[ndim-idim-1] = NAD_SHAPE(axisdata);
strides[ndim-idim-1] = NAD_STRIDES(axisdata)[i];
+
+ NIT_ADVANCE_AXISDATA(axisdata, 1);
}
Py_INCREF(dtype);
@@ -1055,6 +1090,29 @@ NpyIter_GetIterView(NpyIter *iter, npy_intp i)
}
/* Make sure all the flags are good */
PyArray_UpdateFlags(view, NPY_ARRAY_UPDATE_ALL);
+ /*
+ * Add the mask to the view if the operand was NPY_ITER_USE_MASKNA.
+ */
+ if (maskna_indices[i] >= 0) {
+ PyArrayObject_fieldaccess *fview = (PyArrayObject_fieldaccess *)view;
+ int i_maskna = maskna_indices[i];
+ npy_intp *maskna_strides = fview->maskna_strides;
+
+ fview->maskna_dtype = PyArray_MASKNA_DTYPE(obj);
+ Py_INCREF(fview->maskna_dtype);
+ fview->maskna_data = NIT_RESETDATAPTR(iter)[i_maskna];
+
+ axisdata = NIT_AXISDATA(iter);
+ for (idim = 0; idim < ndim; ++idim) {
+ maskna_strides[ndim-idim-1] = NAD_STRIDES(axisdata)[i_maskna];
+
+ NIT_ADVANCE_AXISDATA(axisdata, 1);
+ }
+
+ /* This view doesn't own the mask */
+ fview->flags |= NPY_ARRAY_MASKNA;
+ fview->flags &= ~NPY_ARRAY_OWNMASKNA;
+ }
return view;
}
diff --git a/numpy/core/src/multiarray/nditer_constr.c b/numpy/core/src/multiarray/nditer_constr.c
index 2054a4fdb..5e4401c34 100644
--- a/numpy/core/src/multiarray/nditer_constr.c
+++ b/numpy/core/src/multiarray/nditer_constr.c
@@ -85,6 +85,8 @@ npyiter_allocate_arrays(NpyIter *iter,
PyArray_Descr **op_dtype, PyTypeObject *subtype,
npy_uint32 *op_flags, char *op_itflags,
int **op_axes, int output_scalars);
+static int
+npyiter_fill_maskna_axisdata(NpyIter *iter, int **op_axes);
static void
npyiter_get_priority_subtype(int first_maskna_op, PyArrayObject **op,
char *op_itflags,
@@ -406,6 +408,18 @@ NpyIter_AdvancedNew(int nop, PyArrayObject **op_in, npy_uint32 flags,
return NULL;
}
+ /*
+ * If there were any NA masks added to the iteration, fill in
+ * the strides and other data they need. This is being done
+ * after all the 'allocate' arrays are finished.
+ */
+ if (maskna_nop > 0) {
+ if (!npyiter_fill_maskna_axisdata(iter, op_axes)) {
+ NpyIter_Deallocate(iter);
+ return NULL;
+ }
+ }
+
NPY_IT_TIME_POINT(c_allocate_arrays);
/*
@@ -1149,27 +1163,6 @@ npyiter_prepare_one_operand(PyArrayObject **op,
}
/*
- * Prepares the maskna virtual operand for one of the constructor
- * operands. Assumes that 'op' has already been prepared by
- * npyiter_prepare_one_operand. Fills in 'op_maskna_dataptr',
- * 'op_maskna_dtype', and 'op_maskna_itflags'.
- *
- * This needs to be called after any 'allocate' operands have
- * been allocated.
- *
- * Returns 1 on success, 0 on failure.
- */
-static int
-npyiter_prepare_one_maskna_operand(PyArrayObject *op,
- char **op_maskna_dataptr,
- PyArray_Descr **op_maskna_dtype,
- npy_uint32 flags,
- npy_uint32 op_flags, char *op_maskna_itflags)
-{
- return 1;
-}
-
-/*
* Process all the operands, copying new references so further processing
* can replace the arrays if copying is necessary.
*/
@@ -3128,6 +3121,84 @@ npyiter_allocate_arrays(NpyIter *iter,
}
/*
+ * Prepares the maskna virtual operands for the constructor
+ * operands, and fills in the axisdata. Fills in 'op_maskna_dataptr',
+ * 'op_maskna_dtype', and may modify 'op_maskna_itflags'.
+ *
+ * This needs to be called after any 'allocate' operands have
+ * been allocated. There is no validation of the shape/strides done,
+ * because the shape of a mask exactly matches the shape of the
+ * operand to which it attached.
+ *
+ * Returns 1 on success, 0 on failure.
+ */
+static int
+npyiter_fill_maskna_axisdata(NpyIter *iter, int **op_axes)
+{
+ npy_uint32 itflags = NIT_ITFLAGS(iter);
+ int idim, ndim = NIT_NDIM(iter);
+ int iop, iop_maskna, nop = NIT_NOP(iter);
+ int first_maskna_op = NIT_FIRST_MASKNA_OP(iter);
+
+ npy_int8 *maskna_indices = NIT_MASKNA_INDICES(iter);
+ NpyIter_AxisData *axisdata;
+ npy_intp sizeof_axisdata;
+ PyArrayObject **op = NIT_OPERANDS(iter), *op_cur;
+ char **op_dataptr = NIT_RESETDATAPTR(iter);
+
+ axisdata = NIT_AXISDATA(iter);
+ sizeof_axisdata = NIT_AXISDATA_SIZEOF(itflags, ndim, nop);
+
+ /* Fill in the reset dataptr array with the mask pointers */
+ for (iop = first_maskna_op; iop < nop; ++iop) {
+ op_dataptr[iop] = PyArray_MASKNA_DATA(op[maskna_indices[iop]]);
+ }
+
+ /* Process the maskna operands, filling in the axisdata */
+ for (idim = 0; idim < ndim; ++idim) {
+ npy_intp *strides = NAD_STRIDES(axisdata);
+
+ for (iop = first_maskna_op; iop < nop; ++iop) {
+ /*
+ * iop_maskna is the index of the USE_MASKNA input,
+ * iop is the index of the corresponding mask.
+ */
+ iop_maskna = maskna_indices[iop];
+ op_cur = op[iop_maskna];
+
+ /*
+ * The strides of the mask will be zero exactly
+ * where they're zero for the main data
+ */
+ if (strides[iop_maskna] == 0) {
+ strides[iop] = 0;
+ }
+ else {
+ int i;
+
+ if (op_axes == NULL || op_axes[iop] == NULL) {
+ i = PyArray_NDIM(op_cur) - idim - 1;
+ }
+ else {
+ i = op_axes[iop][ndim-idim-1];
+ }
+
+ strides[iop] = PyArray_MASKNA_STRIDES(op_cur)[i];
+ }
+ }
+
+ /* Initialize the mask data pointers */
+ memcpy(NAD_PTRS(axisdata) + first_maskna_op,
+ op_dataptr + first_maskna_op,
+ NPY_SIZEOF_INTP*(nop - first_maskna_op));
+
+ NIT_ADVANCE_AXISDATA(axisdata, 1);
+ }
+
+ return 1;
+}
+
+/*
* The __array_priority__ attribute of the inputs determines
* the subtype of any output arrays. This function finds the
* subtype of the input array with highest priority.
diff --git a/numpy/core/src/multiarray/nditer_pywrap.c b/numpy/core/src/multiarray/nditer_pywrap.c
index 5ece85eca..0d4da1436 100644
--- a/numpy/core/src/multiarray/nditer_pywrap.c
+++ b/numpy/core/src/multiarray/nditer_pywrap.c
@@ -428,8 +428,17 @@ NpyIter_OpFlagsConverter(PyObject *op_flags_in,
}
break;
case 'u':
- if (strcmp(str, "updateifcopy") == 0) {
- flag = NPY_ITER_UPDATEIFCOPY;
+ switch (str[1]) {
+ case 'p':
+ if (strcmp(str, "updateifcopy") == 0) {
+ flag = NPY_ITER_UPDATEIFCOPY;
+ }
+ break;
+ case 's':
+ if (strcmp(str, "use_maskna") == 0) {
+ flag = NPY_ITER_USE_MASKNA;
+ }
+ break;
}
break;
case 'v':
@@ -659,9 +668,9 @@ npyiter_convert_op_axes(PyObject *op_axes_in, npy_intp nop,
}
/*
- * Converts the operand array and op_flags array into the form NpyIter_AdvancedNew
- * needs. Sets nop, and on success, each op[i] owns a reference
- * to an array object.
+ * Converts the operand array and op_flags array into the form
+ * NpyIter_AdvancedNew needs. Sets nop, and on success, each
+ * op[i] owns a reference to an array object.
*/
static int
npyiter_convert_ops(PyObject *op_in, PyObject *op_flags_in,
@@ -1977,7 +1986,11 @@ static PyObject *npyiter_nop_get(NewNpyArrayIterObject *self)
return NULL;
}
- return PyInt_FromLong(NpyIter_GetNOp(self->iter));
+ /*
+ * We only expose the provided operands, which is everything
+ * before the first MASKNA operand.
+ */
+ return PyInt_FromLong(NpyIter_GetFirstMaskNAOp(self->iter));
}
static PyObject *npyiter_itersize_get(NewNpyArrayIterObject *self)
@@ -2008,7 +2021,11 @@ npyiter_seq_length(NewNpyArrayIterObject *self)
return 0;
}
else {
- return NpyIter_GetNOp(self->iter);
+ /*
+ * We only expose the provided operands, which is everything
+ * before the first MASKNA operand.
+ */
+ return NpyIter_GetFirstMaskNAOp(self->iter);
}
}
@@ -2017,10 +2034,12 @@ npyiter_seq_item(NewNpyArrayIterObject *self, Py_ssize_t i)
{
PyArrayObject *ret;
+ npy_int8 *maskna_indices;
npy_intp ret_ndim;
npy_intp nop, innerloopsize, innerstride;
char *dataptr;
PyArray_Descr *dtype;
+ int has_external_loop;
if (self->iter == NULL || self->finished) {
PyErr_SetString(PyExc_ValueError,
@@ -2035,7 +2054,11 @@ npyiter_seq_item(NewNpyArrayIterObject *self, Py_ssize_t i)
return NULL;
}
- nop = NpyIter_GetNOp(self->iter);
+ /*
+ * We only expose the provided operands, which is everything
+ * before the first MASKNA operand.
+ */
+ nop = NpyIter_GetFirstMaskNAOp(self->iter);
if (i < 0 || i >= nop) {
PyErr_Format(PyExc_IndexError,
"Iterator operand index %d is out of bounds", (int)i);
@@ -2059,8 +2082,10 @@ npyiter_seq_item(NewNpyArrayIterObject *self, Py_ssize_t i)
dataptr = self->dataptrs[i];
dtype = self->dtypes[i];
+ has_external_loop = NpyIter_HasExternalLoop(self->iter);
+ maskna_indices = NpyIter_GetMaskNAIndices(self->iter);
- if (NpyIter_HasExternalLoop(self->iter)) {
+ if (has_external_loop) {
innerloopsize = *self->innerloopsizeptr;
innerstride = self->innerstrides[i];
ret_ndim = 1;
@@ -2085,6 +2110,23 @@ npyiter_seq_item(NewNpyArrayIterObject *self, Py_ssize_t i)
PyArray_UpdateFlags(ret, NPY_ARRAY_UPDATE_ALL);
+ /* If this is a USE_MASKNA operand, include the mask */
+ if (maskna_indices[i] >= 0) {
+ PyArrayObject *obj = NpyIter_GetOperandArray(self->iter)[i];
+ PyArrayObject_fieldaccess *fret = (PyArrayObject_fieldaccess *)ret;
+ int i_maskna = maskna_indices[i];
+
+ fret->maskna_dtype = PyArray_MASKNA_DTYPE(obj);
+ Py_INCREF(fret->maskna_dtype);
+ fret->maskna_data = self->dataptrs[i_maskna];
+ if (has_external_loop) {
+ fret->maskna_strides[0] = self->innerstrides[i_maskna];
+ }
+
+ fret->flags |= NPY_ARRAY_MASKNA;
+ fret->flags &= ~NPY_ARRAY_OWNMASKNA;
+ }
+
return (PyObject *)ret;
}
@@ -2109,7 +2151,11 @@ npyiter_seq_slice(NewNpyArrayIterObject *self,
return NULL;
}
- nop = NpyIter_GetNOp(self->iter);
+ /*
+ * We only expose the provided operands, which is everything
+ * before the first MASKNA operand.
+ */
+ nop = NpyIter_GetFirstMaskNAOp(self->iter);
if (ilow < 0) {
ilow = 0;
}
@@ -2143,10 +2189,11 @@ npyiter_seq_ass_item(NewNpyArrayIterObject *self, Py_ssize_t i, PyObject *v)
{
npy_intp nop, innerloopsize, innerstride;
+ npy_int8 *maskna_indices;
char *dataptr;
PyArray_Descr *dtype;
PyArrayObject *tmp;
- int ret;
+ int ret, has_external_loop;
if (v == NULL) {
PyErr_SetString(PyExc_ValueError,
@@ -2167,7 +2214,11 @@ npyiter_seq_ass_item(NewNpyArrayIterObject *self, Py_ssize_t i, PyObject *v)
return -1;
}
- nop = NpyIter_GetNOp(self->iter);
+ /*
+ * We only expose the provided operands, which is everything
+ * before the first MASKNA operand.
+ */
+ nop = NpyIter_GetFirstMaskNAOp(self->iter);
if (i < 0 || i >= nop) {
PyErr_Format(PyExc_IndexError,
"Iterator operand index %d is out of bounds", (int)i);
@@ -2181,8 +2232,9 @@ npyiter_seq_ass_item(NewNpyArrayIterObject *self, Py_ssize_t i, PyObject *v)
dataptr = self->dataptrs[i];
dtype = self->dtypes[i];
+ has_external_loop = NpyIter_HasExternalLoop(self->iter);
- if (NpyIter_HasExternalLoop(self->iter)) {
+ if (has_external_loop) {
innerloopsize = *self->innerloopsizeptr;
innerstride = self->innerstrides[i];
}
@@ -2191,6 +2243,8 @@ npyiter_seq_ass_item(NewNpyArrayIterObject *self, Py_ssize_t i, PyObject *v)
innerstride = 0;
}
+ maskna_indices = NpyIter_GetMaskNAIndices(self->iter);
+
/* TODO - there should be a better way than this... */
Py_INCREF(dtype);
tmp = (PyArrayObject *)PyArray_NewFromDescr(&PyArray_Type, dtype,
@@ -2200,6 +2254,25 @@ npyiter_seq_ass_item(NewNpyArrayIterObject *self, Py_ssize_t i, PyObject *v)
if (tmp == NULL) {
return -1;
}
+ /* If this is a USE_MASKNA operand, include the mask */
+ if (maskna_indices[i] >= 0) {
+ PyArrayObject *obj = NpyIter_GetOperandArray(self->iter)[i];
+ PyArrayObject_fieldaccess *ftmp = (PyArrayObject_fieldaccess *)tmp;
+ int i_maskna = maskna_indices[i];
+
+ ftmp->maskna_dtype = PyArray_MASKNA_DTYPE(obj);
+ Py_INCREF(ftmp->maskna_dtype);
+ ftmp->maskna_data = self->dataptrs[i_maskna];
+ if (has_external_loop) {
+ ftmp->maskna_strides[0] = self->innerstrides[i_maskna];
+ }
+ else {
+ ftmp->maskna_strides[0] = 0;
+ }
+
+ ftmp->flags |= NPY_ARRAY_MASKNA;
+ ftmp->flags &= ~NPY_ARRAY_OWNMASKNA;
+ }
PyArray_UpdateFlags(tmp, NPY_ARRAY_UPDATE_ALL);
ret = PyArray_CopyObject(tmp, v);
Py_DECREF(tmp);
@@ -2232,7 +2305,11 @@ npyiter_seq_ass_slice(NewNpyArrayIterObject *self, Py_ssize_t ilow,
return -1;
}
- nop = NpyIter_GetNOp(self->iter);
+ /*
+ * We only expose the provided operands, which is everything
+ * before the first MASKNA operand.
+ */
+ nop = NpyIter_GetFirstMaskNAOp(self->iter);
if (ilow < 0) {
ilow = 0;
}
View Lorry Controller Status