Question

I have a list pts containing N points (Python floats). I wish to construct a NumPy array of dimension N*N*N*3 such that the array is equivalent to:

for i in xrange(0, N):
    for j in xrange(0, N):
        for k in xrange(0, N):
            arr[i,j,k,0] = pts[i]
            arr[i,j,k,1] = pts[j]
            arr[i,j,k,2] = pts[k]

I am wondering how I can exploit the array broadcasting rules of NumPy and functions such as tile to simplify this.

Answer 1

I think that the following should work:

pts = np.array(pts)  #Skip if pts is a numpy array already
lp = len(pts)
arr = np.zeros((lp,lp,lp,3))
arr[:,:,:,0] = pts[:,None,None]  #None is the same as np.newaxis
arr[:,:,:,1] = pts[None,:,None]
arr[:,:,:,2] = pts[None,None,:]

A quick test:

import numpy as np
import timeit

def meth1(pts):
   pts = np.array(pts)  #Skip if pts is a numpy array already
   lp = len(pts)
   arr = np.zeros((lp,lp,lp,3))
   arr[:,:,:,0] = pts[:,None,None]  #None is the same as np.newaxis
   arr[:,:,:,1] = pts[None,:,None]
   arr[:,:,:,2] = pts[None,None,:]
   return arr

def meth2(pts):
   lp = len(pts)
   N = lp
   arr = np.zeros((lp,lp,lp,3))
   for i in xrange(0, N):
      for j in xrange(0, N):
         for k in xrange(0, N):
            arr[i,j,k,0] = pts[i]
            arr[i,j,k,1] = pts[j]
            arr[i,j,k,2] = pts[k]

   return arr

pts = range(10)
a1 = meth1(pts)
a2 = meth2(pts)

print np.all(a1 == a2)

NREPEAT = 10000
print timeit.timeit('meth1(pts)','from __main__ import meth1,pts',number=NREPEAT)
print timeit.timeit('meth2(pts)','from __main__ import meth2,pts',number=NREPEAT)

results in:

True
0.873255968094   #my way
11.4249279499    #original

So this new method is an order of magnitude faster as well.