How can I efficiently vectorize interpolation from variable model levels to fixed height levels in a 4D xarray DataArray?

Below is one vectorized approach that uses xarray’s ability to apply a function over entire dimensions via apply_ufunc. The idea is to write a helper function (using, for example, numpy’s interp) that performs 1D interpolation along the vertical profile (the lev dimension) and then have xarray apply it over the (time, lat, lon) dimensions.

Assuming that for each (time, lat, lon) point your “height” coordinate (the geometric height corresponding to lev) is monotonic, you can do the following:

1. Write a 1D Interpolation Function

We define a function that takes a vector of height values, the corresponding data values, and returns interpolated values at the desired target heights. For example:

PYTHONCopied
import numpy as np

def interp_profile(h, data, new_h):
    # h and data must be 1D arrays with the same length
    # np.interp requires h to be sorted in increasing order.
    return np.interp(new_h, h, data)

2. Use xr.apply_ufunc to Vectorize the Interpolation

We then use xarray’s apply_ufunc to apply this function over the lev dimension (which is indicated by the core dimension “lev” in both the coordinate and the data variable). In addition, we add the target height levels via the kwargs:

PYTHONCopied
import xarray as xr

# Your desired new height levels:
height_levels = [1, 2, 3, 4]

# Apply our 1D interpolation function over the lev dimension.
# Here, we assume:
#   - T is the DataArray with dimensions (time, lat, lon, lev)
#   - The 'height' coordinate is stored as a variable named 'height' that has dims (time, lat, lon, lev)
# We use 'input_core_dims' to tell xarray which dimension our function operates over.
result = xr.apply_ufunc(
    interp_profile,                # the function to apply
    T['height'],                   # first input: the array of height values
    T,                             # second input: the data array to interpolate
    input_core_dims=[['lev'], ['lev']],
    output_core_dims=[['height']],
    vectorize=True,                # automatically vectorizes over (time, lat, lon)
    kwargs={'new_h': height_levels},
    dask='parallelized',           # if T is a dask array; otherwise this can be omitted
    output_dtypes=[T.dtype]
)

3. Explanation

• input_core_dims
  Both the height coordinate and the data array have “lev” as the core dimension that the 1D interpolation works over.

• output_core_dims
  We want the new vertical dimension to be named “height” (with the target height level values).

• vectorize=True
  This tells xarray to apply the function to every (time, lat, lon) slice without an explicit Python loop.

• kwargs={'new_h': height_levels}
  Passes the new height levels to our interpolation function.

4. Final Result

The returned DataArray “result” will have dimensions (time, lat, lon, height) where the “height” coordinate is now 1D and corresponds to [1, 2, 3, 4].

This method bypasses the need for nested loops and leverages xarray’s vectorized operations to perform the interpolation in an efficient manner.