使用 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 坐标，我想 return 随着时间的推移最长的 1 或 2 序列。所以我的输入数组是 3D (time, x, y) 而我的输出是 2D (x, y)。 'seq_gufunc' 中的代码受到的启发。我的实际数据集要大得多（landuse 类而不是 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

我会向您指出，直接目标是不一样的，但详细的描述和示例应该可以阐明问题。

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

输出与输入具有相同的形状，但是，如上文 link 中所述，apply_ufunc 期望它与 广播的暗淡[=26] 具有相同的形状=].需要 output_core_dims 才能获得 apply_ufunc 以期望带有暗淡 y, x, time.

的输出

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

Applying a Numba guvectorize function over time dimension of an 3D-Array with Xarray's apply_ufunc

python

numpy

numba

numpy-ufunc

python-xarray