用numpy向量化这个操作

93ze6v8z  于 11个月前  发布在  其他
关注(0)|答案(2)|浏览(201)

为了介绍这个问题,让我们假设我有一个2D数组(Nx3)存储N个点。另一方面,我有一个1D数组(M,3)包含单个3D参考点的坐标,因此M=1
对于这种情况,如果我想计算一个新的1D数组(N),其中包含到参考点的所有距离,我只需要做:

import numpy as np

N = 100
M = 1
points = np.random.random((N, 3))
references = np.random.random((M, 3))

print(np.linalg.norm(points - references, axis=1))

字符串
这是可行的。然而,当我有多个参考点(M>1)时,numpy无法广播操作。现在我只是像之前描述的那样执行相同的操作,但使用for循环逐个遍历参考点,这是非常低效的。
当有M > 1时,我如何以矢量化的方式实现上述行为?我之前考虑过的事情:

  • 使用numba.njit去除python解释for的效果。但我想不使用numba
  • 类似于将点数组广播到3D(MxNx3),将参考数组广播到3D(NxMx3),这样我就可以操作它们了。但是,如果我没有错的话,我会将点数组的大小乘以M,将参考数组的大小乘以N。如果这是解决方案,我也不知道如何做到这一点。
cu6pst1q

cu6pst1q1#

简单地说:

np.linalg.norm(references - points[:, None], axis=-1)

字符串
这给了你一个所有距离的M,N数组。
范例:

np.random.seed(0)  # reproducible setup
N = 4
M = 2
points = np.random.random((N, 3))
references = np.random.random((M, 3))

>>> np.linalg.norm(references - points[:, None], axis=-1)
array([[0.57216693, 0.86543108],
       [0.76350759, 0.63809564],
       [0.90274337, 0.94847268],
       [0.51150232, 0.88049546]])


注意事项:如果你要处理大量的点,并且不一定对完整的距离列表感兴趣,而是对一个子集感兴趣,比如说,最近的邻居,那么你应该考虑使用scipy.spatial.KDTree。这将更快,内存效率更高。

mznpcxlj

mznpcxlj2#

当你取数组的差,然后将np.linalg.norm应用于最后一个轴时,进行广播(我显式地写了2,但你也可以使用-1)。

np.linalg.norm(points[:,None] - references[None,:], axis=2)

字符串
需要注意的是,当M = 1时,你的版本会给你一个1D向量。如果你想要一个行向量,那么交换pointsreferences索引。这也会对最终结果产生转置的效果。我不确定你更喜欢哪个方向。

相关问题