<jupyter_start><jupyter_text>Compute MSD in Python quickly with FFTThe implementation is entirely copied from this SO answer : http://stackoverflow.com/questions/34222272/computing-mean-square-displacement-using-python-and-fft34222273Thanks to [thomasfermi](http://stackoverflow.com/users/2609987/thomasfermi) for the implementation. The original algorithm comes from : ```nMoldyn - Interfacing spectroscopic experiments, molecular dynamics simulations and models for time correlation functions<NAME>, <NAME>, <NAME>, <NAME> and <NAME>Publié en ligne: 9 juin 2011DOI: 10.1051/sfn/201112010```This notebook compares the speed of calculation of MSD with FFT and without with trackpy.<jupyter_code>%load_ext snakeviz
%load_ext autoreload
%autoreload 2
%matplotlib inline

import timeit

import matplotlib.pyplot as plt
import matplotlib
matplotlib.style.use('ggplot')

import numpy as np
import pandas as pd

import trackpy as tp
def autocorr_fft(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

def msd_fft(traj, mpp, fps, max_lagtime=100, detail=False, pos_columns=['x', 'y']):
    """Implementation comes from here : 
    http://stackoverflow.com/questions/34222272/computing-mean-square-displacement-using-python-and-fft#34222273
    """
    
    r = traj[pos_columns].values
    r *= mpp

    t = traj['frame']

    max_lagtime = min(max_lagtime, len(t))  # checking to be safe
    lagtimes = 1 + np.arange(max_lagtime - 1)    

    N = len(r)

    D = np.square(r).sum(axis=1) 
    D = np.append(D, 0)
    S2 = sum([autocorr_fft(r[:, i]) for i in range(len(pos_columns))])

    Q = 2 * D.sum()
    S1 = np.zeros(max_lagtime)

    for m in range(max_lagtime):
        Q = Q - D[m - 1] - D[N - m]
        S1[m] = Q / (N - m)

    msd = S1 - 2 * S2[:max_lagtime]
    msd = msd[1:]

    lagt = lagtimes / fps

    results = pd.DataFrame(np.array([msd, lagt]).T, columns=['msd', 'lagt'])
    results.index = 1 + np.arange(max_lagtime - 1)
    results.index.name = 'lagt'
    
    return results

def get(N, pos_columns):
    traj = np.cumsum(np.random.choice([-1., 0., 1.], size=(N, len(pos_columns))), axis=0)
    traj = pd.DataFrame(traj, columns=pos_columns)
    traj['frame'] = np.arange(N)
    return traj
# Generate a random trajectory

pos_columns = ['x', 'y', 'z']
mpp = 0.3
fps = 1
max_lagtime = 1000

traj = get(N=10000, pos_columns=pos_columns)
# Get MSD with trackpy
msd1 = tp.msd(traj, mpp=mpp, fps=fps, max_lagtime=max_lagtime, pos_columns=pos_columns)
msd1[['msd', 'lagt']].head()
# Get MSD with msd_fft
msd2 = msd_fft(traj, mpp=mpp, fps=fps, max_lagtime=max_lagtime, pos_columns=pos_columns)
msd2[['msd', 'lagt']].head()
# Compate MSD with plots
fig, ax = plt.subplots(figsize=(14, 10))

msd1.plot(x='lagt', y='msd', ax=ax)
msd2.plot(x='lagt', y='msd', ax=ax)
t1 = %timeit -o tp.msd(traj, mpp=mpp, fps=fps, max_lagtime=max_lagtime, pos_columns=pos_columns)
t1.best
t2 = %timeit -o msd_fft(traj, mpp=mpp, fps=fps, max_lagtime=max_lagtime, pos_columns=pos_columns)
t2.best<jupyter_output>100 loops, best of 3: 5.4 ms per loop
{'fraction_non_alphanumeric': 0.14229960972680877, 'fraction_numerical': 0.035124587211047734, 'mean_word_length': 3.9362962962962964, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '77726', 'license': 'MIT', 'n_tokens_mistral': 1403, 'n_tokens_neox': 1308, 'n_words': 321, 'path': 'Code/Quick_MSD/notebook.ipynb', 'repo_name': 'hadim/Public-Notebooks', 'size': '3360'}