用python和FFT计算均方位移

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

---

其中r（m）是m行的位置向量，N是行数。

以下msd的直接方法可行，但它是O（N**2）（我从中改编了代码）

但是，我们可以使用FFT使该代码更快。下面的考虑事项和生成的算法来自，我将只展示如何在python中实现它。 我们可以按以下方式拆分MSD

因此，S_2（m）只是位置的自相关。注意，在一些教科书中，S_2（m）表示为自相关（约定A），而在一些教科书中，S_2（m）*（N-m）表示为自相关（约定B）。 根据维纳-金钦定理，函数的功率谱密度（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节中解释）。 我们定义

并通过

这将产生以下均方位移代码

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

---

使用numpy.cumsum还可以避免S1计算中的范围（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))

---

非常好的实现。例如，您将如何修改

msd_fft（）

以仅计算第100个第一延迟时间？@HadiM您的意思是更有效地计算

msd_fft（r）[：100]

？准确地说，如果可能的话！Mhh，对不起，我认为不可能做得更有效（当然对于S1我们可以，但是对于S2，这是主要的工作，似乎不可能）。。。如果您只需要在前100个时间步中使用MSD，那么使用直接方法可能会更快，并且在循环中只会增加到100。你应该检查你的数据什么更快。好的。非常感谢。我会做更多的测试。为了记录在案，我询问了trackpy开发人员是否对您的算法感兴趣：

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