[Numpy-discussion] Faster array version of ndindex

[Stefan van der Walt]

Hi Sebastian

N.fromiter only works on 1D arrays.  I thought the following may work,
but it doesn't:

 np.fromiter(np.ndindex(10,10,10),N.dtype((int,3)))

This kind of loop is probably best implemented in C, although I think
Jonathan's version is rather clever.

Regards
Stéfan

On Fri, Dec 14, 2007 at 10:31:56AM +0100, Sebastian Haase wrote:
> 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