在Python中加速MSD计算

Question

这是一个呼吁社区的行动，看是否有人有想法改进这个MSD计算实现的速度。这个实现主要基于这篇博客文章：http://damcb.com/mean-square-disp.html。

目前，当前的实现对于一个由5,000个点组成的二维轨迹需要大约9秒的时间。如果您需要计算许多轨迹，那么这显然太慢了...

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

以下是代码：

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

有什么想法吗？

- hadim

---

既然您已经有可工作的代码，这可能是一个很好的代码审查候选项。 - cel

哦，我不知道_codereview_。能否有一位管理员确认一下，我会将其移动到_codereview_吗？ - hadim

我是 Code Review 的版主，我已经标记了这个问题以便迁移到 Code Review。现在我们只能等待看看 Stack Overflow 的版主是否同意。 - Simon Forsberg

我遇到了一个NA问题，在compute_msd的第二行中，floor函数在尝试转换为int时抛出异常。（numpy 1.9.2，Py2.7.10，OSX）还有其他人遇到这个问题吗？ - rll

在我的Ubuntu系统上，使用numpy 1.9.3、pandas 0.16.2和python 3.4，这段代码可以正常运行。 - hadim

@tll - 上述代码需要在dt = max_time / N中使用真正的除法，所以它不能在Python 2.7上工作。 - user2379410

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

- user2379410

Answer 2

补充moarningsun的答案:

• you can speed up using numexpr

• if you plot the MSD in log scale anyway, you don't need to compute it for every time

  import numpy as np
  import numexpr
  
  def logSpaced(L, pointsPerDecade=15):
      """Generate an array of log spaced integers smaller than L"""
      nbdecades = np.log10(L)
      return np.unique(np.logspace(
          start=0, stop=nbdecades, 
          num=nbdecades * pointsPerDecade, 
          base=10, endpoint=False
          ).astype(int))
  
  def compute_msd(xy, pointsPerDecade=15):
      dts = logSpaced(len(xy), pointsPerDecade)
      msd = np.zeros(len(idts))
      msd_std = np.zeros(len(idts))
      for i, dt in enumerate(dts):
          sqdist = numexpr.evaluate(
              '(a-b)**2',
              {'a': xy[:-dt], 'b':xy[dt:]}
              ).sum(axis=-1)
          msd[i] = sqdist.mean()
          msd_std[i] = sqdist.std(ddof=1)
      msds = pd.DataFrame({'msds': msd, 'tau': dt, 'msds_std': msd_std})
      return msds
  

- Mathieu Leocmach

---

谢谢。您是否比较了 numexpr 版本与 moarningsun 版本的速度？ - hadim

Answer 3

到目前为止提到的MSD计算都是O(N**2)，其中N是时间步数。使用FFT可以将其降至O(N*log(N))。参见这个问题和答案以获得Python中的解释和实现。

编辑： 一个小基准测试（我也将此基准测试添加到了这个答案中）：生成一条轨迹。

r = np.cumsum(np.random.choice([-1., 0., 1.], size=(N, 3)), axis=0)

对于N=100,000，我们得到：

$ %timeit msd_straight_forward(r)
1 loops, best of 3: 2min 1s per loop

$ %timeit msd_fft(r)
10 loops, best of 3: 253 ms per loop

- thomasfermi

---

如果能帮到某人，我会感到高兴 :) - thomasfermi

Answer 4

通过注释，我设计了这个函数：

def get_msd(traj, dt, with_nan=True):

    shifts = np.arange(1, len(traj), dtype='int')
    msd = np.empty((len(shifts), 2), dtype='float')
    msd[:] = np.nan

    msd[:, 1] = shifts * dt

    for i, shift in enumerate(shifts):
        diffs = traj[:-shift] - traj[shift:]
        if with_nan:
            diffs = diffs[~np.isnan(diffs).any(axis=1)]
        diffs = np.square(diffs).sum(axis=1)

        if len(diffs) > 0:
            msd[i, 0] = np.mean(diffs)

    msd = pd.DataFrame(msd)
    msd.columns = ["msd", "delay"]

    msd.set_index('delay', drop=True, inplace=True)
    msd.dropna(inplace=True)

    return msd

具有以下功能:

• 以numpy数组为轨迹输入。
• 返回一个几乎没有重叠的pandas.DataFrame。
• with_nan允许处理包含NaN值的轨迹，但会增加很大的开销（超过100%），因此我将其放置在函数参数中。
• 它可以处理多维轨迹（1D、2D、3D等）。

一些分析：

$ print(traj.shape)
(2108, 2)

$ %timeit get_msd(traj, with_nan=True, dt=0.1)
10 loops, best of 3: 143 ms per loop

$ %timeit get_msd(traj, with_nan=False, dt=0.1)
10 loops, best of 3: 68 ms per loop

- hadim

Answer 5

也许这不是主题，但是必须像第37行中计算的那样计算MSD，而不是平均值。

msds[i] = sqdist.mean()

以mean=N为例

你必须除以：

msds[i] = sqdist/N-1 // for lag1

然后：

msds[i] = sqdist/N-2 // for lag2 .... msds[i] = sqdist/N-n // for lag n

等等。

因此，您不会得到标准偏差，只会得到单个轨迹的MSD。

- Jonathan Pacheco