Efficient way to map 3D function to a meshgrid with NumPy

user_7565410 0
Viewed 613

I have a set of data values for a scalar 3D function, arranged as inputs x,y,z in an array of shape (n,3) and the function values f(x,y,z) in an array of shape (n,).

EDIT: For instance, consider the following simple function

data = np.array([np.arange(n)]*3).T
F = np.linalg.norm(data,axis=1)**2

I would like to convolve this function with a spherical kernel in order to perform a 3D smoothing. The easiest way I found to perform this is to map the function values in a 3D spatial grid and then apply a 3D convolution with the kernel I want.

This works fine, however the part that maps the 3D function to the 3D grid is very slow, as I did not find a way to do it with NumPy only. The code below is my actual implementation, where data is the (n,3) array containing the 3D positions in which the function is evaluated, F is the (n,) array containing the corresponding values of the function and M is the (N,N,N) array that contains the 3D space grid.

step = 0.1

# Create meshgrid
xmin = data[:,0].min()
xmax = data[:,0].max()
ymin = data[:,1].min()
ymax = data[:,1].max()
zmin = data[:,2].min()
zmax = data[:,2].max()

x = np.linspace(xmin,xmax,int((xmax-xmin)/step)+1)
y = np.linspace(ymin,ymax,int((ymax-ymin)/step)+1)
z = np.linspace(zmin,zmax,int((zmax-zmin)/step)+1)


# Build image
M = np.zeros((len(x),len(y),len(z)))

for l in range(len(data)):
    for i in range(len(x)-1):
        if x[i] < data[l,0] < x[i+1]:
            for j in range(len(y)-1):
                if y[j] < data[l,1] < y[j+1]:
                    for k in range(len(z)-1):
                        if z[k] < data[l,2] < z[k+1]:
                            M[i,j,k] = F[l]

Is there a more efficient way to fill a 3D spatial grid with the values of a 3D function ?

numpy python performance
3 Answers

For each item of data you're scanning pixels of cuboid to check if it's inside. There is an option to skip this scan. You could calculate corresponding indices of these pixels by yourself, for example:

data = np.array([[1, 2, 3], #14 (corner1)
                 [4, 5, 6], #77 (corner2)
                 [2.5, 3.5, 4.5], #38.75 (duplicated pixel)
                 [2.9, 3.9, 4.9], #47.63 (duplicated pixel)
                 [1.5, 2, 3]]) #15.25 (one step up from [1, 2, 3])
step = 0.5                            
data_idx = ((data - data.min(axis=0))//step).astype(int)
M = np.zeros(np.max(data_idx, axis=0) + 1)
x, y, z = data_idx.T
M[x, y, z] = F

Note that only one value of duplicated pixels is being mapped to M.

user_30448250

All you need is just reshape F[:, 3] (only f(x, y, z)) into a grid. Hard to be more precise without sample data:

If the data is not sorted, you need to sort it:

F_sorted = F[np.lexsort((F[:,0], F[:,1], F[:,2]))]  # sort by x, then y, then z

Choose only f(x, y, z)

F_values = F_sorted[:, 3]

Finally, reshape data into a grid:

M = F_sorted.reshape(N, N, N)
user_46018900

This method is faster than the original (approximatly 20x speed up):

step = 0.1
mins = np.min(data, axis=0)
maxs = np.max(data, axis=0)
ranges = np.floor((maxs - mins) / step + 1).astype(int)
indx = np.zeros(data.shape,dtype=int)
for i in range(3):
    x = np.linspace(mins[i], maxs[i], ranges[i])
    indx[:,i] = np.argmax(data[:,i,np.newaxis] <= (x[np.newaxis,:]), axis=1) -1
    
M = np.zeros(ranges)
M[indx[:,0],indx[:,1],indx[:,2]] = F

The first part sets up the required grid variables. The argmax function provides a simple (and fast) way to find the first true value of the broadcasted array. This produces a set of indices for x, y and z directions for each of the function values.

The resulting array M is not the same as that produced by the original code as the original code loses data. The logic of y[j] < data[l,1] < y[j+1] where y is a vector produced using linspace means the minimum and maximum values for each direction will be missed (data[l,1] might be equal to either y[j] or y[j+1]!). Run it with a dataset of two values each with their own coordinates and the M array will be all zeros.

user_149009350
Related