如何在Xarray中对Dask数组沿维度执行运算且无需回写数据

解决方案 #

Dask数组属于延迟计算的不可变结构，不支持inplace的切片赋值，你可以通过Xarray的apply_ufunc接口实现自定义函数的维度级应用，同时完美兼容Dask的延迟计算逻辑。

完整实现代码 #

import xarray as xr
import numpy as np

# 你的自定义归一化函数
def foo(arr):
    lo, hi = np.percentile(arr, (1,99))
    arr = np.clip(arr, lo, hi)
    arr = (arr - lo) / (hi - lo)
    return arr

# 对time维度的每个二维切片应用foo函数
scaled_data = xr.apply_ufunc(
    foo,
    data,  # 你的xarray DataArray对象
    input_core_dims=[["x", "y"]],  # 每个输入块保留的核心维度
    output_core_dims=[["x", "y"]],  # 每个输出块的核心维度
    vectorize=True,  # 沿非核心维度（即time）自动遍历
    dask="parallelized",  # 启用dask并行计算
    output_dtypes=[np.float64]  # 声明输出数据类型
)

可选替代方案：groupby实现 #

如果你更熟悉groupby语法，也可以按time维度分组后应用函数，写法更直观，但超大分块数据集下性能略低于apply_ufunc方案：

scaled_data = data.groupby("time").map(lambda x: foo(x))

注意事项 #

• 处理后的结果scaled_data仍然是Dask延迟计算的数组，你可以在所有计算逻辑定义完成后调用.compute()方法触发实际计算，或者直接用to_netcdf()等方法写入磁盘，Xarray会自动处理Dask的调度逻辑。
• 如果你的数据分块是沿time维度拆分的，不需要做任何分块调整即可直接运行，无需把全量数据加载到内存。

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