如何根据数据暗淡使用 apply_ufunc？

Question

我想用一个取决于时间和经度的二维数组 newlat (time, lon) 插入 3D 数组 air (time, lat, lon)。

For循环方法

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

air = (
    xr.tutorial.load_dataset("air_temperature")
    .air.sortby("lat")  # np.interp needs coordinate in ascending order
    .isel(time=slice(4), lon=slice(3))
)  # choose a small subset for convenience


newlat = xr.DataArray(np.random.rand(air.sizes['time'], air.sizes['lon'])*75,
                      dims=['time', 'lon'],
                      coords={'time': air.time, 'lon': air.lon}
)


# create empty array to save result
result = np.empty((4, 3))

# loop each dim
for t in range(air.sizes['time']):
    for lon in range(air.sizes['lon']):
        # interpolation relying on time and lon
        f = interp1d(air.lat, air.isel(time=t, lon=lon), kind='linear', fill_value='extrapolate')
        result[t, lon] = f(newlat.isel(time=t, lon=lon))

apply_ufunc方法

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

t_index = 0
lon_index = 0

xr.apply_ufunc(
    interp1d_np,
    air.isel(time=t_index, lon=lon_index),
    air.lat,
    newlat.isel(time=t_index, lon=lon_index),
    input_core_dims=[["lat"], ["lat"], []],
    exclude_dims=set(("lat",)),
    vectorize=True,
)

请注意，t_idnex 和 lon_index 对于输入 air 和 newlat 是相同的。上面的代码仅适用于 air 的一个特定部分。如何将它应用于整个 air DataArray？

临时解决方案

我们可以像这样使用嵌入式函数：

air.interp(lat=newlat, kwargs={"fill_value": None})

但是，我还是很好奇在这种情况下如何使用apply_ufunc，因为用户可能有自己的功能而不是简单的插值。

Answer 1

你基本上就在那里。删除 .isel 调用，它将起作用！ vectorize=True 将自动遍历“非核心维度”，即本例中的 time 和 lon。

xr.apply_ufunc(
    interp1d_np,
    air,
    air.lat,
    newlat,
    input_core_dims=[["lat"], ["lat"], []],
    exclude_dims=set(("lat",)),
    vectorize=True,
)

如何根据数据暗淡使用 apply_ufunc？

How to use apply_ufunc depending on data dims?

python

interpolation

python-xarray

For循环方法

apply_ufunc方法

临时解决方案