11

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.

Stephen Terry
  • 6,169
  • 1
  • 28
  • 31

2 Answers2

8

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
Stephen Terry
  • 6,169
  • 1
  • 28
  • 31
4

Thanks very much for this post, and for having followed it up. I have based myself liberally on your vectorisation in order to give it another speed boost (at least with the data I'm working with)!

I'm working with image correlation, and therefore I to interpolate many sets of different coordinates in the same input_array.

Unfortunately I have made it a little more complicated, but if I can explain what I have done, the extra complication should a) justify itself and b) become clear. Your last line ( output = ) still requires a fair amount of looking up in non-sequential places in the input_array, making it relatively slow.

Suppose my 3D data is NxMxP long. I have decided to do the following thing: If I can get an ( 8 x (NxMxP) ) matrix of pre-computed greyvalues for a point and its nearest neighbours, and I can also calculate a ( (NxMxP) X 8 ) matrix of coefficients ( your first coefficient in the example above is (x-1)(y-1)(z-1) ) then I can just multiply then together and be home free!

A good bonus for me is that I can precalculate the grey matrix, and recycle it!

Here is a sample bit of code (pasted from two different functions, so might not work out of the box, but should serve as a good source of inspiration):

def trilinear_interpolator_speedup( input_array, coords ):
  input_array_precut_2x2x2 = numpy.zeros( (input_array.shape[0]-1, input_array.shape[1]-1, input_array.shape[2]-1, 8 ), dtype=DATA_DTYPE )
  input_array_precut_2x2x2[ :, :, :, 0 ] = input_array[ 0:new_dimension-1, 0:new_dimension-1, 0:new_dimension-1 ]
  input_array_precut_2x2x2[ :, :, :, 1 ] = input_array[ 1:new_dimension  , 0:new_dimension-1, 0:new_dimension-1 ]
  input_array_precut_2x2x2[ :, :, :, 2 ] = input_array[ 0:new_dimension-1, 1:new_dimension  , 0:new_dimension-1 ]
  input_array_precut_2x2x2[ :, :, :, 3 ] = input_array[ 0:new_dimension-1, 0:new_dimension-1, 1:new_dimension   ]
  input_array_precut_2x2x2[ :, :, :, 4 ] = input_array[ 1:new_dimension  , 0:new_dimension-1, 1:new_dimension   ]
  input_array_precut_2x2x2[ :, :, :, 5 ] = input_array[ 0:new_dimension-1, 1:new_dimension  , 1:new_dimension   ]
  input_array_precut_2x2x2[ :, :, :, 6 ] = input_array[ 1:new_dimension  , 1:new_dimension  , 0:new_dimension-1 ]
  input_array_precut_2x2x2[ :, :, :, 7 ] = input_array[ 1:new_dimension  , 1:new_dimension  , 1:new_dimension   ] 
  # adapted from from http://stackoverflow.com/questions/6427276/3d-interpolation-of-numpy-arrays-without-scipy
  # 2012.03.02 - heavy modifications, to vectorise the final calculation... it is now superfast.
  #  - the checks are now removed in order to go faster...

  # IMPORTANT: Input array is a pre-split, 8xNxMxO array.

  # input coords could contain indexes at non-integer values (it's kind of the idea), whereas the coords_0 and coords_1 are integer values.
  if coords.max() > min(input_array.shape[0:3])-1  or coords.min() < 0:
    # do some checks to bring back the extremeties
    # Could check each parameter in x y and z separately, but I know I get cubic data...
    coords[numpy.where(coords>min(input_array.shape[0:3])-1)] = min(input_array.shape[0:3])-1
    coords[numpy.where(coords<0                      )] = 0              

  # for NxNxN data, coords[0].shape = N^3
  output_array = numpy.zeros( coords[0].shape, dtype=DATA_DTYPE )

  # a big array to hold all the coefficients for the trilinear interpolation
  all_coeffs = numpy.zeros( (8,coords.shape[1]), dtype=DATA_DTYPE )

  # the "floored" coordinates x, y, z
  coords_0 = coords.astype(numpy.integer)                  

  # all the above + 1 - these define the top left and bottom right (highest and lowest coordinates)
  coords_1 = coords_0 + 1

  # make the input coordinates "local"
  coords = coords - coords_0

  # Calculate one minus these values, in order to be able to do a one-shot calculation
  #   of the coefficients.
  one_minus_coords = 1 - coords

  # calculate those coefficients.
  all_coeffs[0] = (one_minus_coords[0])*(one_minus_coords[1])*(one_minus_coords[2])
  all_coeffs[1] =      (coords[0])     *(one_minus_coords[1])*(one_minus_coords[2])
  all_coeffs[2] = (one_minus_coords[0])*    (coords[1])      *(one_minus_coords[2])
  all_coeffs[3] = (one_minus_coords[0])*(one_minus_coords[1])*     (coords[2])
  all_coeffs[4] =      (coords[0])     *(one_minus_coords[1])*     (coords[2])      
  all_coeffs[5] = (one_minus_coords[0])*     (coords[1])     *     (coords[2])
  all_coeffs[6] =      (coords[0])     *     (coords[1])     *(one_minus_coords[2])
  all_coeffs[7] =      (coords[0])     *     (coords[1])     *     (coords[2])

  # multiply 8 greyscale values * 8 coefficients, and sum them across the "8 coefficients" direction
  output_array = (  input_array[ coords_0[0], coords_0[1], coords_0[2] ].T * all_coeffs ).sum( axis=0 )

  # and return it...
  return output_array

I didn't split x y and z coordinates as above because it didn't seem that useful to remerge them afterwards. There might be something in the above code which assumed cubic data (N=M=P) but I don't think so...

Let me know what you think!

EdwardAndo
  • 333
  • 1
  • 2
  • 7