Calcolo MSD accelerazione in Python

Question

È una chiamata alla comunità per verificare se qualcuno ha un'idea per migliorare la velocità dell'implementazione del calcolo MSD. Si basa in gran parte sull'implementazione da questo post del blog: http://damcb.com/mean-square-disp.html

Per ora l'implementazione corrente richiede circa 9 secondi per una traiettoria 2D di 5 000 punti. E 'davvero un po' troppo se avete bisogno di calcolare un sacco di traiettorie ...

Non ho provato a parallelizzare esso (con multiprocess o joblib), ma ho la sensazione che la creazione di nuovi processi sarà troppo pesante per questo tipo di algoritmo.

Ecco il codice:

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

E l'output:

  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 

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) 

E l'output:

 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 

E alcuni profiling:

%timeit msd = compute_msd(traj, t_step=dt, coords=['x', 'y']) 

dare a questo:

1 loops, best of 3: 8.53 s per loop 

Qualche idea?

Answer 1

I calcoli MSD menzionati finora sono O (N ** 2) dove N è il numero di passaggi temporali. Usando la FFT questo può essere ridotto a O (N * log (N)). Vedi this question and answer per una spiegazione e un'implementazione in python.

EDIT: Un piccolo punto di riferimento (Ho anche aggiunto questo benchmark to this answer): Genera una traiettoria con

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

Per N = 100.000, otteniamo

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

Ha eseguito il profiling riga per riga e sembra che i panda lo stiano rallentando. Questa versione pura NumPy è di circa 14x più veloce:

def compute_msd_np(xy, t, t_step): 
    shifts = np.floor(t/t_step).astype(np.int) 
    msds = np.zeros(shifts.size) 
    msds_std = np.zeros(shifts.size) 

    for i, shift in enumerate(shifts): 
     diffs = xy[:-shift if shift else None] - xy[shift:] 
     sqdist = np.square(diffs).sum(axis=1) 
     msds[i] = sqdist.mean() 
     msds_std[i] = sqdist.std(ddof=1) 

    msds = pd.DataFrame({'msds': msds, 'tau': t, 'msds_std': msds_std}) 
    return msds 

Answer 3

Aggiungendo a moarningsun risposta di cui sopra:

• è possibile accelerare utilizzando numexpr

• se si traccia il MSD in scala logaritmica in ogni caso, è don 't bisogno di calcolare per ogni volta

  import numpy as np 
  import numexpr 
  
  def logSpaced(L, pointsPerDecade=15): 
      """Generate an array of log spaced integers smaller than L""" 
      nbdecades = np.log10(L) 
      return np.unique(np.logspace(
       start=0, stop=nbdecades, 
       num=nbdecades * pointsPerDecade, 
       base=10, endpoint=False 
       ).astype(int)) 
  
  def compute_msd(xy, pointsPerDecade=15): 
      dts = logSpaced(len(xy), pointsPerDecade) 
      msd = np.zeros(len(idts)) 
      msd_std = np.zeros(len(idts)) 
      for i, dt in enumerate(dts): 
       sqdist = numexpr.evaluate(
        '(a-b)**2', 
        {'a': xy[:-dt], 'b':xy[dt:]} 
        ).sum(axis=-1) 
       msd[i] = sqdist.mean() 
       msd_std[i] = sqdist.std(ddof=1) 
      msds = pd.DataFrame({'msds': msd, 'tau': dt, 'msds_std': msd_std}) 
      return msds 
  

Answer 4

Con i commenti Ho progettato questa funzione:

def get_msd(traj, dt, with_nan=True): 

    shifts = np.arange(1, len(traj), dtype='int') 
    msd = np.empty((len(shifts), 2), dtype='float') 
    msd[:] = np.nan 

    msd[:, 1] = shifts * dt 

    for i, shift in enumerate(shifts): 
     diffs = traj[:-shift] - traj[shift:] 
     if with_nan: 
      diffs = diffs[~np.isnan(diffs).any(axis=1)] 
     diffs = np.square(diffs).sum(axis=1) 

     if len(diffs) > 0: 
      msd[i, 0] = np.mean(diffs) 

    msd = pd.DataFrame(msd) 
    msd.columns = ["msd", "delay"] 

    msd.set_index('delay', drop=True, inplace=True) 
    msd.dropna(inplace=True) 

    return msd 

con le seguenti caratteristiche:

• Prende numpy matrice come input traiettoria.
• Restituisce uno pandas.DataFrame quasi senza sovrapposizione.
• with_nan consente di gestire la traiettoria contenente i valori NaN ma aggiunge un grande overhead (oltre il 100%), quindi l'ho inserito come parametro di funzione.
• Può occuparsi di traiettorie multi dimensionale (1D, 2D, 3D, ecc)

Alcuni profiling:

$ print(traj.shape) 
(2108, 2) 

$ %timeit get_msd(traj, with_nan=True, dt=0.1) 
10 loops, best of 3: 143 ms per loop 

$ %timeit get_msd(traj, with_nan=False, dt=0.1) 
10 loops, best of 3: 68 ms per loop 

Answer 5

Forse non è l'argomento, tuttavia MSD deve essere calcolato non come la linea 37:

msds[i] = sqdist.mean() 

Prendendo come mean=N

È necessario dividere per:

msds[i] = sqdist/N-1 // for lag1 

Poi:

msds[i] = sqdist/N-2 // for lag2 .... msds[i] = sqdist/N-n // for lag n 

E così via.

Di conseguenza non si ottiene la deviazione standard, solo l'MSD per una sola traiettoria