Add option to select MSD algorithm

src/gemdat/plots/__init__.py

@@ -37,6 +37,7 @@
     'displacement_per_atom',
     'displacement_per_element',
     'frequency_vs_occurence',
+    'msd_per_element',
     'jumps_3d',
     'jumps_3d_animation',
     'jumps_vs_distance',

src/gemdat/plots/matplotlib/_displacements.py

@@ -68,7 +68,10 @@ def displacement_per_element(*, trajectory: Trajectory) -> plt.Figure:
     return fig
 
 
-def msd_per_element(*, trajectory: Trajectory) -> plt.Figure:
+def msd_per_element(
+    *,
+    trajectory: Trajectory,
+) -> plt.Figure:
     """Plot mean squared displacement per element.
 
     Parameters
@@ -85,15 +88,26 @@ def msd_per_element(*, trajectory: Trajectory) -> plt.Figure:
 
     fig, ax = plt.subplots()
 
+    # Since we want to plot in picosecond, we convert the time units
+    time_ps = trajectory.time_step * 1e12
+
     for sp in species:
         traj = trajectory.filter(sp.symbol)
-        ax.plot(traj.mean_squared_displacement().mean(axis=0),
-                lw=0.5,
-                label=sp.symbol)
+        msd = traj.mean_squared_displacement()
+        msd_mean = np.mean(msd, axis=0)
+        msd_std = np.std(msd, axis=0)
+        t_values = np.arange(len(msd_mean)) * time_ps
+        ax.plot(t_values, msd_mean, lw=0.5, label=sp.symbol)
+        last_color = ax.lines[-1].get_color()
+        ax.fill_between(t_values,
+                        msd_mean - msd_std,
+                        msd_mean + msd_std,
+                        color=last_color,
+                        alpha=0.2)
 
     ax.legend()
     ax.set(title='Mean squared displacement per element',
-           xlabel='Time step',
+           xlabel='Time lag [ps]',
            ylabel='MSD (Angstrom$^2$)')
 
     return fig

src/gemdat/plots/matplotlib/_rotations.py

@@ -5,7 +5,7 @@
 from scipy.optimize import curve_fit
 from scipy.stats import skewnorm
 
-from gemdat.rotations import Orientations, autocorrelation, calculate_spherical_areas
+from gemdat.rotations import Orientations, calculate_spherical_areas, mean_squared_angular_displacement
 from gemdat.utils import cartesian_to_spherical
 
 
@@ -165,17 +165,19 @@ def unit_vector_autocorrelation(
     # The trajectory is expected to have shape (n_times, n_particles, n_coordinates)
     trajectory = orientations.get_unit_vectors_trajectory()
 
-    ac, std_ac = autocorrelation(trajectory)
+    ac = mean_squared_angular_displacement(trajectory)
+    ac_std = ac.std(axis=0)
+    ac_mean = ac.mean(axis=0)
 
     # Since we want to plot in picosecond, we convert the time units
     time_ps = orientations._time_step * 1e12
-    t_values = np.arange(len(ac)) * time_ps
+    tgrid = np.arange(ac_mean.shape[0]) * time_ps
 
     # and now we can plot the autocorrelation function
     fig, ax = plt.subplots()
 
-    ax.plot(t_values, ac, label='FFT-Autocorrelation')
-    ax.fill_between(t_values, ac - std_ac, ac + std_ac, alpha=0.2)
+    ax.plot(tgrid, ac_mean, label='FFT-Autocorrelation')
+    ax.fill_between(tgrid, ac_mean - ac_std, ac_mean + ac_std, alpha=0.2)
     ax.set_xlabel('Time lag [ps]')
     ax.set_ylabel('Autocorrelation')
 

src/gemdat/plots/plotly/_displacements.py

@@ -117,18 +117,31 @@ def msd_per_element(*, trajectory: Trajectory) -> go.Figure:
 
     species = list(set(trajectory.species))
 
+    # Since we want to plot in picosecond, we convert the time units
+    time_ps = trajectory.time_step * 1e12
+
     for sp in species:
         traj = trajectory.filter(sp.symbol)
 
+        msd = traj.mean_squared_displacement()
+        msd_mean = np.mean(msd, axis=0)
+        msd_std = np.std(msd, axis=0)
+        t_values = np.arange(len(msd_mean)) * time_ps
+
         fig.add_trace(
-            go.Scatter(y=traj.mean_squared_displacement().mean(axis=0),
+            go.Scatter(x=t_values,
+                       y=msd_mean,
+                       error_y=dict(type='data',
+                                    array=msd_std,
+                                    width=0.1,
+                                    thickness=0.1),
                        name=sp.symbol,
                        mode='lines',
                        line={'width': 3},
                        legendgroup=sp.symbol))
 
     fig.update_layout(title='Mean squared displacement per element',
-                      xaxis_title='Time step',
+                      xaxis_title='Time lag [ps]',
                       yaxis_title='MSD (Angstrom<sup>2</sup>)')
 
     return fig

src/gemdat/rotations.py

@@ -298,8 +298,8 @@ def calculate_spherical_areas(shape: tuple[int, int],
     return areas
 
 
-def autocorrelation(trajectory: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
-    """Compute the autocorrelation of the trajectory using FFT.
+def mean_squared_angular_displacement(trajectory: np.ndarray) -> np.ndarray:
+    """Compute the mean squared angular displacement using FFT.
 
     Parameters
     ----------
@@ -309,34 +309,31 @@ def autocorrelation(trajectory: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
 
     Returns
     -------
-    autocorr.mean : np.ndarray
-        The autocorrelation of the signal mediated over the number of particles.
-    autocorr.std : np.ndarray
-        The standard deviation of the autocorrelation of the signal.
+    msad:
+        The mean squared angular displacement
     """
-
     n_times, n_particles, n_coordinates = trajectory.shape
 
-    # Sum the coordinates to get the magnitude of the signal
-    signal = np.sum(trajectory, axis=2)
-
-    # Compute the FFT of the magnitude
-    fft_magnitude = np.fft.fft(signal, n=2 * n_times - 1, axis=0)
+    msad = np.zeros((n_particles, n_times))
+    normalization = np.arange(n_times, 0, -1)
 
-    # Compute the power spectral density
-    psd = np.abs(fft_magnitude)**2
+    for c in range(n_coordinates):
+        signal = trajectory[:, :, c]
 
-    # Compute the inverse FFT of the power spectral density
-    autocorr = np.fft.ifft(psd, axis=0)
+        # Compute the FFT of the signal
+        fft_signal = np.fft.rfft(signal, n=2 * n_times - 1, axis=0)
+        # Compute the power spectral density in-place
+        np.square(np.abs(fft_signal), out=fft_signal)
+        # Compute the inverse FFT of the power spectral density
+        autocorr_c = np.fft.irfft(fft_signal, axis=0)
 
-    # Only keep the positive time lags
-    autocorr = autocorr[:n_times, :]
+        # Only keep the positive times
+        autocorr_c = autocorr_c[:n_times, :]
 
-    # Normalize
-    autocorr = autocorr.T
-    autocorr /= np.arange(n_times, 0, -1)
+        msad += autocorr_c.T / normalization
 
-    # and get the real part
-    autocorr = autocorr.real
+    # Normalize the msad such that it starts from 1
+    # (this makes the normalization independent on the dimensions)
+    msad = msad / msad[:, 0, np.newaxis]
 
-    return autocorr.mean(axis=0), autocorr.std(axis=0)
+    return msad

src/gemdat/trajectory.py

@@ -46,33 +46,6 @@ def _lengths(vectors: np.ndarray, lattice: Lattice) -> np.ndarray:
     return np.sqrt(lengths_sq)
 
 
-def _unwrap_pbcs(coords: np.ndarray) -> np.ndarray:
-    """Unwrap coordinates using periodic boundary conditions.
-
-    Parameters
-    ----------
-    coords : np.ndarray
-        Input coordinates
-
-    Returns
-    -------
-    unwrapped_coords : np.ndarray
-        Output coordinates where periodic boundary conditions have been removed
-    """
-    timesteps, nparticles, _ = coords.shape
-    unwrapped_coords = np.copy(coords)
-
-    for particle in range(nparticles):
-        for t in range(1, timesteps):
-            displacement = coords[t, particle] - coords[t - 1, particle]
-            crossed_boundaries = np.abs(displacement) > 0.5
-
-            unwrapped_coords[
-                t:, particle] -= np.sign(displacement) * crossed_boundaries
-
-    return unwrapped_coords
-
-
 class Trajectory(PymatgenTrajectory):
     """Trajectory of sites from a molecular dynamics simulation."""
 
@@ -474,19 +447,14 @@ def split(self,
 
     def mean_squared_displacement(self) -> np.ndarray:
         """Computes the mean squared displacement using fast Fourier transform.
+
         The algorithm is described in [https://doi.org/10.1051/sfn/201112010].
         See also [https://stackoverflow.com/questions/34222272/computing-mean-square-displacement-using-python-and-fft].
-
-        Returns
-        -------
-        MSD : np.ndarray
-            Output array with mean squared displacement per particle
         """
-        r = self.positions
-        r = _unwrap_pbcs(r)
-
+        r = self.cumulative_displacements
         lattice = self.get_lattice()
         r = lattice.get_cartesian_coords(r)
+
         pos = np.transpose(r, (1, 0, 2))
         n_times = pos.shape[1]
 
@@ -516,8 +484,8 @@ def mean_squared_displacement(self) -> np.ndarray:
         S1 = (double_sum_D - cumsum_D)[:, :-1] / (n_times -
                                                   np.arange(n_times)[None, :])
 
-        MSD = S1 - 2 * S2
-        return MSD
+        msd = S1 - 2 * S2
+        return msd
 
     def transitions_between_sites(
         self,

tests/conftest.py

@@ -4,6 +4,7 @@
 import pytest
 from pymatgen.core import Species
 
+from gemdat.rotations import Orientations
 from gemdat.trajectory import Trajectory
 
 
@@ -26,3 +27,12 @@ def trajectory():
         lattice=np.eye(3),
         metadata={'temperature': 123},
         time_step=1)
+
+
+@pytest.fixture()
+def orientations(trajectory):
+    center_type = 'B'
+    satellite_type = 'Si'
+    nr_central_atoms = 1
+    return Orientations(trajectory, center_type, satellite_type,
+                        nr_central_atoms)

tests/rotations_test.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+from math import isclose
+
+import numpy as np
+from pymatgen.core import Species
+
+from gemdat.rotations import calculate_spherical_areas, mean_squared_angular_displacement
+
+
+def test_orientations(orientations):
+    assert orientations._time_step == 1
+    assert orientations._trajectory_cent.species == [Species('B')]
+    assert orientations._trajectory_sat.species == [Species('Si')]
+
+
+def test_fractional_coordinates(orientations):
+    frac_coord_cent, frac_coord_sat = orientations._fractional_coordinates()
+    assert isinstance(frac_coord_cent, np.ndarray)
+    assert isinstance(frac_coord_sat, np.ndarray)
+
+
+def test_distances(orientations):
+    distances = orientations._distances
+    assert isinstance(distances, np.ndarray)
+
+
+def test_calculate_spherical_areas():
+    shape = (10, 10)
+    areas = calculate_spherical_areas(shape)
+    assert isclose(areas.mean(), 0.00017275712347752164)
+    assert isinstance(areas, np.ndarray)
+    assert areas.shape == shape
+
+
+def test_mean_squared_angular_displacement(trajectory):
+    msad = mean_squared_angular_displacement(trajectory.positions)
+    assert isinstance(msad, np.ndarray)
+    assert isclose(msad.mean(), 0.8142314269325723)
+    assert msad.shape == (trajectory.positions.shape[1],
+                          trajectory.positions.shape[0])

tests/trajectory_test.py

@@ -135,3 +135,11 @@ def test_trajectory_extend(trajectory):
     assert len(trajectory) == 10
     assert_allclose(trajectory.positions[:, 0, 0],
                     [0.2, 0.4, 0.6, 0.8, 0.1, 0.2, 0.4, 0.6, 0.8, 0.1])
+
+
+def test_mean_squared_displacement(trajectory):
+    msd = trajectory.mean_squared_displacement()
+    assert len(msd) == 4
+    assert_allclose(msd[0], [0.0, 0.0525, 0.19, 0.425, 0.81])
+    assert isinstance(msd, np.ndarray)
+    assert_allclose(msd.mean(), 0.073875)