Add post-processing routines to MD

docs/source/apidoc/janus_core.rst

@@ -164,6 +164,16 @@ janus\_core.helpers.descriptors module
    :undoc-members:
    :show-inheritance:
 
+janus\_core.helpers.post_process module
+---------------------------------------
+
+.. automodule:: janus_core.helpers.post_process
+   :members:
+   :special-members:
+   :private-members:
+   :undoc-members:
+   :show-inheritance:
+
 janus\_core.helpers.train module
 --------------------------------
 

janus_core/calculations/md.py

@@ -5,11 +5,11 @@
 from math import isclose
 from pathlib import Path
 import random
-from typing import Any, Optional
+from typing import Any, Optional, Union
 from warnings import warn
 
 from ase import Atoms, units
-from ase.io import write
+from ase.io import read, write
 from ase.md.langevin import Langevin
 from ase.md.npt import NPT as ASE_NPT
 from ase.md.velocitydistribution import (
@@ -18,11 +18,21 @@
     ZeroRotation,
 )
 from ase.md.verlet import VelocityVerlet
+
+try:
+    from ase.geometry.analysis import Analysis
+
+    ASE_GEOMETRY = True
+except ImportError:
+
+    ASE_GEOMETRY = False
+
 import numpy as np
 
 from janus_core.calculations.geom_opt import optimize
-from janus_core.helpers.janus_types import Ensembles, PathLike
+from janus_core.helpers.janus_types import Ensembles, PathLike, PostProcessKwargs
 from janus_core.helpers.log import config_logger
+from janus_core.helpers.post_process import compute_rdf, compute_vaf
 from janus_core.helpers.utils import FileNameMixin
 
 DENS_FACT = (units.m / 1.0e2) ** 3 / units.mol
@@ -97,6 +107,8 @@ class MolecularDynamics(FileNameMixin):  # pylint: disable=too-many-instance-att
         heating.
     temp_time : Optional[float]
         Time between heating steps, in fs. Default is None, which disables heating.
+    post_process_kwargs : Optional[PostProcessKwargs]
+        Keyword arguments to control post-processing operations.
     log_kwargs : Optional[dict[str, Any]]
         Keyword arguments to pass to log config. Default is None.
     seed : Optional[int]
@@ -157,6 +169,7 @@ def __init__(  # pylint: disable=too-many-arguments,too-many-locals,too-many-sta
         temp_end: Optional[float] = None,
         temp_step: Optional[float] = None,
         temp_time: Optional[float] = None,
+        post_process_kwargs: Optional[PostProcessKwargs] = None,
         log_kwargs: Optional[dict[str, Any]] = None,
         seed: Optional[int] = None,
     ) -> None:
@@ -231,6 +244,8 @@ def __init__(  # pylint: disable=too-many-arguments,too-many-locals,too-many-sta
             disables heating.
         temp_time : Optional[float]
             Time between heating steps, in fs. Default is None, which disables heating.
+        post_process_kwargs : Optional[PostProcessKwargs]
+            Keyword arguments to control post-processing operations.
         log_kwargs : Optional[dict[str, Any]]
             Keyword arguments to pass to log config. Default is None.
         seed : Optional[int]
@@ -262,6 +277,9 @@ def __init__(  # pylint: disable=too-many-arguments,too-many-locals,too-many-sta
         self.temp_end = temp_end
         self.temp_step = temp_step
         self.temp_time = temp_time * units.fs if temp_time else None
+        self.post_process_kwargs = (
+            post_process_kwargs if post_process_kwargs is not None else {}
+        )
         self.log_kwargs = log_kwargs
         self.ensemble = ensemble
         self.seed = seed
@@ -315,7 +333,7 @@ def __init__(  # pylint: disable=too-many-arguments,too-many-locals,too-many-sta
 
         self.minimize_kwargs = minimize_kwargs if minimize_kwargs else {}
         self.restart_files = []
-        self.dyn = None
+        self.dyn: Union[Langevin, VelocityVerlet, ASE_NPT]
         self.n_atoms = len(self.struct)
 
         self.stats_file = self._build_filename(
@@ -542,6 +560,76 @@ def _write_final_state(self) -> None:
             columns=["symbols", "positions", "momenta", "masses"],
         )
 
+    def _post_process(self) -> None:
+        """Compute properties after MD run."""
+        # Nothing to do
+        if not any(
+            self.post_process_kwargs.get(kwarg, None)
+            for kwarg in ("rdf_compute", "vaf_compute")
+        ):
+            warn(
+                "Post-processing arguments present, but no computation requested. "
+                "Please set either 'rdf_compute' or 'vaf_compute' "
+                "to do post-processing."
+            )
+
+        data = read(self.traj_file, index=":")
+
+        if ASE_GEOMETRY:
+            ana = Analysis(data)
+        else:
+            ana = None
+
+        param_pref = self._parameter_prefix if self.file_prefix is None else ""
+
+        if self.post_process_kwargs.get("rdf_compute", False):
+            base_name = self.post_process_kwargs.get("rdf_output_file", None)
+            rdf_args = {
+                name: self.post_process_kwargs.get(key, default)
+                for name, (key, default) in (
+                    ("rmax", ("rdf_rmax", 2.5)),
+                    ("nbins", ("rdf_nbins", 50)),
+                    ("elements", ("rdf_elements", None)),
+                )
+            }
+            slice_ = (
+                self.post_process_kwargs.get("rdf_start", len(data) - 1),
+                self.post_process_kwargs.get("rdf_stop", len(data)),
+                self.post_process_kwargs.get("rdf_step", 1),
+            )
+
+            out_paths = [
+                self._build_filename(
+                    "rdf.dat", param_pref, str(ind), prefix_override=base_name
+                )
+                for ind in range(*slice_)
+            ]
+
+            rdf_args["index"] = slice_
+            compute_rdf(data, ana, filename=out_paths, **rdf_args)
+
+        if self.post_process_kwargs.get("vaf_compute", False):
+
+            file_name = self.post_process_kwargs.get("vaf_output_file", None)
+            use_vel = self.post_process_kwargs.get("vaf_velocities", False)
+            fft = self.post_process_kwargs.get("vaf_fft", False)
+
+            out_path = self._build_filename("vaf.dat", param_pref, filename=file_name)
+            slice_ = (
+                self.post_process_kwargs.get("vaf_start", 0),
+                self.post_process_kwargs.get("vaf_stop", None),
+                self.post_process_kwargs.get("vaf_step", 1),
+            )
+
+            compute_vaf(
+                data,
+                out_path,
+                use_velocities=use_vel,
+                fft=fft,
+                index=slice_,
+                filter_atoms=self.post_process_kwargs.get("vaf_atoms", None),
+            )
+
     def _write_restart(self) -> None:
         """Write restart file and (optionally) rotate files saved."""
         step = self.offset + self.dyn.nsteps
@@ -594,6 +682,9 @@ def run(self) -> None:
         self.struct.info["real_time"] = datetime.datetime.now()
         self._run_dynamics()
 
+        if self.post_process_kwargs:
+            self._post_process()
+
     def _run_dynamics(self) -> None:
         """Run dynamics and/or temperature ramp."""
         # Store temperature for final MD

janus_core/cli/md.py

@@ -15,6 +15,7 @@
     EnsembleKwargs,
     LogPath,
     MinimizeKwargs,
+    PostProcessKwargs,
     ReadKwargs,
     StructPath,
     Summary,
@@ -170,6 +171,7 @@ def md(
     temp_time: Annotated[
         float, Option(help="Time between heating steps, in fs.")
     ] = None,
+    post_process_kwargs: PostProcessKwargs = None,
     log: LogPath = "md.log",
     seed: Annotated[
         Optional[int],
@@ -271,6 +273,8 @@ def md(
     temp_time : Optional[float]
         Time between heating steps, in fs. Default is None, which disables
         heating.
+    post_process_kwargs : Optional[PostProcessKwargs]
+        Kwargs to pass to post-processing.
     log : Optional[Path]
         Path to write logs to. Default is "md.log".
     seed : Optional[int]
@@ -284,8 +288,20 @@ def md(
     # Check options from configuration file are all valid
     check_config(ctx)
 
-    [read_kwargs, calc_kwargs, minimize_kwargs, ensemble_kwargs] = parse_typer_dicts(
-        [read_kwargs, calc_kwargs, minimize_kwargs, ensemble_kwargs]
+    [
+        read_kwargs,
+        calc_kwargs,
+        minimize_kwargs,
+        ensemble_kwargs,
+        post_process_kwargs,
+    ] = parse_typer_dicts(
+        [
+            read_kwargs,
+            calc_kwargs,
+            minimize_kwargs,
+            ensemble_kwargs,
+            post_process_kwargs,
+        ]
     )
 
     if not ensemble in get_args(Ensembles):
@@ -338,6 +354,7 @@ def md(
         "temp_end": temp_end,
         "temp_step": temp_step,
         "temp_time": temp_time,
+        "post_process_kwargs": post_process_kwargs,
         "log_kwargs": log_kwargs,
         "seed": seed,
         "ensemble_kwargs": ensemble_kwargs,

janus_core/cli/types.py

@@ -152,6 +152,20 @@ def __str__(self):
     ),
 ]
 
+PostProcessKwargs = Annotated[
+    TyperDict,
+    Option(
+        parser=parse_dict_class,
+        help=(
+            """
+            Keyword arguments to pass to post-processer. Must be passed as a dictionary
+            wrapped in quotes, e.g. "{'key' : value}".
+            """
+        ),
+        metavar="DICT",
+    ),
+]
+
 LogPath = Annotated[Path, Option(help="Path to save logs to.")]
 
 Summary = Annotated[

janus_core/helpers/janus_types.py

@@ -17,7 +17,8 @@
 MaybeList = Union[T, list[T]]
 MaybeSequence = Union[T, Sequence[T]]
 PathLike = Union[str, Path]
-
+StartStopStep = tuple[Optional[int], Optional[int], int]
+SliceLike = Union[slice, range, int, StartStopStep]
 
 # ASE Arg types
 
@@ -43,13 +44,36 @@ class ASEWriteArgs(TypedDict, total=False):
 
 
 class ASEOptArgs(TypedDict, total=False):
-    """Main arugments for ase optimisers."""
+    """Main arguments for ase optimisers."""
 
     restart: Optional[bool]
     logfile: Optional[PathLike]
     trajectory: Optional[str]
 
 
+class PostProcessKwargs(TypedDict, total=False):
+    """Main arguments for MD post-processing."""
+
+    # RDF
+    rdf_compute: bool
+    rdf_rmax: float
+    rdf_nbins: int
+    rdf_elements: MaybeSequence[Union[str, int]]
+    rdf_start: int
+    rdf_stop: Optional[int]
+    rdf_step: int
+    rdf_output_file: Optional[str]
+    # VAF
+    vaf_compute: bool
+    vaf_velocities: bool
+    vaf_fft: bool
+    vaf_atoms: Sequence[Sequence[int]]
+    vaf_start: int
+    vaf_stop: Optional[int]
+    vaf_step: int
+    vaf_output_file: Optional[PathLike]
+
+
 # eos_names from ase.eos
 EoSNames = Literal[
     "sj",