使用 Xarray 的 apply_ufunc 在 3D 数组的时间维度上应用 Numba guvectorize 函数

Question

我在让它正常工作时遇到了一些问题，我也愿意接受其他建议，因为我不能 100% 确定我是否走对了这条路。

这是一些简单的虚拟数据：

times = pd.date_range(start='2012-01-01',freq='1W',periods=25)
x = np.array([range(0,20)]).squeeze()
y = np.array([range(20,40)]).squeeze()

data = np.random.randint(3, size=(25,20,20))

ds = xr.DataArray(data, dims=['time', 'y', 'x'], coords = {'time': times, 'y': y, 'x': x})

对于每个 x,y 坐标，我想随时间返回最长的 1 或 2 序列。所以我的输入数组是 3D（时间，x，y），我的输出是 2D（x，y）。 'seq_gufunc' 中的代码受到this thread 的启发。 我的实际数据集要大得多（使用土地利用类而不是 1s、2s 等），这只是更大工作流程的一小部分，我还使用 dask 进行并行处理。所以最后这应该快速有效地运行，这就是为什么我最终试图弄清楚如何让 numba 的 @guvectorize 和 Xarray 的 apply_ufunc 一起工作：

@guvectorize(
    "(int64[:], int64[:])",
    "(n) -> (n)", target='parallel', nopython=True
)
def seq_gufunc(x, out):

    f_arr = np.array([False])

    bool_stack = np.hstack((f_arr, (x == 1) | (x == 2), f_arr))

    # Get start, stop index pairs for sequences
    idx_pairs = np.where(np.diff(bool_stack))[0].reshape(-1, 2)

    # Get length of longest sequence
    longest_seq = np.max(np.diff(idx_pairs))

    out[:] = longest_seq


## Input for dim would be: 'time' 
def apply_seq_gufunc(data, dim):

    return xr.apply_ufunc(seq_gufunc,
                          data,
                          input_core_dims=[[dim]],
                          exclude_dims=set((dim,)), 
                          dask="allowed") 

可能存在一些非常明显的错误，希望有人能指出。我很难理解后台实际发生了什么，以及我应该如何设置 @guvectorize 的布局字符串和 apply_ufunc 的参数，以便它做我想要的。

---

编辑2： 这是有效的解决方案。有关apply_ufunc 和guvectorize 的参数的更多信息，请参阅@OriolAbril 的答案。还必须实现 if...else... 子句，以防没有匹配的值，并避免引发 ValueError 。

@guvectorize(
    "(int64[:], int64[:])",
    "(n) -> ()", nopython=True
)
def seq_gufunc(x, out):

    f_arr = np.array([False])
    bool_stack = np.hstack((f_arr, (x == 1) | (x == 2), f_arr))

    if np.sum(bool_stack) == 0:
        longest_seq = 0

    else:
        # Get start, stop index pairs for sequences
        idx_pairs = np.where(np.diff(bool_stack))[0].reshape(-1, 2)

        # Get length of longest sequence
        longest_seq = np.max(np.diff(idx_pairs))   

    out[:] = longest_seq


def apply_seq_gufunc(data, dim):

    return xr.apply_ufunc(seq_gufunc,
                          data,
                          input_core_dims=[[dim]],
                          dask="parallelized",
                          output_dtypes=['uint8']
                         )

Answer 1

我会向您指出How to apply a xarray u_function over NetCDF and return a 2D-array (multiple new variables) to the DataSet，近期目标不一样，但详细的描述和示例应该可以澄清问题。

特别是，您使用 time 作为 input_core_dims 是正确的（以确保将其移动到最后一个维度）并且它被正确格式化为列表列表，但是，您不需要 @ 987654325@ 但output_core_dims==[["time"]]。

输出与输入具有相同的形状，但是，正如上面链接中所解释的，apply_ufunc 期望它与广播的暗淡具有相同的形状。需要output_core_dims 才能获得apply_ufunc 以期望输出带有暗淡y, x, time。