使用 Xarray 的 apply_ufunc 在 3D 數組的時間維度上應用 Numba guvectorize function

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）和我的 output 2D（x，y）。 'seq_gufunc' 中的代碼受此線程啟發。 我的實際數據集要大得多（使用土地利用類而不是 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 的參數，以便它做我想要的。

---

EDIT2：這是可行的解決方案。 有關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

我會向您指出如何在 NetCDF 上應用 xarray u_function 並將 2D 數組（多個新變量）返回到 DataSet ，直接目標不一樣，但詳細的描述和示例應該可以澄清這個問題。

特別是，您正確地將time用作input_core_dims （以確保將其移動到最后一個維度）並且它被正確格式化為列表列表，但是，您不需要excluded_dims而是output_core_dims==[["time"]] 。

output 具有與輸入相同的形狀，但是，如上面鏈接中所述， apply_ufunc預計它將具有與廣播的 dims相同的形狀。 需要output_core_dims來獲得apply_ufunc以期望 output 與暗淡y, x, time 。