[Numpy-discussion] Faster array version of ndindex
Sebastian Haase
haase at msg.ucsf.edu
Fri Dec 14 04:31:56 EST 2007
Do you know about
N.fromiter()
?
-Sebastian Haase
On Dec 14, 2007 12:33 AM, 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
>
More information about the NumPy-Discussion
mailing list