Python 中均方位移的矢量化计算

Question

我想计算N个粒子的均方位移，我有粒子位置随时间变化的轨迹。我编写的代码有 3 个 for 循环，这使得它非常慢。你能帮我如何用 numpy 或熊猫的某种矢量化功能替换循环吗？

这是我的代码：

ntime = 10 # number of times represented in data
atom_count = 3 # number of particles 
norigin = 5 # number of origins is half number of time steps
nmin = 2 # minimum number of intervals to contribute to diffusivity
nmax = norigin # maximum number of intervals to contribute to diffusivity
dt = 1.0 # timestep
# creating sample trajectory of particles
traj = pd.DataFrame(np.random.rand(ntime*atom_count,3), columns=['x', 'y', 'z'])
traj['frame_id'] = np.repeat(np.arange(ntime)+1, atom_count)
traj['particle_id'] = np.tile(np.arange(atom_count)+1, ntime)
traj = traj[['frame_id', 'particle_id', 'x', 'y', 'z']]
print(traj.head(6))
ndata = traj.shape[0] # number of rows of data
# store mean square displacements in msd
time_vec= np.arange(dt, norigin*dt+1, dt)
msd_xyz = np.zeros((norigin,3))
# loop over all particles
for i in range(atom_count):
# loop over all time origins
for j in range(norigin):
jstart = j*atom_count + i
# loop over all time windows
for k in range(nmin, nmax):
kend = jstart + k*atom_count
msd_xyz[k, :] += (traj.ix[kend, ['x', 'y', 'z']].values - 
traj.ix[jstart, ['x', 'y', 'z']].values)**2
msd_xyz = msd_xyz / (atom_count * norigin)   
msd = np.mean(msd_xyz, axis=1) # total MSD averaged over x, y, z directions
print()
print("MSD (averaged over all particles and time origins):")
print(msd)

Answer 1

使用 numpy 的索引功能，所有 3 个嵌套循环都可以在网格的帮助下进行矢量化。

要做到这一点，关键是 numpy 数组支持任何形状的列表或数组索引：

a = np.arange(5,10)
b = np.array([[0,2,4],[3,3,0]])
print(a[b])
# Output
[[5 7 9]
[8 8 5]]

因此，我们可以从循环中用作迭代器的数组中定义一个网格网格，以便一次从循环中检索 i、j 和 k 的所有组合，然后对 i 和 j 求和。

重要的是要注意，数组索引已在.values方法之后移动，因为numpy支持这种索引，但pandas仅适用于1D数组。

# define indexing arrays 
k = np.arange(nmin,nmax)
j = np.arange(norigin)
i = np.arange(atom_count)
I,J,K = np.meshgrid(i,j,k) # the meshgrid contains all the combinations of i,j,k, 
# it is equivalent to the 3 nested loops
jstart = J*atom_count + I
kend = jstart + K*atom_count
msd_xyz[k,:] = np.sum(np.sum((traj[['x', 'y', 'z']].values[kend,:] - 
traj[['x', 'y', 'z']].values[jstart,:])**2,axis=0),axis=0)
msd_xyz = msd_xyz / (atom_count * norigin)   
msd = np.mean(msd_xyz, axis=1) # total MSD averaged over x, y, z directions

使用问题中示例数据的维度，这实现了相对于 3 个嵌套循环的 x60 加速。但是，对于大型数据帧，这可能会使用过多内存并变慢，在这种情况下，最好将循环和矢量化结合起来，只矢量化一个或 2 个循环以避免过多的内存使用。