沿 XArray 的时间维度应用函数

Question

我有一个图像堆栈存储在尺寸为时间、x、y 的 XArray DataArray 中，我想在其中沿每个像素的时间轴应用自定义函数，以便输出是尺寸为 x、y 的单个图像。

我尝试过： apply_ufunc 但该函数失败，说明我需要首先将数据加载到 RAM 中（即无法使用 Dask 数组）。理想情况下，我希望在内部将 DataArray 保留为 Dask 数组，因为不可能将整个堆栈加载到 RAM 中。确切的错误消息是：

ValueError: apply_ufunc 在参数上遇到 dask 数组，但尚未启用 dask 数组的处理。设置参数或首先使用或dask将数据加载到内存中.load().compute()

我的代码目前如下所示：

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})

res = xr.apply_ufunc(special_mean, data, input_core_dims=[["time"]], kwargs={'drop_min': True})

如果我确实使用 .compute 将数据加载到 RAM 中，那么我仍然会遇到一个错误，其中指出：

ValueError：应用函数返回的数据具有意外的维度数：0 与 2，对于维度（'y'、'x'）

我完全不确定我错过了什么/做错了什么。

Answer 1

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})

res = xr.apply_ufunc(special_mean, data, input_core_dims=[["time"]], kwargs={'drop_min': True}, dask = 'allowed', vectorize = True)

上面使用参数的代码vectorize应该可以工作。