[Numpy-svn] r3627 - trunk/numpy/core

[numpy-svn at scipy.org]

Author: oliphant
Date: 2007-03-31 18:45:45 -0500 (Sat, 31 Mar 2007)
New Revision: 3627

Modified:
   trunk/numpy/core/numeric.py
Log:
Speed up indicies significantly (about 10x).

Modified: trunk/numpy/core/numeric.py
===================================================================
--- trunk/numpy/core/numeric.py	2007-03-31 22:56:19 UTC (rev 3626)
+++ trunk/numpy/core/numeric.py	2007-03-31 23:45:45 UTC (rev 3627)
@@ -478,26 +478,34 @@
 
 little_endian = (sys.byteorder == 'little')
 
+
 def indices(dimensions, dtype=int):
     """Returns an array representing a grid of indices with row-only, and
     column-only variation.
     """
-    tmp = ones(dimensions, dtype)
-    lst = []
-    for i in range(len(dimensions)):
-        lst.append( add.accumulate(tmp, i, dtype)-1 )
-    return array(lst)
+    dimensions = tuple(dimensions)
+    N = len(dimensions)
+    if N == 0:
+        return array([],dtype=dtype)
+    res = empty((N,)+dimensions, dtype=dtype)
+    for i, dim in enumerate(dimensions):
+        tmp = arange(dim,dtype=dtype)
+        tmp.shape = (1,)*i + (dim,)+(1,)*(N-i-1)
+        newdim = dimensions[:i] + (1,)+ dimensions[i+1:]
+        val = zeros(newdim, dtype)
+        add(tmp, val, res[i])
+    return res
 
 def fromfunction(function, shape, **kwargs):
     """Returns an array constructed by calling a function on a tuple of number
     grids.
 
-    The function should accept as many arguments as the length of shape and work
-    on array inputs.  The shape argument is a sequence of numbers indicating the
-    length of the desired output for each axis.
+    The function should accept as many arguments as the length of shape and
+    work on array inputs.  The shape argument is a sequence of numbers
+    indicating the length of the desired output for each axis.
 
-    The function can also accept keyword arguments (except dtype), which will be
-    passed through fromfunction to the function itself.  The dtype argument
+    The function can also accept keyword arguments (except dtype), which will
+    be passed through fromfunction to the function itself.  The dtype argument
     (default float) determines the data-type of the index grid passed to the
     function.
     """