计算numpy数组的平方距离

Question

我正在脑力衰竭,但我无法让它发挥作用.我有一系列的距离:

import numpy as np
zvals = np.linspace(-5,5,10)
d = np.array([(0,0,z) for z in zvals])

我想计算数组中点的平方距离.实现这项工作的非笨拙方式是:

d2 = np.array([np.dot(d[i,:],d[i,:]) for i in range(d.shape[0])])

但是,我知道必须有一些方法可以通过一次调用dot来做到这一点,对吧？话既然说,也没有

d2 = np.dot(d,d.T)

要么

d2 = np.dot(d.T,d)

给我想要的东西.我意识到,我是傻瓜,但请在这里赐教.谢谢!

Answer 1

编辑:从NumPy 1.9开始,看起来inner1d可能会更快.(感谢Nuno Aniceto指出这一点):

In [9]:  %timeit -n 1000000 inner1d(d,d)
1000000 loops, best of 3: 1.39 µs per loop

In [14]: %timeit -n 1000000 einsum('ij,ij -> i', d, d)
1000000 loops, best of 3: 1.8 µs per loop

PS.始终在类似于您的预期用例的输入上测试您自己的基准.结果可能因各种原因而有所不同,例如输入大小,硬件,操作系统,Python版本,NumPy版本,编译器和库(例如ATLAS,MKL,BLAS).

如果你有NumPy 1.6或更高版本,你可以使用np.einsum:

In [40]: %timeit np.einsum('ij,ij -> i', d, d)
1000000 loops, best of 3: 1.79 us per loop

In [46]: from numpy.core.umath_tests import inner1d

In [48]: %timeit inner1d(d, d)
100000 loops, best of 3: 1.97 us per loop

In [44]: %timeit np.sum(d*d, axis=1)
100000 loops, best of 3: 5.39 us per loop

In [41]: %timeit np.diag(np.dot(d,d.T)) 
100000 loops, best of 3: 7.2 us per loop

In [42]: %timeit np.array([np.dot(d[i,:],d[i,:]) for i in range(d.shape[0])])
10000 loops, best of 3: 26.1 us per loop