interpolate time series of 2d gridded data to a new 2d grid - efficiently

Question

I have an numpy array of 2D gridded data with a temporal axis, so my array has a shape of (nsteps, ny, nx)

For context it's a grid of the Earth (x = 0 to 360 in 1 deg increments and y = -90 to 90 in 1 deg increments). And there's timeseries of data at each point on this grid, 1 array per timestep.

I'm trying to interpolate data from this grid to a very slightly different grid (different resolution and thus node points).

I was able to do this fine via:

import numpy as np
from scipy.interpolate import RectBivariateSpline
from tqdm import tqdm

#interpolate the msl pressures back to the other grid
p_dat = pres[0].values #this is a numpy array (nsteps, ny0, nx0)
w_dat = wind[0].values #this is a numpy array (nsteps, ny1, nx1)

#new array with 2d shape of w_dat
out = np.full((p_dat.shape[0], *w_dat.shape[1:]), np.nan) 

#interpolate one timestep at a time
for i in tqdm(list(range(out.shape[0]))):
    interp = RectBivariateSpline(pres.geometry.y, pres.geometry.x, p_dat[i])
    dat = interp(wind.geometry.y, wind.geometry.x)
    out[i,:,:] = dat

Is there a way I can avoid this loop and vectorize this interpolation over the 0th axis?

Answer 1

Do you have any data you could add to your question? That would enable us to have a closer look.

I think np.apply_along_axis might work here.

Right now i can only provide this untested approach:

def helper(idx: int):
    interp = RectBivariateSpline(pres.geometry.y, pres.geometry.x, p_dat[idx])
    dat = interp(wind.geometry.y, wind.geometry.x)
    return dat

out = np.apply_along_axis(helper, axis=0, arr=np.arange(out.shape[0]))

Edit: I can't comment on speedups due to the lack of test data, but this is very similar to map(helper, np.arange(out.shape[0])) . In essence, it's just hiding the loop.