如何基于XArray Dataset的坐标值应用函数创建数据数组？

XArray中没有直接命名为apply*的函数，但可以通过矢量化运算或xr.apply_ufunc工具实现对坐标的函数映射，以下是具体实现方式：

1. 矢量化运算（推荐） #

利用XArray坐标本身的DataArray特性，结合numpy函数直接进行广播运算，效率最高：

import xarray as xr
import numpy as np

# 初始化Dataset
ds = xr.Dataset(
    coords={
        'x': [1.5, 4, 2],
        'y': [2, -3, -1],
    },
)

# 先将x、y扩展为不同维度，实现笛卡尔积式的组合计算
x_2d = ds.x.expand_dims(y=ds.y)
y_2d = ds.y.expand_dims(x=ds.x)

# 应用函数计算（注意修正你示例中的笔误：应该是x²+y²而非x²+x²）
r = np.sqrt(x_2d**2 + y_2d**2)

运行后r会生成一个二维数据数组，对应所有x与y组合的计算结果：

<xarray.DataArray (x: 3, y: 3)>
array([[2.5       , 3.90512484, 1.80277564],
       [4.47213595, 5.        , 4.12310563],
       [2.82842712, 3.60555128, 2.23606798]])
Coordinates:
  * x        (x) float64 1.5 4.0 2.0
  * y        (y) int64 2 -3 -1

2. 使用xr.apply_ufunc处理复杂函数 #

如果需要应用更复杂的自定义函数，可以用xr.apply_ufunc，它能自动处理维度对齐与广播：

def calculate_distance(x, y):
    return np.sqrt(x**2 + y**2)

# 扩展维度后传入函数
r = xr.apply_ufunc(
    calculate_distance,
    ds.x.expand_dims(y=ds.y),
    ds.y.expand_dims(x=ds.x)
)

关键说明 #

• XArray优先推荐矢量化运算，避免循环式的apply，因为矢量化运算基于numpy，性能远高于循环。
• 你示例中的lambda表达式存在笔误，coords.x**2 + coords.x**2应改为coords.x**2 + coords.y**2，否则无法得到预期的距离计算结果。

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