[Python] Interpolate data onto multidimensional coordinate

My problem involves an xarray Dataarray but I am also interested in vectorised numpy solutions.

I have a 4 dimensional data array T with coordinates (time, lat, lon and lev). The lev coord corresponds to a model coordinate that I have converted into geometric height, and which is time and space dependent. Therefore the Data array has a coordinate 'height' with dimensions (time, lat,lon, lev).

xarray.DataArray
'T'

time: 120 lev: 32 lat: 192 lon: 288

268.5 268.5 268.5 268.5 268.5 268.5 ... 256.1 256.1 256.1 256.1 256.1
Coordinates:
lat
(lat)
float64
-90.0 -89.06 -88.12 ... 89.06 90.0
lon
(lon)
float64
0.0 1.25 2.5 ... 356.2 357.5 358.8
lev
(lev)
float64
3.643 7.595 14.36 ... 976.3 992.6
time
(time)
object
2025-02-01 00:00:00 ... 2035-01-...
height
(lev, time, lat, lon)
float64
<Quantity([[[[2.91215871e+01 2.9...
Indexes: (4)
Attributes:

mdims :
1
units :
K
long_name :
Temperature
cell_methods :
time: mean


I want to interpolate my data from 'lev' onto a given set of height levels (e.g. [1,2,3,4]). The height coordinate will then be one dimensional. However I can not work out how to implement this interpolation in an efficient way. I can only get the naive approach to work but this is of course unacceptably slow:

height_levels=[1,2,3,4]

da_arr1=[]
for i in T.time:
da_arr2=[]
for j in T.lat:
da_arr3=[]
for k in T.lon:
da=T.sel(time=i,lat=j,lon=k)
da=da.swap_dims(lev='height').interp(height_levels)
da_arr3.append(da)
da_arr2.append(xr.concat(da_arr3,'lon'))
da_arr1.append(xr.concat(da_arr2,'lat'))
result=xr.concat(da_arr1,'time'))

Continue reading...