Add in analysis function to calculate MSDs

package/MDAnalysis/analysis/msd.py

@@ -0,0 +1,213 @@
+# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
+# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
+#
+# MDAnalysis --- https://www.mdanalysis.org
+# Copyright (c) 2006-2017 The MDAnalysis Development Team and contributors
+# (see the file AUTHORS for the full list of names)
+#
+# Released under the GNU Public Licence, v2 or any higher version
+#
+# Please cite your use of MDAnalysis in published work:
+#
+# R. J. Gowers, M. Linke, J. Barnoud, T. J. E. Reddy, M. N. Melo, S. L. Seyler,
+# D. L. Dotson, J. Domanski, S. Buchoux, I. M. Kenney, and O. Beckstein.
+# MDAnalysis: A Python package for the rapid analysis of molecular dynamics
+# simulations. In S. Benthall and S. Rostrup editors, Proceedings of the 15th
+# Python in Science Conference, pages 102-109, Austin, TX, 2016. SciPy.
+# doi: 10.25080/majora-629e541a-00e
+#
+# N. Michaud-Agrawal, E. J. Denning, T. B. Woolf, and O. Beckstein.
+# MDAnalysis: A Toolkit for the Analysis of Molecular Dynamics Simulations.
+# J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787
+#
+
+"""
+Mean Squared Displacement --- :mod:`MDAnalysis.analysis.msd`
+==============================================================
+
+This module implements the calculation of Mean Squared Displacmements (MSD).
+MSDs can be used to characterise the speed at which particles move and has its roots
+in the study of Brownian motion.  For a thorough review see XXX et al. MSDs are computed from
+the following expression
+
+Where XX represents an ensemble average over 
+
+The computation of the MSD in this way can be computationally intensive due to it's N^2 scaling. 
+An algorithm to compute the MSD with Nlogn(N) scaling based on a Fast Fourier Transform is known and can be accessed by setting FFT=True.
+
+
+
+
+Classes and Functions
+---------------------
+
+.. autoclass:: MeanSquaredDisplacement
+
+
+"""
+
+from __future__ import division, absolute_import
+from six.moves import zip
+
+import os
+import errno
+import warnings
+import bz2
+import functools
+
+import numpy as np
+from scipy import fft,ifft
+import logging
+import MDAnalysis
+
+
+
+class MeanSquaredDisplacement(object):
+    r"""Class representing a density on a regular cartesian grid.
+
+    Parameters
+    ----------∏
+    u : 
+        An MDAnalysis Universe :class:`Universe`
+    selection : 
+        An MDAnalysis selection string
+
+
+
+    Attributes
+    ----------
+
+
+
+
+    Notes
+    -----
+    Notes
+
+
+    See Also
+    --------
+
+
+    Examples
+    --------
+    Typical use:
+
+
+
+    """
+
+    def __init__(self, u, selection, msd_type='xyz', fft=True):
+
+        #args
+        self.u = u
+        self.selection = selection
+        self.msd_type = msd_type
+        self.fft = fft
+
+        #local
+        self.dim_fac = 0
+        self._dim = None
+        self.atoms = None
+        self.n_frames = len(self.u.trajectory)
+        self._position_array = None
+
+        #result
+        self.timeseries = None
+        #prep
+        self._prepare()
+
+
+    def _prepare(self):
+        self.parse_msd_type()
+        self.select_reference_positions()
+
+    def parse_msd_type(self):
+
+        if self.msd_type == 'xyz': # full 3d
+            self._dim = [0,1,2]
+            self.dim_fac = 3.0
+
+        elif self.msd_type == 'xy': # xy
+            self._dim = [0,1]
+            self.dim_fac = 2.0
+
+        elif self.msd_type == 'xz': # xz
+            self._dim = [0,2]
+            self.dim_fac = 2.0
+
+        elif self.msd_type == 'yz': # yz
+            self._dim = [1,2]
+            self.dim_fac = 2.0
+
+        elif self.msd_type == 'x': # x
+            self._dim = [0]
+            self.dim_fac = 1.0
+
+        elif self.msd_type == 'y': # y
+            self._dim = [1]
+            self.dim_fac = 1.0
+
+        elif self.msd_type == 'z': # z
+            self._dim = [2]
+            self.dim_fac = 1.0
+
+        else:
+            raise ValueError('invalid msd_type specified')
+
+    def select_reference_positions(self):
+        self._position_array = self.u.trajectory.timeseries(self.u.select_atoms(self.selection),order='fac') 
+        self.N_particles = self._position_array.shape[1] 
+
+    def run(self):
+        if self.fft == True:
+            self.timeseries = self._run_fft()
+        else:
+            self.timeseries = self._run_simple()
+
+    def _run_simple(self): # naieve algorithm without FFT
+        # _position_array is shape time, nparticles, 3
+        msds_byparticle = np.zeros([self.n_frames, self.N_particles])
+        lagtimes = np.arange(1,self.n_frames)
+        msds_byparticle[0,:] = np.zeros(self.N_particles) # preset the zero lagtime so we dont have to iterate through
+        for n in range(self.N_particles):
+            for lag in lagtimes:
+                disp = self._position_array[:-lag,n,self._dim if lag else None] - self._position_array[lag:,n,self._dim]
+                sqdist = np.square(disp, dtype=np.float64).sum(axis=1, dtype=np.float64) #accumulation in anything other than f64 is innacurate
+                msds_byparticle[lag,n] = np.mean(sqdist, dtype=np.float64)
+        msds = msds_byparticle.mean(axis=1, dtype=np.float64)
+        return msds
+
+    def _run_fft(self):  #with FFT
+        # _position_array is shape time, nparticles, 3
+        msds_byparticle = []
+        reshape_positions = self._position_array[:,:,self._dim]
+        N=reshape_positions.shape[0]
+        D=np.square(reshape_positions).sum(axis=2, dtype=np.float64) 
+        D=np.append(D,np.zeros(reshape_positions.shape[:2]), axis=0) 
+        Q=2*D.sum(axis=0)
+        S1=np.zeros(reshape_positions.shape[:2])
+        for m in range(N):
+            Q=Q-D[m-1,:]-D[N-m,:]
+            S1[m,:]=Q/(N-m)
+
+        S2accumulate = []
+        for i in range(reshape_positions.shape[2]):
+            S2accumulate.append(self.autocorrFFT(reshape_positions[:,:,i]))
+        S2= np.sum(S2accumulate,axis=0,dtype=np.float64)
+
+        msds_byparticle.append(S1-2*S2)
+        msds = np.concatenate(msds_byparticle,axis=1).mean(axis=-1, dtype=np.float64)
+        return msds
+
+    @staticmethod
+    def autocorrFFT(x):
+        N=(x.shape[0])
+        F = fft(x, n=2*N, axis=0)  #zero pad to get non-cyclic autocorrelation
+        PowerSpectralDensity = F * F.conjugate()
+        inverse = ifft(PowerSpectralDensity,axis=0)
+        autocorr = (inverse[:N]).real   #autocorr convention B
+        n = np.arange(1, N+1)[::-1] 
+        return autocorr/n[:, np.newaxis] #autocorr convention A
+
+

package/doc/sphinx/source/documentation_pages/analysis/msd.rst

@@ -0,0 +1 @@
+.. automodule:: MDAnalysis.analysis.msd

package/doc/sphinx/source/documentation_pages/analysis_modules.rst

@@ -107,7 +107,7 @@ Structure
    analysis/helanal
    analysis/rdf
    analysis/dihedrals
-
+   analysis/msd
 
 Volumetric analysis
 ===================

testsuite/MDAnalysisTests/analysis/test_msd.py

@@ -0,0 +1,128 @@
+# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
+# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 fileencoding=utf-8
+#
+# MDAnalysis --- https://www.mdanalysis.org
+# Copyright (c) 2006-2017 The MDAnalysis Development Team and contributors
+# (see the file AUTHORS for the full list of names)
+#
+# Released under the GNU Public Licence, v2 or any higher version
+#
+# Please cite your use of MDAnalysis in published work:
+#
+# R. J. Gowers, M. Linke, J. Barnoud, T. J. E. Reddy, M. N. Melo, S. L. Seyler,
+# D. L. Dotson, J. Domanski, S. Buchoux, I. M. Kenney, and O. Beckstein.
+# MDAnalysis: A Python package for the rapid analysis of molecular dynamics
+# simulations. In S. Benthall and S. Rostrup editors, Proceedings of the 15th
+# Python in Science Conference, pages 102-109, Austin, TX, 2016. SciPy.
+# doi: 10.25080/majora-629e541a-00e
+#
+# N. Michaud-Agrawal, E. J. Denning, T. B. Woolf, and O. Beckstein.
+# MDAnalysis: A Toolkit for the Analysis of Molecular Dynamics Simulations.
+# J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787
+#
+from __future__ import division, absolute_import, print_function
+
+
+
+import MDAnalysis as mda
+from MDAnalysis.analysis.msd import MeanSquaredDisplacement as MSD
+
+from numpy.testing import (assert_array_less,
+                           assert_almost_equal, assert_equal)
+import numpy as np
+
+from scipy import fft,ifft
+
+from MDAnalysisTests.datafiles import PSF, DCD, DCD
+
+import pytest
+
+SELECTION = 'backbone and name CA and resid 1-10'
+NSTEP = 1000
+
+@pytest.fixture(scope='module')
+def u():
+    return mda.Universe(PSF, DCD)
+
+@pytest.fixture(scope='module')
+def msd(u):
+    m = MSD(u, SELECTION, msd_type='xyz', fft=False)
+    m.run()
+    return m
+
+@pytest.fixture(scope='module')
+def msd_fft(u):
+    m = MSD(u, SELECTION, msd_type='xyz', fft=True)
+    m.run()
+    return m
+
+@pytest.fixture(scope='module')
+def dimension_list():
+    dimensions = ['xyz', 'xy', 'xz', 'yz', 'x', 'y', 'z']
+    return dimensions
+
+@pytest.fixture(scope='module')
+def step_traj():
+    x = np.arange(NSTEP)
+    traj = np.vstack([x,x,x]).T
+    traj_reshape = traj.reshape([NSTEP,1,3])
+    u = mda.Universe.empty(1)
+    u.load_new(traj_reshape)
+    return u
+
+def characteristic_poly(n,d): #polynomial that describes unit step trajectory MSD
+    x = np.arange(1,n+1)
+    y = d*((x-1)*(x-1))
+    return y
+
+def test_fft_vs_simple_default(msd, msd_fft):
+    timeseries_simple = msd.timeseries
+    timeseries_fft = msd_fft.timeseries
+    assert_almost_equal(timeseries_simple, timeseries_fft, decimal=5)
+
+def test_fft_vs_simple_all_dims(dimension_list, u):
+    for dim in dimension_list:
+        m_simple = MSD(u, SELECTION, msd_type=dim, fft=False)
+        m_simple.run()
+        timeseries_simple = m_simple.timeseries
+        m_fft = MSD(u,SELECTION, msd_type=dim, fft=True)
+        m_fft.run()
+        timeseries_fft = m_fft.timeseries
+        assert_almost_equal(timeseries_simple, timeseries_fft, decimal=5)
+
+def test_simple_step_traj_3d(step_traj): # this should fit the polynomial 3(x-1)**2
+    m_simple = MSD(step_traj, 'all' , msd_type='xyz', fft=False)
+    m_simple.run()
+    poly3 = characteristic_poly(NSTEP,3)
+    assert_almost_equal(m_simple.timeseries, poly3)
+
+def test_simple_step_traj_2d(step_traj): # this should fit the polynomial 2(x-1)**2
+    m_simple = MSD(step_traj, 'all' , msd_type='xy',  fft=False)
+    m_simple.run()
+    poly2 = characteristic_poly(NSTEP,2)
+    assert_almost_equal(m_simple.timeseries, poly2)
+
+def test_simple_step_traj_1d(step_traj): # this should fit the polynomial (x-1)**
+    m_simple = MSD(step_traj, 'all' , msd_type='x', fft=False)
+    m_simple.run()
+    poly1 = characteristic_poly(NSTEP,1)
+    assert_almost_equal(m_simple.timeseries, poly1)  
+
+def test_fft_step_traj_3d(step_traj): # this should fit the polynomial 3(x-1)**2
+    m_fft = MSD(step_traj, 'all' , msd_type='xyz', fft=True)
+    m_fft.run()
+    poly3 = characteristic_poly(NSTEP,3)
+    assert_almost_equal(m_fft.timeseries, poly3)
+
+def test_fft_step_traj_2d(step_traj): # this should fit the polynomial 2(x-1)**2
+    m_fft = MSD(step_traj, 'all' , msd_type='xy', fft=True)
+    m_fft.run()
+    poly2 = characteristic_poly(NSTEP,2)
+    assert_almost_equal(m_fft.timeseries, poly2)
+
+def test_fft_step_traj_1d(step_traj): # this should fit the polynomial (x-1)**2
+    m_fft = MSD(step_traj, 'all' , msd_type='x', fft=True)
+    m_fft.run()
+    poly1 = characteristic_poly(NSTEP,1)
+    assert_almost_equal(m_fft.timeseries, poly1)
+