以最有效的方式编写python代码

我正在使用 Python库MDAnalysis编写代码.
我有一个原子位置的数组(502,3)
我想获得一系列债券(Atom(i 1)的位置向量 – Atom(i))
然后我想获得一个平均张量qab =本质上是一个np.outer(ua,ub)
由所有原子平均.

我可以使用fortran子程序重写这段代码,但我认为看到一个美味的numpy切片解决方案会更令人愉快:)这是我编写的代码,这是有效的,我想让它变得更快更漂亮.预先感谢您的帮助.

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非常容易,但是我无法优化double for循环.在我的电脑上,我的方法快了大约26％.同样基于您的问题陈述,我相信您的代码中存在一个错误,我在评论中指出了这个错误.我已在我的答案中修复了这个错误,因此它产生的输出与您的代码略有不同.

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,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)