Question

I am writing a plugin for an application that includes NumPy in the binary distribution, but not SciPy. My plugin needs to interpolate data from one regular 3D grid to another regular 3D grid. Running from source, this can be done very efficiently using scipy.ndimage or, if the user doesn’t have SciPy installed, a weave generated .pyd I’ve written. Unfortunately, neither of those options are available if the user is running the binary.

I’ve written a simple trilinear interpolation routine in python that gives the correct result, but for the array sizes I’m using, takes a long time (~5 minutes). I’m wondering if there is a way to speed it up using only the functionality within NumPy. Just like scipy.ndimage.map_coordinates, It takes a 3D input array and an array with the x, y and z coordinates of each point to be interpolated.

def trilinear_interp(input_array, indices):
    """Evaluate the input_array data at the indices given"""

    output = np.empty(indices[0].shape)
    x_indices = indices[0]
    y_indices = indices[1]
    z_indices = indices[2]
    for i in np.ndindex(x_indices.shape):
        x0 = np.floor(x_indices[i])
        y0 = np.floor(y_indices[i])
        z0 = np.floor(z_indices[i])
        x1 = x0 + 1
        y1 = y0 + 1
        z1 = z0 + 1
        #Check if xyz1 is beyond array boundary:
        if x1 == input_array.shape[0]:
            x1 = x0
        if y1 == input_array.shape[1]:
            y1 = y0
        if z1 == input_array.shape[2]:
            z1 = z0
        x = x_indices[i] - x0
        y = y_indices[i] - y0
        z = z_indices[i] - z0
        output[i] = (input_array[x0,y0,z0]*(1-x)*(1-y)*(1-z) +
                 input_array[x1,y0,z0]*x*(1-y)*(1-z) +
                 input_array[x0,y1,z0]*(1-x)*y*(1-z) +
                 input_array[x0,y0,z1]*(1-x)*(1-y)*z +
                 input_array[x1,y0,z1]*x*(1-y)*z +
                 input_array[x0,y1,z1]*(1-x)*y*z +
                 input_array[x1,y1,z0]*x*y*(1-z) +
                 input_array[x1,y1,z1]*x*y*z)

    return output

Obviously the reason the function is so slow is the for loop over each point in 3D space. Is there any way to perform some sort of slicing or vectorization magic to speed it up? Thanks.

Answer 1

It turns out it’s embarrassingly easy to vectorize it.

output = np.empty(indices[0].shape)
x_indices = indices[0]
y_indices = indices[1]
z_indices = indices[2]

x0 = x_indices.astype(np.integer)
y0 = y_indices.astype(np.integer)
z0 = z_indices.astype(np.integer)
x1 = x0 + 1
y1 = y0 + 1
z1 = z0 + 1

#Check if xyz1 is beyond array boundary:
x1[np.where(x1==input_array.shape[0])] = x0.max()
y1[np.where(y1==input_array.shape[1])] = y0.max()
z1[np.where(z1==input_array.shape[2])] = z0.max()

x = x_indices - x0
y = y_indices - y0
z = z_indices - z0
output = (input_array[x0,y0,z0]*(1-x)*(1-y)*(1-z) +
             input_array[x1,y0,z0]*x*(1-y)*(1-z) +
             input_array[x0,y1,z0]*(1-x)*y*(1-z) +
             input_array[x0,y0,z1]*(1-x)*(1-y)*z +
             input_array[x1,y0,z1]*x*(1-y)*z +
             input_array[x0,y1,z1]*(1-x)*y*z +
             input_array[x1,y1,z0]*x*y*(1-z) +
             input_array[x1,y1,z1]*x*y*z)

return output