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

Question

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

我已经设法实现了一段循环遍历我的 3D 数组像素的代码，但是对于形状为 (70, 550, 350) 的数组，这种方法非常缓慢。有没有办法使用numpy 或xarray 对操作进行矢量化（数组存储在xarray 数据集中）？

这是我想要优化的 sn-p：

# 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 或其他方式在计算之前对索引进行标准化。