Implement circular mode for neighborhood array iterator.

3 files changed, 60 insertions, 0 deletions

diff --git a/numpy/core/include/numpy/_neighborhood_iterator_imp.h b/numpy/core/include/numpy/_neighborhood_iterator_imp.h
index 9b9a70fca..533208229 100644
--- a/numpy/core/include/numpy/_neighborhood_iterator_imp.h
+++ b/numpy/core/include/numpy/_neighborhood_iterator_imp.h
@@ -161,6 +161,41 @@ _PyArrayNeighborhoodIter_SetPtrMirror(PyArrayNeighborhoodIterObject* iter)
 }
 #undef _INF_SET_PTR_MIRROR
 
+/* compute l such as i = k * n + l, 0 <= l < |k| */
+static NPY_INLINE npy_intp _npy_euclidean_division(npy_intp i, npy_intp n)
+{
+        npy_intp l;
+
+        l = i % n;
+
+        if (l < 0) {
+                l += n;
+        }
+        return l;
+}
+#define _INF_SET_PTR_CIRCULAR(c) \
+    bd = iter->coordinates[c] + iter->_internal_iter->coordinates[c]; \
+    truepos = _npy_euclidean_division(bd, iter->dimensions[c]); \
+    offset = (truepos - iter->_internal_iter->coordinates[c]) * iter->strides[c]; \
+    iter->dataptr += offset;
+
+/* set the dataptr from its current coordinates */
+static NPY_INLINE int
+_PyArrayNeighborhoodIter_SetPtrCircular(PyArrayNeighborhoodIterObject* iter)
+{
+    int i;
+    npy_intp offset, bd, truepos;
+
+    iter->dataptr = iter->_internal_iter->dataptr;
+
+    for(i = 0; i < iter->nd; ++i) {
+        _INF_SET_PTR_CIRCULAR(i)
+    }
+
+    return 0;
+}
+#undef _INF_SET_PTR_CIRCULAR
+
 /*
  * Advance to the next neighbour
  */
@@ -214,6 +249,9 @@ static NPY_INLINE int PyArrayNeighborhoodIter_Next(PyArrayNeighborhoodIterObject
         case NPY_NEIGHBORHOOD_ITER_MIRROR_PADDING:
             _PyArrayNeighborhoodIter_SetPtrMirror(iter);
             break;
+        case NPY_NEIGHBORHOOD_ITER_CIRCULAR_PADDING:
+            _PyArrayNeighborhoodIter_SetPtrCircular(iter);
+            break;
     }
 
     return 0;
@@ -271,6 +309,9 @@ PyArrayNeighborhoodIter_Reset(PyArrayNeighborhoodIterObject* iter)
         case NPY_NEIGHBORHOOD_ITER_MIRROR_PADDING:
             _PyArrayNeighborhoodIter_SetPtrMirror(iter);
             break;
+        case NPY_NEIGHBORHOOD_ITER_CIRCULAR_PADDING:
+            _PyArrayNeighborhoodIter_SetPtrCircular(iter);
+            break;
     }
 
     return 0;
diff --git a/numpy/core/src/multiarray/iterators.c b/numpy/core/src/multiarray/iterators.c
index 455ce06ad..729dd3963 100644
--- a/numpy/core/src/multiarray/iterators.c
+++ b/numpy/core/src/multiarray/iterators.c
@@ -1828,6 +1828,7 @@ PyArray_NeighborhoodIterNew(PyArrayIterObject *x, intp *bounds,
                 ret->mode = mode->mode;
                 break;
             case NPY_NEIGHBORHOOD_ITER_MIRROR_PADDING:
+            case NPY_NEIGHBORHOOD_ITER_CIRCULAR_PADDING:
                 ret->mode = mode->mode;
                 ret->constant = NULL;
                 break;
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
index 0e96deab2..7ee4e0c0e 100644
--- a/numpy/core/tests/test_multiarray.py
+++ b/numpy/core/tests/test_multiarray.py
@@ -1168,5 +1168,23 @@ class TestNeighborhoodIter(TestCase):
         from decimal import Decimal
         self._test_mirror(Decimal)
 
+    # Circular mode
+    def _test_circular(self, dt):
+        x = np.linspace(1, 5, 5).astype(dt)
+        r = np.array([[4, 5, 1, 2, 3], [5, 1, 2, 3, 4], [1, 2, 3, 4, 5],
+                [2, 3, 4, 5, 1], [3, 4, 5, 1, 2]], dtype=dt)
+        l = test_neighborhood_iterator(x, [-2, 2], x[0], NEIGH_MODE['circular'])
+        assert_array_equal(l, r)
+
+    def test_circular(self):
+        self._test_circular(np.float)
+
+    @dec.skipif(not can_use_decimal(),
+            "Skip neighborhood iterator tests for decimal objects " \
+            "(decimal module not available")
+    def test_circular_object(self):
+        from decimal import Decimal
+        self._test_circular(Decimal)
+
 if __name__ == "__main__":
     run_module_suite()