Vectorising nested for loops with loop dependency in NumPy

Viewed 334

I have the following function, which generates a series of grid points on a tetrahedron.

def tet_grid(n):

    xv = np.array([
        [-1.,-1.,-1.],
        [ 1.,-1.,-1.],
        [-1., 1.,-1.],
        [-1.,-1., 1.],
        ])

    nsize = int((n+1)*(n+2)*(n+3)/6);
    xg = np.zeros((nsize,3))
    p = 0

    for i in range ( 0, n + 1 ):
        for j in range ( 0, n + 1 - i ):
            for k in range ( 0, n + 1 - i - j ):
                l = n - i - j - k
                xg[p,0]=(i * xv[0,0] + j * xv[1,0] + k * xv[2,0] + l * xv[3,0])/n 
                xg[p,1]=(i * xv[0,1] + j * xv[1,1] + k * xv[2,1] + l * xv[3,1])/n 
                xg[p,2]=(i * xv[0,2] + j * xv[1,2] + k * xv[2,2] + l * xv[3,2])/n 
                p = p + 1

    return xg

Is there an easy way to vectorise this in NumPy?

2 Answers
Related