Added multiter type to expose broadcasting as an iterator.

5 files changed, 252 insertions, 9 deletions

diff --git a/scipy/base/code_generators/generate_array_api.py b/scipy/base/code_generators/generate_array_api.py
index 1d2b9028c..914453306 100644
--- a/scipy/base/code_generators/generate_array_api.py
+++ b/scipy/base/code_generators/generate_array_api.py
@@ -227,6 +227,10 @@ objectapi_list = [
     """,
      'IterNew','PyObject *', 'PyObject *'),
 
+    (r"""Get MultiIterator,
+    """,
+     'MultiIterNew','int, ...', 'PyObject *'),
+
     (r"""
     """,
      'PyIntAsInt', 'PyObject *', 'int'),
@@ -518,7 +522,7 @@ types = ['Generic','Numeric','Integer','SignedInteger','UnsignedInteger', 'Inexa
 
 # API fixes for __arrayobject_api.h
 
-fixed = 4
+fixed = 5
 numtypes = len(types) + fixed
 numobject = len(objectapi_list) + numtypes
 nummulti = len(multiapi_list)
@@ -619,6 +623,7 @@ static PyTypeObject PyBigArray_Type;
 static PyTypeObject PyArray_Type;
 static PyTypeObject PyArrayIter_Type;
 static PyTypeObject PyArrayMapIter_Type;
+static PyTypeObject PyArrayMultiIter_Type;
 static int PyArray_NUMUSERTYPES=0;
 
 %s
@@ -642,7 +647,8 @@ static void **PyArray_API=NULL;
 #define PyBigArray_Type (*(PyTypeObject *)PyArray_API[0])
 #define PyArray_Type (*(PyTypeObject *)PyArray_API[1])
 #define PyArrayIter_Type (*(PyTypeObject *)PyArray_API[2])
-#define PyArray_NUMUSERTYPES (*(int *)PyArray_API[3])
+#define PyArrayMultiIter_Type (*(PyTypeObject *)PyArray_API[3])
+#define PyArray_NUMUSERTYPES (*(int *)PyArray_API[4])
 
 %s
 
@@ -685,6 +691,7 @@ void *PyArray_API[] = {
         (void *) &PyBigArray_Type,
         (void *) &PyArray_Type,
         (void *) &PyArrayIter_Type,
+        (void *) &PyArrayMultiIter_Type,
         (int *) &PyArray_NUMUSERTYPES,
 %s
 };
diff --git a/scipy/base/include/scipy/arrayobject.h b/scipy/base/include/scipy/arrayobject.h
index 520473168..214175677 100644
--- a/scipy/base/include/scipy/arrayobject.h
+++ b/scipy/base/include/scipy/arrayobject.h
@@ -1018,14 +1018,12 @@ typedef struct {
 		}							\
 }
 
-/* Not constructed anywhere.  Just serves as a standard type that
-   PyArray_Broadcast expects.
 
+/*
    Any object passed to PyArray_Broadcast must be binary compatible with 
    this structure.    
 */
 
-
 typedef struct {
 	PyObject_HEAD
 
@@ -1036,7 +1034,39 @@ typedef struct {
 	intp                    dimensions[MAX_DIMS];  /* dimensions */
 	PyArrayIterObject       *iters[MAX_DIMS];      /* iterators */
 } PyArrayMultiIterObject;  
+
+#define PyArray_MultiIter_RESET(multi) {			  \
+		int _mi_;					  \
+		(multi)->index = 0;				  \
+		for (_mi_ = 0; _mi_ < (multi)->numiter; _mi_++) { \
+			PyArray_ITER_RESET((multi)->iters[_mi_]); \
+		}						  \
+	}
+	
+#define PyArray_MultiIter_NEXT(multi) {				 \
+		int _mi_;					 \
+		(multi)->index += 1;				 \
+		for (_mi_=0; _mi_<(multi)->numiter; _mi_++) {	 \
+			PyArray_ITER_NEXT((multi)->iters[_mi_]); \
+		}						 \
+	}
+	
+#define PyArray_MultiIter_GOTO(it, dest) {				\
+		int _mi_;						\
+		for (_mi_=0; _mi_<(multi)->numiter; _mi_++) {		\
+			PyArray_ITER_GOTO((multi)->iters[_mi_], dest);	\
+		}							\
+		(multi)->index = (multi)->iters[0]->index;		\
+	}
 	
+#define PyArray_MultiIter_GOTO1D(it, ind) {				\
+		int _mi_;						\
+		for (_mi_=0; _mi_<(multi)->numiter; _mi_++) {		\
+			PyArray_ITER_GOTO1D((multi)->iters[_mi_], ind);	\
+		}							\
+		(multi)->index = (multi)->iters[0]->index;		\
+	}
+		
 
 /* Store the information needed for fancy-indexing over an array */
 
diff --git a/scipy/base/src/arrayobject.c b/scipy/base/src/arrayobject.c
index 3d00f9e57..0489685a7 100644
--- a/scipy/base/src/arrayobject.c
+++ b/scipy/base/src/arrayobject.c
@@ -4457,7 +4457,7 @@ static PyTypeObject PyBigArray_Type = {
         NULL,                            	  /*tp_as_buffer*/
         (Py_TPFLAGS_DEFAULT 
          | Py_TPFLAGS_BASETYPE
-         | Py_TPFLAGS_CHECKTYPES),           /*tp_flags*/
+         | Py_TPFLAGS_CHECKTYPES),                /*tp_flags*/
         /*Documentation string */
         Arraytype__doc__,			  /*tp_doc*/
 
@@ -6566,13 +6566,14 @@ static PyTypeObject PyArrayIter_Type = {
         0,	                         	/* tp_getattro */
         0,					/* tp_setattro */
         0,					/* tp_as_buffer */
-        Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,/* tp_flags */
+        Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | \
+	Py_TPFLAGS_HAVE_ITER,                   /* tp_flags */
         0,					/* tp_doc */
         (traverseproc)arrayiter_traverse,	/* tp_traverse */
         0,					/* tp_clear */
         0,					/* tp_richcompare */
         0,					/* tp_weaklistoffset */
-        0,		    	                /* tp_iter */
+        PyObject_SelfIter, 	                /* tp_iter */
         (iternextfunc)arrayiter_next,		/* tp_iternext */
         iter_methods,				/* tp_methods */
 };
@@ -7294,6 +7295,206 @@ static PyTypeObject PyArrayMapIter_Type = {
 
 };
 
+static PyObject *
+PyArray_MultiIterNew(int n, ...)
+{
+        va_list va;
+	PyArrayMultiIterObject *multi;
+	PyObject *current;
+	PyObject *arr;
+	
+	int i, err=0;
+	
+	if (n < 2 || n > MAX_DIMS) {
+		PyErr_Format(PyExc_ValueError, 
+			     "Need between 2 and (%d) "			\
+			     "array objects (inclusive).", MAX_DIMS);
+	}
+	
+        multi = PyObject_GC_New(PyArrayMultiIterObject, &PyArrayMultiIter_Type);
+        if (multi == NULL)
+                return NULL;
+	
+	for (i=0; i<n; i++) multi->iters[i] = NULL;
+	multi->numiter = n;
+	multi->index = 0;
+
+        va_start(va, n);
+	for (i=0; i<n; i++) {
+		current = va_arg(va, PyObject *);
+		arr = PyArray_FROM_O(current);
+		if (arr==NULL) {
+			err=1; break;
+		}
+		else {
+			multi->iters[i] = (PyArrayIterObject *)PyArray_IterNew(arr);
+			Py_DECREF(arr);
+		}
+	}
+
+	va_end(va);	
+	
+	if (!err && PyArray_Broadcast(multi) < 0) err=1;
+	
+	if (err) {
+		for (i=0; i<n; i++) Py_XDECREF(multi->iters[i]);
+		PyObject_GC_Del(multi);
+		return NULL;
+	}
+	
+	PyArray_MultiIter_RESET(multi);
+	
+        PyObject_GC_Track(multi);
+        return (PyObject *)multi;
+}
+
+static PyObject *
+arraymultiter_new(PyTypeObject *subtype, PyObject *args, PyObject *kwds)
+{
+	
+	int n, i;
+	PyArrayMultiIterObject *multi;
+	PyObject *arr;
+	
+	if (kwds != NULL) {
+		PyErr_SetString(PyExc_ValueError, 
+				"keyword arguments not accepted.");
+		return NULL;
+	}
+
+	n = PyTuple_Size(args);
+	if (n < 2 || n > MAX_DIMS) {
+		if (PyErr_Occurred()) return NULL;
+		PyErr_Format(PyExc_ValueError, 
+			     "Need at least two and fewer than (%d) "	\
+			     "array objects.", MAX_DIMS);
+		return NULL;
+	}
+	
+        multi = PyObject_GC_New(PyArrayMultiIterObject, &PyArrayMultiIter_Type);
+        if (multi == NULL)
+                return NULL;
+
+	multi->numiter = n;
+	multi->index = 0;
+	for (i=0; i<n; i++) multi->iters[i] = NULL;
+	for (i=0; i<n; i++) {
+		arr = PyArray_FromAny(PyTuple_GET_ITEM(args, i), NULL, 0, 0, 0);
+		if (arr == NULL) goto fail;
+		if ((multi->iters[i] =					\
+		     (PyArrayIterObject *)PyArray_IterNew(arr))==NULL) 
+			goto fail;
+		Py_DECREF(arr);
+	}
+	if (PyArray_Broadcast(multi) < 0) goto fail;
+	PyArray_MultiIter_RESET(multi);
+	
+        PyObject_GC_Track(multi);
+        return (PyObject *)multi;
+	
+ fail:
+	for (i=0; i<n; i++) Py_XDECREF(multi->iters[i]);
+	PyObject_GC_Del(multi);
+	return NULL;
+}
+
+static PyObject *
+arraymultiter_next(PyArrayMultiIterObject *multi)
+{
+	PyObject *ret;
+	int i, n;
+
+	n = multi->numiter;
+	ret = PyTuple_New(n);
+	if (ret == NULL) return NULL;
+	if (multi->index < multi->size) {
+		for (i=0; i < n; i++) {
+			PyArrayIterObject *it=multi->iters[i];
+			PyTuple_SET_ITEM(ret, i, 
+					 PyArray_ToScalar(it->dataptr, it->ao));
+			PyArray_ITER_NEXT(it);
+		}
+		multi->index++;
+		return ret;
+	}
+        return NULL;
+}
+
+static void
+arraymultiter_dealloc(PyArrayMultiIterObject *multi)
+{
+	int i;
+
+        PyObject_GC_UnTrack(multi);
+	for (i=0; i<multi->numiter; i++) 
+		Py_XDECREF(multi->iters[i]);
+        PyObject_GC_Del(multi);
+}
+
+static int
+arraymultiter_traverse(PyArrayMultiIterObject *multi, visitproc visit, void *arg)
+{
+	int err, i;
+
+        for (i=0; i<multi->numiter; i++) 
+		if (multi->iters[i]) {
+			err = visit((PyObject *)(multi->iters[i]), arg);
+			if (err) return err;
+		}
+        return 0;
+}
+
+static PyTypeObject PyArrayMultiIter_Type = {
+        PyObject_HEAD_INIT(NULL)
+        0,					 /* ob_size */
+        "scipy.multiter",		      	 /* tp_name */
+        sizeof(PyArrayMultiIterObject),          /* tp_basicsize */
+        0,					 /* tp_itemsize */
+        /* methods */
+        (destructor)arraymultiter_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 | Py_TPFLAGS_HAVE_GC	\
+	| Py_TPFLAGS_HAVE_ITER,                 /* tp_flags */
+        0,					/* tp_doc */
+        (traverseproc)arraymultiter_traverse,	/* tp_traverse */
+        0,					/* tp_clear */
+        0,					/* tp_richcompare */
+        0,					/* tp_weaklistoffset */
+        PyObject_SelfIter,    	                /* tp_iter */
+        (iternextfunc)arraymultiter_next,	/* 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 */
+        arraymultiter_new,	                  /* tp_new */
+        0,	                                  /* tp_free */
+        0,					  /* tp_is_gc */
+        0,					  /* tp_bases */
+        0,					  /* tp_mro */
+        0,					  /* tp_cache */
+        0,					  /* tp_subclasses */
+        0					  /* tp_weaklist */
+};
+
 /** END of Subscript Iterator **/
 
 
diff --git a/scipy/base/src/multiarraymodule.c b/scipy/base/src/multiarraymodule.c
index 50a529be7..4ac14b005 100644
--- a/scipy/base/src/multiarraymodule.c
+++ b/scipy/base/src/multiarraymodule.c
@@ -4096,6 +4096,9 @@ DL_EXPORT(void) initmultiarray(void) {
 	if (PyType_Ready(&PyArrayMapIter_Type) < 0)
                 return; 
 
+	if (PyType_Ready(&PyArrayMultiIter_Type) < 0)
+		return;
+
 	c_api = PyCObject_FromVoidPtr((void *)PyArray_API, NULL);
 	if (PyErr_Occurred()) goto err;
 	PyDict_SetItemString(d, "_ARRAY_API", c_api);
@@ -4114,6 +4117,8 @@ DL_EXPORT(void) initmultiarray(void) {
 	PyDict_SetItemString(d, "ndarray", (PyObject *)&PyArray_Type);
         Py_INCREF(&PyArrayIter_Type);
 	PyDict_SetItemString(d, "flatiter", (PyObject *)&PyArrayIter_Type);
+        Py_INCREF(&PyArrayMultiIter_Type);
+	PyDict_SetItemString(d, "multiter", (PyObject *)&PyArrayMultiIter_Type);
 
 	/* Doesn't need to be exposed to Python 
         Py_INCREF(&PyArrayMapIter_Type);
diff --git a/scipy/core_version.py b/scipy/core_version.py
index 7763ba426..28b156686 100644
--- a/scipy/core_version.py
+++ b/scipy/core_version.py
@@ -1,4 +1,4 @@
-version='0.6.2'
+version='0.7.0'
 
 import os
 svn_version_file = os.path.join(os.path.dirname(__file__),