support for PLUMED extras in SimulationOutput

Author: armaet

psiflow/functions.py

@@ -22,14 +22,16 @@ def format_output(
     energy: float,
     forces: np.ndarray | None = None,
     stress: np.ndarray | None = None,
+    extras: dict | None = None,
     **kwargs,
 ) -> dict:
     """"""
     forces = np.zeros(shape=(len(geometry), 3)) if forces is None else forces
     stress = np.zeros(shape=(3, 3)) if stress is None else stress
     if stress.size == 6:
         stress = voigt_6_to_full_3x3_stress(stress)
-    return {"energy": energy, "forces": forces, "stress": stress}
+    extras = {} if extras is None else extras
+    return {"energy": energy, "forces": forces, "stress": stress, 'extras': extras}
 
 
 @typeguard.typechecked
@@ -92,6 +94,7 @@ def __call__(
 @dataclass
 class PlumedFunction(EnergyFunction):
     plumed_input: str
+    plumed_extras: list[str]
     external: Optional[Union[str, Path]] = None
 
     def __post_init__(self):
@@ -126,6 +129,20 @@ def __call__(
             plumed_.cmd("init")
             for line in plumed_input.split("\n"):
                 plumed_.cmd("readInputLine", line)
+
+            self.plumed_data = {}
+            plumed_extras = self.plumed_extras
+            for x in plumed_extras:
+                rank = np.zeros(1, dtype=np.int_)
+                plumed_.cmd(f"getDataRank {x}", rank)
+                if rank[0] > 1:
+                    raise ValueError("Cannot retrieve variables with rank > 1")
+                shape = np.zeros(rank[0], dtype=np.int_)
+                plumed_.cmd(f"getDataShape {x}", shape)
+                if shape[0] > 1:
+                    raise ValueError("Cannot retrieve variables with size > 1")
+                self.plumed_data[x] = np.zeros(shape, dtype=np.double)
+                plumed_.cmd(f"setMemoryForData {x}", self.plumed_data[x])   # internal PLUMED variables are piped to arrays
             os.remove(tmp.name)  # remove whatever plumed has created
             self.plumed_instances[key] = plumed_
 
@@ -147,12 +164,19 @@ def __call__(
         plumed_.cmd("setForces", forces)
         plumed_.cmd("setVirial", virial)
         plumed_.cmd("prepareCalc")
-        plumed_.cmd("performCalcNoUpdate")
+        if "METAD" in self.plumed_input:
+            plumed_.cmd("performCalcNoUpdate")
+        else:
+            plumed_.cmd("performCalc")  # TODO: why update? i-PI does not do this
         plumed_.cmd("getBias", energy)
         stress = None
         if geometry.periodic:
             stress = virial / np.linalg.det(geometry.cell)
-        return format_output(geometry, float(energy.item()), forces, stress)
+
+        extras = {}
+        for x in self.plumed_data:
+            extras[str(x)] = self.plumed_data[x].copy()[0]
+        return format_output(geometry, float(energy.item()), forces, stress, extras)
 
     @staticmethod
     def _geometry_to_key(geometry: Geometry) -> tuple:

psiflow/hamiltonians.py

@@ -1,5 +1,6 @@
 import urllib
 from functools import partial
+import re
 from pathlib import Path
 from typing import ClassVar, Optional, Union, Callable, Sequence
 
@@ -282,6 +283,8 @@ def __init__(
     ):
         super().__init__()
 
+        match = re.search(r'^\s*PRINT\s+.*\bARG=(\S+)', plumed_input, re.MULTILINE)  # parse plumed quantities from input file
+        self.plumed_extras = match.group(1).split(",") if match else []
         self.plumed_input = remove_comments_printflush(plumed_input)
         if type(external) in [str, Path]:
             external = File(str(external))
@@ -303,7 +306,7 @@ def parameters(self) -> dict:
             external = copy_app_future(self.external.filepath, inputs=[self.external])
         else:
             external = None
-        return {"plumed_input": self.plumed_input, "external": external}
+        return {"plumed_input": self.plumed_input, "plumed_extras": self.plumed_extras, "external": external}
 
     def __eq__(self, other: Hamiltonian) -> bool:
         if (

psiflow/sampling/driver.py

@@ -2,6 +2,7 @@
 Updated version of the Psiflow driver included in i-Pi
 """
 
+import json
 import numpy as np
 
 from ipi.pes.dummy import Dummy_driver
@@ -78,6 +79,10 @@ def compute_structure(self, cell, pos):
         vir_ipi = np.array(
             unit_to_internal("energy", "electronvolt", vir_calc.T), dtype=np.float64
         )
-        extras = ""
+        extras = outputs["extras"]
+        if extras == {}:
+            extras = ""
+        else:
+            extras = json.dumps(extras)
 
         return pot_ipi, force_ipi, vir_ipi, extras

psiflow/sampling/sampling.py

@@ -322,7 +322,7 @@ def setup_output(
         trajectory = ET.Element(
             "trajectory",
             filename="trajectory",
-            stride=str(step),
+            stride=str(step),   # TODO: separate stride for trajectory and properties?
             format="ase",
             bead="0",
         )
@@ -335,6 +335,21 @@ def setup_output(
     )
     properties.text = create_xml_list(observables)
     output.append(properties)
+    extras_list = []
+    for comp in components:
+        if comp.name.startswith("Plumed"):  # technically other hamiltonians could also have extras
+            extras_list += comp.hamiltonian.plumed_extras    # maybe extras should be a general property of the Hamiltonian class?
+    observables += [extra + "{au}" for extra in extras_list]   # for SimulationOutput  TODO: what to do with units?
+    extras = ",".join(list(set(extras_list)))
+    extras_element = ET.Element(
+        "trajectory",
+        filename=f"extras",
+        stride=str(step),
+        extra_type=extras,
+        bead="0",
+    )
+    extras_element.text = f" extras "
+    output.append(extras_element)
     return output, observables
 
 
@@ -512,7 +527,7 @@ def _sample(
     if step is not None:
         start = math.floor(start / step)  # start is applied on subsampled quantities
     if step is None:
-        keep_trajectory = False
+        keep_trajectory = False # TODO: warning here? 
     # TODO: check whether observables are valid?
     output, observables = setup_output(
         hamiltonian_components,  # for potential components

psiflow/sampling/server.py

@@ -13,6 +13,7 @@
 from ase.units import Bohr, Ha, _e, _hbar
 from ipi.engine.simulation import Simulation
 from ipi.utils.softexit import softexit
+from ipi.utils.parsing import read_output
 
 from psiflow.geometry import Geometry
 from psiflow.sampling.utils import create_xml_list
@@ -93,11 +94,34 @@ def wait_for_clients(input_xml, timeout: int = 60) -> None:
     raise ConnectionError(msg)
 
 
+def add_extras(noutputs: int) -> None:
+    # property headers
+    file_props = next(Path.cwd().glob(f"*0*.properties"))   # properties should be the same for all coupled walkers?
+    _, info_dict = read_output(file_props)
+    props_headers = ["# column   {}     --> {}{} : {}".format(i, key, "{" + info_dict[key][0] + "}", info_dict[key][1]) for i, key in enumerate(info_dict)]
+    # extras headers
+    file_extras = next(Path.cwd().glob(f"*0*.extras*"))   # extras should be the same for all coupled walkers?
+    assert file_extras is not None, "No extras file found."
+    with open(file_extras, "r") as f:
+        line = f.readline()
+        start = line.find("(") + 1
+        end = line.find(")")
+        extra_names = line[start:end].split(",")
+    extras_headers = ["# column   {}     --> {}{} : {}".format(i + len(props_headers), name, "{au}", "PLUMED variable") for i, name in enumerate(extra_names)]
+    # add extras values to properties files
+    for idx in range(noutputs):
+        file_props = next(Path.cwd().glob(f"*{idx}*.properties"))
+        file_extras = next(Path.cwd().glob(f"*{idx}*.extras*"))
+        np.savetxt(file_props, np.hstack((np.loadtxt(file_props), np.loadtxt(file_extras))), fmt='%10.8e', delimiter='     ', header="\n".join(props_headers) + "\n" + "\n".join(extras_headers), comments="")
+        # granted i-Pi properties file has some weird indentations so the files do not perfectly match
+        # but output parser still works so who cares
+
+
 def run(start_xyz: str, input_xml: str):
     # prepare starting geometries from context_dir
     data_start: list[ase.Atoms] = read(start_xyz, index=":")
     for i, at in enumerate(data_start):
-        print(at.pbc)
+        print(at.pbc)   # TODO: why print?
         if not any(at.pbc):  # set fake large cell for i-PI
             at.pbc = True
             at.cell = Cell(NONPERIODIC_CELL)
@@ -134,6 +158,14 @@ def cleanup(output_xyz: str, output_props: str, output_trajs: str) -> None:
     _write_frames(*states, outputs=[output_xyz])
     print("Moved checkpoint geometries")
 
+    output_props = _.split(",") if (_ := output_props) else []
+    output_trajs = _.split(",") if (_ := output_trajs) else []
+
+    # Add collective variables to simulation properties, if they exist
+    if output_props:
+        add_extras(len(output_props))
+        print("Added i-Pi extras output to properties files")
+
     prefix = ""
     if "remd" in content:
         # unshuffle simulation output according to ensemble
@@ -144,9 +176,6 @@ def cleanup(output_xyz: str, output_props: str, output_trajs: str) -> None:
         assert out.returncode == 0  # TODO: what if it isn't?
         print("REMDSORT")
 
-    output_props = _.split(",") if (_ := output_props) else []
-    output_trajs = _.split(",") if (_ := output_trajs) else []
-
     # move recorded simulation observables
     if len(output_props):
         assert len(states) == len(output_props)