Add Correlator CLI and defaults

docs/source/user_guide/command_line.rst

@@ -332,6 +332,33 @@ will compute the RDFs for ``Na`` and ``Cl`` atoms. Seperately for each paring, u
 
 The RDF is also computed from the trajectory file and the options ``rdf_start``, ``rdf_stop``, and ``rdf_step`` may be used to control which trajectory frames are utilised.
 
+On-the-fly Processing
++++++++++++++++++++++
+
+Alongside post-processing, correlations may be calculated during MD. This means that, for example, the VAF may be computed without storing trajectory data. To compute the VAF at runtime the following options can be passed:
+
+.. code-block:: bash
+
+   janus md --ensemble nve --struct tests/data/NaCl.cif --steps 100 --correlation-kwargs "{'vaf': {'a': 'Velocity', 'points': 100, 'correlation_frequency': 2}}"
+
+This would result in the file ```janus_results/NaCl-nve-T300.0-cor.dat``` containing the combined VAF for Na and Cl atoms correlated every other step, meaning 50 correlation lag times.
+
+The option ``a`` specifies the Observable to be correlated. Possible values are ``Velocity``, ``Stress``, ``StressHydrostatic``, and ``StressShear``. The latter two stresses combine the diagonal and off-diagonal components of the stress tensor. When ``b`` is not specified it is set to a copy of ``a`` to form an auto-correlation.
+
+Correlation observables may also specify their own keyword arguments. For example to specify the components of stress to correlated over (with the correlations averaged) the following options may be passed:
+
+.. code-block:: bash
+
+   janus md --ensemble nve --struct tests/data/NaCl.cif --steps 100 --correlation-kwargs "{'saf': {'a': 'Stress', 'points': 100, 'a_kwargs': {'components': ['xy', 'yz', 'zx']}}}"
+
+Resulting in the stress auto-correlation function :math:`\frac{1}{3}(\langle\sigma_{xy}\sigma_{xy}\rangle+\langle\sigma_{yz}\sigma_{yz}\rangle+\langle\sigma_{zx}\sigma_{zx}\rangle)`, calculated every step for 100 correlation lag times.
+
+The Velocity observable may also be computed over specific components (it defaults to all) and atom slices. To compute over odd indexed atoms (Na here) the following options may be passed:
+
+.. code-block:: bash
+
+   janus md --ensemble nve --struct tests/data/NaCl.cif --steps 100 --correlation-kwargs "{'vaf': {'a': 'Velocity', 'points': 100, 'a_kwargs': {'atoms_slice': (0, None, 2)}}}"
+
 Heating
 -------
 

janus_core/calculations/md.py

@@ -156,9 +156,9 @@ class MolecularDynamics(BaseCalculation):
         Keyword arguments to pass to `output_structs` when saving trajectory and final
         files. Default is {}.
     post_process_kwargs
-        Keyword arguments to control post-processing operations.
+        Keyword arguments to control post-processing operations. Default is None.
     correlation_kwargs
-        Keyword arguments to control on-the-fly correlations.
+        Keyword arguments to control on-the-fly correlations. Default is None.
     seed
         Random seed used by numpy.random and random functions, such as in Langevin.
         Default is None.
@@ -324,9 +324,9 @@ def __init__(
             Keyword arguments to pass to `output_structs` when saving trajectory and
             final files. Default is {}.
         post_process_kwargs
-            Keyword arguments to control post-processing operations.
+            Keyword arguments to control post-processing operations. Default is None.
         correlation_kwargs
-            Keyword arguments to control on-the-fly correlations.
+            Keyword arguments to control on-the-fly correlations. Default is None.
         seed
             Random seed used by numpy.random and random functions, such as in Langevin.
             Default is None.

janus_core/cli/md.py

@@ -12,6 +12,7 @@
 from janus_core.cli.types import (
     Architecture,
     CalcKwargs,
+    CorrelationKwargs,
     Device,
     EnsembleKwargs,
     FilePrefix,
@@ -24,7 +25,7 @@
     Summary,
     WriteKwargs,
 )
-from janus_core.cli.utils import yaml_converter_callback
+from janus_core.cli.utils import parse_correlation_kwargs, yaml_converter_callback
 
 app = Typer()
 
@@ -221,6 +222,7 @@ def md(
     ] = None,
     write_kwargs: WriteKwargs = None,
     post_process_kwargs: PostProcessKwargs = None,
+    correlation_kwargs: CorrelationKwargs = None,
     seed: Annotated[
         int | None,
         Option(help="Random seed for numpy.random and random functions."),
@@ -349,7 +351,9 @@ def md(
         Keyword arguments to pass to `output_structs` when saving trajectory and final
         files. Default is {}.
     post_process_kwargs
-        Kwargs to pass to post-processing.
+        Keyword arguments to pass to post-processing. Default is None.
+    correlation_kwargs
+        Keyword arguments to pass for on-the-fly correlations. Default is None.
     seed
         Random seed used by numpy.random and random functions, such as in Langevin.
         Default is None.
@@ -379,7 +383,6 @@ def md(
 
     # Check options from configuration file are all valid
     check_config(ctx)
-
     [
         read_kwargs,
         calc_kwargs,
@@ -409,9 +412,14 @@ def md(
         "ensemble_kwargs": ensemble_kwargs.copy(),
         "write_kwargs": write_kwargs.copy(),
         "post_process_kwargs": post_process_kwargs.copy(),
+        "correlation_kwargs": correlation_kwargs.value.copy(),
     }
     config = get_config(params=ctx.params, all_kwargs=all_kwargs)
 
+    # Handle separately to process short-hands, and Observables.
+    if correlation_kwargs:
+        correlation_kwargs = parse_correlation_kwargs(correlation_kwargs)
+
     # Read only first structure by default and ensure only one image is read
     set_read_kwargs_index(read_kwargs)
 
@@ -482,6 +490,7 @@ def md(
         "temp_time": temp_time,
         "write_kwargs": write_kwargs,
         "post_process_kwargs": post_process_kwargs,
+        "correlation_kwargs": correlation_kwargs,
         "seed": seed,
     }
 

janus_core/cli/types.py

@@ -298,6 +298,21 @@ def __str__(self) -> str:
     ),
 ]
 
+CorrelationKwargs = Annotated[
+    TyperDict | None,
+    Option(
+        parser=parse_dict_class,
+        help=(
+            """
+            Keyword arguments to pass to md for on-the-fly correlations. Must be
+            passed as a list of dictionaries wrapped in quotes, e.g.
+            "[{'key' : values}]".
+            """
+        ),
+        metavar="DICT",
+    ),
+]
+
 LogPath = Annotated[
     Path | None,
     Option(help=("Path to save logs to. Default is inferred from `file_prefix`")),

janus_core/cli/utils.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
+from copy import deepcopy
 import datetime
 import logging
 from pathlib import Path
@@ -17,12 +18,14 @@
     from ase import Atoms
     from typer import Context
 
-    from janus_core.cli.types import TyperDict
+    from janus_core.cli.types import CorrelationKwargs, TyperDict
     from janus_core.helpers.janus_types import (
         MaybeSequence,
         PathLike,
     )
 
+from janus_core.processing import observables
+
 
 def dict_paths_to_strs(dictionary: dict) -> None:
     """
@@ -322,3 +325,82 @@ def check_config(ctx: Context) -> None:
         # Check options individually so can inform user of specific issue
         if option not in ctx.params:
             raise ValueError(f"'{option}' in configuration file is not a valid option")
+
+
+def parse_correlation_kwargs(kwargs: CorrelationKwargs) -> list[dict]:
+    """
+    Parse CLI CorrelationKwargs to md correlation_kwargs.
+
+    Parameters
+    ----------
+    kwargs
+        CLI correlation keyword options.
+
+    Returns
+    -------
+    list[dict]
+        The parsed correlation_kwargs for md.
+    """
+    parsed_kwargs = []
+    for name, cli_kwargs in kwargs.value.items():
+        arguments = {
+            "blocks",
+            "points",
+            "averaging",
+            "update_frequency",
+            "a_kwargs",
+            "b_kwargs",
+            "a",
+            "b",
+        }
+        if not (set(cli_kwargs.keys()) <= arguments):
+            raise ValueError(
+                "correlation_kwargs got unexpected argument(s)"
+                f"{set(cli_kwargs.keys()).difference(arguments)}"
+            )
+
+        if "a" not in cli_kwargs and "b" not in cli_kwargs:
+            raise ValueError("At least one observable must be supplied as 'a' or 'b'")
+
+        if "points" not in cli_kwargs:
+            raise ValueError("Correlation keyword argument 'points' must be specified")
+
+        # Accept an Observable to be replicated.
+        if "b" not in cli_kwargs:
+            a = cli_kwargs["a"]
+            b = deepcopy(a)
+            # Copying Observable, so can copy kwargs as well.
+            if "b_kwargs" not in cli_kwargs and "a_kwargs" in cli_kwargs:
+                cli_kwargs["b_kwargs"] = cli_kwargs["a_kwargs"]
+        elif "a" not in cli_kwargs:
+            b = cli_kwargs["b"]
+            a = deepcopy(b)
+            if "a_kwargs" not in cli_kwargs and "b_kwargs" in cli_kwargs:
+                cli_kwargs["a_kwargs"] = cli_kwargs["b_kwargs"]
+        else:
+            a = cli_kwargs["a"]
+            b = cli_kwargs["b"]
+
+        a_kwargs = cli_kwargs["a_kwargs"] if "a_kwargs" in cli_kwargs else {}
+        b_kwargs = cli_kwargs["b_kwargs"] if "b_kwargs" in cli_kwargs else {}
+
+        # Accept "." in place of one kwargs to repeat.
+        if a_kwargs == "." and b_kwargs == ".":
+            raise ValueError("a_kwargs and b_kwargs cannot 'ditto' each other")
+        if a_kwargs and b_kwargs == ".":
+            b_kwargs = a_kwargs
+        elif b_kwargs and a_kwargs == ".":
+            a_kwargs = b_kwargs
+
+        cor_kwargs = {
+            "name": name,
+            "points": cli_kwargs["points"],
+            "a": getattr(observables, a)(**a_kwargs),
+            "b": getattr(observables, b)(**b_kwargs),
+        }
+
+        for optional in cli_kwargs.keys() & {"blocks", "averaging", "update_frequency"}:
+            cor_kwargs[optional] = cli_kwargs[optional]
+
+        parsed_kwargs.append(cor_kwargs)
+    return parsed_kwargs

janus_core/helpers/janus_types.py

@@ -80,7 +80,7 @@ class PostProcessKwargs(TypedDict, total=False):
     vaf_output_files: Sequence[PathLike] | None
 
 
-class CorrelationKwargs(TypedDict, total=True):
+class Correlation(TypedDict, total=True):
     """Arguments for on-the-fly correlations <ab>."""
 
     #: observable a in <ab>, with optional args and kwargs
@@ -89,16 +89,18 @@ class CorrelationKwargs(TypedDict, total=True):
     b: Observable
     #: name used for correlation in output
     name: str
-    #: blocks used in multi-tau algorithm
-    blocks: int
     #: points per block
     points: int
+    #: blocks used in multi-tau algorithm
+    blocks: int
     #: averaging between blocks
     averaging: int
     #: frequency to update the correlation (steps)
     update_frequency: int
 
 
+CorrelationKwargs = list[Correlation]
+
 # eos_names from ase.eos
 EoSNames = Literal[
     "sj",

janus_core/processing/correlator.py

@@ -253,14 +253,14 @@ class Correlation:
         Observable for b.
     name
         Name of correlation.
-    blocks
-        Number of correlation blocks.
     points
         Number of points per block.
+    blocks
+        Number of correlation blocks. Default is 1.
     averaging
-        Averaging window per block level.
+        Averaging window per block level. Default is 1.
     update_frequency
-        Frequency to update the correlation, md steps.
+        Frequency to update the correlation, md steps. Default is 1.
     """
 
     def __init__(
@@ -269,10 +269,10 @@ def __init__(
         a: Observable,
         b: Observable,
         name: str,
-        blocks: int,
         points: int,
-        averaging: int,
-        update_frequency: int,
+        blocks: int = 1,
+        averaging: int = 1,
+        update_frequency: int = 1,
     ) -> None:
         """
         Initialise a correlation.
@@ -285,14 +285,14 @@ def __init__(
             Observable for b.
         name
             Name of correlation.
-        blocks
-            Number of correlation blocks.
         points
             Number of points per block.
+        blocks
+            Number of correlation blocks. Default is 1.
         averaging
-            Averaging window per block level.
+            Averaging window per block level. Default is 1.
         update_frequency
-            Frequency to update the correlation, md steps.
+            Frequency to update the correlation, md steps. Default is 1.
         """
         self.name = name
         self.blocks = blocks

janus_core/processing/observables.py

@@ -237,7 +237,7 @@ class Velocity(Observable, ComponentMixin):
     def __init__(
         self,
         *,
-        components: list[str],
+        components: list[str] | None = None,
         atoms_slice: list[int] | SliceLike | None = None,
     ):
         """
@@ -246,12 +246,12 @@ def __init__(
         Parameters
         ----------
         components
-            Symbols for tensor components, x, y, and z.
+            Symbols for returned velocity components, x, y, and z (default is all).
         atoms_slice
             List or slice of atoms to observe velocities from.
         """
         ComponentMixin.__init__(self, components={"x": 0, "y": 1, "z": 2})
-        self.components = components
+        self.components = components if components else ["x", "y", "z"]
 
         Observable.__init__(self, atoms_slice)
 

tests/test_correlator.py

@@ -113,6 +113,12 @@ def test_vaf(tmp_path):
                 "averaging": 1,
                 "update_frequency": 1,
             },
+            {
+                "a": Velocity(),
+                "b": Velocity(),
+                "points": 1,
+                "name": "vaf_default",
+            },
         ],
         write_kwargs={"invalidate_calc": False},
     )
@@ -133,6 +139,12 @@ def test_vaf(tmp_path):
     assert vaf_na * 3 == approx(vaf_post[1][0], rel=1e-5)
     assert vaf_cl * 3 == approx(vaf_post[1][1], rel=1e-5)
 
+    # Default arguments are equivalent to mean square velocities.
+    v = np.mean([np.mean(atoms.get_velocities() ** 2) for atoms in traj])
+    vaf_default = vaf["vaf_default"]
+    assert len(vaf_default["value"]) == 1
+    assert v == approx(vaf_default["value"][0], rel=1e-5)
+
 
 def test_md_correlations(tmp_path):
     """Test correlations as part of MD cycle."""