使用python和FFT计算均方位移 [英] Computing mean square displacement using python and FFT

问题描述

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

Given a 2 dimensional array, where each row represents the position vector of a particle, how do I compute the mean square displacement efficiently (using FFT)? The mean square displacement is defined as

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

where r(m) is the position vector of row m, and N is the number of rows.

推荐答案

适用于msd的以下直接方法有效，但它是O(N ** 2)(我从

The following straight forward method for the msd works, but it is O(N**2) (I adapted the code from this 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来使此代码更快.以下考虑因素和产生的算法来自本文，我将展示如何在python中实现它. 我们可以通过以下方式拆分MSD

However, we can make this code way faster using the FFT. The following consideration and the resulting algorithm are from this paper, I will just show how to implement it in python. We can split the MSD in the following way

因此，S_2(m)仅仅是位置的自相关.注意，在某些教科书中，S_2(m)表示为自相关(约定A)，在某些S_2(m)*(N-m)中表示为自相关(约定B). 根据维纳-欣钦定理，函数的功率谱密度(PSD)是自相关的傅立叶变换. 这意味着我们可以计算信号的PSD并对其进行傅立叶反转，以获得自相关(按照惯例B).对于离散信号，我们得到了循环自相关. 但是，通过对数据进行零填充，我们可以获得非循环自相关.算法看起来像这样

Thereby, S_2(m) is just the autocorrelation of the position. Note that in some textbooks S_2(m) is denoted as autocorrelation (convention A) and in some S_2(m)*(N-m) is denoted as autocorrelation (convention B). By the Wiener–Khinchin theorem, the power spectral density (PSD) of a function is the Fourier transform of the autocorrelation. This means we can compute a PSD of a signal and Fourier-invert it, to get the autocorrelation (in convention B). For discrete signals we get the cyclic autocorrelation. However, by zero-padding the data, we can get the non-cyclic autocorrelation. The algorithm looks like this

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)的递归关系(这在

For the term S_1(m), we exploit the fact, that a recursive relation for (N-m)*S_1(m) can be found (This is explained in this paper in section 4.2). We define

并通过

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

This yields the following code for the mean square displacement

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而言速度更快

You can compare msd_straight_forward() and msd_fft() and will find that they yield the same results, though msd_fft() is way faster for large N

一个小基准:使用

r = np.cumsum(np.random.choice([-1., 0., 1.], size=(N, 3)), axis=0)

对于N = 100.000，我们得到

For N=100.000, we get

$ %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

这篇关于使用python和FFT计算均方位移的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！