python-3.x - 在 Python 中加速 MSD 计算

Question

这是对社区的呼吁，看看是否有人有提高此 MSD 计算实现速度的想法。它主要基于此博客文章中的实现:http://damcb.com/mean-square-disp.html

目前，对于 5000 点的 2D 轨迹，当前的实现需要大约 9 秒。如果你需要计算很多轨迹，那真的太多了......

我没有尝试将它并行化(使用 multiprocess 或 joblib )，但我觉得创建新进程对于这种算法来说太重了。

这是代码:

import os

import matplotlib
import matplotlib.pyplot as plt

import pandas as pd
import numpy as np

# Parameters
N = 5000
max_time = 100
dt = max_time / N

# Generate 2D brownian motion

t = np.linspace(0, max_time, N)
xy = np.cumsum(np.random.choice([-1, 0, 1], size=(N, 2)), axis=0)
traj = pd.DataFrame({'t': t, 'x': xy[:,0], 'y': xy[:,1]})
print(traj.head())

# Draw motion
ax = traj.plot(x='x', y='y', alpha=0.6, legend=False)

# Set limits
ax.set_xlim(traj['x'].min(), traj['x'].max())
ax.set_ylim(traj['y'].min(), traj['y'].max())

和输出:

          t  x  y
0  0.000000 -1 -1
1  0.020004 -1  0
2  0.040008 -1 -1
3  0.060012 -2 -2
4  0.080016 -2 -2

def compute_msd(trajectory, t_step, coords=['x', 'y']):

    tau = trajectory['t'].copy()
    shifts = np.floor(tau / t_step).astype(np.int)
    msds = np.zeros(shifts.size)
    msds_std = np.zeros(shifts.size)

    for i, shift in enumerate(shifts):
        diffs = trajectory[coords] - trajectory[coords].shift(-shift)
        sqdist = np.square(diffs).sum(axis=1)
        msds[i] = sqdist.mean()
        msds_std[i] = sqdist.std()

    msds = pd.DataFrame({'msds': msds, 'tau': tau, 'msds_std': msds_std})
    return msds

# Compute MSD
msd = compute_msd(traj, t_step=dt, coords=['x', 'y'])
print(msd.head())

# Plot MSD
ax = msd.plot(x="tau", y="msds", logx=True, logy=True, legend=False)
ax.fill_between(msd['tau'], msd['msds'] - msd['msds_std'], msd['msds'] + msd['msds_std'], alpha=0.2)

和输出:

       msds  msds_std       tau
0  0.000000  0.000000  0.000000
1  1.316463  0.668169  0.020004
2  2.607243  2.078604  0.040008
3  3.891935  3.368651  0.060012
4  5.200761  4.685497  0.080016

还有一些分析:

%timeit msd = compute_msd(traj, t_step=dt, coords=['x', 'y'])

给这个:

1 loops, best of 3: 8.53 s per loop

任何的想法 ？

Answer 1

它逐行进行了一些分析，看来 Pandas 正在使这一点变慢。这个纯 numpy 版本快了大约 14 倍:

def compute_msd_np(xy, t, t_step):
    shifts = np.floor(t / t_step).astype(np.int)
    msds = np.zeros(shifts.size)
    msds_std = np.zeros(shifts.size)

    for i, shift in enumerate(shifts):
        diffs = xy[:-shift if shift else None] - xy[shift:]
        sqdist = np.square(diffs).sum(axis=1)
        msds[i] = sqdist.mean()
        msds_std[i] = sqdist.std(ddof=1)

    msds = pd.DataFrame({'msds': msds, 'tau': t, 'msds_std': msds_std})
    return msds