使用python和FFT計算均方位移

Question

給定一個二維數組，其中每一行代表一個粒子的位置矢量，如何有效地計算均方位移（使用FFT）？ 均方位移定義為

其中r（m）是行m的位置矢量，N是行數。

Answer 1

以下為msd的直接方法有效，但它是O（N ** 2）（我使用用戶morningsun修改了此stackoverflow應答的代碼）

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更快地編寫代碼。 下面的考慮和由此產生的算法來自本文 ，我將展示如何在python中實現它。 我們可以通過以下方式拆分MSD

因此，S_2（m）只是位置的自相關。 注意，在一些教科書中，S_2（m）表示為自相關（約定A），而在一些S_2（m）*（Nm）中表示為自相關（約定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），我們利用了這樣一個事實：可以找到（Nm）* S_1（m）的遞歸關系（ 本文在4.2節中對此進行了解釋）。 我們定義

並找到S_1（m）via

這產生以下均方位移代碼

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

Answer 2

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