diff options
| -rw-r--r-- | numpy/core/src/multiarray/arraytypes.c.src | 12 | ||||
| -rw-r--r-- | numpy/core/src/multiarray/item_selection.c | 39 |
2 files changed, 31 insertions, 20 deletions
diff --git a/numpy/core/src/multiarray/arraytypes.c.src b/numpy/core/src/multiarray/arraytypes.c.src index 685aba542..05b843fc4 100644 --- a/numpy/core/src/multiarray/arraytypes.c.src +++ b/numpy/core/src/multiarray/arraytypes.c.src @@ -3680,11 +3680,10 @@ static void @name@_fastputmask(@type@ *in, npy_bool *mask, npy_intp ni, @type@ *vals, npy_intp nv) { - npy_intp i; - @type@ s_val; + npy_intp i, j; if (nv == 1) { - s_val = *vals; + @type@ s_val = *vals; for (i = 0; i < ni; i++) { if (mask[i]) { in[i] = s_val; @@ -3692,9 +3691,12 @@ static void } } else { - for (i = 0; i < ni; i++) { + for (i = 0, j = 0; i < ni; i++, j++) { + if (j >= nv) { + j = 0; + } if (mask[i]) { - in[i] = vals[i%nv]; + in[i] = vals[j]; } } } diff --git a/numpy/core/src/multiarray/item_selection.c b/numpy/core/src/multiarray/item_selection.c index 7016760b3..11075cb81 100644 --- a/numpy/core/src/multiarray/item_selection.c +++ b/numpy/core/src/multiarray/item_selection.c @@ -427,10 +427,11 @@ NPY_NO_EXPORT PyObject * PyArray_PutMask(PyArrayObject *self, PyObject* values0, PyObject* mask0) { PyArray_FastPutmaskFunc *func; - PyArrayObject *mask, *values; + PyArrayObject *mask, *values; PyArray_Descr *dtype; - npy_intp i, chunk, ni, max_item, nv, tmp; + npy_intp i, j, chunk, ni, max_item, nv; char *src, *dest; + npy_bool *mask_data; int copied = 0; mask = NULL; @@ -469,6 +470,7 @@ PyArray_PutMask(PyArrayObject *self, PyObject* values0, PyObject* mask0) "the same size"); goto fail; } + mask_data = PyArray_DATA(mask); dtype = PyArray_DESCR(self); Py_INCREF(dtype); values = (PyArrayObject *)PyArray_FromAny(values0, dtype, @@ -483,14 +485,20 @@ PyArray_PutMask(PyArrayObject *self, PyObject* values0, PyObject* mask0) Py_INCREF(Py_None); return Py_None; } + src = PyArray_DATA(values); + if (PyDataType_REFCHK(PyArray_DESCR(self))) { - for (i = 0; i < ni; i++) { - tmp = ((npy_bool *)(PyArray_DATA(mask)))[i]; - if (tmp) { - src = PyArray_BYTES(values) + chunk * (i % nv); - PyArray_Item_INCREF(src, PyArray_DESCR(self)); - PyArray_Item_XDECREF(dest+i*chunk, PyArray_DESCR(self)); - memmove(dest + i * chunk, src, chunk); + for (i = 0, j = 0; i < ni; i++, j++) { + if (j >= nv) { + j = 0; + } + if (mask_data[i]) { + char *src_ptr = src + j*chunk; + char *dest_ptr = dest + i*chunk; + + PyArray_Item_INCREF(src_ptr, PyArray_DESCR(self)); + PyArray_Item_XDECREF(dest_ptr, PyArray_DESCR(self)); + memmove(dest_ptr, src_ptr, chunk); } } } @@ -499,16 +507,17 @@ PyArray_PutMask(PyArrayObject *self, PyObject* values0, PyObject* mask0) NPY_BEGIN_THREADS_DESCR(PyArray_DESCR(self)); func = PyArray_DESCR(self)->f->fastputmask; if (func == NULL) { - for (i = 0; i < ni; i++) { - tmp = ((npy_bool *)(PyArray_DATA(mask)))[i]; - if (tmp) { - src = PyArray_BYTES(values) + chunk*(i % nv); - memmove(dest + i*chunk, src, chunk); + for (i = 0, j = 0; i < ni; i++, j++) { + if (j >= nv) { + j = 0; + } + if (mask_data[i]) { + memmove(dest + i*chunk, src + j*chunk, chunk); } } } else { - func(dest, PyArray_DATA(mask), ni, PyArray_DATA(values), nv); + func(dest, mask_data, ni, src, nv); } NPY_END_THREADS; } |