Simplify the ndenumerate class by exposing coords and index from flat iterator. Simple cosmetic changes to ffts.

3 files changed, 33 insertions, 24 deletions

diff --git a/numpy/core/src/arrayobject.c b/numpy/core/src/arrayobject.c
index 033eef26d..0cc35efc6 100644
--- a/numpy/core/src/arrayobject.c
+++ b/numpy/core/src/arrayobject.c
@@ -7017,9 +7017,36 @@ static PyMethodDef iter_methods[] = {
 
 static PyMemberDef iter_members[] = {
 	{"base", T_OBJECT, offsetof(PyArrayIterObject, ao), RO, NULL},
+        {"index", T_INT, offsetof(PyArrayIterObject, index), RO, NULL},
 	{NULL},
 };
 
+static PyObject *
+iter_coords_get(PyArrayIterObject *self)
+{
+        int nd;
+        nd = self->ao->nd;
+        if (self->contiguous) { /* coordinates not kept track of --- need to generate
+                                   from index */
+                intp val;
+                int i;
+                val = self->index;
+                for (i=0;i<nd; i++) {
+                        self->coordinates[i] = val / self->factors[i];
+                        val = val % self->factors[i];
+                }
+        }
+        return PyArray_IntTupleFromIntp(nd, self->coordinates);
+}
+
+static PyGetSetDef iter_getsets[] = {
+	{"coords", 
+	 (getter)iter_coords_get,
+	 NULL,
+	 "An N-d tuple of current coordinates."},
+	{NULL, NULL, NULL, NULL},
+};
+
 static PyTypeObject PyArrayIter_Type = {
         PyObject_HEAD_INIT(NULL)
         0,					 /* ob_size */
@@ -7052,7 +7079,7 @@ static PyTypeObject PyArrayIter_Type = {
         (iternextfunc)arrayiter_next,		/* tp_iternext */
         iter_methods,				/* tp_methods */
         iter_members,	             	        /* tp_members */
-        0,                                      /* tp_getset */
+        iter_getsets,                           /* tp_getset */
 
 };
 
diff --git a/numpy/dft/fftpack.py b/numpy/dft/fftpack.py
index 441aae5f5..05efc759c 100644
--- a/numpy/dft/fftpack.py
+++ b/numpy/dft/fftpack.py
@@ -108,7 +108,7 @@ def inverse_fft(a, n=None, axis=-1):
     run into memory problems if you call this too many times with too many
     different n's."""
 
-    a = asarray(a).astype(Complex)
+    a = asarray(a).astype(complex)
     if n == None:
         n = shape(a)[axis]
     return _raw_fft(a, n, axis, fftpack.cffti, fftpack.cfftb, _fft_cache) / n
@@ -131,7 +131,7 @@ def real_fft(a, n=None, axis=-1):
 
     This is most efficient for n a power of two."""
 
-    a = asarray(a).astype(Float)
+    a = asarray(a).astype(float)
     return _raw_fft(a, n, axis, fftpack.rffti, fftpack.rfftf, _real_fft_cache)
 
 
@@ -153,7 +153,7 @@ def inverse_real_fft(a, n=None, axis=-1):
     inverse_real_fft(real_fft(a), len(a)) == a
     within numerical accuracy."""
 
-    a = asarray(a).astype(Complex)
+    a = asarray(a).astype(complex)
     if n == None:
         n = (shape(a)[axis] - 1) * 2
     return _raw_fft(a, n, axis, fftpack.rffti, fftpack.rfftb,
@@ -283,7 +283,6 @@ def real_fftnd(a, s=None, axes=None):
         a = fft(a, s[ii], axes[ii])
     return a
 
-
 def real_fft2d(a, s=None, axes=(-2,-1)):
     """real_fft2d(a, s=None, axes=(-2,-1)) 
 
diff --git a/numpy/lib/index_tricks.py b/numpy/lib/index_tricks.py
index 0e625d3fa..604bd6059 100644
--- a/numpy/lib/index_tricks.py
+++ b/numpy/lib/index_tricks.py
@@ -243,7 +243,6 @@ c_.__doc__ = """Translates slice objects to concatenation along the second axis.
 #col = concatenator(-1,1)
 
 
-
 class ndenumerate(object):
     """
     A simple nd index iterator over an array.
@@ -258,31 +257,15 @@ class ndenumerate(object):
     (1, 1) 4
     """
     def __init__(self, arr):
-        arr = asarray(arr)
-        self.iter = enumerate(arr.flat)
-        self.ashape = arr.shape
-        self.nd = arr.ndim
-        self.factors = [None]*(self.nd-1)
-        val = self.ashape[-1]
-        for i in range(self.nd-1,0,-1):
-            self.factors[i-1] = val
-            val *= self.ashape[i-1]
+        self.iter = asarray(arr).flat
 
     def next(self):
-        res = self.iter.next()
-        indxs = [None]*self.nd
-        val = res[0]
-        for i in range(self.nd-1):
-            indxs[i] = val / self.factors[i]
-            val = val % self.factors[i]
-        indxs[self.nd-1] = val
-        return tuple(indxs), res[1]
+        return self.iter.coords, self.iter.next()
 
     def __iter__(self):
         return self
 
 
-
 # You can do all this with slice() plus a few special objects,
 # but there's a lot to remember. This version is simpler because
 # it uses the standard array indexing syntax.