Add neighborhood iterator.

6 files changed, 305 insertions, 0 deletions

diff --git a/numpy/core/code_generators/generate_numpy_api.py b/numpy/core/code_generators/generate_numpy_api.py
index 4d639c3e5..33e3aeaa4 100644
--- a/numpy/core/code_generators/generate_numpy_api.py
+++ b/numpy/core/code_generators/generate_numpy_api.py
@@ -28,6 +28,7 @@ extern NPY_NO_EXPORT PyTypeObject PyArrayFlags_Type;
 extern NPY_NO_EXPORT PyTypeObject PyArrayIter_Type;
 extern NPY_NO_EXPORT PyTypeObject PyArrayMapIter_Type;
 extern NPY_NO_EXPORT PyTypeObject PyArrayMultiIter_Type;
+extern NPY_NO_EXPORT PyTypeObject PyArrayNeighborhoodIter_Type;
 extern NPY_NO_EXPORT PyTypeObject PyBoolArrType_Type;
 extern NPY_NO_EXPORT PyBoolScalarObject _PyArrayScalar_BoolValues[2];
 #else
@@ -39,6 +40,7 @@ NPY_NO_EXPORT PyTypeObject PyArrayFlags_Type;
 NPY_NO_EXPORT PyTypeObject PyArrayIter_Type;
 NPY_NO_EXPORT PyTypeObject PyArrayMapIter_Type;
 NPY_NO_EXPORT PyTypeObject PyArrayMultiIter_Type;
+NPY_NO_EXPORT PyTypeObject PyArrayNeighborhoodIter_Type;
 NPY_NO_EXPORT PyTypeObject PyBoolArrType_Type;
 NPY_NO_EXPORT PyBoolScalarObject _PyArrayScalar_BoolValues[2];
 #endif
diff --git a/numpy/core/code_generators/numpy_api_order.txt b/numpy/core/code_generators/numpy_api_order.txt
index e111f0fd3..3d1e40add 100644
--- a/numpy/core/code_generators/numpy_api_order.txt
+++ b/numpy/core/code_generators/numpy_api_order.txt
@@ -174,3 +174,4 @@ PyArray_MultiIterFromObjects
 PyArray_GetEndianness
 PyArray_GetNDArrayCFeatureVersion
 PyArray_Acorrelate
+PyArray_NeighborhoodIterNew
diff --git a/numpy/core/include/numpy/_neighborhood_iterator_imp.h b/numpy/core/include/numpy/_neighborhood_iterator_imp.h
new file mode 100644
index 000000000..b110f1c15
--- /dev/null
+++ b/numpy/core/include/numpy/_neighborhood_iterator_imp.h
@@ -0,0 +1,133 @@
+#ifndef _NPY_INCLUDE_NEIGHBORHOOD_IMP
+#error You should not include this header directly
+#endif
+/*
+ * Private API (here for inline)
+ */
+static NPY_INLINE int
+_PyArrayNeighborhoodIter_IncrCoord(PyArrayNeighborhoodIterObject* iter);
+static NPY_INLINE int
+_PyArrayNeighborhoodIter_SetPtr(PyArrayNeighborhoodIterObject* iter);
+
+/*
+ * Inline implementations
+ */
+static NPY_INLINE int PyArrayNeighborhoodIter_Reset(PyArrayNeighborhoodIterObject* iter)
+{
+    int i;
+
+    for (i = 0; i < iter->nd; ++i) {
+        iter->coordinates[i] = iter->bounds[i][0];
+    }
+    _PyArrayNeighborhoodIter_SetPtr(iter);
+
+    return 0;
+}
+
+/*
+ * Update to next item of the iterator
+ *
+ * Note: this simply increment the coordinates vector, last dimension
+ * incremented first , i.e, for dimension 3
+ * ...
+ * -1, -1, -1
+ * -1, -1,  0
+ * -1, -1,  1
+ *  ....
+ * -1,  0, -1
+ * -1,  0,  0
+ *  ....
+ * 0,  -1, -1
+ * 0,  -1,  0
+ *  ....
+ */
+#define _UPDATE_COORD_ITER(c) \
+    wb = iter->coordinates[c] < iter->bounds[c][1]; \
+    if (wb) { \
+        iter->coordinates[c] += 1; \
+        return 0; \
+    } \
+    else { \
+        iter->coordinates[c] = iter->bounds[c][0]; \
+    }
+
+static NPY_INLINE int _PyArrayNeighborhoodIter_IncrCoord(PyArrayNeighborhoodIterObject* iter)
+{
+    int i, wb;
+
+    for (i = iter->nd - 1; i >= 0; --i) {
+        _UPDATE_COORD_ITER(i)
+    }
+
+    return 0;
+}
+
+/*
+ * Version optimized for 2d arrays, manual loop unrolling
+ */
+static NPY_INLINE int _PyArrayNeighborhoodIter_IncrCoord2D(PyArrayNeighborhoodIterObject* iter)
+{
+    int wb;
+
+    _UPDATE_COORD_ITER(1)
+    _UPDATE_COORD_ITER(0)
+
+    return 0;
+}
+#undef _UPDATE_COORD_ITER
+
+#define _INF_SET_PTR(c) \
+    bd = iter->coordinates[c] + iter->_internal_iter->coordinates[c]; \
+    if (bd < 0 || bd > iter->dimensions[c]) { \
+        iter->dataptr = iter->zero; \
+        return 1; \
+    } \
+    offset = iter->coordinates[c] * iter->strides[c]; \
+    iter->dataptr += offset;
+
+/* set the dataptr from its current coordinates */
+static NPY_INLINE int _PyArrayNeighborhoodIter_SetPtr(PyArrayNeighborhoodIterObject* iter)
+{
+    int i;
+    npy_intp offset, bd;
+
+    iter->dataptr = iter->_internal_iter->dataptr;
+
+    for(i = 0; i < iter->nd; ++i) {
+        _INF_SET_PTR(i)
+    }
+
+    return 0;
+}
+
+static NPY_INLINE int _PyArrayNeighborhoodIter_SetPtr2D(PyArrayNeighborhoodIterObject* iter)
+{
+    npy_intp offset, bd;
+
+    iter->dataptr = iter->_internal_iter->dataptr;
+
+    _INF_SET_PTR(0)
+    _INF_SET_PTR(1)
+
+    return 0;
+}
+#undef _INF_SET_PTR
+
+/*
+ * Advance to the next neighbour
+ */
+static NPY_INLINE int PyArrayNeighborhoodIter_Next2D(PyArrayNeighborhoodIterObject* iter)
+{
+    _PyArrayNeighborhoodIter_IncrCoord2D(iter);
+    _PyArrayNeighborhoodIter_SetPtr2D(iter);
+
+    return 0;
+}
+
+static NPY_INLINE int PyArrayNeighborhoodIter_Next(PyArrayNeighborhoodIterObject* iter)
+{
+    _PyArrayNeighborhoodIter_IncrCoord (iter);
+    _PyArrayNeighborhoodIter_SetPtr(iter);
+
+    return 0;
+}
diff --git a/numpy/core/include/numpy/ndarrayobject.h b/numpy/core/include/numpy/ndarrayobject.h
index 0826f59a6..57c4daab5 100644
--- a/numpy/core/include/numpy/ndarrayobject.h
+++ b/numpy/core/include/numpy/ndarrayobject.h
@@ -909,6 +909,57 @@ typedef struct {
 
 } PyArrayMapIterObject;
 
+typedef struct {
+    PyObject_HEAD
+
+    /* 
+     * PyArrayIterObject part: keep this in this exact order 
+     */
+    int               nd_m1;            /* number of dimensions - 1 */
+    npy_intp          index, size;
+    npy_intp          coordinates[NPY_MAXDIMS];/* N-dimensional loop */
+    npy_intp          dims_m1[NPY_MAXDIMS];    /* ao->dimensions - 1 */
+    npy_intp          strides[NPY_MAXDIMS];    /* ao->strides or fake */
+    npy_intp          backstrides[NPY_MAXDIMS];/* how far to jump back */
+    npy_intp          factors[NPY_MAXDIMS];     /* shape factors */
+    PyArrayObject     *ao;
+    char              *dataptr;        /* pointer to current item*/
+    npy_bool          contiguous;
+
+    /* 
+     * New members
+     */
+    npy_intp nd;
+
+    /* Dimensions is the dimension of the array */
+    npy_intp dimensions[NPY_MAXDIMS];
+    /* Bounds of the neighborhood to iterate over */
+    npy_intp bounds[NPY_MAXDIMS][2];
+
+    /* Neighborhood points coordinates are computed relatively to the point pointed
+     * by _internal_iter */
+    PyArrayIterObject* _internal_iter;
+    /* To keep a reference to the zero representation correponding to the dtype
+     * of the array we iterate over */
+    char* zero;
+} PyArrayNeighborhoodIterObject;
+
+/*
+ * Neighborhood iterator API
+ */
+static NPY_INLINE int
+PyArrayNeighborhoodIter_Reset(PyArrayNeighborhoodIterObject* iter);
+static NPY_INLINE int
+PyArrayNeighborhoodIter_Next(PyArrayNeighborhoodIterObject* iter);
+static NPY_INLINE int
+PyArrayNeighborhoodIter_Next2D(PyArrayNeighborhoodIterObject* iter);
+
+/* Include inline implementations - functions defined there are not considered
+ * public API */
+#define _NPY_INCLUDE_NEIGHBORHOOD_IMP
+#include "_neighborhood_iterator_imp.h"
+#undef _NPY_INCLUDE_NEIGHBORHOOD_IMP
+
 /* The default array type
  */
 #define NPY_DEFAULT_TYPE NPY_DOUBLE
diff --git a/numpy/core/src/multiarray/iterators.c b/numpy/core/src/multiarray/iterators.c
index ca5518de5..842522d73 100644
--- a/numpy/core/src/multiarray/iterators.c
+++ b/numpy/core/src/multiarray/iterators.c
@@ -1731,3 +1731,116 @@ NPY_NO_EXPORT PyTypeObject PyArrayMultiIter_Type = {
     0,                                           /* *tp_next */
 #endif
 };
+
+/*========================= Neighborhood iterator ======================*/
+
+/*NUMPY_API*/
+NPY_NO_EXPORT PyObject*
+PyArray_NeighborhoodIterNew(PyArrayIterObject *x, intp *bounds)
+{
+    int i;
+    PyArrayNeighborhoodIterObject *ret;
+
+    ret = _pya_malloc(sizeof(*ret));
+    if (ret == NULL) {
+        return NULL;
+    }
+    PyObject_Init((PyObject *)ret, &PyArrayNeighborhoodIter_Type);
+
+    array_iter_base_init((PyArrayIterObject*)ret, x->ao);
+    Py_INCREF(x);
+    ret->_internal_iter = x;
+
+    ret->nd = x->ao->nd;
+
+    /* Compute the neighborhood size and copy the shape */
+    ret->size = 1;
+    for (i = 0; i < ret->nd; ++i) {
+        ret->bounds[i][0] = bounds[2 * i];
+        ret->bounds[i][1] = bounds[2 * i + 1];
+        ret->size *= (bounds[2*i+1] - bounds[2*i]) + 1;
+    }
+
+    for (i = 0; i < ret->nd; ++i) {
+        ret->dimensions[i] = x->ao->dimensions[i];
+    }
+    ret->zero = PyArray_Zero(x->ao);
+
+    /*
+     * XXX: we force x iterator to be non contiguous because we need
+     * coordinates... Modifying the iterator here is not great
+     */
+    x->contiguous = 0;
+
+    PyArrayNeighborhoodIter_Reset(ret);
+
+    return (PyObject*)ret;
+}
+
+static void neighiter_dealloc(PyArrayNeighborhoodIterObject* iter)
+{
+    PyDataMem_FREE(iter->zero);
+    Py_DECREF(iter->_internal_iter);
+
+    array_iter_base_dealloc((PyArrayIterObject*)iter);
+    _pya_free((PyArrayObject*)iter);
+}
+
+NPY_NO_EXPORT PyTypeObject PyArrayNeighborhoodIter_Type = {
+    PyObject_HEAD_INIT(NULL)
+    0,                         /*ob_size*/
+    "numpy.neigh_internal_iter",          /*tp_name*/
+    sizeof(PyArrayNeighborhoodIterObject), /*tp_basicsize*/
+    0,                         /*tp_itemsize*/
+    (destructor)neighiter_dealloc,         /*tp_dealloc*/
+    0,                         /*tp_print*/
+    0,                         /*tp_getattr*/
+    0,                         /*tp_setattr*/
+    0,                         /*tp_compare*/
+    0,                         /*tp_repr*/
+    0,                         /*tp_as_number*/
+    0,                         /*tp_as_sequence*/
+    0,                         /*tp_as_mapping*/
+    0,                         /*tp_hash */
+    0,                         /*tp_call*/
+    0,                         /*tp_str*/
+    0,                         /*tp_getattro*/
+    0,                         /*tp_setattro*/
+    0,                         /*tp_as_buffer*/
+    Py_TPFLAGS_DEFAULT,        /*tp_flags*/
+    0,                         /* tp_doc */
+    0,                                           /* tp_traverse */
+    0,                                           /* tp_clear */
+    0,                                           /* tp_richcompare */
+    0,                                           /* tp_weaklistoffset */
+    0,                                           /* tp_iter */
+    (iternextfunc)0,            /* tp_iternext */
+    0,                    /* tp_methods */
+    0,                    /* tp_members */
+    0,                    /* tp_getset */
+    0,                                           /* tp_base */
+    0,                                           /* tp_dict */
+    0,                                           /* tp_descr_get */
+    0,                                           /* tp_descr_set */
+    0,                                           /* tp_dictoffset */
+    (initproc)0,                                 /* tp_init */
+    0,                                           /* tp_alloc */
+    0,                           /* tp_new */
+    0,                                           /* tp_free */
+    0,                                           /* tp_is_gc */
+    0,                                           /* tp_bases */
+    0,                                           /* tp_mro */
+    0,                                           /* tp_cache */
+    0,                                           /* tp_subclasses */
+    0,                                           /* tp_weaklist */
+    0,                                           /* tp_del */
+
+#ifdef COUNT_ALLOCS
+    /* these must be last and never explicitly initialized */
+    0,                                           /* tp_allocs */
+    0,                                           /* tp_frees */
+    0,                                           /* tp_maxalloc */
+    0,                                           /* tp_prev */
+    0,                                           /* *tp_next */
+#endif
+};
diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c
index a10636f16..084a3dd02 100644
--- a/numpy/core/src/multiarray/multiarraymodule.c
+++ b/numpy/core/src/multiarray/multiarraymodule.c
@@ -2674,6 +2674,11 @@ PyMODINIT_FUNC initmultiarray(void) {
     if (PyType_Ready(&PyArrayMultiIter_Type) < 0) {
         return;
     }
+    PyArrayNeighborhoodIter_Type.tp_new = PyType_GenericNew;
+    if (PyType_Ready(&PyArrayNeighborhoodIter_Type) < 0) {
+        return;
+    }
+
     PyArrayDescr_Type.tp_hash = PyArray_DescrHash;
     if (PyType_Ready(&PyArrayDescr_Type) < 0) {
         return;