使用 python 和 FFT 计算均方位移
Computing mean square displacement using python and FFT
给定一个二维数组,其中每一行代表一个粒子的位置向量,我如何有效地计算均方位移(使用 FFT)?
均方位移定义为
其中r(m)是第m行的位置向量,N是行数
以下用于 msd 的直接方法有效,但它是 O(N**2)(我改编了此 中的代码)
def msd_straight_forward(r):
shifts = np.arange(len(r))
msds = np.zeros(shifts.size)
for i, shift in enumerate(shifts):
diffs = r[:-shift if shift else None] - r[shift:]
sqdist = np.square(diffs).sum(axis=1)
msds[i] = sqdist.mean()
return msds
但是,我们可以使用 FFT 使这段代码更快。下面的考虑和最终的算法来自this paper,我将在python中展示如何实现它。
我们可以通过以下方式拆分MSD
因此,S_2(m)就是位置的自相关。请注意,在某些教科书中 S_2(m) 表示为自相关(约定 A),而在某些教科书中 S_2(m)*(N-m) 表示为自相关(约定 B)。
根据 Wiener-Khinchin 定理,函数的功率谱密度 (PSD) 是自相关的傅里叶变换。
这意味着我们可以计算信号的 PSD 并对它进行傅立叶反演,以获得自相关(在约定 B 中)。对于离散信号,我们得到循环自相关。
然而,通过对数据进行零填充,我们可以获得非循环自相关。算法看起来像这样
def autocorrFFT(x):
N=len(x)
F = np.fft.fft(x, n=2*N) #2*N because of zero-padding
PSD = F * F.conjugate()
res = np.fft.ifft(PSD)
res= (res[:N]).real #now we have the autocorrelation in convention B
n=N*np.ones(N)-np.arange(0,N) #divide res(m) by (N-m)
return res/n #this is the autocorrelation in convention A
对于术语 S_1(m),我们利用了这样一个事实,即可以找到 (N-m)*S_1(m) 的递归关系(这在 paper 在第 4.2 节)。
我们定义
并通过
找到 S_1(m)
这会产生以下用于均方位移的代码
def msd_fft(r):
N=len(r)
D=np.square(r).sum(axis=1)
D=np.append(D,0)
S2=sum([autocorrFFT(r[:, i]) for i in range(r.shape[1])])
Q=2*D.sum()
S1=np.zeros(N)
for m in range(N):
Q=Q-D[m-1]-D[N-m]
S1[m]=Q/(N-m)
return S1-2*S2
您可以比较 msd_straight_forward() 和 msd_fft() 并会发现它们产生相同的结果,尽管 msd_fft() 对于大 N[=21 更快=]
一个小基准:用
生成轨迹
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
使用numpy.cumsum您还可以避免在 S1 计算中循环 range(N):
sq = map(sum, map(np.square, r))
s1 = 2 * sum(sq) - np.cumsum(np.insert(sq[0:-1], 0, 0) + np.flip(np.append(sq[1:], 0), 0))
给定一个二维数组,其中每一行代表一个粒子的位置向量,我如何有效地计算均方位移(使用 FFT)? 均方位移定义为
其中r(m)是第m行的位置向量,N是行数
以下用于 msd 的直接方法有效,但它是 O(N**2)(我改编了此
def msd_straight_forward(r):
shifts = np.arange(len(r))
msds = np.zeros(shifts.size)
for i, shift in enumerate(shifts):
diffs = r[:-shift if shift else None] - r[shift:]
sqdist = np.square(diffs).sum(axis=1)
msds[i] = sqdist.mean()
return msds
但是,我们可以使用 FFT 使这段代码更快。下面的考虑和最终的算法来自this paper,我将在python中展示如何实现它。 我们可以通过以下方式拆分MSD
因此,S_2(m)就是位置的自相关。请注意,在某些教科书中 S_2(m) 表示为自相关(约定 A),而在某些教科书中 S_2(m)*(N-m) 表示为自相关(约定 B)。 根据 Wiener-Khinchin 定理,函数的功率谱密度 (PSD) 是自相关的傅里叶变换。 这意味着我们可以计算信号的 PSD 并对它进行傅立叶反演,以获得自相关(在约定 B 中)。对于离散信号,我们得到循环自相关。 然而,通过对数据进行零填充,我们可以获得非循环自相关。算法看起来像这样
def autocorrFFT(x):
N=len(x)
F = np.fft.fft(x, n=2*N) #2*N because of zero-padding
PSD = F * F.conjugate()
res = np.fft.ifft(PSD)
res= (res[:N]).real #now we have the autocorrelation in convention B
n=N*np.ones(N)-np.arange(0,N) #divide res(m) by (N-m)
return res/n #this is the autocorrelation in convention A
对于术语 S_1(m),我们利用了这样一个事实,即可以找到 (N-m)*S_1(m) 的递归关系(这在 paper 在第 4.2 节)。 我们定义
并通过
找到 S_1(m)这会产生以下用于均方位移的代码
def msd_fft(r):
N=len(r)
D=np.square(r).sum(axis=1)
D=np.append(D,0)
S2=sum([autocorrFFT(r[:, i]) for i in range(r.shape[1])])
Q=2*D.sum()
S1=np.zeros(N)
for m in range(N):
Q=Q-D[m-1]-D[N-m]
S1[m]=Q/(N-m)
return S1-2*S2
您可以比较 msd_straight_forward() 和 msd_fft() 并会发现它们产生相同的结果,尽管 msd_fft() 对于大 N[=21 更快=]
一个小基准:用
生成轨迹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
使用numpy.cumsum您还可以避免在 S1 计算中循环 range(N):
sq = map(sum, map(np.square, r))
s1 = 2 * sum(sq) - np.cumsum(np.insert(sq[0:-1], 0, 0) + np.flip(np.append(sq[1:], 0), 0))