基于 2D 数组的 3D numpy 切片的平均值

Question

我正在尝试计算第一个轴上两个索引之间的 3D 数组的平均值。开始和结束索引因单元格而异，由两个单独的二维数组表示，这两个数组的形状与 3D 数组的切片相同。

我已经成功地实现了一段循环遍历我的 3D 数组像素的代码，但是对于形状为 (70, 550, 350) 的数组，这种方法非常慢。有没有办法使用 numpy 或 xarray 对操作进行矢量化（数组存储在 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 或其他方式在计算前对索引进行标准化。

基于 2D 数组的 3D numpy 切片的平均值

Average of a 3D numpy slice based on 2D arrays

python

arrays

numpy

mean

slice