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

Posted 2023-03-24

【中文标题】基于 2D 数组的 3D numpy 切片的平均值

【英文标题】：Average of a 3D numpy slice based on 2D arrays

【发布时间】：2021-04-22 23:42:16

【问题描述】：

我正在尝试计算第一个轴上两个索引之间的 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)

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

【讨论】：

我的值通常在 -10 到 +10 之间。如果我理解正确，我将不得不将它们从 0 标准化为 values.shape[0]，进行操作并将平均值转换回初始范围？

不，值可以是任何值。开始和结束索引必须在范围内，这两种方法才能产生相同的结果。

感谢您的信息，我测试了您的代码，它运行良好。它比循环（600x）快得多。有趣的是，它比使用 numba jit 循环慢 2 倍。但话虽这么说，它是同一个数量级，正如我指出的那样，nuba 在抛出错误时更难调试。