[Numpy-tickets] [NumPy] #636: Faster array version of ndindex

[NumPy]

#636: Faster array version of ndindex
----------------------------+-----------------------------------------------
 Reporter:  jarrod.millman  |       Owner:  jarrod.millman
     Type:  enhancement     |      Status:  new           
 Priority:  normal          |   Milestone:  1.0.5         
Component:  Other           |     Version:  none          
 Severity:  normal          |    Keywords:                
----------------------------+-----------------------------------------------
 The following is from an email from Jonathan Taylor:

 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
 }}}

-- 
Ticket URL: <http://scipy.org/scipy/numpy/ticket/636>
NumPy <http://projects.scipy.org/scipy/numpy>
The fundamental package needed for scientific computing with Python.