[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          |   Resolution:                
 Keywords:                  |  
----------------------------+-----------------------------------------------
Comment (by jamartinh):

 == ndlinspace ==


 Hi there.
 Thanks for the ndtuples, now I have built what I really need which is the
 n-dimensional version of linspace: '''ndlinspace''', which of course with
 ndtuples is a lot of faster than using python list or other approaches.

 So my little contribution:

 {{{
 # Jose Antonio Martin H.
 # jamartinh <AT> fdi.ucm.es
 # 29-Dec-2007

 from numpy import *

 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 = arange(N)
    cur = cur[:,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 = kron(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 = arange(d).repeat(N)[:,newaxis]

        # This puts these two together.
        cur = column_stack((front,cur))
        N *= d

    return cur

 def ndlinspace(input_ranges,nelems):
     """ ndlinspace: n-dimensional linspace function
         input_ranges = [[-1,1],[-0.07,0.07]]
         nelems = (5,5)
     """

     #x = array(list(ndindex(*nelems)),float)+1.0
     x = ndtuples(*nelems)+1.0

     lbounds = []
     ubounds = []
     from_b  = array(nelems,float)
     for r in input_ranges:
         lbounds.append(r[0])
         ubounds.append(r[1])


     lbounds = array(lbounds,float)
     ubounds = array(ubounds,float)

     y = (lbounds) + (((x-1)/(from_b-1))*((ubounds)-(lbounds)))


     return y

 input_ranges = [[-1,1],[-0.07,0.07]]
 nelems = (5,5)
 b = ndlinspace(input_ranges,nelems)
 print b

 # result
 ##[[-1.    -0.07 ]
 ## [-1.    -0.035]
 ## [-1.     0.   ]
 ## [-1.     0.035]
 ## [-1.     0.07 ]
 ## [-0.5   -0.07 ]
 ## [-0.5   -0.035]
 ## [-0.5    0.   ]
 ## [-0.5    0.035]
 ## [-0.5    0.07 ]
 ## [ 0.    -0.07 ]
 ## [ 0.    -0.035]
 ## [ 0.     0.   ]
 ## [ 0.     0.035]
 ## [ 0.     0.07 ]
 ## [ 0.5   -0.07 ]
 ## [ 0.5   -0.035]
 ## [ 0.5    0.   ]
 ## [ 0.5    0.035]
 ## [ 0.5    0.07 ]
 ## [ 1.    -0.07 ]
 ## [ 1.    -0.035]
 ## [ 1.     0.   ]
 ## [ 1.     0.035]
 ## [ 1.     0.07 ]]

 }}}

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