Add a few basic analysis functions

CHANGELOG.md

@@ -15,6 +15,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   such as in Monte Carlo simulations.
 - The `ForceField` object is extended with `try_move` and `accept_move`
   to support (relatively) efficient Monte Carlo algorithms with TinyFF.
+- Basic analysis routines for radial distribution functions and autocorrelation functions.
 
 
 ### Changed

README.md

@@ -39,18 +39,18 @@ pip install tinyff
 
 TinyFF is a Python package with the following modules:
 
+- `tinyff.analysis`: helper functions to compute a radial distribution function
+  and an autocorrelation function.
 - `tinyff.atomsmithy`: functions for creating initial atomic positions.
 - `tinyff.forcefield`: implements a general force field interface: energy, atomic forces and pressure
 - `tinyff.pairwise`: pairwise potentials to be used in force fields.
 - `tinyff.neighborlist`: used by the `forcefield` module to compute pairwise interactions
-   with a real-space cut-off.
+  with a real-space cut-off.
 - `tinyff.trajectory`: tools for writing molecular dynamics trajectories to file(s).
 - `tinyff.utils`: utility functions used in `tinyff`.
 
 Most relevant functions and classes can be imported directly from the top-level `tinyff` package.
-
-None of these modules implement any molecular dynamics integrators,
-nor any post-processing of molecular dynamics trajectories.
+None of these modules implement any molecular dynamics integrators.
 This is the part that you are expected to write.
 
 

src/tinyff/__init__.py

@@ -18,13 +18,16 @@
 # --
 """The TinyFF package."""
 
+from .analysis import compute_acf, compute_rdf
 from .atomsmithy import build_bcc_lattice, build_cubic_lattice, build_fcc_lattice, build_random_cell
 from .forcefield import ForceField
 from .neighborlist import NBuildCellLists, NBuildSimple
 from .pairwise import CutOffWrapper, LennardJones, PairwiseTerm
 from .trajectory import NPYWriter, PDBWriter
 
 __all__ = (
+    "compute_rdf",
+    "compute_acf",
     "build_bcc_lattice",
     "build_fcc_lattice",
     "build_cubic_lattice",

src/tinyff/analysis.py

@@ -0,0 +1,89 @@
+# TinyFF is a minimalistic Force Field evaluator.
+# Copyright (C) 2024 Toon Verstraelen
+#
+# This file is part of TinyFF.
+#
+# TinyFF is free software; you can redistribute it and/or
+# modify it under the terms of the GNU General Public License
+# as published by the Free Software Foundation; either version 3
+# of the License, or (at your option) any later version.
+#
+# TinyFF is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program; if not, see <http://www.gnu.org/licenses/>
+# --
+"""Trajectory analysis functions."""
+
+import numpy as np
+from numpy.typing import ArrayLike, NDArray
+from scipy.signal import correlate
+
+from .neighborlist import NBuild
+
+__all__ = ("compute_rdf", "compute_acf")
+
+
+def compute_rdf(
+    traj_atpos: ArrayLike, cell_lengths: ArrayLike, spacing: float, nbuild: NBuild
+) -> tuple[NDArray[float], NDArray[float]]:
+    """Compute an RDF in post-processing, given a trajectory of atomic positions.
+
+    Parameters
+    ----------
+    traj_atpos
+        Configurational snaphots in a 3D array with shape `(nstep, natom, 3)`.
+    cell_length
+        The length of the edge of a cubic simulation cell.
+    spacing
+        The with of the bins of the histogram.
+    nbuild
+        The neighborlist build algorithm for the computation of pairwise distances.
+
+    Returns
+    -------
+    bin_mids
+        The midpoints of the histogram bins used to count the number of pairs.
+    rdf
+        The radial distribution function at each bin midpoint.
+    """
+    bins = np.arange(int(np.floor(nbuild.rmax / spacing)) + 1) * spacing
+    counts = 0
+    for atpos in traj_atpos:
+        cell_lengths = nbuild.update(atpos, cell_lengths)
+        counts += np.histogram(nbuild.nlist["dist"], bins)[0]
+    sphere_vols = (4 * np.pi / 3) * bins**3
+    delta_vols = sphere_vols[1:] - sphere_vols[:-1]
+    rho_pair = counts / (traj_atpos.shape[0] * delta_vols)
+    natom = traj_atpos.shape[1]
+    rho_pair0 = ((natom - 1) * natom) / (2 * np.prod(cell_lengths))
+    bin_mids = (bins[1:] + bins[:-1]) / 2
+    return bin_mids, rho_pair / rho_pair0
+
+
+def compute_acf(traj_data):
+    """Compute the autocorrelation function of time-dependent data.
+
+    Parameters
+    ----------
+    traj_data
+        An array of which the first index corresponds to an equidistant time step.
+        The autocorrelation function is averaged over all remaining indexes.
+
+    Returns
+    -------
+    acf
+        The autocorrelation function, as a function of time lag,
+        at the same equidistant time steps of the input.
+    """
+    traj_data = traj_data.reshape((traj_data.shape[0], -1))
+    acf = 0
+    for column in traj_data.T:
+        acf += correlate(column, column, mode="full")
+    acf = acf[traj_data.shape[0] - 1 :]
+    acf /= traj_data.shape[1]
+    acf /= np.arange(traj_data.shape[0], 0, -1)
+    return acf

tests/test_analysis.py

@@ -0,0 +1,47 @@
+# TinyFF is a minimalistic Force Field evaluator.
+# Copyright (C) 2024 Toon Verstraelen
+#
+# This file is part of TinyFF.
+#
+# TinyFF is free software; you can redistribute it and/or
+# modify it under the terms of the GNU General Public License
+# as published by the Free Software Foundation; either version 3
+# of the License, or (at your option) any later version.
+#
+# TinyFF is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program; if not, see <http://www.gnu.org/licenses/>
+# --
+"""Unit tests for tinyff.analysis."""
+
+import numpy as np
+import pytest
+
+from tinyff.analysis import compute_acf, compute_rdf
+from tinyff.neighborlist import NBuildSimple
+
+
+def test_compute_rdf_simple():
+    traj_atpos = np.array([[[0.0, 0.0, 0.0], [0.0, 0.0, 3.0], [0.0, 4.0, 0.0]]])
+    r, g = compute_rdf(traj_atpos, cell_lengths=12.1, spacing=0.7, nbuild=NBuildSimple(rmax=6.0))
+    assert r == pytest.approx(np.arange(0.35, 5.5, 0.7))
+    mask = np.zeros(r.shape, dtype=bool)
+    mask[abs(r - 3.0).argmin()] = True
+    mask[abs(r - 4.0).argmin()] = True
+    mask[abs(r - 5.0).argmin()] = True
+    assert (g[mask] != 0).all()
+    assert (g[~mask] == 0).all()
+    nz = g[mask]
+    assert (nz[1:] < nz[:-1]).all()
+    assert (abs((g * r**2)[mask] - 67) < 0.5).all()
+
+
+def test_compute_vaf_simple():
+    rng = np.random.default_rng(42)
+    traj_atvel = rng.uniform(-1, 1, (100, 10, 3))
+    acf = compute_acf(traj_atvel)
+    assert acf.shape == (100,)