Fast interpolation over 3D array

Question

I have a 3D array that I need to interpolate over one axis (the last dimension). Let's say y.shape = (nx, ny, nz), I want to interpolate in nz for every (nx, ny). However, I want to interpolate for a different value in each [i, j].

Here's some code to exemplify. If I wanted to interpolate to a single value, say new_z, I'd use scipy.interpolate.interp1d like this

# y is a 3D ndarray
# x is a 1D ndarray with the abcissa values
# new_z is a number
f = scipy.interpolate.interp1d(x, y, axis=-1, kind='linear')
result = f(new_z)

However, for this problem what I actually want is to interpolate to a different new_z for each y[i, j]. So I do this:

# y is a 3D ndarray
# x is a 1D ndarray with the abcissa values
# new_z is a 2D array
result = numpy.empty(y.shape[:-1])
for i in range(nx):
    for j in range(ny):
        f = scipy.interpolate.interp1d(x, y[i, j], axis=-1, kind='linear')
        result[i, j] = f(new_z[i, j])

Unfortunately, with multiple loops this becomes inefficient and slow. Is there a better way to do this kind of interpolation? Linear interpolation is sufficient. A possibility is to implement this in Cython, but I was trying to avoid that because I want to have the flexibility of changing to cubic interpolation and don't want to do it by hand in Cython.

Answer 1

To speedup high order interpolate, you can call interp1d() only once, and then use the _spline attribute and the low level function _bspleval() in the _fitpack module. Here is the code:

from scipy.interpolate import interp1d
import numpy as np

nx, ny, nz = 30, 40, 50
x = np.arange(0, nz, 1.0)
y = np.random.randn(nx, ny, nz)
new_x = np.random.random_integers(1, (nz-1)*10, size=(nx, ny))/10.0

def original_interpolation(x, y, new_x):
    result = np.empty(y.shape[:-1])
    for i in xrange(nx):
        for j in xrange(ny):
            f = interp1d(x, y[i, j], axis=-1, kind=3)
            result[i, j] = f(new_x[i, j])
    return result

def fast_interpolation(x, y, new_x):
    from scipy.interpolate._fitpack import _bspleval
    f = interp1d(x, y, axis=-1, kind=3)
    xj,cvals,k = f._spline
    result = np.empty_like(new_x)
    for (i, j), value in np.ndenumerate(new_x):
        result[i, j] = _bspleval(value, x, cvals[:, i, j], k, 0)
    return result

r1 = original_interpolation(x, y, new_x)
r2 = fast_interpolation(x, y, new_x)

>>> np.allclose(r1, r2)
True

%timeit original_interpolation(x, y, new_x)
%timeit fast_interpolation(x, y, new_x)
1 loops, best of 3: 3.78 s per loop
100 loops, best of 3: 15.4 ms per loop

Answer 2

Building on @pv.'s answer, and vectorising the inner loop, the following gives a substantial speedup (EDIT: changed the expensive numpy.tile to using numpy.lib.stride_tricks.as_strided):

import numpy
from scipy import interpolate

nx = 30
ny = 40
nz = 50

y = numpy.random.randn(nx, ny, nz)
x = numpy.float64(numpy.arange(0, nz))

# We select some locations in the range [0.1, nz-0.1]
new_z = numpy.random.random_integers(1, (nz-1)*10, size=(nx, ny))/10.0

# y is a 3D ndarray
# x is a 1D ndarray with the abcissa values
# new_z is a 2D array

def original_interpolation():
    result = numpy.empty(y.shape[:-1])
    for i in range(nx):
        for j in range(ny):
            f = interpolate.interp1d(x, y[i, j], axis=-1, kind='linear')
            result[i, j] = f(new_z[i, j])

    return result

grid_x, grid_y = numpy.mgrid[0:nx, 0:ny]
def faster_interpolation():
    flat_new_z = new_z.ravel()
    k = numpy.searchsorted(x, flat_new_z)
    k = k.reshape(nx, ny)

    lower_index = [grid_x, grid_y, k-1]
    upper_index = [grid_x, grid_y, k]

    tiled_x = numpy.lib.stride_tricks.as_strided(x, shape=(nx, ny, nz), 
        strides=(0, 0, x.itemsize))

    z_upper = tiled_x[upper_index]
    z_lower = tiled_x[lower_index]

    z_step = z_upper - z_lower
    z_delta = new_z - z_lower

    y_lower = y[lower_index]
    result = y_lower + z_delta * (y[upper_index] - y_lower)/z_step

    return result

# both should be the same (giving a small difference)
print numpy.max(
        numpy.abs(original_interpolation() - faster_interpolation()))

That gives the following times on my machine:

In [8]: timeit foo.original_interpolation()
10 loops, best of 3: 102 ms per loop

In [9]: timeit foo.faster_interpolation()
1000 loops, best of 3: 564 us per loop

Going to nx = 300, ny = 300 and nz = 500, gives a 130x speedup:

In [2]: timeit original_interpolation()
1 loops, best of 3: 8.27 s per loop

In [3]: timeit faster_interpolation()
10 loops, best of 3: 60.1 ms per loop

You'd need a write your own algorithm for cubic interpolation, but it shouldn't be so hard.