XArray时间维度应用自定义函数报错，如何保留Dask Array解决？

让我们一步步拆解你遇到的两个问题，然后给出针对性的修正方案：

第一个错误：Dask数组未启用处理的问题 #

你已经将数据分块为Dask数组（data.chunk(...)），但apply_ufunc默认不会自动处理Dask分块数据，因此需要显式开启Dask支持。你需要在apply_ufunc中添加两个关键参数：

• dask='parallelized'：允许函数在Dask分块上并行执行
• output_dtypes=[float]：告诉Dask输出的数据类型（你的函数返回浮点数，所以指定为float）

第二个错误：维度不匹配的问题 #

当你把数据加载到内存后，报错的核心原因是apply_ufunc没有明确知晓你的函数返回结果的维度结构。你的special_mean函数接收一个一维的时间序列（每个像素的所有时间点数据），返回一个标量值，但apply_ufunc需要你明确指定输出的核心维度——这里因为每个时间轴处理后返回0维的标量，所以要设置output_core_dims=[[]]，告诉它处理后去掉time维度，保留y和x维度。

修正后的完整代码 #

情况1：使用Dask分块处理（不加载到内存） #

import numpy as np
import xarray as xr
import pandas as pd

def special_mean(x, drop_min=False):
    s = np.sum(x)
    n = len(x)
    if drop_min:
        s = s - x.min()
        n -= 1
    return s/n

times = pd.date_range('2019-01-01', '2019-01-10', name='time')
data = xr.DataArray(np.random.rand(10, 8, 8), dims=["time", "y", "x"], coords={'time': times})
data = data.chunk({'time':10, 'x':1, 'y':1})

# 添加必要参数修正Dask和维度问题
res = xr.apply_ufunc(
    special_mean,
    data,
    input_core_dims=[["time"]],
    output_core_dims=[[]],  # 声明每个time序列处理后返回0维标量
    kwargs={'drop_min': True},
    dask='parallelized',    # 启用Dask并行处理
    output_dtypes=[float]   # 指定输出数据类型
)

# 验证结果维度，应为('y', 'x')
print(res.dims)

情况2：加载到内存后处理 #

如果你想先将数据加载到RAM，只需去掉Dask相关参数，但保留output_core_dims=[[]]：

data_loaded = data.load()  # 加载数据到内存
res_loaded = xr.apply_ufunc(
    special_mean,
    data_loaded,
    input_core_dims=[["time"]],
    output_core_dims=[[]],
    kwargs={'drop_min': True}
)

# 验证结果维度
print(res_loaded.dims)

额外说明 #

output_core_dims的作用是明确告知apply_ufunc函数处理后的维度变化：input_core_dims=[["time"]]指定要沿time轴逐个处理每个(y,x)位置的时间序列，而output_core_dims=[[]]则说明每个处理返回无维度的标量，最终组合成(y,x)维度的结果。如果不指定这个参数，apply_ufunc会误解输出的维度结构，从而触发维度不匹配的错误。

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