Set defaults, test in test_correlator

Author: harveydevereux

janus_core/cli/utils.py

@@ -328,23 +328,23 @@ def parse_correlation_kwargs(kwargs: CorrelationKwargs) -> list[dict]:
         The parsed correlation_kwargs for md.
     """
     parsed_kwargs = []
-    for name, kwarg in kwargs.value.items():
-        if "a" not in kwarg and "b" not in kwarg:
+    for name, cli_kwargs in kwargs.value.items():
+        if "a" not in cli_kwargs and "b" not in cli_kwargs:
             raise ValueError("At least on observable must be supplied as 'a' or 'b'")
 
         # Accept on Observable to be replicated.
-        if "b" not in kwarg:
-            a = kwarg["a"]
+        if "b" not in cli_kwargs:
+            a = cli_kwargs["a"]
             b = a
-        elif "a" not in kwarg:
-            a = kwarg["b"]
+        elif "a" not in cli_kwargs:
+            a = cli_kwargs["b"]
             b = a
         else:
-            a = kwarg["a"]
-            b = kwarg["b"]
+            a = cli_kwargs["a"]
+            b = cli_kwargs["b"]
 
-        a_kwargs = kwarg["a_kwargs"] if "a_kwargs" in kwarg else {}
-        b_kwargs = kwarg["b_kwargs"] if "b_kwargs" in kwarg else {}
+        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 == ".":
@@ -354,22 +354,15 @@ def parse_correlation_kwargs(kwargs: CorrelationKwargs) -> list[dict]:
         elif b_kwargs and a_kwargs == ".":
             a_kwargs = b_kwargs
 
-        blocks = kwarg["blocks"] if "blocks" in kwarg else 1
-        points = kwarg["points"] if "points" in kwarg else 1
-        averaging = kwarg["averaging"] if "averaging" in kwarg else 1
-        update_frequency = (
-            kwarg["update_frequency"] if "update_frequency" in kwarg else 1
-        )
+        cor_kwargs = {
+            "name": name,
+            "a": getattr(observables, a)(**a_kwargs),
+            "b": getattr(observables, b)(**b_kwargs),
+        }
 
-        parsed_kwargs.append(
-            {
-                "name": name,
-                "a": getattr(observables, a)(**a_kwargs),
-                "b": getattr(observables, b)(**b_kwargs),
-                "blocks": blocks,
-                "points": points,
-                "averaging": averaging,
-                "update_frequency": update_frequency,
-            }
-        )
+        for optional in ["blocks", "points", "averaging", "update_frequency"]:
+            if optional in cli_kwargs:
+                cor_kwargs[optional] = cli_kwargs[optional]
+
+        parsed_kwargs.append(cor_kwargs)
     return parsed_kwargs

janus_core/processing/correlator.py

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

tests/test_correlator.py

@@ -113,6 +113,11 @@ def test_vaf(tmp_path):
                 "averaging": 1,
                 "update_frequency": 1,
             },
+            {
+                "a": Velocity(),
+                "b": Velocity(),
+                "name": "vaf_default",
+            },
         ],
         write_kwargs={"invalidate_calc": False},
     )
@@ -133,6 +138,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."""

tests/test_md_cli.py

@@ -104,7 +104,7 @@ def test_md(ensemble):
                 "--correlation-kwargs",
                 (
                     "{'vaf': {'a': 'Velocity'},"
-                    " 'vaf_x': {'a': 'velocity',"
+                    " 'vaf_x': {'a': 'Velocity',"
                     "'a_kwargs': {'components': ['x']}, 'b_kwargs': '.'}}"
                 ),
             ],