Allow .item(<index>) for fast selection of Python scalars from a numpy array. Re-name a few confusing constants in ufuncobject.c

2 files changed, 94 insertions, 22 deletions

diff --git a/numpy/core/src/arraymethods.c b/numpy/core/src/arraymethods.c
index 610a6a372..d98f0b00d 100644
--- a/numpy/core/src/arraymethods.c
+++ b/numpy/core/src/arraymethods.c
@@ -452,14 +452,86 @@ array_tofile(PyArrayObject *self, PyObject *args, PyObject *kwds)
 
 static PyObject *
 array_toscalar(PyArrayObject *self, PyObject *args) {
-        if (!PyArg_ParseTuple(args, "")) return NULL;
-	if (self->nd == 0 || PyArray_SIZE(self) == 1)
-		return self->descr->f->getitem(self->data, self);
-	else {
-		PyErr_SetString(PyExc_ValueError, "can only convert an"	\
-				" array of size 1 to Python scalar.");
-		return NULL;
-	}
+        int n, nd;
+        n = PyTuple_GET_SIZE(args);
+        
+        if (n==1) {
+                PyObject *obj;
+                obj = PyTuple_GET_ITEM(args, 0);
+                if (PyTuple_Check(obj)) {
+                        args = obj;
+                        n = PyTuple_GET_SIZE(args);
+                }
+        }
+        
+        if (n==0) {
+                if (self->nd == 0 || PyArray_SIZE(self) == 1)
+                        return self->descr->f->getitem(self->data, self);
+                else {
+                        PyErr_SetString(PyExc_ValueError, 
+                                        "can only convert an array "    \
+                                        " of size 1 to a Python scalar");
+                        return NULL;
+                }
+        }
+        else if (n != self->nd && (n > 1 || self->nd==0)) {
+                PyErr_SetString(PyExc_ValueError, 
+                                "incorrect number of indices for "      \
+                                "array");
+                return NULL;
+        }
+        else if (n==1) { /* allows for flat getting as well as 1-d case */
+                intp value, loc, index, factor;
+                intp factors[MAX_DIMS];
+                value = PyArray_PyIntAsIntp(PyTuple_GET_ITEM(args, 0));
+                if (error_converting(value)) {
+                        PyErr_SetString(PyExc_ValueError, "invalid integer");
+                        return NULL;
+                }
+                if (value >= PyArray_SIZE(self)) {
+                        PyErr_SetString(PyExc_ValueError, 
+                                        "index out of bounds");
+                        return NULL;
+                }
+                if (self->nd == 1) {
+                        return self->descr->f->getitem(self->data + value, 
+                                                       self);
+                }
+                nd = self->nd;
+                factor = 1;
+                while (nd--) {
+                        factors[nd] = factor;
+                        factor *= self->dimensions[nd];
+                }
+                loc = 0;
+                for (nd=0; nd < self->nd; nd++) {
+                        index = value / factors[nd];
+                        value = value % factors[nd];
+                        loc += self->strides[nd]*index;
+                }
+                
+                return self->descr->f->getitem(self->data + loc,
+                                               self);
+                
+        }
+        else {
+                intp loc, index[MAX_DIMS];
+                nd = PyArray_IntpFromSequence(args, index, MAX_DIMS);
+                if (nd < n) return NULL;
+                loc = 0;
+                while (nd--) {
+                        if (index[nd] < 0) 
+                                index[nd] += self->dimensions[nd];
+                        if (index[nd] < 0 || 
+                            index[nd] >= self->dimensions[nd]) {
+                                PyErr_SetString(PyExc_ValueError, 
+                                                "index out of bounds");
+                                return NULL;
+                        }
+                        loc += self->strides[nd]*index[nd];
+                }
+                return self->descr->f->getitem(self->data + loc, self);
+        }
 }
 
 
diff --git a/numpy/core/src/ufuncobject.c b/numpy/core/src/ufuncobject.c
index 2b4016b09..b08eb64af 100644
--- a/numpy/core/src/ufuncobject.c
+++ b/numpy/core/src/ufuncobject.c
@@ -578,9 +578,9 @@ PyUFunc_clearfperr()
 
 
 #define NO_UFUNCLOOP        0
-#define ZERODIM_REDUCELOOP  0
+#define ZERO_EL_REDUCELOOP  0
 #define ONE_UFUNCLOOP       1
-#define ONEDIM_REDUCELOOP   1
+#define ONE_EL_REDUCELOOP   1
 #define NOBUFFER_UFUNCLOOP  2
 #define NOBUFFER_REDUCELOOP 2
 #define BUFFER_UFUNCLOOP  3
@@ -1832,12 +1832,12 @@ construct_reduce(PyUFuncObject *self, PyArrayObject **arr, PyArrayObject *out,
 	aar = *arr;
 
         if (loop->N == 0) {
-                loop->meth = ZERODIM_REDUCELOOP;
+                loop->meth = ZERO_EL_REDUCELOOP;
         }
         else if (PyArray_ISBEHAVED_RO(aar) &&		\
                  otype == (aar)->descr->type_num) {
 		if (loop->N == 1) {
-			loop->meth = ONEDIM_REDUCELOOP;
+			loop->meth = ONE_EL_REDUCELOOP;
 		}
 		else {
 			loop->meth = NOBUFFER_UFUNCLOOP;
@@ -1856,7 +1856,7 @@ construct_reduce(PyUFuncObject *self, PyArrayObject **arr, PyArrayObject *out,
         else
                 loop->obj = 0;
 
-        if (loop->meth == ZERODIM_REDUCELOOP) {
+        if (loop->meth == ZERO_EL_REDUCELOOP) {
                 idarr = _getidentity(self, otype, str);
                 if (idarr == NULL) goto fail;
                 if (idarr->descr->elsize > UFUNC_MAXIDENTITY) {
@@ -1913,10 +1913,10 @@ construct_reduce(PyUFuncObject *self, PyArrayObject **arr, PyArrayObject *out,
                         outsize = PyArray_MultiplyList(loop_i, aar->nd);
                 }
 		if (ind_size == 0) {
-			loop->meth = ZERODIM_REDUCELOOP;
+			loop->meth = ZERO_EL_REDUCELOOP;
 			return loop;
 		}
-		if (loop->meth == ONEDIM_REDUCELOOP)
+		if (loop->meth == ONE_EL_REDUCELOOP)
 			loop->meth = NOBUFFER_REDUCELOOP;
 		break;
 	}
@@ -1938,7 +1938,7 @@ construct_reduce(PyUFuncObject *self, PyArrayObject **arr, PyArrayObject *out,
         loop->outsize = loop->ret->descr->elsize;
         loop->bufptr[1] = loop->ret->data;
 
-	if (loop->meth == ZERODIM_REDUCELOOP) {
+	if (loop->meth == ZERO_EL_REDUCELOOP) {
 		loop->size = PyArray_SIZE(loop->ret);
 		return loop;
 	}
@@ -1946,7 +1946,7 @@ construct_reduce(PyUFuncObject *self, PyArrayObject **arr, PyArrayObject *out,
 	loop->it = (PyArrayIterObject *)PyArray_IterNew((PyObject *)aar);
         if (loop->it == NULL) return NULL;
 
-	if (loop->meth == ONEDIM_REDUCELOOP) {
+	if (loop->meth == ONE_EL_REDUCELOOP) {
 		loop->size = loop->it->size;
 		return loop;
 	}
@@ -2046,7 +2046,7 @@ PyUFunc_Reduce(PyUFuncObject *self, PyArrayObject *arr, PyArrayObject *out,
 
         NPY_LOOP_BEGIN_THREADS
         switch(loop->meth) {
-        case ZERODIM_REDUCELOOP:
+        case ZERO_EL_REDUCELOOP:
 		/* fprintf(stderr, "ZERO..%d\n", loop->size); */
 		for(i=0; i<loop->size; i++) {
 			if (loop->obj) Py_INCREF(*((PyObject **)loop->idptr));
@@ -2054,7 +2054,7 @@ PyUFunc_Reduce(PyUFuncObject *self, PyArrayObject *arr, PyArrayObject *out,
 			loop->bufptr[1] += loop->outsize;
 		}
                 break;
-        case ONEDIM_REDUCELOOP:
+        case ONE_EL_REDUCELOOP:
 		/*fprintf(stderr, "ONEDIM..%d\n", loop->size); */
                 while(loop->index < loop->size) {
 			if (loop->obj)
@@ -2189,7 +2189,7 @@ PyUFunc_Accumulate(PyUFuncObject *self, PyArrayObject *arr, PyArrayObject *out,
 
 	NPY_LOOP_BEGIN_THREADS
         switch(loop->meth) {
-        case ZERODIM_REDUCELOOP: /* Accumulate */
+        case ZERO_EL_REDUCELOOP: /* Accumulate */
 		/* fprintf(stderr, "ZERO..%d\n", loop->size); */
 		for(i=0; i<loop->size; i++) {
 			if (loop->obj)
@@ -2198,7 +2198,7 @@ PyUFunc_Accumulate(PyUFuncObject *self, PyArrayObject *arr, PyArrayObject *out,
 			loop->bufptr[1] += loop->outsize;
 		}
                 break;
-        case ONEDIM_REDUCELOOP: /* Accumulate */
+        case ONE_EL_REDUCELOOP: /* Accumulate */
 		/* fprintf(stderr, "ONEDIM..%d\n", loop->size); */
                 while(loop->index < loop->size) {
 			if (loop->obj)
@@ -2369,7 +2369,7 @@ PyUFunc_Reduceat(PyUFuncObject *self, PyArrayObject *arr, PyArrayObject *ind,
 	NPY_LOOP_BEGIN_THREADS
 	switch(loop->meth) {
 	/* zero-length index -- return array immediately */
-	case ZERODIM_REDUCELOOP:
+	case ZERO_EL_REDUCELOOP:
 		/* fprintf(stderr, "ZERO..\n"); */
 		break;
 	/* NOBUFFER -- behaved array and same type */