Tính toán MSD tăng tốc bằng Python

Question

Đó là một cuộc gọi đến cộng đồng để xem liệu có ai có ý tưởng cải thiện tốc độ thực hiện tính toán MSD này hay không. Phần lớn dựa trên việc triển khai từ bài đăng trên blog này: http://damcb.com/mean-square-disp.html

Hiện tại, việc triển khai hiện tại mất khoảng 9 giây cho quỹ đạo 2D là 5 000 điểm. Nó thực sự là quá nhiều nếu bạn cần tính toán rất nhiều quỹ đạo ...

Tôi không cố gắng song song nó (với multiprocess hoặc joblib) nhưng tôi có cảm giác rằng việc tạo quy trình mới sẽ quá nặng nề loại thuật toán.

Đây là mã:

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()) 

Và kết quả:

  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) 

Và kết quả:

 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 

Và một số hồ sơ:

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

Give này:

1 loops, best of 3: 8.53 s per loop 

Bất kỳ ý tưởng?

Answer 1

Các tính toán MSD đề cập đến nay đều là O (N ** 2) trong đó N là số bước thời gian. Sử dụng FFT này có thể được giảm xuống O (N * log (N)). Xem this question and answer để được giải thích và triển khai trong python.

EDIT: Một benchmark nhỏ (Tôi cũng đã thêm điểm chuẩn này to this answer): Tạo một quỹ đạo với

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

Đối với N = 100.000, chúng tôi nhận

$ %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

Nó đã làm một số dòng profiling bởi dòng và nó xuất hiện rằng gấu trúc đang làm chậm này. phiên bản numpy tinh khiết này là khoảng 14x nhanh hơn:

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

Thêm vào moarningsun câu trả lời ở trên:

• bạn có thể tăng tốc độ sử dụng numexpr

• nếu bạn lô MSD ở dạng thang loga dù sao, bạn don 't cần phải tính toán nó cho mỗi lần

  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

Với những ý kiến ​​tôi thiết kế chức năng này:

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 

với các tính năng sau:

• Phải mất numpy mảng như là đầu vào quỹ đạo.
• Nó trả về một pandas.DataFrame với hầu như không có lớp phủ.
• with_nan cho phép xử lý quỹ đạo có chứa các giá trị NaN nhưng nó thêm một chi phí lớn (trên 100%) vì vậy tôi đặt nó làm thông số hàm.
• Nó có thể đối phó với quỹ đạo đa chiều (1D, 2D, 3D, vv)

Một số hồ sơ:

$ 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

Có lẽ không phải là chủ đề, tuy nhiên MSD phải được tính không phải là giá trị trung bình như trong dòng 37:

msds[i] = sqdist.mean() 

Lấy như mean=N

Bạn phải phân chia theo:

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

Sau đó:

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

Và vân vân. Vì vậy, bạn không có độ lệch chuẩn, chỉ cần MSD cho một quỹ đạo duy nhất