如何并行写入 numpy 数组中的同一单元格？

Question

背景：我在 2D 空间中有数百万个点，(x_position, y_position, value)每个点都与之相关。我试图通过创建图像来总结这些点，其中每个像素可以包含多个点。总而言之，每个像素存储图像中value该位置处的 s之和。(x_pixel, y_pixel)问题：我怎样才能有效地做到这一点？目前，我的代码做了这样的事情：image = np.zeros((4096,4096)) for each point in data:     x_pixel, y_pixel = convertPointPos2PixelPos(point)     image[x_pixel, y_pixel] += point.getValue()但完成此代码的预计时间为 450 小时，这是不可接受的。有没有办法并行化这个？image[x,y]该代码多次写入同一索引。我发现 StackOverflow 帖子建议使用multiprocessing，但我认为需要lock防止竞争条件意味着这将花费与没有并行化时一样多的时间。

Answer 1

假设您想要常规网格上的某些内容，您可以使用简单的除法来对数据进行分类。这是一个例子：

size = (4096, 4096)

data = np.random.rand(100000000, 3)

image = np.zeros(size)

coords = data[:, :2]

min = coords.min(0)

max = coords.max(0)

index = np.floor_divide(coords - min, (max - min) / np.subtract(size, 1), out=np.empty(coords.shape, dtype=int), casting='unsafe')

index现在是一个索引数组，image您要在其中添加相应的值。您可以使用以下命令进行无缓冲添加np.add.at：

np.add.at(image, tuple(index.T), data[:, -1])

如果您的数据范围比仅坐标的边界框定义得更好，则可以通过不计算coord.max()和来节省一些时间coord.min()。

结果是这样的：

在我的 10M 点的中等功率机器上，整个操作需要 6.4 秒，包括调用plt.imshow和plt.colorbar运行前的垃圾收集。

使用 IPython 中的单元魔法收集时间%%timeit。

不管怎样，你的工作时间远低于 450 小时。即使您的坐标变换不是线性分箱，我希望只要正确地对其进行矢量化，您就可以在合理的时间内运行。此外，多重处理不太可能给您带来巨大的提升，因为它需要复制数据。