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

Question

I am writing a code using a Python library MDAnalysis. And I have an array (502,3) of atom positions I want to get an array of bonds (vectors of position of Atom(i+1) - Atom(i)) And then I want to obtain an average tensor qab = which is essentially a np.outer(ua,ub) averaged by all atoms.

I can rewrite this code using fortran subroutines but I think it would be more pleasing to see a delicious numpy slicing solution:) Here is my code that I wrote and that works, I want to make it faster and prettier. Thank you in advance for your help.

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.        ]]

Answer 1

I've done my best below. It's pretty easy to generate your bb_res more efficiently, but I was unable to optimize the double for loop. On my computer, my method is about 26% faster. Also based on the statement of your question I believe there is a bug in your code which I pointed out in a comment. I have fixed this bug in my answer, so it produces a slightly different output than your code.

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)

Answer 2

Assuming you meant nq = bb_res.shape[1] in your code snippet, I can reproduce the output with the following vectorized code:

d = np.diff(coor, axis=0)
I, J = np.triu_indices(d.shape[0])
print np.dot(d[J].T, d[I]) / len(J)