以最有效的方式编写python代码
我正在使用Python库MDAnalysis编写代码。 我有一个原子位置的数组(502,3) 我想得到一组键(原子(I+1)的位置向量,原子(I)) 然后我想得到一个平均张量qab=它本质上是一个np.outer(ua,ub) 所有原子的平均值 我可以用fortran子程序重写这段代码,但我认为看到一个美味的numpy切片解决方案会更令人高兴:)这是我写的代码,我想让它更快、更漂亮。提前感谢您的帮助以最有效的方式编写python代码,python,numpy,scipy,Python,Numpy,Scipy,我正在使用Python库MDAnalysis编写代码。 我有一个原子位置的数组(502,3) 我想得到一组键(原子(I+1)的位置向量,原子(I)) 然后我想得到一个平均张量qab=它本质上是一个np.outer(ua,ub) 所有原子的平均值 我可以用fortran子程序重写这段代码,但我认为看到一个美味的numpy切片解决方案会更令人高兴:)这是我写的代码,我想让它更快、更漂亮。提前感谢您的帮助 bb = u.selectAtoms("all") coor = bb.positions pr
bb = u.selectAtoms("all")
coor = bb.positions
print coor.shape # return is (502,3)
# the coordinates of the atoms that we split on 3 dimensions
bbx = coor[:,0]
bby = coor[:,1]
bbz = coor[:,2]
#the bond vectors we obtain by subtructing atoms positions
bbx_ave = bbx[:-1] - bbx[1:]
bby_ave = bby[:-1] - bby[1:]
bbz_ave = bbz[:-1] - bbz[1:]
#now we concentrate everything back so we have one big array
bb_res = np.vstack((bbx_ave, bby_ave, bbz_ave))
# print bb_res.shape # the return is (3,501)
# now we have an array of bonds
# qab - the result tensor which we want to get
nq = len(bb_res)
qab = np.zeros((3,3))
count = 0.
for i in xrange(0, nq):
for j in xrange(0, i+1):
qab = qab + np.outer(bb_res[:,i], bb_res[:,j])
count += 1.
qab = qab/count
print qab
[[ 0.21333394 0.5333341 0. ]
[ 0.5333341 4. 0. ]
[ 0. 0. 0. ]]
我在下面已经尽力了。更有效地生成
bb_res
非常容易,但我无法优化双for
循环。在我的电脑上,我的方法大约快26%。同样基于您的问题陈述,我相信您的代码中有一个bug,我在评论中指出了这个bug。我已经在我的答案中修复了这个bug,所以它产生的输出与您的代码略有不同
import numpy as np
from numpy.random import rand
# create some mock data
coor = rand(502,3)
def qab(coor):
# this is the transpose of your bb_res
# transposing is as simple as bb_res.T
bb_res = coor[:-1] - coor[1:]
nq = bb_res.shape[1]
out = np.zeros((3,3))
for i in xrange(0, nq):
for j in xrange(0, i):
tmp = np.outer(bb_res[i], bb_res[j])
out += tmp + tmp.T
out += np.outer(bb_res[i], bb_res[i])
return out / nq**2
print qab(coor)
假设您在代码片段中的意思是
nq=bb_res.shape[1]
,我可以使用以下矢量化代码再现输出:
d = np.diff(coor, axis=0)
I, J = np.triu_indices(d.shape[0])
print np.dot(d[J].T, d[I]) / len(J)
这应该是开着的,我认为你的双for循环没有达到你想要的效果。如果在两个循环上都有
xrange(0,nq)
,则会产生不同的结果(但所需时间大约是原来的两倍)。似乎你正试图利用对称性,通过对张量的一半求和来节省时间。但是当你这样做的时候,你应该包括这样一个事实:非对角项的重要性是对角项的两倍。为了更清楚地理解这一点,请尝试手工制作一个简单的例子,其中nq=2
。现在你有了nq==3
,而如果它是nq==501
则更有意义。我认为你不应该添加2*
外积,而是添加外积及其转置,以获得与两个循环运行到nq
时相同的输出。而nq
应该等于bb_res.shape[0]
?@moarningsun,我在利用问题的对称性。这种对称性是np.outer(bb_res[i],bb_res[j])=np.outer(bb_res[j],bb_res[i])
请验证,对称性应该是:np.outer(bb_res[i],bb_res[j])==np转置(np.outer(bb res[j],bb_res[i])
Oh derp,我怎么忘了这个?谢谢,我已经更新了答案。现在你可以用out=np.einsum('ij,kl->jl',bb_res,bb_res)
替换嵌套循环,它没有利用对称性,但却提供了很好的加速效果:)