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