How to leverage xarray in unique interp use case

[marstonsward]

import numpy as np
import xarray as xr

def create_da(arr, time, lev, lat, lon, units, description):
    ts = xr.DataArray(arr,
        coords={
            "time": time,
            "lev": xr.DataArray(lev, dims=('lev',), attrs={'units': 'Pa', 'axis': 'Z'}),
            "lat": xr.DataArray(lat, dims=('lat',), attrs={'units': 'degrees', 'axis': 'Y'}),
            "lon": xr.DataArray(lon, dims=('lon',), attrs={'units': 'degrees', 'axis': 'X'}),
        },
    )
    ts.attrs['units'] = units
    ts.attrs['description'] = description
    return ts

levs = np.array([101300, 100000,  97500,  95000,  92500,  90000,  87500,  85000,
                         82500,  80000,  77500,  75000,  72500,  70000,  65000,  60000,
                         55000,  50000,  45000,  40000,  35000,  30000,  25000,  20000,
                         15000,  10000,   7000,   5000,   3000,   2000,   1000])

msllevs = np.arange(-1, 80, 0.5)
res = 0.25
lon = np.arange(-180, 179.75+res, res)
lat = np.arange(90, -90-res, -res)
temp = np.random.uniform(low=210, high=310, size=(1, 31, 721, 1440))  
gh = np.random.uniform(low=-630, high=71000, size=(1, 31, 721, 1440))
time = pd.date_range(start='1/1/2018', periods=1, freq='H')
    
gtmp = create_da(temp, time, levs, lat, lon, "temp", "K", "Temperature")
gmh = create_da(gh, time, levs, lat, lon, 'gmh', 'm', "geometric height")

I would like to interpolate gtemp using gmh from being expressed on pressure levels to MSL levels msllevs, i.e., change the vertical coordinate. In a brute force kind of what of a way, one could do this:

mslArr = np.zeros( (nlev, nlat, nlon), dtype=np.float32 )
for i in range(nlat):
    for j in range(nlon):
        iGrid = gmh[:, i, j] 
        iValues = iarr[:, i, j]
        if all( [np.isnan(i) for i in iValues] ): 
            mslArr[:, i, j] = np.nan
            continue                            
        else:
            f = interpolate.interp1d(iGrid, iValues, kind='linear', bounds_error=False)
            nVals = f(newz)
            mslArr[:, i, j] = nVals

I've been looking for a way to do this in a very fast way, maybe with dask, but none of the examples I've seen provide a clear way to achieve this with xarray. Is this use case easily done with xarray?

[dcherian]

Does https://xgcm.readthedocs.io/en/latest/transform.html help?

[marstonsward]

This looks like a good tool. I've never seen it before, though I do recognize some of the models.
I gave it a try but there was an error in the results. I'll move there to try and sort it out.
Thanks for the tip!

[TomNicholas]

Please raise an issue on the xgcm repo if you have any problems!