加快 numpy 寻找最佳索引的速度

Question

我有一个 numpy 数组，它将 x-y 坐标映射到适当的 z 坐标。为此，我使用一个 2D 数组，它表示 x 和 y 作为其轴并包含相应的 z 值：

import numpy as np
x_size = 2000
y_size = 2500
z_size = 400
rng = np.random.default_rng(123)
z_coordinates = np.linspace(0, z_size, y_size) + rng.laplace(0, 1, (x_size, y_size))

因此，2000*2500 个 x-y 点中的每一个都分配有一个 z 值（0 到 400 之间的浮点数）。现在我想查找每个整数 z 和整数 x（最接近的 y 值），本质上是创建一个形状为

(x_size, z_size)

并保存最佳 y 值的地图。

最简单的方法是创建目标形状的空数组并迭代每个 z 值：

y_coordinates = np.empty((x_size, z_size), dtype=np.uint16)
for i in range(z_size):
    y_coordinates[:, i] = np.argmin(
        np.abs(z_coordinates - i),
        axis=1,
    )

然而，这在我的机器上大约需要 11 秒，不幸的是速度太慢了。

使用更加矢量化的方法当然会更快，例如：

y_coordinates = np.argmin(
    np.abs(
        z_coordinates[..., np.newaxis] - np.arange(z_size)
    ),
    axis=1,
)

令人惊讶的是，它的运行速度比上面的版本慢了大约 60%（在 1/10 大小下测试，因为在全尺寸下它使用了过多的内存）。

还将代码块包装在函数中并用 numba 的

@jit(nopython=True)

装饰它们并没有帮助。

如何加快计算速度？

Answer 1

提高性能的一种方法是添加以下代码行：

z_coordinates = z_coordinates.astype(np.float32).copy(order='C')

确保

C-contiguous

内存布局并使用

float32

代替 float64，通过优化内存访问模式和减少数据大小，显着加快计算速度。

• C-contiguous

  数组按行优先顺序存储元素，通过最大限度地减少缓存未命中并利用 CPU 优化来提高缓存效率并启用矢量化操作。

• 切换到

  float32

  可将内存使用量减半，减少内存带宽需求，并允许更多数据放入CPU缓存，这对于大型阵列特别有利。

这些变化共同使数据与 CPU 的期望保持一致，从而实现更快、更高效的处理，并大约将您的案例中的执行时间减半。