How to interpolate two 3D arrays by 1D array?

[zxdawn]

Hi, I have downloaded model-level data and want to convert it to pressure level by myself. Here's the example:

import xarray as xr
import numpy as np
from scipy.interpolate import interp1d

x = 10
y = 10
z = 4
values = np.arange(x*y*z).reshape((z, y, x))

#  data
model_data = xr.DataArray(values,
                    coords={'x': np.arange(x),
                            'y': np.arange(y),
                            'z': np.arange(z)},
                    dims=['z', 'y', 'x']
                    )

# corresponding pressure data
pressure = xr.DataArray(model_data+1,
                        coords={'x': np.arange(x),
                                'y': np.arange(y),
                                'z': np.arange(z)},
                        dims=['z', 'y', 'x']
                        )

# pressure levels which data are interpolated to
interp_p = np.array([50, 100])


def interp1d_np(data, x, xi):
    f = interp1d(x, data, kind='linear', fill_value='extrapolate')
    return f(xi)

xr.apply_ufunc(
    interp1d_np,
    model_data,
    pressure,
    interp_p,
    input_core_dims=[["z"], ["z"], []],
    exclude_dims=set(("z",)),
    # vectorize=True,
)

I got this error:

-> 1204     return apply_dataarray_vfunc(
   1205         variables_vfunc,
   1206         *args,
   1207         signature=signature,
   1208         join=join,
   1209         exclude_dims=exclude_dims,
   1210         keep_attrs=keep_attrs,
...
    125                      "interpolation axis.")
    127 self._y_axis = (axis % yi.ndim)
    128 self._y_extra_shape = yi.shape[:self._y_axis]+yi.shape[self._y_axis+1:]

ValueError: x and y arrays must be equal in length along interpolation axis.

It seems I have to extend the 1D pressure data to the same 3D data. Is there a more efficient way? Thanks!

[dcherian]

You need vectorize=True with numpy's interp. See https://tutorial.xarray.dev/advanced/apply_ufunc/vectorize_1d.html

[zxdawn]

@dcherian Thanks, but vectorize=True raised this error

ValueError: shape mismatch: objects cannot be broadcast to a single shape.  Mismatch is between arg 0 with shape (10, 10) and arg 2 with shape (2,).

[zxdawn]

Aha, I have to change the dimension name:

    input_core_dims=[["z"], ["z"], ['new_z']],
    output_core_dims=[["new_z"]],