xarray apply_ufunc新增维度实现及核心参数语义解析

一、核心需求 #

从两个具有(y, x)维度的变量，生成一个(y, x, xy)的三维变量，其中xy维度长度为2，对应原两个变量在相同(y,x)位置的元素堆叠。

二、常见认知误区 #

很多人对apply_ufunc的**核心维度（core dims）**概念理解颠倒：

• 错误认知：核心维度是需要保留、不被函数处理的维度。
• 正确逻辑：核心维度是自定义函数直接操作的维度，而非核心维度会被作为「批量维度」进行广播对齐，apply_ufunc会对每个批量维度的切片调用函数。

简单总结：

• 批量维度：apply_ufunc自动循环处理的维度（无需函数关心）
• 核心维度：函数要直接接收、处理的维度（函数内部需负责这些维度的计算）

三、关键参数语义拆解 #

1. input_core_dims #

• 格式：列表的列表，每个子列表对应一个输入变量的核心维度名称。
• 作用：告知apply_ufunc，每个输入变量中哪些维度是给函数直接处理的，剩余维度作为批量维度。

针对本需求：我们希望函数对每个(y,x)位置的标量值进行堆叠（而非处理整个(y,x)数组），因此每个输入的核心维度为空（函数接收标量），即：

input_core_dims=[[], []]

2. output_core_dims #

• 格式：列表的列表，子列表对应函数输出的核心维度名称。
• 作用：告知apply_ufunc，函数返回的数组中哪些维度是核心维度（即函数新增或处理后的维度），这些维度会和批量维度组合成最终结果的维度。

针对本需求：函数接收两个标量，返回一个长度为2的一维数组，这个新增的维度就是我们要的xy，因此输出核心维度为["xy"]，即：

output_core_dims=[["xy"]]

3. output_sizes #

• 格式：字典，键为维度名称，值为维度长度。
• 作用：当函数新增了xarray未知的维度时，需显式指定该维度的长度，否则xarray无法推断形状（尤其是使用dask并行计算时）。

针对本需求：新增的xy维度长度为2，因此：

output_sizes={"xy": 2}

无需合并原数据集的sizes，因为批量维度y和x的大小会自动从输入变量中继承。

四、修正后的完整代码 #

import numpy as np
import xarray as xr

a1 = np.array([[11,12,13,14],[15,16,17,18],[19,110,111,112],[113,114,115,116]], dtype=np.int16)
a2 = np.array([[21,22,23,24],[25,26,27,28],[29,210,211,212],[213,214,215,216]], dtype=np.int16)

ds = xr.Dataset(
    data_vars= {
        "d1": (("y", "x"), a1),
        "d2": (("y", "x"), a2) 
    }
).chunk({"y":2,"x":2})

def stackuf(xa1, xa2):
    # 此时xa1、xa2是标量，返回长度为2的数组
    return np.stack((xa1, xa2), axis=-1)

out = xr.apply_ufunc(
    stackuf,
    ds.d1,
    ds.d2,
    input_core_dims=[[], []],
    output_core_dims=[["xy"]],
    output_sizes={"xy": 2},
    dask="parallelized",
    vectorize=False,
    output_dtypes=[np.int16],
)
out.compute()

运行后out的维度为(y:4, x:4, xy:2)，完全符合需求。

五、补充验证逻辑 #

如果修改函数为接收整个(y,x)数组（而非标量），参数设置会完全不同：

• input_core_dims=[["y","x"], ["y","x"]]（核心维度是y和x，无批量维度）
• output_core_dims=[["y","x","xy"]]（函数直接返回三维数组，核心维度包含所有三个维度）
• 此时无需output_sizes，因为函数返回的数组形状可被xarray直接推断

但这种方式会失去dask并行处理的优势（无法按chunk拆分计算），因此更推荐第一种参数设置。

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