Merge pull request #19016 from ReallyNiceGuy/main

BUG: Update coordinates on PyArray_ITER_GOTO1D

3 files changed, 21 insertions, 3 deletions

diff --git a/numpy/core/include/numpy/ndarraytypes.h b/numpy/core/include/numpy/ndarraytypes.h
index d1acfdf26..3ea66b049 100644
--- a/numpy/core/include/numpy/ndarraytypes.h
+++ b/numpy/core/include/numpy/ndarraytypes.h
@@ -1236,6 +1236,8 @@ struct PyArrayIterObject_tag {
                 _PyAIT(it)->dataptr = PyArray_BYTES(_PyAIT(it)->ao); \
                 for (__npy_i = 0; __npy_i<=_PyAIT(it)->nd_m1; \
                      __npy_i++) { \
+                        _PyAIT(it)->coordinates[__npy_i] = \
+                                (__npy_ind / _PyAIT(it)->factors[__npy_i]); \
                         _PyAIT(it)->dataptr += \
                                 (__npy_ind / _PyAIT(it)->factors[__npy_i]) \
                                 * _PyAIT(it)->strides[__npy_i]; \
diff --git a/numpy/core/src/multiarray/_multiarray_tests.c.src b/numpy/core/src/multiarray/_multiarray_tests.c.src
index bfdeae079..79140bdb7 100644
--- a/numpy/core/src/multiarray/_multiarray_tests.c.src
+++ b/numpy/core/src/multiarray/_multiarray_tests.c.src
@@ -87,7 +87,7 @@ static int copy_@name@(PyArrayIterObject *itx, PyArrayNeighborhoodIterObject *ni
      * For each point in itx, copy the current neighborhood into an array which
      * is appended at the output list
      */
-    for (i = 0; i < itx->size; ++i) {
+    for (i = itx->index; i < itx->size; ++i) {
         PyArrayNeighborhoodIter_Reset(niterx);
 
         for (j = 0; j < PyArray_NDIM(itx->ao); ++j) {
@@ -130,7 +130,7 @@ static int copy_object(PyArrayIterObject *itx, PyArrayNeighborhoodIterObject *ni
      * For each point in itx, copy the current neighborhood into an array which
      * is appended at the output list
      */
-    for (i = 0; i < itx->size; ++i) {
+    for (i = itx->index; i < itx->size; ++i) {
         PyArrayNeighborhoodIter_Reset(niterx);
 
         for (j = 0; j < PyArray_NDIM(itx->ao); ++j) {
@@ -161,10 +161,11 @@ test_neighborhood_iterator(PyObject* NPY_UNUSED(self), PyObject* args)
     PyArrayObject *ax, *afill;
     PyArrayIterObject *itx;
     int i, typenum, mode, st;
+    Py_ssize_t idxstart = 0;
     npy_intp bounds[NPY_MAXDIMS*2];
     PyArrayNeighborhoodIterObject *niterx;
 
-    if (!PyArg_ParseTuple(args, "OOOi", &x, &b, &fill, &mode)) {
+    if (!PyArg_ParseTuple(args, "OOOi|n", &x, &b, &fill, &mode, &idxstart)) {
         return NULL;
     }
 
@@ -224,12 +225,20 @@ test_neighborhood_iterator(PyObject* NPY_UNUSED(self), PyObject* args)
         }
     }
 
+    if (idxstart >= itx->size) {
+        PyErr_SetString(PyExc_ValueError,
+                "start index not compatible with x input");
+        goto clean_itx;
+    }
+
     niterx = (PyArrayNeighborhoodIterObject*)PyArray_NeighborhoodIterNew(
                     (PyArrayIterObject*)itx, bounds, mode, afill);
     if (niterx == NULL) {
         goto clean_afill;
     }
 
+    PyArray_ITER_GOTO1D((PyArrayIterObject*)itx, idxstart);
+
     switch (typenum) {
         case NPY_OBJECT:
             st = copy_object(itx, niterx, bounds, &out);
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
index d567653f5..25dd76256 100644
--- a/numpy/core/tests/test_multiarray.py
+++ b/numpy/core/tests/test_multiarray.py
@@ -6947,6 +6947,13 @@ class TestNeighborhoodIter:
                 x, [-1, 0, -1, 1], 4, NEIGH_MODE['constant'])
         assert_array_equal(l, r)
 
+        # Test with start in the middle
+        r = [np.array([[4, 0, 1], [4, 2, 3]], dtype=dt),
+             np.array([[0, 1, 4], [2, 3, 4]], dtype=dt)]
+        l = _multiarray_tests.test_neighborhood_iterator(
+                x, [-1, 0, -1, 1], 4, NEIGH_MODE['constant'], 2)
+        assert_array_equal(l, r)
+
     def test_mirror2d(self, dt):
         x = np.array([[0, 1], [2, 3]], dtype=dt)
         r = [np.array([[0, 0, 1], [0, 0, 1]], dtype=dt),