More tests for neighborhood iterator.

Mirror mode now works for object and basic types.

5 files changed, 279 insertions, 70 deletions

diff --git a/numpy/core/SConscript b/numpy/core/SConscript
index 0be3c996d..972a1f5cd 100644
--- a/numpy/core/SConscript
+++ b/numpy/core/SConscript
@@ -306,6 +306,8 @@ umathmodule_src = env.GenerateFromTemplate(pjoin('src', 'umath',
     'umathmodule.c.src'))
 umath_tests_src = env.GenerateFromTemplate(pjoin('src', 'umath',
     'umath_tests.c.src'))
+multiarray_tests_src = env.GenerateFromTemplate(pjoin('src', 'multiarray',
+    'multiarray_tests.c.src'))
 scalarmathmodule_src = env.GenerateFromTemplate(
                             pjoin('src', 'scalarmathmodule.c.src'))
 
@@ -366,6 +368,7 @@ if ENABLE_SEPARATE_COMPILATION:
 else:
     multiarray_src = [pjoin('src', 'multiarray', 'multiarraymodule_onefile.c')]
 multiarray = env.DistutilsPythonExtension('multiarray', source = multiarray_src)
+env.DistutilsPythonExtension('multiarray_tests', source=multiarray_tests_src)
 
 #------------------
 # Build sort module
diff --git a/numpy/core/include/numpy/_neighborhood_iterator_imp.h b/numpy/core/include/numpy/_neighborhood_iterator_imp.h
index 6773a9a40..9b9a70fca 100644
--- a/numpy/core/include/numpy/_neighborhood_iterator_imp.h
+++ b/numpy/core/include/numpy/_neighborhood_iterator_imp.h
@@ -7,40 +7,11 @@
 static NPY_INLINE int
 _PyArrayNeighborhoodIter_IncrCoord(PyArrayNeighborhoodIterObject* iter);
 static NPY_INLINE int
-_PyArrayNeighborhoodIter_SetPtr(PyArrayNeighborhoodIterObject* iter);
+_PyArrayNeighborhoodIter_SetPtrConstant(PyArrayNeighborhoodIterObject* iter);
 static NPY_INLINE int
 _PyArrayNeighborhoodIter_SetPtrMirror(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;
-}
-
-static NPY_INLINE int PyArrayNeighborhoodIter_ResetMirror(
-                PyArrayNeighborhoodIterObject* iter)
-{
-    int i;
-
-    assert(iter->mode == NPY_NEIGHBORHOOD_ITER_MIRROR_PADDING);
-    for (i = 0; i < iter->nd; ++i) {
-        iter->coordinates[i] = iter->bounds[i][0];
-    }
-    _PyArrayNeighborhoodIter_SetPtrMirror(iter);
-
-    return 0;
-}
-
-/*
  * Update to next item of the iterator
  *
  * Note: this simply increment the coordinates vector, last dimension
@@ -102,11 +73,16 @@ static NPY_INLINE int _PyArrayNeighborhoodIter_IncrCoord2D(PyArrayNeighborhoodIt
     iter->dataptr += offset;
 
 /* set the dataptr from its current coordinates */
-static NPY_INLINE int _PyArrayNeighborhoodIter_SetPtr(PyArrayNeighborhoodIterObject* iter)
+static NPY_INLINE int
+_PyArrayNeighborhoodIter_SetPtrConstant(PyArrayNeighborhoodIterObject* iter)
 {
     int i;
     npy_intp offset, bd;
 
+    assert((iter->mode == NPY_NEIGHBORHOOD_ITER_ONE_PADDING) 
+           | (iter->mode == NPY_NEIGHBORHOOD_ITER_ZERO_PADDING)
+           | (iter->mode == NPY_NEIGHBORHOOD_ITER_CONSTANT_PADDING));
+
     iter->dataptr = iter->_internal_iter->dataptr;
 
     for(i = 0; i < iter->nd; ++i) {
@@ -116,7 +92,8 @@ static NPY_INLINE int _PyArrayNeighborhoodIter_SetPtr(PyArrayNeighborhoodIterObj
     return 0;
 }
 
-static NPY_INLINE int _PyArrayNeighborhoodIter_SetPtr2D(PyArrayNeighborhoodIterObject* iter)
+static NPY_INLINE int
+_PyArrayNeighborhoodIter_SetPtrConstant2D(PyArrayNeighborhoodIterObject* iter)
 {
     npy_intp offset, bd;
 
@@ -168,8 +145,8 @@ static NPY_INLINE npy_intp _npy_pos_remainder(npy_intp i, npy_intp n)
     iter->dataptr += offset;
 
 /* set the dataptr from its current coordinates */
-static NPY_INLINE int _PyArrayNeighborhoodIter_SetPtrMirror(
-                PyArrayNeighborhoodIterObject* iter)
+static NPY_INLINE int
+_PyArrayNeighborhoodIter_SetPtrMirror(PyArrayNeighborhoodIterObject* iter)
 {
     int i;
     npy_intp offset, bd, truepos;
@@ -187,18 +164,29 @@ static NPY_INLINE int _PyArrayNeighborhoodIter_SetPtrMirror(
 /*
  * Advance to the next neighbour
  */
-static NPY_INLINE int PyArrayNeighborhoodIter_Next2D(PyArrayNeighborhoodIterObject* iter)
+static NPY_INLINE int
+PyArrayNeighborhoodIter_Next2D(PyArrayNeighborhoodIterObject* iter)
 {
+    assert((iter->mode == NPY_NEIGHBORHOOD_ITER_ONE_PADDING) 
+           | (iter->mode == NPY_NEIGHBORHOOD_ITER_ZERO_PADDING)
+           | (iter->mode == NPY_NEIGHBORHOOD_ITER_CONSTANT_PADDING));
+    assert(iter->nd == 2);
+
     _PyArrayNeighborhoodIter_IncrCoord2D(iter);
-    _PyArrayNeighborhoodIter_SetPtr2D(iter);
+    _PyArrayNeighborhoodIter_SetPtrConstant2D(iter);
 
     return 0;
 }
 
-static NPY_INLINE int PyArrayNeighborhoodIter_Next(PyArrayNeighborhoodIterObject* iter)
+static NPY_INLINE int
+PyArrayNeighborhoodIter_NextConstant(PyArrayNeighborhoodIterObject* iter)
 {
-    _PyArrayNeighborhoodIter_IncrCoord (iter);
-    _PyArrayNeighborhoodIter_SetPtr(iter);
+    assert((iter->mode == NPY_NEIGHBORHOOD_ITER_ONE_PADDING) 
+           | (iter->mode == NPY_NEIGHBORHOOD_ITER_ZERO_PADDING)
+           | (iter->mode == NPY_NEIGHBORHOOD_ITER_CONSTANT_PADDING));
+
+    _PyArrayNeighborhoodIter_IncrCoord(iter);
+    _PyArrayNeighborhoodIter_SetPtrConstant(iter);
 
     return 0;
 }
@@ -206,8 +194,84 @@ static NPY_INLINE int PyArrayNeighborhoodIter_Next(PyArrayNeighborhoodIterObject
 static NPY_INLINE
 int PyArrayNeighborhoodIter_NextMirror(PyArrayNeighborhoodIterObject* iter)
 {
+    assert(iter->mode == NPY_NEIGHBORHOOD_ITER_MIRROR_PADDING);
+
     _PyArrayNeighborhoodIter_IncrCoord(iter);
     _PyArrayNeighborhoodIter_SetPtrMirror(iter);
 
     return 0;
 }
+
+static NPY_INLINE int PyArrayNeighborhoodIter_Next(PyArrayNeighborhoodIterObject* iter)
+{
+    _PyArrayNeighborhoodIter_IncrCoord (iter);
+    switch (iter->mode) {
+        case NPY_NEIGHBORHOOD_ITER_ZERO_PADDING:
+        case NPY_NEIGHBORHOOD_ITER_ONE_PADDING:
+        case NPY_NEIGHBORHOOD_ITER_CONSTANT_PADDING:
+            _PyArrayNeighborhoodIter_SetPtrConstant(iter);
+            break;
+        case NPY_NEIGHBORHOOD_ITER_MIRROR_PADDING:
+            _PyArrayNeighborhoodIter_SetPtrMirror(iter);
+            break;
+    }
+
+    return 0;
+}
+
+/*
+ * Reset functions
+ */
+static NPY_INLINE int
+PyArrayNeighborhoodIter_ResetConstant(PyArrayNeighborhoodIterObject* iter)
+{
+    int i;
+
+    assert((iter->mode == NPY_NEIGHBORHOOD_ITER_ONE_PADDING) 
+           | (iter->mode == NPY_NEIGHBORHOOD_ITER_ZERO_PADDING)
+           | (iter->mode == NPY_NEIGHBORHOOD_ITER_CONSTANT_PADDING));
+
+    for (i = 0; i < iter->nd; ++i) {
+        iter->coordinates[i] = iter->bounds[i][0];
+    }
+    _PyArrayNeighborhoodIter_SetPtrConstant(iter);
+
+    return 0;
+}
+
+static NPY_INLINE int
+PyArrayNeighborhoodIter_ResetMirror(PyArrayNeighborhoodIterObject* iter)
+{
+    int i;
+
+    assert(iter->mode == NPY_NEIGHBORHOOD_ITER_MIRROR_PADDING);
+
+    for (i = 0; i < iter->nd; ++i) {
+        iter->coordinates[i] = iter->bounds[i][0];
+    }
+    _PyArrayNeighborhoodIter_SetPtrMirror(iter);
+
+    return 0;
+}
+
+static NPY_INLINE int
+PyArrayNeighborhoodIter_Reset(PyArrayNeighborhoodIterObject* iter)
+{
+    int i;
+
+    for (i = 0; i < iter->nd; ++i) {
+        iter->coordinates[i] = iter->bounds[i][0];
+    }
+    switch (iter->mode) {
+        case NPY_NEIGHBORHOOD_ITER_ZERO_PADDING:
+        case NPY_NEIGHBORHOOD_ITER_ONE_PADDING:
+        case NPY_NEIGHBORHOOD_ITER_CONSTANT_PADDING:
+            _PyArrayNeighborhoodIter_SetPtrConstant(iter);
+            break;
+        case NPY_NEIGHBORHOOD_ITER_MIRROR_PADDING:
+            _PyArrayNeighborhoodIter_SetPtrMirror(iter);
+            break;
+    }
+
+    return 0;
+}
diff --git a/numpy/core/include/numpy/ndarrayobject.h b/numpy/core/include/numpy/ndarrayobject.h
index f151c2f77..2a4e21495 100644
--- a/numpy/core/include/numpy/ndarrayobject.h
+++ b/numpy/core/include/numpy/ndarrayobject.h
@@ -962,6 +962,8 @@ typedef struct {
 /*
  * Neighborhood iterator API
  */
+
+/* General: those work for any mode */
 static NPY_INLINE int
 PyArrayNeighborhoodIter_Reset(PyArrayNeighborhoodIterObject* iter);
 static NPY_INLINE int
@@ -969,6 +971,13 @@ PyArrayNeighborhoodIter_Next(PyArrayNeighborhoodIterObject* iter);
 static NPY_INLINE int
 PyArrayNeighborhoodIter_Next2D(PyArrayNeighborhoodIterObject* iter);
 
+/* Mode specific: those are faster, but have undefined behavior if the mode
+ * does not match. Sanity checks are enabled in debug mode. */
+static NPY_INLINE int
+PyArrayNeighborhoodIter_ResetConstant(PyArrayNeighborhoodIterObject* iter);
+static NPY_INLINE int
+PyArrayNeighborhoodIter_NextConstant(PyArrayNeighborhoodIterObject* iter);
+
 static NPY_INLINE int
 PyArrayNeighborhoodIter_ResetMirror(PyArrayNeighborhoodIterObject* iter);
 static NPY_INLINE int
diff --git a/numpy/core/src/multiarray/multiarray_tests.c b/numpy/core/src/multiarray/multiarray_tests.c.src
index fb4bd998c..1f225d5ec 100644
--- a/numpy/core/src/multiarray/multiarray_tests.c
+++ b/numpy/core/src/multiarray/multiarray_tests.c.src
@@ -1,7 +1,4 @@
-/*
- *****************************************************************************
- **                            INCLUDES                                     **
- *****************************************************************************
+/* vim:ft=c
  */
 #include <Python.h>
 
@@ -10,36 +7,39 @@
 /*
  * TODO:
  *  - Handle mode
- *  - Handle any dtype
- *  - Handle Object type
  */
-static int copy_double(PyArrayIterObject *itx, PyArrayNeighborhoodIterObject *niterx,
+
+/**begin repeat
+ * #type = double, int#
+ * #typenum = NPY_DOUBLE, NPY_INT#
+ */
+static int copy_@type@(PyArrayIterObject *itx, PyArrayNeighborhoodIterObject *niterx,
         npy_intp *bounds,
         PyObject **out)
 {
     npy_intp i, j;
-    double *ptr;
+    @type@ *ptr;
     npy_intp odims[NPY_MAXDIMS];
     PyArrayObject *aout;
 
     /* 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) {
-        PyArrayNeighborhoodIter_ResetMirror(niterx);
+        PyArrayNeighborhoodIter_Reset(niterx);
 
         for(j = 0; j < itx->ao->nd; ++j) {
             odims[0] = bounds[2 * j + 1] - bounds[2 * j] + 1;
         }
-        aout = (PyArrayObject*)PyArray_SimpleNew(itx->ao->nd, odims, NPY_DOUBLE);
+        aout = (PyArrayObject*)PyArray_SimpleNew(itx->ao->nd, odims, @typenum@);
         if (aout == NULL) {
             return -1;
         }
 
-        ptr = (double*)aout->data;
+        ptr = (@type@*)aout->data;
 
         for(j = 0; j < niterx->size; ++j) {
-            *ptr = *((double*)niterx->dataptr);
-            PyArrayNeighborhoodIter_NextMirror(niterx);
+            *ptr = *((@type@*)niterx->dataptr);
+            PyArrayNeighborhoodIter_Next(niterx);
             ptr += 1;
         }
 
@@ -51,19 +51,56 @@ static int copy_double(PyArrayIterObject *itx, PyArrayNeighborhoodIterObject *ni
 
     return 0;
 }
+/**end repeat**/
+
+static int copy_object(PyArrayIterObject *itx, PyArrayNeighborhoodIterObject *niterx,
+        npy_intp *bounds,
+        PyObject **out)
+{
+    npy_intp i, j;
+    npy_intp odims[NPY_MAXDIMS];
+    PyArrayObject *aout;
+    PyArray_CopySwapFunc *copyswap = itx->ao->descr->f->copyswap;
+    npy_int itemsize = PyArray_ITEMSIZE(itx->ao);
+
+    /* 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) {
+        PyArrayNeighborhoodIter_Reset(niterx);
+
+        for(j = 0; j < itx->ao->nd; ++j) {
+            odims[0] = bounds[2 * j + 1] - bounds[2 * j] + 1;
+        }
+        aout = (PyArrayObject*)PyArray_SimpleNew(itx->ao->nd, odims, NPY_OBJECT);
+        if (aout == NULL) {
+            return -1;
+        }
+
+        for(j = 0; j < niterx->size; ++j) {
+            copyswap(aout->data + j * itemsize, niterx->dataptr, 0, NULL);
+            PyArrayNeighborhoodIter_Next(niterx);
+        }
 
+        Py_INCREF(aout);
+        PyList_Append(*out, (PyObject*)aout);
+        Py_DECREF(aout);
+        PyArray_ITER_NEXT(itx);
+    }
+
+    return 0;
+}
 static PyObject*
 test_neighborhood_iterator(PyObject* NPY_UNUSED(self), PyObject* args)
 {
-    PyObject *x, *c, *out, *b;
-    PyArrayObject *ax;
+    PyObject *x, *fill, *out, *b;
+    PyArrayObject *ax, *afill;
     PyArrayIterObject *itx;
-    int i, typenum;
+    int i, typenum, imode;
     npy_intp bounds[NPY_MAXDIMS*2];
     PyArrayNeighborhoodIterObject *niterx;
     PyArrayNeighborhoodIterMode mode;
 
-    if (!PyArg_ParseTuple(args, "OOO", &x, &b, &c)) {
+    if (!PyArg_ParseTuple(args, "OOOi", &x, &b, &fill, &imode)) {
         return NULL;
     }
 
@@ -72,68 +109,96 @@ test_neighborhood_iterator(PyObject* NPY_UNUSED(self), PyObject* args)
     }
 
     typenum = PyArray_ObjectType(x, 0);
-    typenum = PyArray_ObjectType(c, typenum);
+    typenum = PyArray_ObjectType(fill, typenum);
 
     ax = (PyArrayObject*)PyArray_FromObject(x, typenum, 1, 10);
     if (ax == NULL) {
-        printf("Bleh\n");
         return NULL;
     }
     if (PySequence_Size(b) != 2 * ax->nd) {
         PyErr_SetString(PyExc_ValueError,
                 "bounds sequence size not compatible with x input");
-        return NULL;
+        goto clean_ax;
     }
 
     out = PyList_New(0);
     if (out == NULL) {
-        printf("Bleh\n");
-        return NULL;
+        goto clean_ax;
     }
 
     itx = (PyArrayIterObject*)PyArray_IterNew(x);
     if (itx == NULL) {
-        printf("bleh\n");
-        return NULL;
+        goto clean_out;
     }
 
     /* Compute boundaries for the neighborhood iterator */
     for(i = 0; i < 2 * ax->nd; ++i) {
         PyObject* bound;
         bound = PySequence_GetItem(b, i);
+        if (bounds == NULL) {
+            goto clean_itx;
+        }
         if (!PyInt_Check(bound)) {
             PyErr_SetString(PyExc_ValueError, "bound not long");
-            return NULL;
+            Py_DECREF(bound);
+            goto clean_itx;
         }
         bounds[i] = PyInt_AsLong(bound);
         Py_DECREF(bound);
     }
 
     /* Create the neighborhood iterator */
-    mode.mode = NPY_NEIGHBORHOOD_ITER_MIRROR_PADDING;
+    mode.mode = imode;
     mode.constant = NULL;
+    afill = NULL;
+    if (imode == NPY_NEIGHBORHOOD_ITER_CONSTANT_PADDING) {
+	    afill = (PyArrayObject *)PyArray_FromObject(fill, typenum, 0, 0);
+	    if (afill == NULL) {
+            goto clean_itx;
+        }
+        mode.constant = (PyObject*)afill;
+    }
+
     niterx = (PyArrayNeighborhoodIterObject*)PyArray_NeighborhoodIterNew(
                     (PyArrayIterObject*)itx, bounds, &mode);
     if (niterx == NULL) {
-        printf("bleh\n");
-        return NULL;
+        goto clean_afill;
     }
 
     switch (typenum) {
+        case NPY_OBJECT:
+            copy_object(itx, niterx, bounds, &out);
+            break;
+        case NPY_INT:
+            copy_int(itx, niterx, bounds, &out);
+            break;
         case NPY_DOUBLE:
             copy_double(itx, niterx, bounds, &out);
             break;
         default:
             PyErr_SetString(PyExc_ValueError, "Type not supported");
-            return NULL;
+            goto clean_niterx;
     }
 
-    Py_DECREF((PyArrayIterObject*)niterx);
-    Py_DECREF((PyArrayIterObject*)itx);
+    Py_DECREF(niterx);
+    Py_XDECREF(afill);
+    Py_DECREF(itx);
 
     Py_DECREF(ax);
 
     return out;
+
+clean_niterx:
+    Py_DECREF(niterx);
+clean_afill:
+    Py_XDECREF(afill);
+clean_itx:
+    Py_DECREF(itx);
+clean_out:
+    Py_DECREF(out);
+clean_ax:
+    Py_DECREF(ax);
+    return NULL;
 }
 
 static PyMethodDef Multiarray_TestsMethods[] = {
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
index 7022ef14d..f6102bef2 100644
--- a/numpy/core/tests/test_multiarray.py
+++ b/numpy/core/tests/test_multiarray.py
@@ -4,6 +4,7 @@ import os
 import numpy as np
 from numpy.testing import *
 from numpy.core import *
+from numpy.core.multiarray_tests import test_neighborhood_iterator
 
 from test_print import in_foreign_locale
 
@@ -1058,6 +1059,73 @@ class TestChoose(TestCase):
         A = np.choose(self.ind, (self.x, self.y2))
         assert_equal(A, [[2,2,3],[2,2,3]])
 
+def can_use_decimal():
+    try:
+        from decimal import Decimal
+        return True
+    except ImportError:
+        return False
+
+# TODO: test for multidimensional
+NEIGH_MODE = {'zero': 0, 'one': 1, 'constant': 2, 'circular': 3, 'mirror': 4}
+class TestNeighborhoodIter(TestCase):
+    ## Simple, 2d tests
+    #def test_simple2d(self):
+    #    # Test zero and one padding for simple data type
+    #    x = np.array([[0, 1], [2, 3]], dtype=np.float)
+    #    r = [np.array([[0, 0], [0, 0]], dtype=np.float), 
+    #         np.array([[0, 0], [0, 1]], dtype=np.float), 
+    #         np.array([[0, 0], [1, 0]], dtype=np.float), 
+    #         np.array([[0, 0], [0, 2]], dtype=np.float), 
+    #         np.array([[0, 1], [2, 3]], dtype=np.float),]
+    #    l = test_neighborhood_iterator(x, [-1, 1, -1, 1], x[0], NEIGH_MODE['zero'])
+    #    print l
+    #    #assert_array_equal(l, r)
+
+    # Simple, 1d tests
+    def _test_simple(self, dt):
+        # Test padding with constant values
+        x = np.linspace(1, 5, 5).astype(dt)
+        r = [[0, 1, 2], [1, 2, 3], [2, 3, 4], [3, 4, 5], [4, 5, 0]]
+        l = test_neighborhood_iterator(x, [-1, 1], x[0], NEIGH_MODE['zero'])
+        assert_array_equal(l, r)
+
+        r = [[1, 1, 2], [1, 2, 3], [2, 3, 4], [3, 4, 5], [4, 5, 1]]
+        l = test_neighborhood_iterator(x, [-1, 1], x[0], NEIGH_MODE['one'])
+        assert_array_equal(l, r)
+
+        r = [[x[4], 1, 2], [1, 2, 3], [2, 3, 4], [3, 4, 5], [4, 5, x[4]]]
+        l = test_neighborhood_iterator(x, [-1, 1], x[4], NEIGH_MODE['constant'])
+        assert_array_equal(l, r)
+
+    def test_simple_float(self):
+        self._test_simple(np.float)
+
+    @dec.skipif(not can_use_decimal(),
+            "Skip neighborhood iterator tests for decimal objects " \
+            "(decimal module not available")
+    def test_simple_object(self):
+        from decimal import Decimal
+        self._test_simple(Decimal)
+
+    # Test mirror modes
+    def _test_mirror(self, dt):
+        x = np.linspace(1, 5, 5).astype(dt)
+        r = np.array([[2, 1, 1, 2, 3], [1, 1, 2, 3, 4], [1, 2, 3, 4, 5],
+                [2, 3, 4, 5, 5], [3, 4, 5, 5, 4]], dtype=dt)
+        l = test_neighborhood_iterator(x, [-2, 2], x[1], NEIGH_MODE['mirror'])
+        self.failUnless([i.dtype == dt for i in l])
+        assert_array_equal(l, r)
+
+    def test_mirror(self):
+        self._test_mirror(np.float)
+
+    @dec.skipif(not can_use_decimal(),
+            "Skip neighborhood iterator tests for decimal objects " \
+            "(decimal module not available")
+    def test_mirror_object(self):
+        from decimal import Decimal
+        self._test_mirror(Decimal)
 
 if __name__ == "__main__":
     run_module_suite()