在 Xarray DataArray 上使用 sp.ndimage.label 和 apply_ufunc

Question

假设我有一个简单的数组：

data = [1,1,0,0,1,1,1]

我可以使用 scipy ndimage 模块对这些数据进行标记：

groups, _ = sp.ndimage.label(data)

导致

In [68]: print(groups)
[1 1 0 0 2 2 2]

现在，我想在 xarray DataArray 上执行相同的标记功能。

xr_data = xr.DataArray([1,1,0,0,1,1,1], coords = [("x", [0,1,2,3,4,5,6])])

我知道我可以在 xr_data 上调用与以前相同的函数，但执行此调用的输出是一个 numpy 数组，在我的实际数据集中，它太大而无法放入内存。

看来xr.apply_ufunc 函数是我需要的。但是，我无法让它工作。

def xr_label(arr):
    return xr.apply_ufunc(sp.ndimage.label, arr)

xr_groups, _ = xr_label(xr_data)

这会导致： “ValueError：应用函数返回的数据包含意外的维度数。收到 0 个维度，但预期 1 个维度的名称为：('x',)”

我发现有关 apply_ufunc 方法的文档难以解释。有人可以帮我解决这个问题吗？

Answer 1

您必须将input_core_dims 和output_core_dims 定义为apply_ufunc 的参数。请参阅文档：http://xarray.pydata.org/en/stable/generated/xarray.apply_ufunc.html

在你的情况下，我认为这将是：

xr.apply_ufunc(sp.ndimage.label, arr, input_core_dims=[['x']], output_core_dims=[['x']])

我最近也在努力理解apply_ufunc（公平地说，我仍然没有完全理解），但是http://xarray.pydata.org/en/stable/examples/apply_ufunc_vectorize_1d.html 的示例对我帮助很大。