如何基于坐标值在xarray维度上优雅应用函数？能否用apply_ufunc实现？

可以用xarray.apply_ufunc直接实现这个需求，无需转成numpy再重构。核心是通过指定核心维度和启用向量化，让函数自动遍历目标维度并使用对应坐标作为参数。

实现代码 #

import xarray as xr
import numpy as np

# 创建测试数据
data = xr.DataArray(np.random.randn(4, 3), dims=("x", "y"), coords={"x": [0, 1, 2, 3]})
data = data.assign_coords(shift=('x', [1,1,2,0]))

# 定义单组数据的处理函数：接收y维度数组和对应shift值，返回滚动结果
def roll_with_shift(arr, shift_val):
    return np.roll(arr, shift_val)

# 用apply_ufunc执行映射
result = xr.apply_ufunc(
    roll_with_shift,
    data,
    data.shift,
    input_core_dims=[['y'], []],  # 第一个输入的核心处理维度是y，第二个输入为标量（无核心维度）
    vectorize=True,  # 自动遍历非核心维度（x），逐个处理每组数据
    output_core_dims=[['y']],  # 输出保留y维度
    dask="allowed"  # 可选：支持dask数组并行处理
)

print(result)

结果说明 #

运行后得到的结果和你之前的变通方案完全一致：

<xarray.DataArray (x: 4, y: 3)>
array([[-2.30255789,  1.66474014,  0.59244195],
       [ 1.21868344, -0.0252584 ,  0.24441252],
       [ 0.03254358,  0.3009372 ,  1.10716318],
       [ 0.15047633, -1.22922668,  2.56393195]])
Coordinates:
  * x        (x) int64 0 1 2 3
    shift    (x) int64 1 1 2 0
Dimensions without coordinates: y

关键参数解释 #

• input_core_dims：明确每个输入的核心处理维度，告诉apply_ufunc需要将哪些维度作为一个整体传入函数。这里把y作为核心维度，意味着函数会接收每个x对应的一维y数组，以及对应的shift标量。
• vectorize=True：启用向量化遍历，自动循环处理x维度的每个元素，无需手动写循环。
• output_core_dims：指定输出的核心维度，确保结果的维度结构和原数据一致。

内容的提问来源于stack exchange，提问作者efpeel