[Numpy-discussion] Faster array version of ndindex

[Jarrod Millman]

Could you provide more details about this to the ticket I created
based on your email:
http://projects.scipy.org/scipy/numpy/ticket/636

Thanks,

On Thu, Dec 13, 2007 at 3:33 PM, Jonathan Taylor
<jonathan.taylor at utoronto.ca> wrote:
> I was needing an array representation of ndindex since ndindex only
>  gives an iterator but array(list(ndindex)) takes too long.  There is
>  prob some obvious way to do this I am missing but if not feel free to
>  include this code which is much faster.
>
>  In [252]: time a=np.array(list(np.ndindex(10,10,10,10,10,10)))
>  CPU times: user 11.61 s, sys: 0.09 s, total: 11.70 s
>  Wall time: 11.82
>
>  In [253]: time a=ndtuples(10,10,10,10,10,10)
>  CPU times: user 0.32 s, sys: 0.21 s, total: 0.53 s
>  Wall time: 0.60
>
>  def ndtuples(*dims):
>     """Fast implementation of array(list(ndindex(*dims)))."""
>
>     # Need a list because we will go through it in reverse popping
>     # off the size of the last dimension.
>     dims = list(dims)
>
>     # N will keep track of the current length of the indices.
>     N = dims.pop()
>
>     # At the beginning the current list of indices just ranges over the
>     # last dimension.
>     cur = np.arange(N)
>     cur = cur[:,np.newaxis]
>
>     while dims != []:
>
>         d = dims.pop()
>
>         # This repeats the current set of indices d times.
>         # e.g. [0,1,2] -> [0,1,2,0,1,2,...,0,1,2]
>         cur = np.kron(np.ones((d,1)),cur)
>
>         # This ranges over the new dimension and 'stretches' it by N.
>         # e.g. [0,1,2] -> [0,0,...,0,1,1,...,1,2,2,...,2]
>         front = np.arange(d).repeat(N)[:,np.newaxis]
>
>         # This puts these two together.
>         cur = np.column_stack((front,cur))
>         N *= d
>
>     return cur
>  _______________________________________________
>  Numpy-discussion mailing list
>  Numpy-discussion at scipy.org
>  http://projects.scipy.org/mailman/listinfo/numpy-discussion
>



-- 
Jarrod Millman
Computational Infrastructure for Research Labs
10 Giannini Hall, UC Berkeley
phone: 510.643.4014
http://cirl.berkeley.edu/