Python中的加速MSD計算

Question

這是一個致社區的調查，看看是否有人有想法提高此MSD計算實現的速度。這主要是基於這篇博文的實施：http://damcb.com/mean-square-disp.html

目前，當前的實施需要大約9秒的2D軌跡5 000點。這真是太多，如果你需要計算大量的軌跡......

我沒有嘗試並行它（與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

提到到目前爲止是所有的O的MSD計算（N ** 2）其中N是時間步數。使用FFT可以將其減少到O（N * log（N））。請參閱this question and answer以獲取python中的解釋和實現。

編輯： 小基準（我也加入這一基準to this answer）：生成具有

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 

Answer 2

它一行一行地做了一些分析，看起來熊貓正在使這個緩慢。這純粹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 

Answer 3

添加到moarningsun上面的回答：

• ，如果你對數尺度繪製MSD可以加快使用numexpr

• 反正，你不「T需要計算它每次

  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 
  

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 

Answer 5

也許不是話題，但是MSD必須被計算爲不像37行中的平均值：

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單個軌跡