Fixed bug in ma.py and added fast quicksorts.

5 files changed, 144 insertions, 11 deletions

diff --git a/scipy/base/__init__.py b/scipy/base/__init__.py
index c8c28efe9..0c0c158df 100644
--- a/scipy/base/__init__.py
+++ b/scipy/base/__init__.py
@@ -6,6 +6,7 @@ import multiarray
 import umath
 import numerictypes as nt
 multiarray.set_typeDict(nt.typeDict)
+import _sort
 from numeric import *
 from oldnumeric import *
 from matrix import *
diff --git a/scipy/base/include/scipy/arrayobject.h b/scipy/base/include/scipy/arrayobject.h
index a0aec4326..b6f3e330d 100644
--- a/scipy/base/include/scipy/arrayobject.h
+++ b/scipy/base/include/scipy/arrayobject.h
@@ -781,7 +781,7 @@ typedef int (PyArray_ScanFunc)(FILE *, void *, void *, void *);
 typedef int (PyArray_FillFunc)(void *, intp, void *);
 
 typedef void (PyArray_SortFunc)(void *, intp, int, void *);
-typedef int (PyArray_ArgSortFunc)(void *, void *);
+typedef void (PyArray_ArgSortFunc)(void *, intp *, intp, int, void *);
 
 typedef struct {
         intp *ptr;
diff --git a/scipy/base/ma.py b/scipy/base/ma.py
index 49b953179..e0c81d8c8 100644
--- a/scipy/base/ma.py
+++ b/scipy/base/ma.py
@@ -610,7 +610,7 @@ class MaskedArray (object):
     def __array__ (self, t = None):
         "Special hook for numeric. Converts to numeric if possible."
         if self._mask is not None:
-            if umath.sometrue(oldnumeric.ravel(self._mask)):
+            if oldnumeric.sometrue(oldnumeric.ravel(self._mask)):
                 raise MAError, \
                 """Cannot automatically convert masked array to numeric because data
                    is masked in one or more locations.
diff --git a/scipy/base/src/_sortmodule.c.src b/scipy/base/src/_sortmodule.c.src
index 00098938f..d3959097d 100644
--- a/scipy/base/src/_sortmodule.c.src
+++ b/scipy/base/src/_sortmodule.c.src
@@ -1,8 +1,11 @@
 /* The purpose of this module is to add faster sort functions
    that are type-specific.  This is done by altering the
    function table for the builtin descriptors. 
+
+   These sorting functions are copied from numarray 
 */
 
+
 #include "Python.h"
 #include "scipy/arrayobject.h"
 
@@ -11,13 +14,17 @@
 #define STDC_EQ(a,b) ((a) == (b))
 #define SWAP(a,b) {SWAP_temp = (b); (b)=(a); (a) = SWAP_temp;}
 
+/**begin repeat
+#TYPE=BOOL,BYTE,UBYTE,SHORT,USHORT,INT,UINT,LONG,ULONG,LONGLONG,ULONGLONG,FLOAT,DOUBLE,LONGDOUBLE#
+#type=Bool,byte,ubyte,short,ushort,int,uint,long,ulong,longlong,ulonglong,float,double,longdouble#
+ **/
 static void 
-LONG_quicksort(long *start, intp num, int elsize, void *unused)
+@TYPE@_quicksort(@type@ *start, intp num, int elsize, void *unused)
 {
-	long *pl = start;
-	long *pr = start + num - 1;
-	long vp, SWAP_temp;
-	long *stack[100], **sptr = stack, *pm, *pi, *pj, *pt;
+	@type@ *pl = start;
+	@type@ *pr = start + num - 1;
+	@type@ vp, SWAP_temp;
+	@type@ *stack[100], **sptr = stack, *pm, *pi, *pj, *pt;
 
         for(;;) {
                 while ((pr - pl) > 15) {
@@ -61,14 +68,75 @@ LONG_quicksort(long *start, intp num, int elsize, void *unused)
                 pl = *(--sptr);
         }
 }
+
+static void 
+@TYPE@_aquicksort(@type@ *v, intp* tosort, intp num, int elsize, void *unused)
+{
+	@type@ vp;
+	intp *pl, *pr, SWAP_temp;
+	intp *stack[100], **sptr=stack, *pm, *pi, *pj, *pt, vi;
+
+	pl = tosort;
+	pr = tosort + num - 1;
+
+        for(;;) {
+                while ((pr - pl) > 15) {
+                        /* quicksort partition */
+                        pm = pl + ((pr - pl) >> 1);
+                        if (STDC_LT(v[*pm],v[*pl])) SWAP(*pm,*pl);
+                        if (STDC_LT(v[*pr],v[*pm])) SWAP(*pr,*pm);
+                        if (STDC_LT(v[*pm],v[*pl])) SWAP(*pm,*pl);
+                        vp = v[*pm];
+                        pi = pl;
+                        pj = pr - 1;
+                        SWAP(*pm,*pj);
+                        for(;;) {
+                                do ++pi; while (STDC_LT(v[*pi],vp));
+                                do --pj; while (STDC_LT(vp,v[*pj]));
+                                if (pi >= pj)  break;
+                                SWAP(*pi,*pj);
+                        }
+                        SWAP(*pi,*(pr-1));
+                        /* push largest partition on stack */
+                        if (pi - pl < pr - pi) {
+                                *sptr++ = pi + 1;
+                                *sptr++ = pr;
+                                pr = pi - 1;
+                        }else{
+                                *sptr++ = pl;
+                                *sptr++ = pi - 1;
+                                pl = pi + 1;
+                        }
+                }
+                /* insertion sort */
+                for(pi = pl + 1; pi <= pr; ++pi) {
+                        vi = *pi;
+                        vp = v[vi];
+                        for(pj = pi, pt = pi - 1; pj > pl && STDC_LT(vp, v[*pt]);)
+				{
+					*pj-- = *pt--;
+				}
+                        *pj = vi;
+                }
+                if (sptr == stack) break;
+                pr = *(--sptr);
+                pl = *(--sptr);
+        }
+}
+/**end repeat**/
 	
 static void 
 add_sortfuncs(void)
 {
 	PyArray_Descr *descr;
 
-	descr = PyArray_DescrFromType(PyArray_LONG);
-	descr->f->sort[PyArray_QUICKSORT] = (PyArray_SortFunc *)LONG_quicksort;
+/**begin repeat
+#TYPE=BOOL,BYTE,UBYTE,SHORT,USHORT,INT,UINT,LONG,ULONG,LONGLONG,ULONGLONG,FLOAT,DOUBLE,LONGDOUBLE#
+**/
+	descr = PyArray_DescrFromType(PyArray_@TYPE@);
+	descr->f->sort[PyArray_QUICKSORT] = (PyArray_SortFunc *)@TYPE@_quicksort;
+	descr->f->argsort[PyArray_QUICKSORT] = (PyArray_ArgSortFunc *)@TYPE@_aquicksort;
+/**end repeat**/
 }
 
 static struct PyMethodDef methods[] = {
diff --git a/scipy/base/src/multiarraymodule.c b/scipy/base/src/multiarraymodule.c
index bb8195929..af6acfc34 100644
--- a/scipy/base/src/multiarraymodule.c
+++ b/scipy/base/src/multiarraymodule.c
@@ -1613,6 +1613,7 @@ _new_sort(PyArrayObject *op, int axis, PyArray_SORTKIND which)
 				      N, elsize);
 			PyArray_ITER_NEXT(it);
 		}
+		PyDataMem_FREE(buffer);
 	}
 	else {
 		while (size--) {
@@ -1620,13 +1621,76 @@ _new_sort(PyArrayObject *op, int axis, PyArray_SORTKIND which)
 			PyArray_ITER_NEXT(it);
 		}
 	}	
+	Py_DECREF(it);
 	return 0;
 }
 
 static PyObject*
 _new_argsort(PyArrayObject *op, int axis, PyArray_SORTKIND which) 
 {
-	PyErr_SetNone(PyExc_NotImplementedError);
+
+	PyArrayIterObject *it=NULL;
+	PyArrayIterObject *rit=NULL;
+	PyObject *ret;
+	int needcopy=0, i;
+	intp N, size;
+	int elsize;
+	intp astride, rstride, *iptr;
+	PyArray_ArgSortFunc *argsort;
+
+	ret = PyArray_New(op->ob_type, op->nd,
+			  op->dimensions, PyArray_INTP,
+			  NULL, NULL, 0, 0, (PyObject *)op);
+	if (ret == NULL) return NULL;
+
+	it = (PyArrayIterObject *)PyArray_IterAllButAxis((PyObject *)op, axis);
+	rit = (PyArrayIterObject *)PyArray_IterAllButAxis(ret, axis);
+	if (rit == NULL || it == NULL) goto fail;
+	argsort = op->descr->f->argsort[which];
+	size = it->size;
+	N = op->dimensions[axis];
+	elsize = op->descr->elsize;
+	astride = op->strides[axis];
+	rstride = PyArray_STRIDE(ret,axis);
+
+	needcopy = !(op->flags & ALIGNED) || (astride != (intp) elsize) || \
+		(rstride != sizeof(intp));
+	
+	if (needcopy) {
+		char *valbuffer, *indbuffer;
+		valbuffer = PyDataMem_NEW(N*(elsize+sizeof(intp)));
+		indbuffer = valbuffer + (N*elsize);
+		while (size--) {
+			_strided_copy(valbuffer, (intp) elsize, it->dataptr, astride, 
+				      N, elsize);
+			iptr = (intp *)indbuffer;
+			for (i=0; i<N; i++) *iptr++ = i;
+			argsort(valbuffer, (intp *)indbuffer, N, elsize, op);
+			_strided_copy(rit->dataptr, rstride, indbuffer, sizeof(intp),
+				      N, sizeof(intp));
+			PyArray_ITER_NEXT(it);
+			PyArray_ITER_NEXT(rit);
+		}
+		PyDataMem_FREE(valbuffer);
+	}
+	else {
+		while (size--) {
+			iptr = (intp *)rit->dataptr;
+			for (i=0; i<N; i++) *iptr++ = i;
+			argsort(it->dataptr, (intp *)rit->dataptr, N, elsize, op);
+			PyArray_ITER_NEXT(it);
+			PyArray_ITER_NEXT(rit);
+		}
+	}
+
+	Py_DECREF(it);
+	Py_DECREF(rit);
+	return ret;
+
+ fail:
+	Py_DECREF(ret);
+	Py_XDECREF(it);
+	Py_XDECREF(rit);
 	return NULL;
 }
 
@@ -1713,7 +1777,7 @@ PyArray_Sort(PyArrayObject *op, int axis, PyArray_SORTKIND which)
 		return -1;
 	}
 
-	/* Determine if we should use new algorithm or not */
+	/* Determine if we should use type-specific algorithm or not */
 	if (op->descr->f->sort[which] != NULL) {
 		return _new_sort(op, axis, which);
 	}