Computing the root mean square displacement using python and FFT

Question

For a two-dimensional array, where each row represents a particle position vector, how can I efficiently calculate the mean square displacement (using FFT)? The root mean square displacement is defined as

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

Answer 1

The following direct method is used for the msd method, but it is O (N ** 2) (I adapted the code from this fooobar.com/questions/1233128 / ... )

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 

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

Thus, S_2 (m) is simply an autocorrelation of a position. Note that in some textbooks, S_2 (m) is denoted as autocorrelation (agreement A), and in some S_2 (m) * (Nm) is denoted as autocorrelation (agreement B). By the Wiener-Khinchin theorem, the power spectral density (PSD) of a function is the Fourier transform of autocorrelation. This means that we can calculate the PSD of the signal and Fourier transform to obtain autocorrelation (in agreement B). For discrete signals, we obtain cyclic autocorrelation. However, with zero data filling, we can get non-cyclic autocorrelation. The algorithm is as follows:

 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 (Nm) return res/n #this is the autocorrelation in convention A 

For the term S_1 (m), we use the fact that we can find a recursive relation for (Nm) * S_1 (m) (this is explained in this in Section 4.2). Define

And find S_1 (m) through

This gives the following code for the mean square offset

 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=QD[m-1]-D[Nm] S1[m]=Q/(Nm) return S1-2*S2 

You can compare msd_straight_forward () and msd_fft () and find that they give the same results, although msd_fft () works faster with large N

Small benchmark: create a path with

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

At 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 

Answer 2

Using numpy.cumsum, you can also avoid looping over range (N) in S1 calculation:

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