Python中的加速MSD计算

时间:2015-10-07 09:12:04

标签: python-3.x numpy pandas physics

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

目前,对于500个点的2D轨迹,当前实施需要大约9秒。如果你需要计算很多轨迹,这真的太过分了......

我没有尝试将其并行化(使用multiprocessjoblib),但我觉得创建新流程对于这种算法来说太重了。

以下是代码:

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

enter image description here

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

enter image description here

还有一些分析:

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

给这个:

1 loops, best of 3: 8.53 s per loop

有什么想法吗?

5 个答案:

答案 0 :(得分:4)

它逐行进行了一些分析,看起来大熊猫正在慢慢进行。这个纯粹的numpy版本快了大约14倍:

props

答案 1 :(得分:3)

添加到上面的moarningsun答案:

  • 你可以加快使用numexpr
  • 如果您无论如何都以对数比例绘制MSD,您不需要每次都计算它

    Link Binary with Libraries

答案 2 :(得分:2)

到目前为止提到的MSD计算都是O(N ** 2),其中N是时间步数。使用FFT可以将其减少到O(N * log(N))。有关python中的解释和实现,请参阅this question and answer

编辑: 一个小基准(我也添加了这个基准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

答案 3 :(得分:1)

根据评论我设计了这个功能:

ggplot(dat, aes(x = id, y = val, fill = as.factor(group))) + 
  geom_bar(stat = "identity", position = "dodge") +
  scale_fill_manual(values = brewer.pal(n, name = "BrBG")) # BrBG as an example

具有以下功能:

  • 它需要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值的轨迹,但会增加很大的开销(超过100%),所以我把它作为函数参数。
  • 它可以处理多维轨迹(1D,2D,3D等)

一些分析:

NaN

答案 4 :(得分:0)

可能不是主题,但是必须计算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