Best way to interpolate a numpy.ndarray along an axis

Question

I have 4-dimensional data, say for the temperature, in an numpy.ndarray. The shape of the array is (ntime, nheight_in, nlat, nlon).

I have corresponding 1D arrays for each of the dimensions that tell me which time, height, latitude, and longitude a certain value corresponds to, for this example I need height_in giving the height in metres.

Now I need to bring it onto a different height dimension, height_out, with a different length.

The following seems to do what I want:

ntime, nheight_in, nlat, nlon = t_in.shape

nheight_out = len(height_out)
t_out = np.empty((ntime, nheight_out, nlat, nlon))

for time in range(ntime):
    for lat in range(nlat):
        for lon in range(nlon):
            t_out[time, :, lat, lon] = np.interp(
                height_out, height_in, t[time, :, lat, lon]
            )

But with 3 nested loops, and lots of switching between python and numpy, I don't think this is the best way to do it.

Any suggestions on how to improve this? Thanks

Answer 1

scipy's interp1d can help:

import numpy as np
from scipy.interpolate import interp1d

ntime, nheight_in, nlat, nlon = (10, 20, 30, 40)

heights = np.linspace(0, 1, nheight_in)

t_in = np.random.normal(size=(ntime, nheight_in, nlat, nlon))
f_out = interp1d(heights, t_in, axis=1)

nheight_out = 50
new_heights = np.linspace(0, 1, nheight_out)
t_out = f_out(new_heights)