title: Speedup MSD calculation in Python
body: <p>It's a call to the community to see if anyone has an idea to improve the speed of this MSD calculation implementation. It is largely based on the implementation from this blog post : <a href="http://damcb.com/mean-square-disp.html" rel="noreferrer">http://damcb.com/mean-square-disp.html</a></p>
  
  <p>For now the current implementation takes about 9s for a 2D trajectory of 5 000 points. It's really way too much if you need to compute a lot of trajectories...</p>
  
  <p>I didn't try to parallelize it (with <code>multiprocess</code> or <code>joblib</code>) but I have the feeling that creating new processes will be too heavy for this kind of algorithm.</p>
  
  <p>Here is the code : </p>
  
  <pre class="lang-py prettyprint-override"><code>import os
  
  import matplotlib
  import matplotlib.pyplot as plt
  
  import pandas as pd
  import numpy as np
  
  # Parameters
  N = 5000
  max_time = 100
  dt = max_time / N
  
  # Generate 2D brownian motion
  
  t = np.linspace(0, max_time, N)
  xy = np.cumsum(np.random.choice([-1, 0, 1], size=(N, 2)), axis=0)
  traj = pd.DataFrame({'t': t, 'x': xy[:,0], 'y': xy[:,1]})
  print(traj.head())
  
  # Draw motion
  ax = traj.plot(x='x', y='y', alpha=0.6, legend=False)
  
  # Set limits
  ax.set_xlim(traj['x'].min(), traj['x'].max())
  ax.set_ylim(traj['y'].min(), traj['y'].max())
  </code></pre>
  
  <p>And the output : </p>
  
  <pre><code>          t  x  y
  0  0.000000 -1 -1
  1  0.020004 -1  0
  2  0.040008 -1 -1
  3  0.060012 -2 -2
  4  0.080016 -2 -2
  </code></pre>
  
  <p><a href="https://i.stack.imgur.com/yyP8x.png" rel="noreferrer"><img src="https://i.stack.imgur.com/yyP8x.png" alt="enter image description here"></a></p>
  
  <pre class="lang-py prettyprint-override"><code>def compute_msd(trajectory, t_step, coords=['x', 'y']):
  
      tau = trajectory['t'].copy()
      shifts = np.floor(tau / t_step).astype(np.int)
      msds = np.zeros(shifts.size)
      msds_std = np.zeros(shifts.size)
  
      for i, shift in enumerate(shifts):
          diffs = trajectory[coords] - trajectory[coords].shift(-shift)
          sqdist = np.square(diffs).sum(axis=1)
          msds[i] = sqdist.mean()
          msds_std[i] = sqdist.std()
  
      msds = pd.DataFrame({'msds': msds, 'tau': tau, 'msds_std': msds_std})
      return msds
  
  # Compute MSD
  msd = compute_msd(traj, t_step=dt, coords=['x', 'y'])
  print(msd.head())
  
  # Plot MSD
  ax = msd.plot(x="tau", y="msds", logx=True, logy=True, legend=False)
  ax.fill_between(msd['tau'], msd['msds'] - msd['msds_std'], msd['msds'] + msd['msds_std'], alpha=0.2)
  </code></pre>
  
  <p>And the output : </p>
  
  <pre><code>       msds  msds_std       tau
  0  0.000000  0.000000  0.000000
  1  1.316463  0.668169  0.020004
  2  2.607243  2.078604  0.040008
  3  3.891935  3.368651  0.060012
  4  5.200761  4.685497  0.080016
  </code></pre>
  
  <p><a href="https://i.stack.imgur.com/vK7Sc.png" rel="noreferrer"><img src="https://i.stack.imgur.com/vK7Sc.png" alt="enter image description here"></a></p>
  
  <p>And some profiling : </p>
  
  <pre class="lang-py prettyprint-override"><code>%timeit msd = compute_msd(traj, t_step=dt, coords=['x', 'y'])
  </code></pre>
  
  <p>Give this : </p>
  
  <pre><code>1 loops, best of 3: 8.53 s per loop
  </code></pre>
  
  <p>Any idea ?</p>