基於二維 arrays 的 3D numpy 切片的平均值

Question

我正在嘗試計算第一個軸上兩個索引之間的 3D 數組的平均值。 開始和結束索引因單元格而異，由兩個單獨的 2D arrays 表示，它們的形狀與 3D 陣列的切片相同。

我已經設法實現了一段代碼，它循環遍歷我的 3D 數組的像素，但是對於形狀為(70, 550, 350)的數組，這種方法非常緩慢。 有沒有辦法使用numpy或xarray對操作進行矢量化（arrays 存儲在xarray數據集中）？

這是我想要優化的片段：

# My 3D raster containing values; shape = (time, x, y)
values = np.random.rand(10, 55, 60)

# A 2D raster containing start indices for the averaging
start_index = np.random.randint(0, 4, size=(values.shape[1], values.shape[2]))

# A 2D raster containing end indices for the averaging
end_index = np.random.randint(5, 9, size=(values.shape[1], values.shape[2]))

# Initialise an array that will contain results
mean_array = np.zeros_like(values[0, :, :])

# Loop over 3D raster to calculate the average between indices on axis 0
for i in range(0, values.shape[1]):
    for j in range(0, values.shape[2]):
        mean_array[i, j] = np.mean(values[start_index[i, j]: end_index[i, j], i, j], axis=0)

Answer 1

在沒有循環的情況下執行此操作的一種方法是將您不想使用的條目清零，計算剩余項目的總和，然后除以非零條目的數量。 例如：

i = np.arange(values.shape[0])[:, None, None]
mean_array_2 = np.where((i >= start_index) & (i < end_index), values, 0).sum(0) / (end_index - start_index)
np.allclose(mean_array, mean_array_2)
# True

請注意，這假設索引在0 <= i < values.shape[0]范圍內； 如果不是這種情況，您可以使用np.clip或其他方法在計算之前標准化索引。