Support for PLUMED output

psiflow/functions.py

@@ -147,7 +147,10 @@ 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")
         plumed_.cmd("getBias", energy)
         stress = None
         if geometry.periodic:

psiflow/hamiltonians.py

@@ -288,7 +288,7 @@ def __init__(
     ):
         super().__init__()
 
-        self.plumed_input = remove_comments_printflush(plumed_input)
+        self.plumed_input = plumed_input
         if type(external) in [str, Path]:
             external = File(str(external))
         if external is not None:

psiflow/sampling/output.py

@@ -224,6 +224,7 @@ class SimulationOutput:
     time: Union[float, AppFuture, None]
     temperature: Union[float, AppFuture, None]
     trajectory: Optional[Dataset]
+    plumed_output: Union[dict[str, np.ndarray], AppFuture, None]
 
     def __init__(self, fields: list[str]):
         self._data = {key: None for key in fields}
@@ -234,6 +235,7 @@ def __init__(self, fields: list[str]):
         self.time = None
         self.temperature = None
         self.trajectory = None
+        self.plumed_output = None
         self.hamiltonians = None
 
     def __getitem__(self, key: str) -> AppFuture:

psiflow/sampling/sampling.py

@@ -5,6 +5,8 @@
 from dataclasses import dataclass
 from typing import Optional, Union, Iterable
 
+from warnings import warn
+
 import parsl
 import typeguard
 import numpy as np
@@ -14,14 +16,15 @@
 
 import psiflow
 from psiflow.data import Dataset
-from psiflow.hamiltonians import Hamiltonian, MixtureHamiltonian, Zero
+from psiflow.hamiltonians import Hamiltonian, MixtureHamiltonian, Zero, PlumedHamiltonian
 from psiflow.sampling.output import (
     DEFAULT_OBSERVABLES,
     SimulationOutput,
     potential_component_names,
 )
-from psiflow.sampling.walker import Coupling, Walker, partition
+from psiflow.sampling.walker import Coupling, Walker, partition, set_plumed_print_files
 from psiflow.utils.io import save_xml
+from psiflow.utils._plumed import parse_plumed_output
 from psiflow.utils import TMP_COMMAND, CD_COMMAND
 
 
@@ -400,6 +403,7 @@ def _execute_ipi(
     hamiltonian_names: list[str],
     client_args: list[str],
     keep_trajectory: bool,
+    print_plumed: bool,
     max_force: Optional[float],
     coupling_command: Optional[str],
     command_server: str,
@@ -429,6 +433,8 @@ def _execute_ipi(
         commands_copy += f'cp walker-{i}_output.properties {outputs[i + 1].filepath}',
         if keep_trajectory:
             commands_copy += f'cp walker-{i}_output.trajectory_0.extxyz {outputs[i + nwalkers + 1].filepath}',
+        if print_plumed:
+            commands_copy += f'cp walker-{i}_plumed.log {outputs[i + (2 if keep_trajectory else 1)*nwalkers + 1].filepath}',
     if coupling_command is not None:
         commands_copy += coupling_command,
     command_copy = '; '.join(commands_copy)
@@ -457,6 +463,7 @@ def _sample(
     step: Optional[int] = None,
     start: int = 0,
     keep_trajectory: bool = True,
+    print_plumed: bool = False,
     max_force: Optional[float] = None,
     observables: Optional[list[str]] = None,
     motion_defaults: Union[None, str, ET.Element] = None,
@@ -536,6 +543,14 @@ def _sample(
         simulation,
         outputs=[context.new_file("input_", ".xml")],
     ).outputs[0]
+
+    if print_plumed:
+        set_plumed_print_files(walkers)
+        if walkers[0].nbeads > 1:
+            warn("Beware: PLUMED output files are currently not bead-specific in i-PI simulations!")
+        if walkers[0].coupling is not None:
+            warn("Beware: PLUMED output files may be inconsistent when doing replica exchange or when using coupled walkers!")
+
     inputs = [
         input_future,
         Dataset([w.state for w in walkers]).extxyz,
@@ -550,11 +565,14 @@ def _sample(
         client_args.append(args)
     outputs = [context.new_file("data_", ".xyz")]
     outputs += [context.new_file("simulation_", ".txt") for _ in walkers]
+    n = 1
     if keep_trajectory:
         outputs += [context.new_file("data_", ".xyz") for _ in walkers]
-        assert len(outputs) == 2 * len(walkers) + 1
-    else:
-        assert len(outputs) == len(walkers) + 1
+        n += 1
+    if print_plumed:
+        outputs += [context.new_file("plumed_", ".log") for _ in walkers]
+        n += 1
+    assert len(outputs) == n * len(walkers) + 1
 
     # add coupling inputs after all other ones;
     # these are updated again with the corresponding outputs from execute_ipi
@@ -575,6 +593,7 @@ def _sample(
         [comp.name for comp in hamiltonian_components],
         client_args,
         keep_trajectory,
+        print_plumed,
         max_force,
         coupling_copy_command,
         command_server,
@@ -600,12 +619,19 @@ def _sample(
             result.outputs[i + 1],
             hamiltonians=[comp.hamiltonian for comp in hamiltonian_components],
         )
+        n = 1
         if keep_trajectory:
             j = len(walkers) + 1 + i
             trajectory = Dataset(None, result.outputs[j])
             if start > 0:
                 trajectory = trajectory[start:]
             simulation_output.trajectory = trajectory
+            n += 1
+        if print_plumed:
+            j = n * len(walkers) + 1 + i
+            plumed_output = result.outputs[j]
+            plumed_output = parse_plumed_output(plumed_output)
+            simulation_output.plumed_output = plumed_output
         if walkers[i].metadynamics is not None:
             walkers[i].metadynamics.wait_for(result)
         simulation_output.update_walker(walkers[i])
@@ -623,6 +649,7 @@ def sample(
     step: Optional[int] = None,
     start: int = 0,
     keep_trajectory: bool = True,
+    print_plumed: bool = False,
     max_force: Optional[float] = None,
     observables: Optional[list[str]] = None,
     motion_defaults: Union[None, str, ET.Element] = None,
@@ -646,6 +673,7 @@ def sample(
             step=step,
             start=start,
             keep_trajectory=keep_trajectory,
+            print_plumed=print_plumed,
             max_force=max_force,
             observables=observables,
             motion_defaults=motion_defaults,

psiflow/sampling/walker.py

@@ -3,6 +3,8 @@
 import xml.etree.ElementTree as ET
 from typing import Optional, Union
 
+from warnings import warn
+
 import numpy as np
 import typeguard
 from parsl.app.app import python_app
@@ -13,10 +15,11 @@
 import psiflow
 from psiflow.data import Dataset
 from psiflow.geometry import Geometry, check_equality
-from psiflow.hamiltonians import Hamiltonian, Zero
+from psiflow.hamiltonians import Hamiltonian, Zero, MixtureHamiltonian
 from psiflow.order_parameters import OrderParameter
 from psiflow.sampling.metadynamics import Metadynamics
 from psiflow.utils.apps import copy_app_future
+from psiflow.utils._plumed import set_path_in_plumed
 
 
 @typeguard.typechecked
@@ -240,6 +243,28 @@ def validate_coupling(walkers: list[Walker]):
         assert coupling.nwalkers == counts[i]
 
 
+# this does not make sense if walkers have multiple PlumedHamiltonians
+# however there is no reason to do this
+@typeguard.typechecked
+def set_plumed_print_files(walkers: list[Walker]) -> None:
+    for i, walker in enumerate(walkers):
+        ham = walker.hamiltonian
+        if not isinstance(ham, MixtureHamiltonian):
+            ham *= 1.0
+        plumed_ids = [j for j, name in enumerate(ham.get_named_components()) if name.startswith("PlumedHamiltonian")]
+        if not plumed_ids:
+            continue
+        if len(plumed_ids) > 1: # this will spam the output if there are many walkers :/
+            warn(f"Walker {i} has multiple PlumedHamiltonians, only setting output file for the first one. Consider combining PLUMED input files.")
+        old_str = ham.hamiltonians[plumed_ids[0]].plumed_input
+        assert "PRINT" in old_str, f"Walker {i} PLUMED input does not contain PRINT directive"
+        # skip if user already set file
+        if "FILE" in old_str:
+            warn(f"Walker {i} PLUMED input already contains FILE directive, modifying to walker-specific file.")
+        new_str = set_path_in_plumed(old_str, "PRINT", f"walker-{i}_plumed.log")
+        ham.hamiltonians[plumed_ids[0]].plumed_input = new_str
+
+
 @typeguard.typechecked
 @psiflow.serializable
 class ReplicaExchange(Coupling):

psiflow/utils/_plumed.py

@@ -2,6 +2,13 @@
 import os
 
 import typeguard
+from typing import Union
+
+import numpy as np
+
+from parsl.app.app import python_app
+from parsl.dataflow.futures import AppFuture
+from parsl.data_provider.files import File
 
 
 @typeguard.typechecked
@@ -52,3 +59,34 @@ def set_path_in_plumed(plumed_input: str, keyword: str, path_to_set: str) -> str
                 line_before + "FILE={} ".format(path_to_set) + " ".join(line_after)
             )
     return "\n".join(lines)
+
+
+@typeguard.typechecked
+def _parse_plumed_output(plumed_output: File) -> dict[str, np.ndarray]:
+    with open(plumed_output.filepath, "r") as f:
+        first = f.readlines()[0]
+    fields = first.strip().split(" ")[3:]
+    data = np.loadtxt(plumed_output.filepath, skiprows=11)  # skip warmup values
+    output_dict = {}
+    for i, field in enumerate(fields):
+        output_dict[field] = data[:, i+1]
+    return output_dict
+
+parse_plumed_output = python_app(_parse_plumed_output, executors=["default_threads"])
+
+@typeguard.typechecked
+def _write_plumed_output(plumed_output_dict: Union[dict[str, np.ndarray], AppFuture], path: str, keys: Union[str, list[str], None] = None) -> None:
+    if keys is None:
+        keys = list(plumed_output_dict.keys())
+    if isinstance(keys, str):
+        keys = [keys]
+    with open(path, "w") as f:
+        f.write("# " + " ".join(keys) + "\n")
+        n_rows = len(plumed_output_dict[keys[0]])
+        for i in range(n_rows):
+            row = []
+            for key in keys:
+                row.append(str(plumed_output_dict[key][i])) if key in keys else None
+            f.write(" ".join(row) + "\n")
+
+write_plumed_output = python_app(_write_plumed_output, executors=["default_threads"])