如何在xarray Dataset中按坐标逐点应用自定义函数？

假设我有一个通过以下代码生成的xarray Dataset，命名为ds：

import pandas as pd
import pymc as pm
import numpy as np

data = pd.DataFrame(dict(
    r = pm.Uniform.dist(shape=100).eval(),
    h = pm.Normal.dist(3, 0.5, shape=100).eval(),
    t = range(100)
)).set_index('t')
    
ds = data.to_xarray()

此时，t序列成为坐标，数据集包含两个数据变量（数据数组）r和h。

现在我有如下自定义函数：

def vol(r, h):
    return h * r ** 2 * np.pi

请问如何在该xarray数据集中按坐标逐点应用此函数？最终我希望得到一个与原坐标t一致的DataArray。

在原始的pandas DataFrame中，我可以通过以下代码实现相同效果：

data.apply(lambda row: vol(row.r, row.h), axis=1)

解决方案 #

方法1：直接向量化调用（最简单） #

xarray的DataArray原生支持向量化运算，直接将对应变量传入函数即可自动按t维度逐点计算：

vol_da = vol(ds['r'], ds['h'])

返回的vol_da是与原数据集t坐标完全一致的DataArray。

方法2：使用xarray.apply_ufunc（通用场景） #

如果函数本身非向量化，或需要更精细的维度控制，可使用apply_ufunc处理：

import xarray as xr

vol_da = xr.apply_ufunc(
    vol,
    ds['r'],
    ds['h'],
    vectorize=True
)

该方法会自动处理维度对齐，保留xarray的元数据信息。

方法3：用Dataset.apply对齐pandas逻辑 #

若想和pandas的apply逻辑保持一致，可指定按t维度逐行处理：

vol_da = ds.apply(lambda row: vol(row['r'], row['h']), dim='t')

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