在 Python 中加速 MSD 计算
Speedup MSD calculation in Python
向社区发出呼吁,看看是否有人有想法提高此 MSD 计算实现的速度。它主要基于此博客的实现 post : 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
有什么想法吗?
它逐行进行了一些分析,看来 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
添加到上面的 moarningsun 回答:
- 你可以使用 numexpr 来加速
如果你以对数刻度绘制 MSD,则不需要每次都计算它
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
根据评论我设计了这个功能:
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
数组
- 它returns
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
到目前为止提到的MSD计算都是O(N**2),其中N是时间步数。使用 FFT,这可以减少到 O(N*log(N))。有关 python 中的解释和实现,请参见 。
编辑:
一个小benchmark(我也加了这个benchmark):Generate a trajectory with
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
也许不是主题,但是必须计算 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
向社区发出呼吁,看看是否有人有想法提高此 MSD 计算实现的速度。它主要基于此博客的实现 post : 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
有什么想法吗?
它逐行进行了一些分析,看来 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
添加到上面的 moarningsun 回答:
- 你可以使用 numexpr 来加速
如果你以对数刻度绘制 MSD,则不需要每次都计算它
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
根据评论我设计了这个功能:
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
数组 - 它returns
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
到目前为止提到的MSD计算都是O(N**2),其中N是时间步数。使用 FFT,这可以减少到 O(N*log(N))。有关 python 中的解释和实现,请参见
编辑:
一个小benchmark(我也加了这个benchmark
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
也许不是主题,但是必须计算 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