diff --git a/psiflow/functions.py b/psiflow/functions.py index e12bec9..1115812 100644 --- a/psiflow/functions.py +++ b/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: diff --git a/psiflow/hamiltonians.py b/psiflow/hamiltonians.py index df1e39d..c211939 100644 --- a/psiflow/hamiltonians.py +++ b/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: diff --git a/psiflow/sampling/output.py b/psiflow/sampling/output.py index 7f73458..459008b 100644 --- a/psiflow/sampling/output.py +++ b/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: diff --git a/psiflow/sampling/sampling.py b/psiflow/sampling/sampling.py index b2f2ea9..1c4aac2 100644 --- a/psiflow/sampling/sampling.py +++ b/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, diff --git a/psiflow/sampling/walker.py b/psiflow/sampling/walker.py index c98cf5c..a88b344 100644 --- a/psiflow/sampling/walker.py +++ b/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): diff --git a/psiflow/utils/_plumed.py b/psiflow/utils/_plumed.py index 89b6789..7f30bd7 100644 --- a/psiflow/utils/_plumed.py +++ b/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"]) \ No newline at end of file