diff --git a/psiflow/functions.py b/psiflow/functions.py index a461e7d..189c35e 100644 --- a/psiflow/functions.py +++ b/psiflow/functions.py @@ -22,6 +22,7 @@ def format_output( energy: float, forces: np.ndarray | None = None, stress: np.ndarray | None = None, + extras: dict | None = None, **kwargs, ) -> dict: """""" @@ -29,7 +30,8 @@ def format_output( 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} class Function: @@ -47,6 +49,7 @@ def compute( ) -> dict[str, float | np.ndarray]: """Evaluate multiple geometries and merge data into single arrays""" value, grad_pos, grad_cell = create_outputs(self.outputs, geometries) + extras = {} for i, geometry in enumerate(geometries): if geometry == NullState: # TODO: remove continue @@ -54,7 +57,12 @@ def compute( value[i] = out["energy"] grad_pos[i, : len(geometry)] = out["forces"] grad_cell[i] = out["stress"] - return {"energy": value, "forces": grad_pos, "stress": grad_cell} + extras_out = out["extras"] + for key in extras_out: + if key not in extras: + extras[key] = np.empty(shape=(len(geometries),), dtype=type(np.float64)) + extras[key][i] = extras_out[key] + return {"energy": value, "forces": grad_pos, "stress": grad_cell, "extras": extras} @dataclass(frozen=True) @@ -80,6 +88,7 @@ def __call__( @dataclass(frozen=True) class PlumedFunction(Function): plumed_input: str + plumed_extras: Optional[list[str]] = None external: Optional[Union[str, Path]] = None plumed_instances: dict[tuple, "plumed.Plumed"] = field(default_factory=dict) @@ -115,6 +124,20 @@ def __call__( plumed_.cmd("init") for line in self.plumed_input.split("\n"): plumed_.cmd("readInputLine", line) + + self.plumed_data = {} + plumed_extras = self.plumed_extras if self.plumed_extras is not None else [] + 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_ @@ -136,12 +159,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: diff --git a/psiflow/hamiltonians.py b/psiflow/hamiltonians.py index f101188..5a98052 100644 --- a/psiflow/hamiltonians.py +++ b/psiflow/hamiltonians.py @@ -1,5 +1,6 @@ import logging from functools import partial +import re from pathlib import Path from typing import Optional, Union, Callable, Sequence, Any, ClassVar from dataclasses import dataclass, field @@ -287,6 +288,25 @@ def from_geometry( @psiflow.register_serializable @dataclass class PlumedHamiltonian(Hamiltonian): + plumed_input: str # TODO: or future? + external: Optional[psiflow._DataFuture] + function_name: str = "PlumedFunction" + + def __init__( + self, + plumed_input: str, + external: Union[None, str, Path, File, DataFuture] = None, + ): + 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)) + if external is not None: + assert external.filepath in self.plumed_input + self.external = external plumed_input: str | AppFuture external: Optional[psiflow._DataFuture] = None function_name: ClassVar[str] = "PlumedFunction" @@ -316,8 +336,12 @@ def get_app(self) -> Callable: ) def parameters(self) -> dict: + if self.external is not None: # ensure parameters depends on self.external + external = copy_app_future(self.external.filepath, inputs=[self.external]) + else: + external = None path = self.external.filepath if self.external is not None else None - return {"plumed_input": self.plumed_input, "external": path} + return {"plumed_input": self.plumed_input, "plumed_extras": self.plumed_extras, "external": path} def __eq__(self, other: Hamiltonian) -> bool: if self is other: diff --git a/psiflow/sampling/driver.py b/psiflow/sampling/driver.py index b56b852..16c1654 100644 --- a/psiflow/sampling/driver.py +++ b/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 diff --git a/psiflow/sampling/sampling.py b/psiflow/sampling/sampling.py index 117f8b7..3e31a90 100644 --- a/psiflow/sampling/sampling.py +++ b/psiflow/sampling/sampling.py @@ -290,6 +290,7 @@ def setup_output( components: list[HamiltonianComponent], observables: Optional[list[str]], step: Optional[int], + step_properties: Optional[int], keep_trajectory: bool, checkpoint_step: int, ) -> tuple[ET.Element, list]: @@ -307,9 +308,8 @@ def setup_output( full_list.append("ensemble_bias{electronvolt}") observables = list(set(full_list)) - if step is None: - # TODO: this logic should be elsewhere - step = checkpoint_step + step = step or step_properties or checkpoint_step # TODO: this logic should be elsewhere + step_properties = step_properties or step output = ET.Element("output", prefix="output") checkpoint = ET.Element( @@ -323,7 +323,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", ) @@ -332,10 +332,26 @@ def setup_output( properties = ET.Element( "properties", filename="properties", - stride=str(step), + stride=str(step_properties), ) 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? + if extras_list: + extras = ",".join(list(set(extras_list))) + extras_element = ET.Element( + "trajectory", + filename="extras", + stride=str(step_properties), + extra_type=extras, + bead="0", + ) + extras_element.text = " extras " + output.append(extras_element) return output, observables @@ -509,6 +525,7 @@ def _sample( walkers: list[Walker], steps: int, step: Optional[int] = None, + step_properties: Optional[int] = None, start: int = 0, keep_trajectory: bool = True, max_force: Optional[float] = None, @@ -550,12 +567,13 @@ 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 observables, step, + step_properties, keep_trajectory, checkpoint_step, ) @@ -666,6 +684,7 @@ def sample( walkers: list[Walker], steps: int, step: Optional[int] = None, + step_properties: Optional[int] = None, start: int = 0, keep_trajectory: bool = True, max_force: Optional[float] = None, @@ -689,6 +708,7 @@ def sample( _walkers, steps, step=step, + step_properties=step_properties, start=start, keep_trajectory=keep_trajectory, max_force=max_force, diff --git a/psiflow/sampling/server.py b/psiflow/sampling/server.py index d0bc4e5..091938a 100644 --- a/psiflow/sampling/server.py +++ b/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,6 +94,28 @@ 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("*0*.properties")) # properties should be the same for all coupled walkers? + _, info_dict = read_output(file_props) + props_headers = ["# column {} --> {}{} : {}".format(i + 1, key, "{" + info_dict[key][0] + "}", info_dict[key][1]) for i, key in enumerate(info_dict)] + # extras headers + file_extras = next(Path.cwd().glob("*0*.extras*")) # extras should be the same for all coupled walkers? + 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 + 1 + 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=":") @@ -133,6 +156,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 extras to simulation properties, if they exist + if len(output_props) and any(Path.cwd().glob("*.extras*")): + 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 @@ -143,9 +174,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) diff --git a/tests/test_function.py b/tests/test_function.py index 39ec1ca..41d1c29 100644 --- a/tests/test_function.py +++ b/tests/test_function.py @@ -130,6 +130,15 @@ def test_plumed_function(tmp_path, dataset, dataset_h2): RESTRAINT ARG=CV AT=50 KAPPA=1 """ function = PlumedFunction(plumed_input) + energy, forces, stress, _ = function.compute(dataset.geometries().result()).values() + + volumes = np.linalg.det(dataset.get("cell").result()) + energy_ = (volumes - 50) ** 2 * (kJ / mol) / 2 + assert np.allclose( + energy, + energy_, + ) + hamiltonian = PlumedHamiltonian(plumed_input) energy, forces, stress = function.compute(geometries).values() @@ -175,6 +184,21 @@ def test_plumed_function(tmp_path, dataset, dataset_h2): with open(path_hills, "r") as f: assert f.read() == hills + # check extraction of PLUMED extras + plumed_input = """ +UNITS LENGTH=A ENERGY=kj/mol TIME=fs +CV: VOLUME +PRINT ARG=CV +""" + ham = PlumedHamiltonian(plumed_input) + plumed_extras = ham.plumed_extras + assert plumed_extras == ["CV"] + function = PlumedFunction(plumed_input, plumed_extras=plumed_extras) + _, _, _, extras = function.compute(dataset.geometries().result()).values() + assert "CV" in extras + + volumes = np.linalg.det(dataset.get("cell").result()) + assert np.allclose(extras["CV"].astype(np.float64), volumes) def test_harmonic_function(dataset): test_set = dataset[:10] diff --git a/tests/test_sampling.py b/tests/test_sampling.py index 0f3ccf8..d236913 100644 --- a/tests/test_sampling.py +++ b/tests/test_sampling.py @@ -245,6 +245,19 @@ def test_sample(dataset, mace_foundation): temp1 = output["temperature{kelvin}"].result()[-1] assert np.allclose(temp0, temp1) + # check extraction of PLUMED extras + plumed_str = """ +UNITS LENGTH=A +CV: VOLUME +PRINT ARG=CV +""" + walker = Walker(start=geom_start, temperature=300, hamiltonian=plumed+einstein) + [output] = sample( + [walker], + steps=10, + ) + assert np.allclose(output["CV{au}"].result(), output["volume{angstrom3}"].result()) + return