沿 XArray 的时间维度应用函数

Question

我在 XArray DataArray 中存储了一个图像堆栈，其维度为时间、x、y，我想在其上沿每个像素的时间轴应用自定义函数，以便输出为维度为单个图像x,y.

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

ValueError: apply_ufunc encountered a dask array on an argument, but handling for dask arrays has not been enabled. Either set the dask argument or load your data into memory first with .load() or .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: applied function returned data with unexpected number of dimensions: 0 vs 2, for dimensions ('y', 'x')

我不太确定我是什么missing/doing错了。

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 参数的代码应该可以工作。

Answer 2

我的目标也是从 Xarray 实现 apply_ufunc，这样它就可以计算 x 和 y 的特殊平均值。

我喜欢 Ales 的例子；当然是通过省略与块相关的行。否则：

ValueError: applied function returned data with unexpected number of dimensions. Received 0 dimension(s) but expected 2 dimensions with names: ('y', 'x')

沿 XArray 的时间维度应用函数

Apply function along time dimension of XArray

dask

python-xarray