python - 使用 python 和 FFT 计算均方位移

Question

给定一个二维数组，其中每一行代表一个粒子的位置向量，我如何有效地计算均方位移(使用 FFT)？均方位移定义为

其中r(m)为第m行的位置向量，N为行数。

Answer 1

以下用于 msd 的直接方法有效，但它是 O(N**2)(我改编自此 stackoverflow answer by user morningsun 的代码)

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) 的递归关系(这在第 4.2 节的 paper 中有解释)。我们定义

通过

找到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 更快

一个小的基准:生成一个轨迹

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