A 1D numpy/scipy interpolation that's not quite 1D

Question

Numpy array data has .... data. Numpy array z has distances. data and z have the same shape, and each point of z is the distance where the corresponding point of data was measured. To complicate matters, users will be providing data/z arrays with 3, 4, or 5 dimensions.

I want to interpolate data to a set of distances in the 1D numpy array dists. Because of the data structure, the interpolation axis is always two axes from the end, i.e. if the arrays have 3 dimensions, interpolation axis is 0; if arrays have 4 dimensions, interpolation axis is 1, etc. Since, AFAICT, all of the numpy/scipy interpolation routines want the original distances to be given in a 1D array, interpolating data and z over dists seems to be a somewhat complex task. This is what I have:

def dist_interp(data, z, dists):
    # construct array to hold interpolation results
    num_dims = len(data.shape)
    interp_axis = num_dims-3
    interp_shape = list(data.shape)
    interp_shape[interp_axis] = dists.shape[0]
    interp_res = np.zeros(shape=interp_shape)
    # depending on usage, data could have between 3 and five dimensions.
    # add dims to generalize.  I hate doing it this way.  Must be
    # some other way.
    for n in range(num_dims, 5) :
        data = np.expand_dims(data, axis=0)
        z = np.expand_dims(z, axis=0)
        interp_res = np.expand_dims(interp_res, axis=0)
    for m in range(data.shape[0]):
        for l in range(data.shape[1]):
            for j in range(data.shape[3]):
                for i in range(data.shape[4]):
                    interp_res[m,l,:,j,i]=(
                                  np.interp(dists,z[m,l,:,j,i],
                                            data[m,l,:,j,i]))
    # now remove extra "wrapping" dimensions
    for n in range(0,5-num_dims):
        interp_res = interp_res[0]
    return(interp_res)

I think this will work, but adding and removing extra "wrapping" dummy dimensions is extremely inelegant and makes for code that is not at all compact. Any better ideas? Thanks.

Answer 1

As a general rule of thumb for this kind of situations:

1. move the axis you want to do stuff on to the end of the shape tuple
2. reshape the resulting array to be 2D
3. create your new array from this 2D one
4. undo steps 2 and 1 on the new array

For your case, it could look something like:

# Create some random test data
axis = -2
shape = np.random.randint(10, size=(5,))
data = np.random.rand(*shape)
data = np.sort(data, axis=axis)
z = np.random.rand(*shape)
dists = np.linspace(0,1, num=100)

data = np.rollaxis(data, axis, data.ndim)
new_shape = data.shape
data = data.reshape(-1, data.shape[-1])
z = np.rollaxis(z, axis, z.ndim)
z = z.reshape(-1, z.shape[-1])

out = np.empty(z.shape[:1]+dists.shape)
for o, x, f in zip(out, data, z):
    o[:] = np.interp(dists, x, f)

out = out.reshape(new_shape[:-1]+dists.shape)
out = np.rollaxis(out, -1, axis)