diff --git a/.gitignore b/.gitignore
index ebc1d1f0..50231d5c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,6 +3,7 @@ __pycache__/
*.py[cod]
*$py.class
*.swp
+*.pyc
# C extensions
*.so
diff --git a/psiflow/__init__.py b/psiflow/__init__.py
index 10816e2c..e8de3e30 100644
--- a/psiflow/__init__.py
+++ b/psiflow/__init__.py
@@ -34,41 +34,3 @@ def resolve_and_check(path: Path) -> Path:
context = ExecutionContextLoader.context
wait = ExecutionContextLoader.wait
-
-# TODO: EXECUTION
-# - max_runtime is in seconds, max_simulation_time (and others) is in minutes
-# - ExecutionDefinition gpu argument?
-# - ExecutionDefinition wrap_in_timeout functionality
-# - ExecutionDefinition wrap_in_srun functionality? Actually for MD this is more iffy right
-# - ExecutionDefinition why properties?
-# - ExecutionDefinition centralize wq_resources
-# - configuration file with all options
-# - timeout -s 9 or -s 15?
-# - executor keys are hardcoded strings..
-# - define a 'format_env_variables' util
-# - update GPAW + ORCA + Default containers
-# include s-dftd3 in modelevaluation + install s-dftd3 with openmp
-# - what with mem_per_node / mem_per_worker
-# - always GPU for training?
-# - currently reference mpi_args have to be tuple according to typeguard..
-# - cores_per_block has to be specified even when exclusive..?
-# - can we do something with WQ priority?
-# - see chatgpt convo for process memory limits and such
-# - make /tmp for app workdirs an option?
-# - what is scaling_cores_per_worker in WQ
-# - can we clean up psiflow_internal slightly?
-# -
-# TODO: REFERENCE
-# - reference MPI args not really checked
-# - mpi flags are very finicky across implementations --> use ENV args?
-# OMPI_MCA_orte_report_bindings=1 I_MPI_DEBUG=4
-# - commands ends with 'exit 0' - what if we do not want to exit yet?
-# - some actual logging?
-# - safe_compute_dataset functionality?
-# - ReferenceDummy is a bit sloppy
-# -
-# TODO: MISC
-# - think about test efficiency
-# - some imports take a very long time
-# - fix serialisation
-# -
diff --git a/psiflow/execution.py b/psiflow/execution.py
index 9653031d..c632935a 100644
--- a/psiflow/execution.py
+++ b/psiflow/execution.py
@@ -10,7 +10,7 @@
from threading import Lock
# see https://stackoverflow.com/questions/59904631/python-class-constants-in-dataclasses
-from typing import Any, Optional, Union, ClassVar, Protocol
+from typing import Any, Optional, Union, ClassVar, Protocol, Iterable
import parsl
import psutil
@@ -155,7 +155,7 @@ def create_executor(self, path: Path) -> ParslExecutor:
# only ModelEvaluation / ModelTraining / default_htex run in containers
if not isinstance(self, ReferenceEvaluation) and self.container:
- prepend = self.container.launch_command()
+ prepend = self.container.launch_command(self.gpu)
worker_executable = f"{prepend} work_queue_worker"
else:
worker_executable = "work_queue_worker"
@@ -261,27 +261,37 @@ def client_command(self):
command_list = ["psiflow-client"]
return " ".join(command_list)
- def get_client_args(
- self,
- hamiltonian_name: str,
- nwalkers: int,
- motion: str,
- ) -> list[str]:
- if "MACE" in hamiltonian_name:
- if motion in ["minimize", "vibrations"]:
- dtype = "float64"
- else:
- dtype = "float32"
- nclients = min(nwalkers, self.max_workers)
- if self.gpu:
- template = "--dtype={} --device=cuda:{}"
- args = [template.format(dtype, i) for i in range(nclients)]
- else:
- template = "--dtype={} --device=cpu"
- args = [template.format(dtype) for i in range(nclients)]
- return args
+ # def get_client_args(
+ # self,
+ # hamiltonian_name: str,
+ # nwalkers: int,
+ # motion: str,
+ # ) -> list[str]:
+ # # TODO: redo this
+ # if "MACE" in hamiltonian_name:
+ # if motion in ["minimize", "vibrations"]:
+ # dtype = "float64"
+ # else:
+ # dtype = "float32"
+ # nclients = min(nwalkers, self.max_workers)
+ # if self.gpu:
+ # template = "--dtype={} --device=cuda:{}"
+ # args = [template.format(dtype, i) for i in range(nclients)]
+ # else:
+ # template = "--dtype={} --device=cpu"
+ # args = [template.format(dtype) for i in range(nclients)]
+ # return args
+ # else:
+ # return [""]
+
+ def get_driver_devices(self, nwalkers: int) -> list[dict]:
+ # assumes driver is GPU capable
+ # TODO: what if only 1 gpu is available?
+ nclients = min(nwalkers, self.max_workers)
+ if self.gpu:
+ return [{'device': f'cuda:{i}'} for i in range(nclients)]
else:
- return [""]
+ return [{'device': 'cpu'} for _ in range(nclients)]
def wq_resources(self, nwalkers):
if self.use_threadpool:
@@ -433,7 +443,7 @@ class ReferenceSpec(Protocol):
name: ClassVar[str]
reference_args: ClassVar[tuple[str, ...]]
mpi_command: str
- mpi_args: tuple[str, ...]
+ mpi_args: Iterable[str]
executable: str
def launch_command(self) -> str:
diff --git a/psiflow/free_energy/phonons.py b/psiflow/free_energy/phonons.py
index 8b60223e..c93dd47d 100644
--- a/psiflow/free_energy/phonons.py
+++ b/psiflow/free_energy/phonons.py
@@ -1,11 +1,8 @@
-from __future__ import annotations # necessary for type-guarding class methods
-
import xml.etree.ElementTree as ET
from typing import Optional, Union
import numpy as np
import parsl
-import typeguard
from ase.units import Bohr, Ha, J, _c, _hplanck, _k, kB, second
from parsl.app.app import bash_app, python_app
from parsl.dataflow.futures import AppFuture
@@ -13,19 +10,19 @@
import psiflow
from psiflow.data import Dataset
from psiflow.geometry import Geometry, mass_weight
-from psiflow.hamiltonians import Hamiltonian, MixtureHamiltonian
+from psiflow.hamiltonians import Hamiltonian, MACEHamiltonian
from psiflow.sampling.sampling import (
setup_sockets,
- make_force_xml,
- make_start_command,
- make_client_command
+ make_server_command,
+ make_driver_commands,
+ make_wait_for_sockets_command,
)
+from psiflow.sampling.optimize import setup_forces, export_env_command
from psiflow.utils.apps import multiply
from psiflow.utils.io import load_numpy, save_xml
from psiflow.utils import TMP_COMMAND, CD_COMMAND
-@typeguard.typechecked
def _compute_frequencies(hessian: np.ndarray, geometry: Geometry) -> np.ndarray:
assert hessian.shape[0] == hessian.shape[1]
assert len(geometry) * 3 == hessian.shape[0]
@@ -35,7 +32,6 @@ def _compute_frequencies(hessian: np.ndarray, geometry: Geometry) -> np.ndarray:
compute_frequencies = python_app(_compute_frequencies, executors=["default_threads"])
-@typeguard.typechecked
def _harmonic_free_energy(
frequencies: Union[float, np.ndarray],
temperature: float,
@@ -65,9 +61,7 @@ def _harmonic_free_energy(
harmonic_free_energy = python_app(_harmonic_free_energy, executors=["default_threads"])
-@typeguard.typechecked
def setup_motion(
- mode: str,
asr: str,
pos_shift: float,
energy_shift: float,
@@ -91,35 +85,33 @@ def setup_motion(
def _execute_ipi(
- hamiltonian_names: list[str],
- client_args: list[list[str]],
+ driver_kwargs: list[dict],
command_server: str,
- command_client: str,
- stdout: str = "",
- stderr: str = "",
+ env_vars: dict = {},
+ stdout: str = parsl.AUTO_LOGNAME,
+ stderr: str = parsl.AUTO_LOGNAME,
inputs: list = [],
outputs: list = [],
parsl_resource_specification: Optional[dict] = None,
) -> str:
- command_start = make_start_command(command_server, inputs[0], inputs[1])
- commands_client = []
- for i, name in enumerate(hamiltonian_names):
- args = client_args[i]
- assert len(args) == 1 # only have one client per hamiltonian
- for arg in args:
- commands_client += make_client_command(command_client, name, inputs[2 + i], inputs[1], arg),
-
- command_end = f'{command_server} --cleanup'
- command_copy = f'cp i-pi.output_full.hess {outputs[0]}'
+ file_xml, file_xyz_in, *files_in = inputs
+ command_start = make_server_command(
+ command_server, file_xml, file_xyz_in, outputs[0], [], []
+ )
+ command_wait = make_wait_for_sockets_command(
+ set(d["address"] for d in driver_kwargs)
+ )
+ commands_driver = make_driver_commands(driver_kwargs, file_xyz_in, files_in)
command_list = [
TMP_COMMAND,
CD_COMMAND,
+ export_env_command(env_vars),
command_start,
- *commands_client,
+ command_wait,
+ *commands_driver,
"wait",
- command_end,
- command_copy,
+ f"cp i-pi.output_full.hess {outputs[0]}",
]
return "\n".join(command_list)
@@ -127,7 +119,6 @@ def _execute_ipi(
execute_ipi = bash_app(_execute_ipi, executors=["ModelEvaluation"])
-@typeguard.typechecked
def compute_harmonic(
state: Union[Geometry, AppFuture],
hamiltonian: Hamiltonian,
@@ -136,26 +127,23 @@ def compute_harmonic(
pos_shift: float = 0.01,
energy_shift: float = 0.00095,
) -> AppFuture:
- hamiltonian: MixtureHamiltonian = 1 * hamiltonian
- names = hamiltonian.get_named_components()
- sockets = setup_sockets(names)
- forces = make_force_xml(hamiltonian, names)
+ components, force_xml = setup_forces(hamiltonian)
+ sockets = setup_sockets(components)
initialize = ET.Element("initialize", nbeads="1")
start = ET.Element("file", mode="ase", cell_units="angstrom")
start.text = " start_0.xyz "
initialize.append(start)
- motion = setup_motion(mode, asr, pos_shift, energy_shift)
+ motion = setup_motion(asr, pos_shift, energy_shift)
+ if mode != "fd":
+ raise NotImplementedError
system = ET.Element("system")
system.append(initialize)
system.append(motion)
- system.append(forces)
-
- # output = setup_output(keep_trajectory)
+ system.append(force_xml)
simulation = ET.Element("simulation", mode="static")
- # simulation.append(output)
for socket in sockets:
simulation.append(socket)
simulation.append(system)
@@ -169,33 +157,22 @@ def compute_harmonic(
simulation,
outputs=[context.new_file("input_", ".xml")],
).outputs[0]
- inputs = [
- input_future,
- Dataset([state]).extxyz,
- ]
- inputs += hamiltonian.serialize(dtype="float64")
-
- client_args = []
- for name in names:
- args = definition.get_client_args(name, 1, "vibrations")
- client_args.append(args)
- outputs = [
- context.new_file("hess_", ".txt"),
- ]
+ inputs = [input_future, Dataset([state]).extxyz]
- command_server = definition.server_command()
- command_client = definition.client_command()
- resources = definition.wq_resources(1)
+ driver_kwargs = []
+ for i, comp in enumerate(components):
+ inputs.append(comp.hamiltonian.serialize_function(dtype="float64"))
+ kwargs = {"idx": i, "address": comp.address}
+ if isinstance(comp.hamiltonian, MACEHamiltonian):
+ kwargs |= definition.get_driver_devices(1)[0]
+ driver_kwargs.append(kwargs)
result = execute_ipi(
- names,
- client_args,
- command_server,
- command_client,
- stdout=parsl.AUTO_LOGNAME,
- stderr=parsl.AUTO_LOGNAME,
+ driver_kwargs,
+ definition.server_command(),
+ env_vars=definition.env_vars,
inputs=inputs,
- outputs=outputs,
- parsl_resource_specification=resources,
+ outputs=[context.new_file("hess_", ".txt")],
+ parsl_resource_specification=definition.wq_resources(1),
)
- return multiply(load_numpy(inputs=[result.outputs[0]]), Ha / Bohr**2)
\ No newline at end of file
+ return multiply(load_numpy(inputs=[result.outputs[0]]), Ha / Bohr**2)
diff --git a/psiflow/geometry.py b/psiflow/geometry.py
index 4e00d4cc..c9400c41 100644
--- a/psiflow/geometry.py
+++ b/psiflow/geometry.py
@@ -304,9 +304,16 @@ def from_string(cls, s: str, natoms: Optional[int] = None) -> Optional[Geometry]
if key.startswith("order_"):
order[key.replace("order_", "")] = value
+ # TODO: this needs a better solution
+ cell = comment_dict.pop("Lattice", np.zeros((3, 3)))
+ if isinstance(cell, str):
+ # most likely an array of NaNs
+ cell = np.array(cell.split()).reshape((3, 3))
+ cell = cell.T # transposed!
+ pbc = sum(comment_dict.pop("pbc", [])) == 3
geometry = cls(
per_atom=per_atom,
- cell=comment_dict.pop("Lattice", np.zeros((3, 3))).T, # transposed!
+ cell=cell if pbc else np.zeros((3, 3)),
energy=comment_dict.pop("energy", None),
stress=comment_dict.pop("stress", None),
delta=comment_dict.pop("delta", None),
diff --git a/psiflow/hamiltonians.py b/psiflow/hamiltonians.py
index df1e39d2..a065484d 100644
--- a/psiflow/hamiltonians.py
+++ b/psiflow/hamiltonians.py
@@ -1,12 +1,9 @@
-from __future__ import annotations # necessary for type-guarding class methods
-
import urllib
from functools import partial
from pathlib import Path
-from typing import ClassVar, Optional, Union, Callable
+from typing import ClassVar, Optional, Union, Callable, Sequence
import numpy as np
-import typeguard
from parsl.app.app import python_app
from parsl.app.futures import DataFuture
from parsl.data_provider.files import File
@@ -29,7 +26,6 @@
from psiflow.utils.io import dump_json
-@typeguard.typechecked
@psiflow.serializable
class Hamiltonian(Computable):
# TODO: app is actually an instance variable, but serialization complains..
@@ -55,24 +51,24 @@ def compute(
batch_size=batch_size,
)
- def __eq__(self, hamiltonian: Hamiltonian) -> bool:
+ def __eq__(self, hamiltonian: "Hamiltonian") -> bool:
raise NotImplementedError
- def __mul__(self, a: float) -> MixtureHamiltonian:
+ def __mul__(self, a: float) -> "MixtureHamiltonian":
return MixtureHamiltonian([self], [a])
- def __add__(self, hamiltonian: Hamiltonian) -> Hamiltonian:
+ def __add__(self, hamiltonian: "Hamiltonian") -> "MixtureHamiltonian | Hamiltonian":
if type(hamiltonian) is Zero:
return self
mixture = MixtureHamiltonian([self], [1.0])
return mixture.__add__(hamiltonian)
- def __sub__(self, hamiltonian: Hamiltonian) -> Hamiltonian:
- return self + (-1.0) * hamiltonian
+ def __sub__(self, hamiltonian: "Hamiltonian") -> "MixtureHamiltonian":
+ return self + hamiltonian * (-1.0)
__rmul__ = __mul__ # handle float * Hamiltonian
- def serialize_function(self, **kwargs):
+ def serialize_function(self, **kwargs) -> DataFuture:
parameters = self.parameters()
for key, value in kwargs.items():
if key in parameters:
@@ -87,7 +83,6 @@ def parameters(self) -> dict:
return {}
-@typeguard.typechecked
@psiflow.serializable
class Zero(Hamiltonian):
@@ -100,7 +95,7 @@ def __eq__(self, hamiltonian: Hamiltonian) -> bool:
return True
return False
- def __mul__(self, a: float) -> Hamiltonian:
+ def __mul__(self, a: float) -> "Zero":
return Zero()
def __add__(self, hamiltonian: Hamiltonian) -> Hamiltonian:
@@ -110,7 +105,6 @@ def __add__(self, hamiltonian: Hamiltonian) -> Hamiltonian:
__rmul__ = __mul__ # handle float * Zero
-@typeguard.typechecked
@psiflow.serializable
class MixtureHamiltonian(Hamiltonian):
hamiltonians: list[Hamiltonian]
@@ -118,8 +112,8 @@ class MixtureHamiltonian(Hamiltonian):
def __init__(
self,
- hamiltonians: Union[tuple, list][Hamiltonian],
- coefficients: Union[tuple, list][float],
+ hamiltonians: Sequence[Hamiltonian],
+ coefficients: Sequence[float],
) -> None:
assert len(hamiltonians) == len(coefficients)
self.hamiltonians = list(hamiltonians)
@@ -146,9 +140,10 @@ def compute( # override compute for efficient batching
batch_size=batch_size,
)
- def __eq__(self, hamiltonian: Hamiltonian | MixtureHamiltonian) -> bool:
+ def __eq__(self, hamiltonian: Hamiltonian) -> bool:
if type(hamiltonian) is not MixtureHamiltonian:
return False
+ hamiltonian: MixtureHamiltonian
if len(self.coefficients) != len(hamiltonian.coefficients):
return False
for c, h in zip(self.coefficients, self.hamiltonians):
@@ -158,7 +153,7 @@ def __eq__(self, hamiltonian: Hamiltonian | MixtureHamiltonian) -> bool:
return False
return True
- def __mul__(self, a: float) -> MixtureHamiltonian:
+ def __mul__(self, a: float) -> "MixtureHamiltonian":
return MixtureHamiltonian(
self.hamiltonians,
[c * a for c in self.coefficients],
@@ -169,7 +164,7 @@ def __mul__(self, a: float) -> MixtureHamiltonian:
def __len__(self) -> int:
return len(self.coefficients)
- def __add__(self, hamiltonian: Hamiltonian) -> MixtureHamiltonian:
+ def __add__(self, hamiltonian: Hamiltonian) -> "MixtureHamiltonian":
if type(hamiltonian) is Zero:
return self
if type(hamiltonian) is not MixtureHamiltonian:
@@ -192,7 +187,7 @@ def get_index(self, hamiltonian: Hamiltonian) -> Optional[int]:
except ValueError:
return None
- def get_indices(self, mixture: MixtureHamiltonian) -> Optional[tuple[int, ...]]:
+ def get_indices(self, mixture: "MixtureHamiltonian") -> Optional[tuple[int, ...]]:
# TODO: why do we not just return None for the missing components?
assert type(mixture) is MixtureHamiltonian
indices = []
@@ -209,7 +204,7 @@ def get_coefficient(self, hamiltonian: Hamiltonian) -> Optional[float]:
return self.coefficients[idx]
def get_coefficients(
- self, mixture: MixtureHamiltonian
+ self, mixture: "MixtureHamiltonian"
) -> Optional[tuple[float, ...]]:
assert type(mixture) is MixtureHamiltonian
for h in mixture.hamiltonians:
@@ -223,6 +218,7 @@ def get_coefficients(
def get_named_components(self) -> list[str]:
"""Create unique string name for every hamiltonian component to be used in i-Pi sampling"""
+ # TODO: iter_named_components?
names, counts = [], {}
for h in self.hamiltonians:
name = h.__class__.__name__
@@ -235,7 +231,6 @@ def serialize(self, **kwargs) -> list[DataFuture]:
return [h.serialize_function(**kwargs) for h in self.hamiltonians]
-@typeguard.typechecked
@psiflow.serializable
class EinsteinCrystal(Hamiltonian):
reference_geometry: Union[Geometry, AppFuture]
@@ -264,7 +259,7 @@ def parameters(self) -> dict:
"volume": get_attribute(self.reference_geometry, "volume"),
}
- def __eq__(self, hamiltonian: Hamiltonian | EinsteinCrystal) -> bool:
+ def __eq__(self, hamiltonian: Hamiltonian) -> bool:
if (
not isinstance(hamiltonian, EinsteinCrystal)
or not np.allclose(self.force_constant, hamiltonian.force_constant)
@@ -274,10 +269,9 @@ def __eq__(self, hamiltonian: Hamiltonian | EinsteinCrystal) -> bool:
return True
-@typeguard.typechecked
@psiflow.serializable
class PlumedHamiltonian(Hamiltonian):
- plumed_input: str
+ plumed_input: str # TODO: or future?
external: Optional[psiflow._DataFuture]
function_name: ClassVar[str] = "PlumedFunction"
@@ -311,7 +305,7 @@ def parameters(self) -> dict:
external = None
return {"plumed_input": self.plumed_input, "external": external}
- def __eq__(self, other: Hamiltonian | PlumedHamiltonian) -> bool:
+ def __eq__(self, other: Hamiltonian) -> bool:
if (
not isinstance(other, PlumedHamiltonian)
or self.plumed_input != other.plumed_input
@@ -320,7 +314,6 @@ def __eq__(self, other: Hamiltonian | PlumedHamiltonian) -> bool:
return True
-@typeguard.typechecked
@psiflow.serializable
class Harmonic(Hamiltonian):
reference_geometry: Union[Geometry, AppFuture[Geometry]]
@@ -354,9 +347,10 @@ def parameters(self) -> dict:
"hessian": self.hessian,
}
- def __eq__(self, hamiltonian: Hamiltonian | Harmonic) -> bool:
+ def __eq__(self, hamiltonian: Hamiltonian) -> bool:
if type(hamiltonian) is not Harmonic:
return False
+ hamiltonian: Harmonic
if hamiltonian.reference_geometry != self.reference_geometry:
return False
@@ -365,10 +359,7 @@ def __eq__(self, hamiltonian: Hamiltonian | Harmonic) -> bool:
is_array0 = type(hamiltonian.hessian) is np.ndarray
is_array1 = type(self.hessian) is np.ndarray
if is_array0 and is_array1:
- equal = np.allclose(
- hamiltonian.hessian,
- self.hessian,
- )
+ equal = np.allclose(hamiltonian.hessian, self.hessian)
else:
equal = hamiltonian.hessian == self.hessian
if not equal:
@@ -376,7 +367,6 @@ def __eq__(self, hamiltonian: Hamiltonian | Harmonic) -> bool:
return True
-@typeguard.typechecked
@psiflow.serializable
class D3Hamiltonian(Hamiltonian):
method: str
@@ -406,7 +396,7 @@ def _create_apps(self):
def parameters(self) -> dict:
return {"method": self.method, "damping": self.damping}
- def __eq__(self, hamiltonian: Hamiltonian | D3Hamiltonian) -> bool:
+ def __eq__(self, hamiltonian: Hamiltonian) -> bool:
if (
not isinstance(hamiltonian, D3Hamiltonian)
or self.method != hamiltonian.method
@@ -416,7 +406,6 @@ def __eq__(self, hamiltonian: Hamiltonian | D3Hamiltonian) -> bool:
return True
-@typeguard.typechecked
@psiflow.serializable
class MACEHamiltonian(Hamiltonian):
external: psiflow._DataFuture
@@ -449,6 +438,7 @@ def _create_apps(self):
)
def parameters(self) -> dict:
+ # TODO: Why is the future copy needed? Can we not pass the File/DataFuture directly?
model_path = copy_app_future(self.external.filepath, inputs=[self.external])
evaluation = psiflow.context().definitions["ModelEvaluation"]
return {
@@ -460,9 +450,10 @@ def parameters(self) -> dict:
"env_vars": evaluation.env_vars,
}
- def __eq__(self, hamiltonian: Hamiltonian | MACEHamiltonian) -> bool:
+ def __eq__(self, hamiltonian: Hamiltonian) -> bool:
if type(hamiltonian) is not MACEHamiltonian:
return False
+ hamiltonian: MACEHamiltonian
if self.external.filepath != hamiltonian.external.filepath:
return False
if len(self.atomic_energies) != len(hamiltonian.atomic_energies):
@@ -478,7 +469,7 @@ def __eq__(self, hamiltonian: Hamiltonian | MACEHamiltonian) -> bool:
# TODO: the methods below are outdated..
@classmethod
- def mace_mp0(cls, size: str = "small") -> MACEHamiltonian:
+ def mace_mp0(cls, size: str = "small") -> 'MACEHamiltonian':
urls = dict(
small="https://github.com/ACEsuit/mace-mp/releases/download/mace_mp_0/2023-12-10-mace-128-L0_energy_epoch-249.model", # 2023-12-10-mace-128-L0_energy_epoch-249.model
large="https://github.com/ACEsuit/mace-mp/releases/download/mace_mp_0/2023-12-03-mace-128-L1_epoch-199.model",
@@ -492,7 +483,7 @@ def mace_mp0(cls, size: str = "small") -> MACEHamiltonian:
return cls(parsl_file, {})
@classmethod
- def mace_cc(cls) -> MACEHamiltonian:
+ def mace_cc(cls) -> 'MACEHamiltonian':
url = "https://github.com/molmod/psiflow/raw/main/examples/data/ani500k_cc_cpu.model"
parsl_file = psiflow.context().new_file("mace_mp_", ".pth")
urllib.request.urlretrieve(
@@ -500,3 +491,7 @@ def mace_cc(cls) -> MACEHamiltonian:
parsl_file.filepath,
)
return cls(parsl_file, {})
+
+
+def combine_hamiltonians(hamiltonians: list[Hamiltonian]) -> MixtureHamiltonian:
+ return sum(hamiltonians, start=Zero()) # mostly for type hinting
diff --git a/psiflow/reference/cp2k_.py b/psiflow/reference/cp2k_.py
index cec01f0e..6af60952 100644
--- a/psiflow/reference/cp2k_.py
+++ b/psiflow/reference/cp2k_.py
@@ -1,5 +1,3 @@
-from __future__ import annotations # necessary for type-guarding class methods
-
import copy
import io
import warnings
diff --git a/psiflow/reference/dummy.py b/psiflow/reference/dummy.py
index d175c9d3..f947efbf 100644
--- a/psiflow/reference/dummy.py
+++ b/psiflow/reference/dummy.py
@@ -1,5 +1,3 @@
-from __future__ import annotations # necessary for type-guarding class methods
-
from typing import Optional, Union
from functools import partial
diff --git a/psiflow/reference/gpaw_.py b/psiflow/reference/gpaw_.py
index 41e4c58b..5c6d5bdd 100644
--- a/psiflow/reference/gpaw_.py
+++ b/psiflow/reference/gpaw_.py
@@ -1,5 +1,3 @@
-from __future__ import annotations # necessary for type-guarding class methods
-
import json
from pathlib import Path
from typing import Optional, Union
diff --git a/psiflow/reference/orca_.py b/psiflow/reference/orca_.py
index 71f4338f..bac4a447 100644
--- a/psiflow/reference/orca_.py
+++ b/psiflow/reference/orca_.py
@@ -1,5 +1,3 @@
-from __future__ import annotations # necessary for type-guarding class methods
-
import warnings
import re
from functools import partial
diff --git a/psiflow/reference/reference.py b/psiflow/reference/reference.py
index b6ed9321..96ff842e 100644
--- a/psiflow/reference/reference.py
+++ b/psiflow/reference/reference.py
@@ -17,7 +17,7 @@
from psiflow.data.utils import extract_quantities
from psiflow.geometry import Geometry, NullState
from psiflow.utils.apps import copy_app_future, setup_logger
-from psiflow.utils.parse import LineNotFoundError
+from psiflow.utils.parse import LineNotFoundError, get_task_name_id
logger = setup_logger(__name__) # logging per module
@@ -44,7 +44,7 @@ class SinglePointResult:
def update_geometry(geom: Geometry, data: dict) -> Geometry:
""""""
- _, task_id, task_name = data["stdout"].stem.split("_", maxsplit=2)
+ task_name, task_id = get_task_name_id(data["stdout"])
logger.info(f'Task "{task_name}" (ID {task_id}): {data["status"].name}')
geom = geom.copy()
diff --git a/psiflow/sampling/__init__.py b/psiflow/sampling/__init__.py
index 3e95469c..c374efd5 100644
--- a/psiflow/sampling/__init__.py
+++ b/psiflow/sampling/__init__.py
@@ -1,7 +1,7 @@
from .metadynamics import Metadynamics # noqa: F401
-# from .optimize import optimize, optimize_dataset # noqa: F401
from .output import SimulationOutput # noqa: F401
from .sampling import sample # noqa: F401
+from .ase import optimize, optimize_dataset
from .walker import ReplicaExchange # noqa: F401
from .walker import Walker # noqa: F401
from .walker import quench # noqa: F401
diff --git a/psiflow/sampling/_ase.py b/psiflow/sampling/_ase.py
index 37c4a035..41495f81 100644
--- a/psiflow/sampling/_ase.py
+++ b/psiflow/sampling/_ase.py
@@ -1,9 +1,6 @@
"""
Structure optimisation through ASE
TODO: do we need to check for very large forces?
-TODO: what units are pressure?
-TODO: what to do when max_steps is reached before converging?
-TODO: timeout is duplicated code
"""
import os
@@ -11,16 +8,17 @@
import warnings
import signal
import argparse
+import shutil
+import time
from pathlib import Path
from types import SimpleNamespace
import ase
import ase.io
import numpy as np
-from ase.io.extxyz import save_calc_results
from ase.calculators.calculator import Calculator, all_properties
from ase.calculators.mixing import LinearCombinationCalculator
-from ase.optimize.precon import PreconLBFGS
+from ase.optimize import BFGS
from ase.filters import FrechetCellFilter
from psiflow.geometry import Geometry
@@ -28,13 +26,13 @@
from psiflow.sampling.utils import TimeoutException, timeout_handler
-ALLOWED_MODES: tuple[str, ...] = ('full', 'fix_volume', 'fix_shape', 'fix_cell')
-FILE_OUT: str = 'out.xyz'
-FILE_TRAJ: str = 'out.traj'
+ALLOWED_MODES: tuple[str, ...] = ("full", "fix_volume", "fix_shape", "fix_cell")
+FILE_OUT: str = "out.xyz"
+FILE_TRAJ: str = "out.traj"
class FunctionCalculator(Calculator):
- implemented_properties = ['energy', 'free_energy', 'forces', 'stress']
+ implemented_properties = ["energy", "free_energy", "forces", "stress"]
def __init__(self, function: EnergyFunction, **kwargs):
super().__init__(**kwargs)
@@ -49,145 +47,153 @@ def calculate(
super().calculate(atoms, properties, system_changes)
geometry = Geometry.from_atoms(self.atoms)
self.results = self.function(geometry)
- self.results['free_energy'] = self.results['energy'] # required by optimiser
+ self.results["free_energy"] = self.results["energy"] # required by optimiser
def log_state(atoms: ase.Atoms) -> None:
""""""
+
def make_log(data: list[tuple[str]]):
""""""
- txt = ['', 'Current atoms state:']
- txt += [f'{_[0]:<15}: {_[1]:<25}[{_[2]}]' for _ in data]
- txt += 'End', ''
- print(*txt, sep='\n')
+ txt = ["", "Current atoms state:"]
+ txt += [f"{_[0]:<15}: {_[1]:<25}[{_[2]}]" for _ in data]
+ txt += "End", ""
+ print(*txt, sep="\n")
data = []
+ energy, max_force = [np.nan] * 2
if atoms.calc:
- energy, max_force = atoms.get_potential_energy(), np.linalg.norm(atoms.get_forces(), axis=0).max()
- else:
- energy, max_force = [np.nan] * 2
- data += ('Energy', f'{energy:.2f}', 'eV'), ('Max. force', f'{max_force:.2E}', 'eV/A')
+ energy = atoms.get_potential_energy()
+ max_force = np.linalg.norm(atoms.get_forces(), axis=1).max()
+ data.append(("Energy", f"{energy:.2f}", "eV"))
+ data.append(("Max. force", f"{max_force:.2E}", "eV/A"))
if not all(atoms.pbc):
make_log(data)
return
volume, cell = atoms.get_volume(), atoms.get_cell().cellpar().round(3)
- data += ('Cell volume', f'{atoms.get_volume():.2f}', 'A^3'),
- data += ('Box norms', str(cell[:3])[1:-1], 'A'), ('Box angles', str(cell[3:])[1:-1], 'degrees')
+ data.append(("Cell volume", f"{atoms.get_volume():.2f}", "A^3"))
+ data.append(("Box norms", str(cell[:3])[1:-1], "A"))
+ data.append(("Box angles", str(cell[3:])[1:-1], "degrees"))
make_log(data)
return
-def get_dof_filter(atoms: ase.Atoms, mode: str, pressure: float) -> ase.Atoms | FrechetCellFilter:
+def get_dof_filter(
+ atoms: ase.Atoms, mode: str, pressure: float
+) -> ase.Atoms | FrechetCellFilter:
""""""
- if mode == 'fix_cell':
+ if mode == "fix_cell":
if pressure:
- warnings.warn('Ignoring external pressure..')
+ warnings.warn("Ignoring external pressure..")
return atoms
- kwargs = {'mask': [True] * 6, 'scalar_pressure': pressure} # enable cell DOFs
- if mode == 'fix_shape':
- kwargs['hydrostatic_strain'] = True
- if mode == 'fix_volume':
- kwargs['constant_volume'] = True
+ kwargs = {"mask": [True] * 6, "scalar_pressure": pressure} # enable cell DOFs
+ if mode == "fix_shape":
+ kwargs["hydrostatic_strain"] = True
+ if mode == "fix_volume":
+ kwargs["constant_volume"] = True
if pressure:
- warnings.warn('Ignoring applied pressure during fixed volume optimisation..')
+ warnings.warn(
+ "Ignoring applied pressure during fixed volume optimisation.."
+ )
return FrechetCellFilter(atoms, **kwargs)
-def run(args: SimpleNamespace):
+def optimise(
+ dof: ase.Atoms | FrechetCellFilter,
+ f_max: float,
+ max_steps: int,
+ trajectory: Path | str = None,
+ **kwargs,
+) -> bool:
+ """"""
+ steps = max_steps // 2
+ opt = BFGS(dof, trajectory=trajectory)
+
+ converged = opt.run(fmax=f_max, steps=steps)
+ if converged:
+ return True
+
+ print(f"No convergence after {steps} steps. Resetting optimiser..")
+ atoms = opt.atoms.atoms if isinstance(opt.atoms, FrechetCellFilter) else opt.atoms
+ atoms.rattle(stdev=1e-4)
+
+ opt = BFGS(dof, trajectory=trajectory, append_trajectory=True)
+ return opt.run(fmax=f_max, steps=steps)
+
+
+def main(args: SimpleNamespace):
""""""
- config = json.load(Path(args.input_config).open('r'))
+ config = json.load(Path(args.input_config).open("r"))
+ config["trajectory"] = FILE_TRAJ if args.output_traj else None
atoms = ase.io.read(args.start_xyz)
if not any(atoms.pbc):
- atoms.center(vacuum=0) # optimiser mysteriously requires a nonzero unit cell
- if config['mode'] != 'fix_cell':
- config['mode'] = 'fix_cell'
- warnings.warn('Molecular structure is not periodic. Ignoring cell..')
+ atoms.center(vacuum=0) # optimiser mysteriously requires a nonzero unit cell
+ if config["mode"] != "fix_cell":
+ config["mode"] = "fix_cell"
+ warnings.warn("Molecular structure is not periodic. Ignoring cell..")
# construct calculator by combining hamiltonians
assert args.path_hamiltonian is not None
- print('Making calculator from:', *config['forces'], sep='\n')
+ print("Making calculator from:", *config["forces"], sep="\n")
functions = [function_from_json(p) for p in args.path_hamiltonian]
calc = LinearCombinationCalculator(
[FunctionCalculator(f) for f in functions],
- [float(h['weight']) for h in config['forces']]
+ [float(h["weight"]) for h in config["forces"]],
)
atoms.calc = calc
- dof = get_dof_filter(atoms, config['mode'], config['pressure'])
- opt = PreconLBFGS(dof, trajectory=FILE_TRAJ if config['keep_trajectory'] else None)
+ dof = get_dof_filter(atoms, config["mode"], config["pressure"])
+ log_state(atoms)
print(f"pid: {os.getpid()}")
print(f"CPU affinity: {os.sched_getaffinity(os.getpid())}")
- log_state(atoms)
+ converged = False
try:
- opt.run(fmax=config['f_max'], steps=config['max_steps'])
+ converged = optimise(dof, **config)
+ assert converged
+ print("OPTIMISATION SUCCESSFUL")
+ except RuntimeError as e:
+ print("OPTIMISATION FAILED", e, sep="\n")
except TimeoutException:
- print('OPTIMISATION TIMEOUT')
- # TODO: what to do here?
- return
+ print("OPTIMISATION TIMEOUT")
+ except AssertionError:
+ print("OPTIMISATION DID NOT CONVERGE")
log_state(atoms)
- save_calc_results(atoms, calc_prefix='', remove_atoms_calc=True)
+ if not converged:
+ # touch output file so Parsl thinks the task completed regardless
+ Path(args.output_xyz).touch()
+ return
+
+ atoms.set_constraint() # usually noop
if not any(atoms.pbc):
- atoms.cell = None # remove meaningless cell
- ase.io.write(FILE_OUT, atoms)
- print('OPTIMISATION SUCCESSFUL')
- return
+ atoms.cell = None # remove meaningless cell
+ ase.io.write(FILE_OUT, atoms, format="extxyz")
+ shutil.copy(FILE_OUT, args.output_xyz)
+ if args.output_traj is not None:
+ atoms = [at for at in ase.io.trajectory.Trajectory(FILE_TRAJ)]
+ ase.io.write(args.output_traj, atoms, format="extxyz")
+ print("FILES MOVED")
-def clean(args: SimpleNamespace):
- """"""
- from psiflow.data.utils import _write_frames
-
- geometry = Geometry.load(FILE_OUT)
- _write_frames(geometry, outputs=[args.output_xyz])
- if Path(FILE_TRAJ).is_file():
- traj = [at for at in ase.io.trajectory.Trajectory(FILE_TRAJ)]
- geometries = [Geometry.from_atoms(at) for at in traj]
- _write_frames(*geometries, outputs=[args.output_traj])
- print('FILES MOVED')
return
-def main():
+if __name__ == "__main__":
signal.signal(signal.SIGTERM, timeout_handler)
parser = argparse.ArgumentParser()
- subparsers = parser.add_subparsers(help='what to do', dest='action')
- run_parser = subparsers.add_parser("run")
- run_parser.set_defaults(func=run)
- run_parser.add_argument(
- "--path_hamiltonian",
- action='extend',
- nargs='*',
- type=str,
- )
- run_parser.add_argument(
- "--input_config",
- type=str,
- default=None,
- )
- run_parser.add_argument(
- "--start_xyz",
- type=str,
- default=None,
- )
- clean_parser = subparsers.add_parser("clean")
- clean_parser.set_defaults(func=clean)
- clean_parser.add_argument(
- "--output_xyz",
- type=str,
- default=None,
- )
- clean_parser.add_argument(
- "--output_traj",
- type=str,
- default=None,
- )
+ parser.add_argument("--path_hamiltonian", action="extend", nargs="*")
+ parser.add_argument("--input_config", required=True)
+ parser.add_argument("--start_xyz", required=True)
+ parser.add_argument("--output_xyz", required=True)
+ parser.add_argument("--output_traj")
args = parser.parse_args()
- args.func(args)
-
+ time_start = time.time()
+ main(args)
+ time_stop = time.time()
+ print(f"Total elapsed time: {time_stop - time_start: .3f} seconds")
diff --git a/psiflow/sampling/ase.py b/psiflow/sampling/ase.py
index dfa63a6d..87434fea 100644
--- a/psiflow/sampling/ase.py
+++ b/psiflow/sampling/ase.py
@@ -1,10 +1,7 @@
-from __future__ import annotations # necessary for type-guarding class methods
-
from typing import Optional, Union
import parsl
-import typeguard
-from parsl.app.app import bash_app, join_app
+from parsl import bash_app, join_app, python_app, File
from parsl.dataflow.futures import AppFuture, DataFuture
import psiflow
@@ -12,12 +9,38 @@
from psiflow.data.utils import write_frames
from psiflow.geometry import Geometry
from psiflow.hamiltonians import Hamiltonian
-from psiflow.utils.io import dump_json
-from psiflow.utils import TMP_COMMAND, CD_COMMAND
+from psiflow.utils.apps import setup_logger
+from psiflow.utils.io import _dump_json
+from psiflow.utils import TMP_COMMAND, CD_COMMAND, export_env_command
+from psiflow.utils.parse import get_task_name_id
+
+from ._ase import ALLOWED_MODES, __file__ as file_ase
+
+DEFAULT_EXECUTABLE = 'script.py'
+logger = setup_logger(__name__) # logging per module
+
+
+class OptimisationFailedError(Exception):
+ pass
-from ._ase import ALLOWED_MODES
-EXECUTABLE = 'psiflow-ase-opt' # not stored in ModelEvaluation (yet?)
+@python_app(executors=["default_threads"])
+def check_task_output(
+ file: DataFuture, file_stdout: str, file_stderr: str
+) -> AppFuture:
+ # TODO: find actual reason for fail?
+ task_name, task_id = get_task_name_id(file_stdout)
+ try:
+ geom = Geometry.load(file.filepath)
+ status = "SUCCESS"
+ except Exception as e:
+ # output file is empty because the optimisation failed
+ status = "FAILED"
+
+ logger.info(f'Task "{task_name}" (ID {task_id}): {status}')
+ if status == "SUCCESS":
+ return geom
+ raise OptimisationFailedError(f"Task {task_id}")
def _execute_ase(
@@ -25,26 +48,27 @@ def _execute_ase(
inputs: list[DataFuture],
outputs: list[DataFuture],
env_vars: dict = {},
- stdout: str = "",
- stderr: str = "",
+ stdout: str = parsl.AUTO_LOGNAME,
+ stderr: str = parsl.AUTO_LOGNAME,
parsl_resource_specification: Optional[dict] = None,
) -> str:
- env_command = 'export ' + ' '.join([f"{name}={value}" for name, value in env_vars.items()])
- command_start = ' '.join([
- f'{command_launch} run --input_config={inputs[0].filepath} --start_xyz={inputs[1].filepath}',
- *[f'--path_hamiltonian={future.filepath}' for future in inputs[2:]], '&'
- ])
- command_end = f'{command_launch} clean --output_xyz={outputs[0].filepath}'
+ command_opt_args = [
+ command_launch,
+ f"--input_config={inputs[1].filepath}",
+ f"--start_xyz={inputs[2].filepath}",
+ *[f"--path_hamiltonian={future.filepath}" for future in inputs[3:]],
+ f"--output_xyz={outputs[0].filepath}",
+ ]
if len(outputs) == 2:
- command_end += f' --output_traj={outputs[1].filepath}'
+ command_opt_args.append(f"--output_traj={outputs[1].filepath}")
command_list = [
TMP_COMMAND,
CD_COMMAND,
- env_command,
- command_start,
- "wait",
- command_end,
+ export_env_command(env_vars),
+ f"cp {inputs[0].filepath} {DEFAULT_EXECUTABLE}",
+ " ".join(command_opt_args),
+ "exit 0", # ignore timeout exitcode
]
return "\n".join(command_list)
@@ -52,15 +76,15 @@ def _execute_ase(
execute_ase = bash_app(_execute_ase, executors=["ModelEvaluation"])
-@typeguard.typechecked
def optimize(
state: Union[Geometry, AppFuture],
hamiltonian: Hamiltonian,
- mode: str = 'full',
+ mode: str = "full",
steps: int = int(1e12),
keep_trajectory: bool = False,
pressure: float = 0,
f_max: float = 1e-3,
+ script: str = file_ase,
) -> Union[AppFuture, tuple[AppFuture, Dataset]]:
assert mode in ALLOWED_MODES
@@ -70,22 +94,22 @@ def optimize(
context = psiflow.context()
definition = context.definitions["ModelEvaluation"]
- command_list = [EXECUTABLE]
+ command_list = [f"python -u {DEFAULT_EXECUTABLE}"]
if definition.max_simulation_time is not None:
max_time = 0.9 * (60 * definition.max_simulation_time)
command_list = ["timeout -s 15 {}s".format(max_time), *command_list]
command_launch = " ".join(command_list)
- input_geometry = Dataset([state]).extxyz
+ input_geometry = Dataset([state]).extxyz # state can be future
hamiltonian = 1.0 * hamiltonian # convert to mixture
names, coeffs = hamiltonian.get_named_components(), hamiltonian.coefficients
- input_forces = hamiltonian.serialize(dtype="float64") # double precision for MLPs
+ input_forces = hamiltonian.serialize(dtype="float64") # double precision for MLPs
forces = [
- dict(forcefield=n, weight=str(c), file=f.filename) for n, c, f in zip(names, coeffs, input_forces)
+ dict(forcefield=n, weight=str(c), file=f.filename)
+ for n, c, f in zip(names, coeffs, input_forces)
]
-
config = dict(
- task='ASE optimisation',
+ task="ASE optimisation",
forces=forces,
mode=mode,
f_max=f_max,
@@ -93,12 +117,10 @@ def optimize(
max_steps=steps,
keep_trajectory=keep_trajectory,
)
- input_future = dump_json(
- outputs=[context.new_file("input_", ".json")],
- **config,
- ).outputs[0]
- inputs = [input_future, input_geometry, *input_forces]
+ file_config = context.new_file("input_", ".json")
+ _dump_json(outputs=[file_config], **config)
+ inputs = [File(script), file_config, input_geometry, *input_forces]
outputs = [context.new_file("data_", ".xyz")]
if keep_trajectory:
outputs.append(context.new_file("opt_", ".xyz"))
@@ -108,32 +130,27 @@ def optimize(
env_vars=definition.env_vars,
inputs=inputs,
outputs=outputs,
- stdout=parsl.AUTO_LOGNAME,
- stderr=parsl.AUTO_LOGNAME,
parsl_resource_specification=definition.wq_resources(1),
)
- final = Dataset(None, result.outputs[0])[-1]
+ geom = check_task_output(result.outputs[0], result.stdout, result.stderr)
if keep_trajectory:
trajectory = Dataset(None, result.outputs[1])
- return final, trajectory
+ return geom, trajectory
else:
- return final
+ return geom
@join_app
-@typeguard.typechecked
def _optimize_dataset(
geometries: list[Geometry], *args, outputs: list = [], **kwargs
) -> AppFuture:
assert not kwargs.get("keep_trajectory", False)
- optimized = []
- for geometry in geometries:
- optimized.append(optimize(geometry, *args, **kwargs))
+ logger.info(f"Performing {len(geometries)} structure optimisations.")
+ optimized = [optimize(geometry, *args, **kwargs) for geometry in geometries]
return write_frames(*optimized, outputs=[outputs[0]])
-@typeguard.typechecked
def optimize_dataset(dataset: Dataset, *args, **kwargs) -> Dataset:
extxyz = _optimize_dataset(
dataset.geometries(),
diff --git a/psiflow/sampling/client.py b/psiflow/sampling/client.py
deleted file mode 100644
index c720d8f8..00000000
--- a/psiflow/sampling/client.py
+++ /dev/null
@@ -1,110 +0,0 @@
-# top level imports should be lightweight!
-import os
-
-
-class SocketNotFoundException(Exception):
- pass
-
-
-def wait_for_socket(address: 'Path', timeout: float = 10, interval: float = 0.1) -> None:
- """"""
- import time
- while not address.exists():
- time.sleep(interval)
- timeout -= interval
- if timeout < 0:
- raise SocketNotFoundException(f'Could not find socket "{address}" to connect to..')
- return
-
-
-def main():
- import argparse
- import time
- from pathlib import Path
-
- from ase.io import read
- from ipi._driver.driver import run_driver
-
- from psiflow.functions import function_from_json
- from psiflow.geometry import Geometry
- from psiflow.sampling.utils import ForceMagnitudeException, FunctionDriver
-
- print("OS environment values:")
- for key, value in os.environ.items():
- print(key, value)
- parser = argparse.ArgumentParser()
- parser.add_argument(
- "--path_hamiltonian",
- type=str,
- default=None,
- )
- parser.add_argument(
- "--device",
- type=str,
- default=None,
- )
- parser.add_argument(
- "--dtype",
- type=str,
- default=None,
- )
- parser.add_argument(
- "--address",
- type=str,
- default=None,
- )
- parser.add_argument(
- "--start",
- type=str,
- default=None,
- )
- parser.add_argument(
- "--max_force",
- type=float,
- default=None,
- )
- args = parser.parse_args()
- assert args.path_hamiltonian is not None
- assert args.address is not None
- assert args.start is not None
-
- print("pid: {}".format(os.getpid()))
- affinity = os.sched_getaffinity(os.getpid())
- print("CPU affinity before function init: {}".format(affinity))
-
- template = Geometry.from_atoms(read(args.start))
- function = function_from_json(
- args.path_hamiltonian,
- device=args.device,
- dtype=args.dtype,
- )
-
- driver = FunctionDriver(
- template=template,
- function=function,
- max_force=args.max_force,
- verbose=True,
- )
-
- affinity = os.sched_getaffinity(os.getpid())
- print("CPU affinity after function init: {}".format(affinity))
- try:
- t0 = time.time()
- for _ in range(10):
- function(template) # torch warm-up before simulation
- print("time for 10 evaluations: {}".format(time.time() - t0))
- socket_address = Path.cwd() / args.address
- wait_for_socket(socket_address)
- run_driver(
- unix=True,
- address=str(socket_address),
- driver=driver,
- sockets_prefix="",
- )
- except ForceMagnitudeException as e:
- print(e) # induce timeout in server
- except ConnectionResetError as e: # some other client induced a timeout
- print(e)
- except SocketNotFoundException as e:
- print(e, *list(Path.cwd().iterdir()), sep='\n') # server-side socket not found
-
diff --git a/psiflow/sampling/driver.py b/psiflow/sampling/driver.py
new file mode 100644
index 00000000..b56b852e
--- /dev/null
+++ b/psiflow/sampling/driver.py
@@ -0,0 +1,83 @@
+"""
+Updated version of the Psiflow driver included in i-Pi
+"""
+
+import numpy as np
+
+from ipi.pes.dummy import Dummy_driver
+from ipi.utils.units import unit_to_internal, unit_to_user
+from ipi.utils.messages import warning
+
+try:
+ from ase.io import read
+ from psiflow.geometry import Geometry
+ from psiflow.functions import function_from_json, Function
+ from psiflow.sampling.utils import initialise_driver, check_output
+except ImportError as e:
+ message = (
+ "Could not find the Psiflow driver dependencies, "
+ "make sure they are installed with "
+ "`pip install git+https://github.com/molmod/psiflow.git`."
+ )
+ warning(f"{message}: {e}")
+ raise ImportError(message) from e
+
+
+__DRIVER_NAME__ = "psiflow"
+__DRIVER_CLASS__ = "Psiflow_driver"
+
+
+class Psiflow_driver(Dummy_driver):
+ """
+ Driver for Psiflow functions.
+ The driver requires a template file (in ASE readable format) and a JSON file
+ containing the function definition. Requires Psiflow to be installed.
+
+ Command-line:
+ i-pi-driver -m psiflow -o template=template.xyz,hamiltonian=hamiltonian.json
+
+ Parameters:
+ :param template: str, filename of an ASE-readable structure file
+ :param hamiltonian: str, filename of a JSON file defining the Psiflow function
+ """
+
+ template: str
+ hamiltonian: str
+ geometry: Geometry
+ function: Function
+
+ def __init__(self, template, hamiltonian, *args, **kwargs):
+ self.template = template
+ self.hamiltonian = hamiltonian
+ super().__init__(*args, **kwargs)
+
+ def check_parameters(self):
+ self.geometry = Geometry.from_atoms(read(self.template))
+ self.function = function_from_json(self.hamiltonian, **self.kwargs)
+ initialise_driver(self)
+
+ def compute_structure(self, cell, pos):
+ pos = unit_to_user("length", "angstrom", pos)
+ cell = unit_to_user("length", "angstrom", cell.T)
+ self.geometry.per_atom.positions[:] = pos
+ if self.geometry.periodic:
+ self.geometry.cell[:] = cell
+
+ outputs = self.function(self.geometry)
+ check_output(self, outputs)
+
+ # converts to internal quantities
+ energy = outputs["energy"]
+ forces = outputs["forces"]
+ stress = outputs["stress"]
+ pot_ipi = np.asarray(
+ unit_to_internal("energy", "electronvolt", energy), np.float64
+ )
+ force_ipi = np.asarray(unit_to_internal("force", "ev/ang", forces), np.float64)
+ vir_calc = -stress * self.geometry.volume
+ vir_ipi = np.array(
+ unit_to_internal("energy", "electronvolt", vir_calc.T), dtype=np.float64
+ )
+ extras = ""
+
+ return pot_ipi, force_ipi, vir_ipi, extras
diff --git a/psiflow/sampling/metadynamics.py b/psiflow/sampling/metadynamics.py
index c26087a2..36352642 100644
--- a/psiflow/sampling/metadynamics.py
+++ b/psiflow/sampling/metadynamics.py
@@ -1,9 +1,6 @@
-from __future__ import annotations # necessary for type-guarding class methods
-
from pathlib import Path
from typing import Optional, Union
-import typeguard
from parsl.data_provider.files import File
from parsl.dataflow.futures import AppFuture
@@ -12,7 +9,6 @@
from psiflow.utils.apps import copy_app_future, copy_data_future
-@typeguard.typechecked
@psiflow.serializable
class Metadynamics:
_plumed_input: str
@@ -33,13 +29,11 @@ def __init__(
# PLUMED + WQ cannot deal with nonexisting hills files!
if type(external) in [str, Path]:
external = File(str(external))
- Path(external).touch()
if external is None:
external = psiflow.context().new_file("hills_", ".txt")
- Path(external.filepath).touch()
else:
assert external.filepath in _plumed_input
- Path(external.filepath).touch()
+ Path(external.filepath).touch()
_plumed_input = set_path_in_plumed(
_plumed_input,
"METAD",
@@ -49,14 +43,13 @@ def __init__(
self.external = external
def plumed_input(self):
- plumed_input = self._plumed_input
- plumed_input = plumed_input.replace("PLACEHOLDER", self.external.filepath)
- return plumed_input
+ return self._plumed_input.replace("PLACEHOLDER", self.external.filepath)
def input(self) -> AppFuture:
return copy_app_future(self.plumed_input(), inputs=[self.external])
def wait_for(self, result: AppFuture) -> None:
+ # TODO: what does this do?
self.external = copy_app_future(
0,
inputs=[result, self.external],
@@ -71,7 +64,7 @@ def __eq__(self, other) -> bool:
return False
return self.plumed_input() == other.plumed_input()
- def copy(self) -> Metadynamics:
+ def copy(self) -> "Metadynamics":
new_external = copy_data_future(
inputs=[self.external],
outputs=[psiflow.context().new_file("hills_", ".txt")],
diff --git a/psiflow/sampling/optimize.py b/psiflow/sampling/optimize.py
index 61e472b7..74d7d54e 100644
--- a/psiflow/sampling/optimize.py
+++ b/psiflow/sampling/optimize.py
@@ -1,10 +1,8 @@
-from __future__ import annotations # necessary for type-guarding class methods
-
+import warnings
import xml.etree.ElementTree as ET
from typing import Optional, Union
import parsl
-import typeguard
from ase.units import Bohr, Ha
from parsl.app.app import bash_app, join_app
from parsl.dataflow.futures import AppFuture
@@ -13,31 +11,42 @@
from psiflow.data import Dataset
from psiflow.data.utils import write_frames
from psiflow.geometry import Geometry
-from psiflow.hamiltonians import Hamiltonian
-from psiflow.sampling.sampling import setup_sockets, make_start_command, make_client_command
+from psiflow.hamiltonians import Hamiltonian, MACEHamiltonian
+from psiflow.sampling.sampling import (
+ setup_sockets,
+ make_server_command,
+ make_driver_commands,
+ make_wait_for_sockets_command,
+)
+from psiflow.sampling.output import HamiltonianComponent
from psiflow.utils.io import save_xml
-from psiflow.utils import TMP_COMMAND, CD_COMMAND
+from psiflow.utils import TMP_COMMAND, CD_COMMAND, export_env_command
+
+warnings.warn(
+ "The 'optimize' module will likely be removed in a future release. "
+ "Consider using the 'ase' module for structure optimisations instead.",
+ FutureWarning,
+)
-@typeguard.typechecked
-def setup_forces(hamiltonian: Hamiltonian) -> tuple[dict[str, Hamiltonian], ET.Element]:
+
+def setup_forces(
+ hamiltonian: Hamiltonian,
+) -> tuple[list[HamiltonianComponent], ET.Element]:
hamiltonian = 1.0 * hamiltonian # convert to mixture
- counts = {}
- hamiltonians_map = {}
forces = ET.Element("forces")
- for h, c in zip(hamiltonian.hamiltonians, hamiltonian.coefficients):
- name = h.__class__.__name__
- if name not in counts:
- counts[name] = 0
- count = counts.get(name)
- counts[name] += 1
- force = ET.Element("force", forcefield=name + str(count), weight=str(c))
+ components = []
+ for n, h, c in zip(
+ hamiltonian.get_named_components(),
+ hamiltonian.hamiltonians,
+ hamiltonian.coefficients,
+ ):
+ components.append(comp := HamiltonianComponent(n, h, True))
+ force = ET.Element("force", forcefield=comp.name, weight=str(c))
forces.append(force)
- hamiltonians_map[name + str(count)] = h
- return hamiltonians_map, forces
+ return components, forces
-@typeguard.typechecked
def setup_motion(
mode: str,
etol: float,
@@ -62,7 +71,6 @@ def setup_motion(
return motion
-@typeguard.typechecked
def setup_output(keep_trajectory: bool) -> ET.Element:
output = ET.Element("output", prefix="output")
checkpoint = ET.Element(
@@ -86,38 +94,32 @@ def setup_output(keep_trajectory: bool) -> ET.Element:
def _execute_ipi(
- hamiltonian_names: list[str],
- client_args: list[list[str]],
- keep_trajectory: bool,
+ driver_kwargs: list[dict],
command_server: str,
- command_client: str,
env_vars: dict = {},
- stdout: str = "",
- stderr: str = "",
+ stdout: str = parsl.AUTO_LOGNAME,
+ stderr: str = parsl.AUTO_LOGNAME,
inputs: list = [],
outputs: list = [],
parsl_resource_specification: Optional[dict] = None,
) -> str:
- env_command = 'export ' + ' '.join([f"{name}={value}" for name, value in env_vars.items()])
- command_start = make_start_command(command_server, inputs[0], inputs[1])
- commands_client = []
- for i, name in enumerate(hamiltonian_names):
- args = client_args[i]
- assert len(args) == 1 # only have one client per hamiltonian
- for arg in args:
- commands_client += make_client_command(command_client, name, inputs[2 + i], inputs[1], arg),
-
- command_end = f'{command_server} --cleanup --output_xyz={outputs[0].filepath}'
- command_copy = f'cp walker-0_output.trajectory_0.ase {outputs[1].filepath}' if keep_trajectory else ''
+ file_xml, file_xyz_in, *files_in = inputs
+ command_start = make_server_command(
+ command_server, file_xml, file_xyz_in, outputs[0], [], outputs[1:]
+ )
+ command_wait = make_wait_for_sockets_command(
+ set(d["address"] for d in driver_kwargs)
+ )
+ commands_driver = make_driver_commands(driver_kwargs, file_xyz_in, files_in)
+
command_list = [
TMP_COMMAND,
CD_COMMAND,
- env_command,
+ export_env_command(env_vars),
command_start,
- *commands_client,
+ command_wait,
+ *commands_driver,
"wait",
- command_end,
- command_copy,
]
return "\n".join(command_list)
@@ -125,7 +127,6 @@ def _execute_ipi(
execute_ipi = bash_app(_execute_ipi, executors=["ModelEvaluation"])
-@typeguard.typechecked
def optimize(
state: Union[Geometry, AppFuture],
hamiltonian: Hamiltonian,
@@ -136,8 +137,8 @@ def optimize(
ptol: float = 1e-5,
ftol: float = 1e-3,
) -> Union[AppFuture, tuple[AppFuture, Dataset]]:
- hamiltonians_map, forces = setup_forces(hamiltonian)
- sockets = setup_sockets(hamiltonians_map)
+ hamiltonian_components, forces = setup_forces(hamiltonian)
+ sockets = setup_sockets(hamiltonian_components)
initialize = ET.Element("initialize", nbeads="1")
start = ET.Element("file", mode="ase", cell_units="angstrom")
@@ -171,33 +172,25 @@ def optimize(
input_future,
Dataset([state]).extxyz,
]
- inputs += [h.serialize_function(dtype="float64") for h in hamiltonians_map.values()]
-
- hamiltonian_names = list(hamiltonians_map.keys())
- client_args = []
- for name in hamiltonian_names:
- args = definition.get_client_args(name, 1, "minimize")
- client_args.append(args)
outputs = [context.new_file("data_", ".xyz")]
if keep_trajectory:
outputs.append(context.new_file("opt_", ".xyz"))
- command_server = definition.server_command()
- command_client = definition.client_command()
- resources = definition.wq_resources(1)
+ driver_kwargs = []
+ for i, comp in enumerate(hamiltonian_components):
+ inputs.append(comp.hamiltonian.serialize_function(dtype="float64"))
+ kwargs = {"idx": i, "address": comp.address}
+ if isinstance(comp.hamiltonian, MACEHamiltonian):
+ kwargs |= definition.get_driver_devices(1)[0]
+ driver_kwargs.append(kwargs)
result = execute_ipi(
- hamiltonian_names,
- client_args,
- keep_trajectory,
- command_server,
- command_client,
+ driver_kwargs,
+ definition.server_command(),
env_vars=definition.env_vars,
- stdout=parsl.AUTO_LOGNAME,
- stderr=parsl.AUTO_LOGNAME,
inputs=inputs,
outputs=outputs,
- parsl_resource_specification=resources,
+ parsl_resource_specification=definition.wq_resources(1),
)
final = Dataset(None, result.outputs[0]).evaluate(hamiltonian)[-1]
@@ -209,7 +202,6 @@ def optimize(
@join_app
-@typeguard.typechecked
def _optimize_dataset(
geometries: list[Geometry], *args, outputs: list = [], **kwargs
) -> AppFuture:
@@ -220,7 +212,6 @@ def _optimize_dataset(
return write_frames(*optimized, outputs=[outputs[0]])
-@typeguard.typechecked
def optimize_dataset(dataset: Dataset, *args, **kwargs) -> Dataset:
extxyz = _optimize_dataset(
dataset.geometries(),
diff --git a/psiflow/sampling/order.py b/psiflow/sampling/order.py
index fb8a4c9b..c5fa96d3 100644
--- a/psiflow/sampling/order.py
+++ b/psiflow/sampling/order.py
@@ -1,16 +1,21 @@
"""
TODO: these imports are outdated.. Is this module still used?
"""
-from __future__ import annotations # necessary for type-guarding class methods
+import warnings
from functools import partial
from typing import Optional, Union
-import typeguard
from ase.units import kJ, mol
from parsl.app.app import python_app
from parsl.dataflow.futures import AppFuture
+warnings.warn(
+ "The 'order' module is severely outdated and no longer works. "
+ "Please reach out on GitHub if you require this functionality.",
+)
+
+
import psiflow
from psiflow.data import Dataset, batch_apply
from psiflow.geometry import Geometry
@@ -18,7 +23,6 @@
from psiflow.hamiltonians.hamiltonian import Hamiltonian
-@typeguard.typechecked
def insert_in_state(
state: Geometry,
name: str,
@@ -29,7 +33,6 @@ def insert_in_state(
return state
-@typeguard.typechecked
def _insert(
state_or_states: Union[Geometry, list[Geometry]],
name: str,
@@ -45,7 +48,6 @@ def _insert(
insert = python_app(_insert, executors=["default_threads"])
-@typeguard.typechecked
def insert_in_dataset(
data: Dataset,
name: str,
@@ -57,7 +59,6 @@ def insert_in_dataset(
return Dataset(geometries)
-@typeguard.typechecked
class OrderParameter:
# TODO: batched evaluation
@@ -71,7 +72,6 @@ def __eq__(self, other):
raise NotImplementedError
-@typeguard.typechecked
@psiflow.serializable
class HamiltonianOrderParameter(OrderParameter):
name: str
@@ -113,7 +113,7 @@ def __eq__(self, other):
@classmethod
def from_plumed(
cls, name: str, hamiltonian: PlumedHamiltonian
- ) -> HamiltonianOrderParameter:
+ ) -> "HamiltonianOrderParameter":
assert name in hamiltonian.plumed_input()
action_prefixes = [
"ABMD",
diff --git a/psiflow/sampling/output.py b/psiflow/sampling/output.py
index 7f734586..318a86c7 100644
--- a/psiflow/sampling/output.py
+++ b/psiflow/sampling/output.py
@@ -1,19 +1,28 @@
import copy
import re
+from enum import Enum
from pathlib import Path
+from dataclasses import dataclass, field, InitVar
from typing import Optional, Union
import numpy as np
-import typeguard
+from ipi.utils.parsing import read_output
from parsl.app.app import python_app
-from parsl.app.futures import DataFuture
+from parsl.app.futures import DataFuture, File
from parsl.dataflow.futures import AppFuture
import psiflow
from psiflow.data import Dataset
from psiflow.geometry import Geometry, NullState
from psiflow.hamiltonians import Hamiltonian, MixtureHamiltonian, Zero
-from psiflow.utils.apps import unpack_i
+from psiflow.sampling.walker import Walker
+from psiflow.utils.io import save_npz
+from psiflow.utils.apps import setup_logger
+from psiflow.utils.parse import get_task_name_id
+
+
+logger = setup_logger(__name__) # logging per module
+
DEFAULT_OBSERVABLES = [
"time{picosecond}",
@@ -23,281 +32,230 @@
]
-@typeguard.typechecked
-def potential_component_names(n: int):
- return [f'pot_component_raw({i}){{electronvolt}}' for i in range(n)]
+class Status(Enum):
+ UNKNOWN = -1
+ DONE = 0
+ TIMEOUT = 1
+ FORCE_EXCEEDED = 2
+ OOM = 3
+ BROKEN_PIPE = 4
+ EXPLODED = 5
-@typeguard.typechecked
-def _get_final_temperature(temperatures: np.ndarray) -> float:
- if not len(temperatures) > 0:
- return np.nan
- return temperatures[-1]
+@dataclass
+class HamiltonianComponent:
+ name: str
+ hamiltonian: Hamiltonian
+ shared: bool
+ def __post_init__(self):
+ self.address = f"psiflow_{self.name.lower()}"
-get_final_temperature = python_app(
- _get_final_temperature, executors=["default_threads"]
-)
-
-
-@typeguard.typechecked
-def read_output(filename): # from i-PI
- # Regex pattern to match header lines and capture relevant parts
- header_pattern = re.compile(
- r"#\s*(column|cols\.)\s+(\d+)(?:-(\d+))?\s*-->\s*([^\s\{]+)(?:\{([^\}]+)\})?\s*:\s*(.*)"
- )
-
- # Reading the file
- with open(filename, "r") as file:
- lines = file.readlines()
-
- header_lines = [line for line in lines if line.startswith("#")]
- data_lines = [line for line in lines if not line.startswith("#") and line.strip()]
-
- # Interprets properties
- properties = {}
- for line in header_lines:
- match = header_pattern.match(line)
- if match:
- # Extracting matched groups
- (
- col_type,
- start_col,
- end_col,
- property_name,
- units,
- description,
- ) = match.groups()
- col_info = f"{start_col}-{end_col}" if end_col else start_col
- properties[col_info] = {
- "name": property_name,
- "units": units,
- "description": description,
- }
-
- # Parse data
- values_dict = {}
- info_dict = {}
- for prop_info in properties.values():
- # Initialize list to hold values for each property
- values_dict[prop_info["name"]] = []
- # Save units and description
- info_dict[prop_info["name"]] = (prop_info["units"], prop_info["description"])
-
- for line in data_lines:
- values = line.split()
- for column_info, prop_info in properties.items():
- if "-" in column_info: # Multi-column property
- start_col, end_col = map(
- int, column_info.split("-")
- ) # 1-based indexing
- prop_values = values[
- start_col - 1 : end_col
- ] # Adjust to 0-based indexing
- else: # Single column property
- col_index = int(column_info) - 1 # Adjust to 0-based indexing
- prop_values = [values[col_index]]
-
- values_dict[prop_info["name"]].append([float(val) for val in prop_values])
-
- for prop_name, prop_values in values_dict.items():
- values_dict[prop_name] = np.array(
- prop_values
- ).squeeze() # make 1-col into a flat array
-
- return values_dict, info_dict
-
-
-@typeguard.typechecked
-def _parse_data(
- status: int,
- start: int,
- keys: list[str],
- inputs: list = [],
-) -> list[np.ndarray]:
- from psiflow.sampling.output import read_output
- bare_keys = []
- for key in keys:
- if "{" in key:
- bare_key = key.split("{")[0]
- else:
- bare_key = key
- bare_keys.append(bare_key)
- if status not in [0, 1]:
- return [np.zeros(0) for key in bare_keys] # read_output might fail
- values, _ = read_output(inputs[0].filepath)
- return [values[key][start:] for key in bare_keys]
+def split_units(key: str) -> tuple[str, str]:
+ name, unit = key.split("{")
+ return name, unit[:-1]
-parse_data = python_app(_parse_data, executors=["default_threads"])
+def potential_component_name(n: int) -> str:
+ return f"pot_component_raw({n}){{electronvolt}}"
-@typeguard.typechecked
-def _add_contributions(
- coefficients: tuple[float, ...],
- *values: np.ndarray,
-) -> np.ndarray:
- assert len(coefficients) == len(values)
- total = np.zeros(len(values[0]))
- for i, c in enumerate(coefficients):
- total += c * values[i]
- return total
-
-
-add_contributions = python_app(_add_contributions, executors=["default_threads"])
-
-
-@typeguard.typechecked
-def _parse(
- state: Geometry,
- inputs: list = [],
-) -> tuple[float, int]:
- time = state.order.pop("time")
-
- # determine status based on stdout
- with open(inputs[0], "r") as f:
- content = f.read()
- if "force exceeded" in content:
- status = 2 # max_force exception
+def get_simulation_status(stdout: str, stderr: str) -> Status:
+ content = Path(stdout).read_text()
+ if "@PSIFLOW: We are done here" in content:
+ return Status.DONE
elif "@SOFTEXIT: Kill signal received" in content:
- status = 1 # timeout
- elif "@ SIMULATION: Exiting cleanly" in content:
- status = 0 # everything OK
- else:
- status = -1
- return time, status
-
+ return Status.TIMEOUT # i-Pi intercepts SIG_INT and SIG_TERM by default
+ elif "@PSIFLOW: Simulation went boom" in content:
+ return Status.EXPLODED
-parse = python_app(_parse, executors=["default_threads"])
+ content = Path(stderr).read_text()
+ if "force exceeded" in content:
+ return Status.FORCE_EXCEEDED
+ elif "BrokenPipeError: [Errno 32] Broken pipe" in content:
+ return Status.BROKEN_PIPE
+ elif "Killed" in content:
+ return Status.OOM
+ return Status.UNKNOWN
-@typeguard.typechecked
-def _update_walker(
- status: int,
+def _parse_simulation_data(
state: Geometry,
- start: Geometry,
-) -> Geometry:
- from psiflow.geometry import NullState
+ observables: list[str],
+ start: int,
+ file_props: File,
+ simulation_stdout: str,
+ simulation_stderr: str,
+) -> tuple[Status, dict, float, float]:
+ """"""
+ status = get_simulation_status(simulation_stdout, simulation_stderr)
+ try:
+ # read_output strips unit information from keys
+ values, _ = read_output(file_props.filepath)
+ data = {k: values[split_units(k)[0]][start:] for k in observables}
+ temperature = data["temperature{kelvin}"][-1]
+ except IndexError as e:
+ # nothing was written
+ data = {k: np.array([]) for k in observables}
+ temperature = np.nan
+
+ task_id = get_task_name_id(simulation_stdout)[-1]
+ logger.info(f"Simulation [ID {task_id}]: {status}")
+ return status, data, temperature, state.order.pop("time")
+
+
+parse_simulation_data = python_app(
+ _parse_simulation_data, executors=["default_threads"]
+)
- if (status in [0, 1]) and state != NullState:
- return state
- else:
- return copy.deepcopy(start)
+def _log_status(
+ task_id: str,
+ status: Status,
+ observables: list[str],
+ time: float,
+ temperature: float,
+ component_map: dict[str, HamiltonianComponent],
+ bias_components: list[HamiltonianComponent],
+) -> None:
+ """Print out some simulation stats"""
+ info = [
+ f"Simulation [ID {task_id}]",
+ f"Status: {status.name}",
+ f"Elapsed time: {time*1000:.1f} fs",
+ f"Final temperature: {temperature:.1f} K",
+ f"Observables: {observables}",
+ f"Force components: {[f'{c.name} -> {k}' for k, c in component_map.items()]}",
+ f"Bias components: {[c.name for c in bias_components]}",
+ ]
+ print(*info, sep="\n")
-update_walker = python_app(_update_walker, executors=["default_threads"])
+log_status = python_app(_log_status, executors=["default_threads"])
-@typeguard.typechecked
-def _get_state(
- geometry: Geometry,
- status: int,
-) -> Geometry:
- if status in [0, 1]:
- return geometry
- else:
- return NullState
+def _add_contributions(
+ coefficients: tuple[float, ...],
+ *values: np.ndarray,
+) -> np.ndarray:
+ assert len(coefficients) == len(values) and values[0] is not None
+ total = np.zeros(len(values[0]))
+ for c, v in zip(coefficients, values):
+ total += c * v
+ return total
-get_state = python_app(_get_state, executors=["default_threads"])
+
+add_contributions = python_app(_add_contributions, executors=["default_threads"])
-@typeguard.typechecked
@psiflow.serializable
+@dataclass
class SimulationOutput:
- """Gathers simulation output
-
- status is an integer which represents an exit code of the run:
-
- -1: unknown error
- 0: run completed successfully
- 1: run terminated early due to time limit
- 2: run terminated early due to max force exception
-
- """
-
- _data: dict[str, Optional[AppFuture]]
- state: Union[Geometry, AppFuture, None]
- stdout: Optional[str]
- status: Union[int, AppFuture, None]
- time: Union[float, AppFuture, None]
- temperature: Union[float, AppFuture, None]
+ task_id: str
+ walker: Walker
+ status: Union[Status, AppFuture]
+ state: Union[Geometry, AppFuture]
+ data: AppFuture
+ time: Union[float, AppFuture]
+ temperature: Union[float, AppFuture]
trajectory: Optional[Dataset]
+ observables: list[str]
- def __init__(self, fields: list[str]):
- self._data = {key: None for key in fields}
+ hamiltonian_components: InitVar[list[HamiltonianComponent]]
+ component_map: dict[str, HamiltonianComponent] = field(init=False)
+ force_comps: list[HamiltonianComponent] = field(init=False)
+ bias_comps: list[HamiltonianComponent] = field(init=False)
+ _data: dict[str, AppFuture] = field(init=False)
- self.state = None
- self.stdout = None
- self.status = None
- self.time = None
- self.temperature = None
- self.trajectory = None
- self.hamiltonians = None
+ def __post_init__(self, hamiltonian_components):
+ self._data = {} # to cache deferred_getitem calls
+ self.force_comps = [comp for comp in hamiltonian_components if comp.shared]
+ self.bias_comps = [comp for comp in hamiltonian_components if not comp.shared]
+
+ # map generic i-Pi keys to hamiltonian force components
+ self.component_map = {
+ potential_component_name(i): comp for i, comp in enumerate(self.force_comps)
+ }
def __getitem__(self, key: str) -> AppFuture:
- if self._data.get(key, None) is None:
- raise ValueError("output {} not available".format(key))
+ assert key in self.observables
+ if key not in self._data:
+ self._data[key] = self.data[key]
return self._data[key]
- def parse(
- self,
- result: AppFuture, # result from ipi execution
- state: AppFuture,
- ):
- self.state = state
- self.stdout = result.stdout
- parsed = parse(state, inputs=[result.stdout, result.stderr])
- self.time = unpack_i(parsed, 0)
- self.status = unpack_i(parsed, 1)
-
- def parse_data(
- self,
- start: int,
- data_future: DataFuture,
- hamiltonians: list[Hamiltonian],
- ):
- data = parse_data(
- self.status,
- start,
- list(self._data.keys()),
- inputs=[data_future],
- )
- for i, key in enumerate(self._data.keys()):
- self._data[key] = unpack_i(data, i)
- self.hamiltonians = hamiltonians
- assert "temperature{kelvin}" in self._data.keys()
- self.temperature = get_final_temperature(self._data["temperature{kelvin}"])
-
- def update_walker(self, walker):
- walker.state = update_walker(
+ def update_walker(self):
+ # TODO: when do we want to reset?
+ self.walker.state = self.state
+
+ def log_status(self):
+ log_status(
+ self.task_id,
self.status,
- self.state,
- walker.start,
+ self.observables,
+ self.time,
+ self.temperature,
+ self.component_map,
+ self.bias_comps,
)
def get_energy(self, hamiltonian: Hamiltonian) -> AppFuture:
- all_h = MixtureHamiltonian(
- list(self.hamiltonians),
- [1.0 for h in self.hamiltonians],
- )
+ """Use stored energy contributions from i-Pi to evaluate a custom (Mixture)Hamiltonian."""
if hamiltonian == Zero():
- name = potential_component_names(1)[0]
- return add_contributions((0.0,), self._data[name])
- coefficients = all_h.get_coefficients(1.0 * hamiltonian)
- names = potential_component_names(len(self.hamiltonians))
- values = [self._data[name] for name in names]
+ # future because array length not known
+ return add_contributions((0.0,), self.data[DEFAULT_OBSERVABLES[0]])
+
+ f_comps = self.force_comps
+ f_hamiltonian = MixtureHamiltonian(
+ [comp.hamiltonian for comp in f_comps], [1] * len(f_comps)
+ )
+ coefficients = f_hamiltonian.get_coefficients(1.0 * hamiltonian)
+ if coefficients is None:
+ raise ValueError(
+ f"Provided hamiltonian is not fully in i-Pi forces. "
+ f"Make sure you do not ask for bias contributions."
+ )
return add_contributions(
- coefficients,
- *values,
+ coefficients, *[self[name] for name in self.component_map]
)
- def get_state(self):
- return get_state(self.state, self.status)
-
- def save(self, path: Union[str, Path]):
- if type(path) is str:
- path = Path(path)
- assert not path.exists()
- raise NotImplementedError
+ def save_data(self, path: Union[str, Path], **kwargs: np.ndarray) -> DataFuture:
+ component_map = {c.name: k for k, c in self.component_map.items()}
+ future = save_npz(self.data, outputs=[File(path)], **component_map, **kwargs)
+ return future.outputs[0]
+
+ @classmethod
+ def from_md(
+ cls,
+ walker: Walker,
+ state: AppFuture[Geometry],
+ observables: list[str],
+ hamiltonian_components: list[HamiltonianComponent],
+ start: int,
+ file_props: DataFuture,
+ file_traj: Optional[DataFuture],
+ output_log: str,
+ error_log: str,
+ ) -> "SimulationOutput":
+ """"""
+ future = parse_simulation_data(
+ state, observables, start, file_props, output_log, error_log
+ )
+ status, data, temperature, time = future[0], future[1], future[2], future[3]
+ trajectory = None
+ if file_traj:
+ trajectory = Dataset(None, file_traj)
+ if start > 0:
+ trajectory = trajectory[start:]
+ return cls(
+ get_task_name_id(output_log)[-1],
+ walker,
+ status,
+ state,
+ data,
+ time,
+ temperature,
+ trajectory,
+ observables,
+ hamiltonian_components,
+ )
diff --git a/psiflow/sampling/sampling.py b/psiflow/sampling/sampling.py
index b2f2ea98..4328ea3d 100644
--- a/psiflow/sampling/sampling.py
+++ b/psiflow/sampling/sampling.py
@@ -1,12 +1,9 @@
-from __future__ import annotations # necessary for type-guarding class methods
-
import math
import xml.etree.ElementTree as ET
from dataclasses import dataclass
from typing import Optional, Union, Iterable
import parsl
-import typeguard
import numpy as np
from parsl.app.app import bash_app
from parsl.data_provider.files import File
@@ -14,22 +11,19 @@
import psiflow
from psiflow.data import Dataset
-from psiflow.hamiltonians import Hamiltonian, MixtureHamiltonian, Zero
+from psiflow.data.utils import read_frames
+from psiflow.hamiltonians import Hamiltonian, MixtureHamiltonian, Zero, MACEHamiltonian
from psiflow.sampling.output import (
DEFAULT_OBSERVABLES,
SimulationOutput,
- potential_component_names,
+ potential_component_name,
+ HamiltonianComponent,
)
-from psiflow.sampling.walker import Coupling, Walker, partition
-from psiflow.utils.io import save_xml
-from psiflow.utils import TMP_COMMAND, CD_COMMAND
-
-
-@dataclass
-class HamiltonianComponent:
- name: str
- hamiltonian: Hamiltonian
- shared: bool
+from psiflow.sampling.utils import create_xml_list
+from psiflow.sampling.walker import Coupling, Walker, partition, Ensemble
+from psiflow.utils.io import _save_xml
+from psiflow.utils import TMP_COMMAND, CD_COMMAND, export_env_command
+from psiflow.sampling.driver import __file__ as PATH_DRIVER
@dataclass
@@ -46,24 +40,18 @@ def get_key(self, key: str) -> np.ndarray:
def __len__(self) -> int:
return self.weights.shape[0]
- def __eq__(self, other: EnsembleTable) -> bool:
+ def __eq__(self, other: "EnsembleTable") -> bool:
if (
- not isinstance(other, EnsembleTable) or
- len(self) != len(other) or
- set(self.keys) != set(other.keys) or
- self.weights.shape != other.weights.shape
+ not isinstance(other, EnsembleTable)
+ or len(self) != len(other)
+ or set(self.keys) != set(other.keys)
+ or self.weights.shape != other.weights.shape
):
return False
key_order = [other.keys.index(key) for key in self.keys]
return np.all(self.weights == other.weights[:, key_order])
-def create_xml_list(items: list[str]) -> str:
- """Pure helper"""
- return ' [ {} ] '.format(', '.join(items))
-
-
-@typeguard.typechecked
def template(
walkers: list[Walker],
) -> tuple[list[HamiltonianComponent], EnsembleTable, list[AppFuture]]:
@@ -71,30 +59,17 @@ def template(
total_hamiltonian = 1.0 * sum([w.hamiltonian for w in walkers], start=Zero())
assert not total_hamiltonian == Zero()
- # create string names for hamiltonians and sort TODO: why sort?
- # names = label_forces(total_hamiltonian)
- # names, hamiltonians, coefficients = zip(
- # *sorted(zip(names, total_hamiltonian.hamiltonians, total_hamiltonian.coefficients))
- # )
- # assert MixtureHamiltonian(hamiltonians, coefficients) == total_hamiltonian
- # total_hamiltonian = MixtureHamiltonian(hamiltonians, coefficients)
-
# construct the table of potential / ensemble / bias weights for every system instance
names = total_hamiltonian.get_named_components()
- weights_hamiltonian = np.zeros(shape=(len(walkers), len(names)))
+ weights_h = np.zeros(shape=(len(walkers), len(names)))
for idx, walker in enumerate(walkers):
- weights_hamiltonian[idx] = total_hamiltonian.get_coefficients(walker.hamiltonian * 1)
-
- components = []
- for idx, (name, hamiltonian) in enumerate(zip(names, total_hamiltonian.hamiltonians)):
- weights = weights_hamiltonian[:, idx]
- components.append(HamiltonianComponent(name, hamiltonian, all(weights)))
+ weights_h[idx] = total_hamiltonian.get_coefficients(walker.hamiltonian * 1)
weights_dict = {}
- if walkers[0].nvt or walkers[0].npt:
- weights_dict['TEMP'] = [w.temperature for w in walkers]
- if walkers[0].npt:
- weights_dict['PRESSURE'] = [w.pressure for w in walkers]
+ if walkers[0].ensemble != Ensemble.NVE:
+ weights_dict["TEMP"] = [w.temperature for w in walkers]
+ if walkers[0].ensemble in (Ensemble.NPT, Ensemble.NVST):
+ weights_dict["PRESSURE"] = [w.pressure for w in walkers]
# inspect metadynamics attributes and allocate additional weights per MTD
metad_objects = []
@@ -108,85 +83,87 @@ def template(
"psiflow.sampling.multiple_walker_metadynamics"
)
metad_objects.append(mtd)
- weights_dict[f'METAD{len(metad_objects) - 1}'] = [(i == j) for j in range(len(walkers))]
+ weights_dict[f"METAD{len(metad_objects) - 1}"] = [
+ (i == j) for j in range(len(walkers))
+ ]
plumed_list = [mtd.input() for mtd in metad_objects]
- weights_ensemble = np.array([*weights_dict.values()]).T
weights_name = tuple(names + list(weights_dict.keys()))
- weights_table = np.concatenate([weights_hamiltonian, weights_ensemble], axis=-1)
-
- return components, EnsembleTable(weights_name, weights_table), plumed_list
+ if len(weights_dict) == 0: # NVE without metadynamics
+ weights_table = weights_h
+ else:
+ weights_ensemble = np.array([*weights_dict.values()]).T
+ weights_table = np.concatenate([weights_h, weights_ensemble], axis=-1)
+ ensemble_table = EnsembleTable(weights_name, weights_table)
+
+ # make list of all (shared) hamiltonian components
+ components = [
+ HamiltonianComponent(n, h, all(weights_h[:, idx]))
+ for idx, (n, h) in enumerate(zip(names, total_hamiltonian.hamiltonians))
+ ]
+ return components, ensemble_table, plumed_list
-def make_force_xml(hamiltonian: MixtureHamiltonian, names: list[str]) -> ET.Element:
- # TODO: move to relevant module
- forces = ET.Element("forces")
- for n, c in zip(names, hamiltonian.coefficients):
- forces.append(ET.Element("force", forcefield=n, weight=str(c)))
- return forces
+def setup_sockets(components: list[HamiltonianComponent]) -> list[ET.Element]:
+ xml_str = """
+
+ {timeout}
+ TRUE
+ {address}
+
+ """
+ timeout = 60 * psiflow.context().definitions["ModelEvaluation"].timeout
-@typeguard.typechecked
-def setup_sockets(
- hamiltonian_labels: Iterable[str],
-) -> list[ET.Element]:
sockets = []
- for name in hamiltonian_labels:
- ffsocket = ET.Element("ffsocket", mode="unix", name=name, pbc="False")
- timeout = ET.Element("timeout")
- timeout.text = str(
- 60 * psiflow.context().definitions["ModelEvaluation"].timeout
- )
- ffsocket.append(timeout)
- exit_on = ET.Element("exit_on_disconnect")
- exit_on.text = " TRUE "
- ffsocket.append(exit_on)
- address = ET.Element("address") # placeholder
- address.text = name.lower()
- ffsocket.append(address)
-
- sockets.append(ffsocket)
+ for comp in components:
+ xml = xml_str.format(name=comp.name, address=comp.address, timeout=str(timeout))
+ sockets.append(ET.fromstring(xml))
return sockets
-@typeguard.typechecked
def setup_motion(walker: Walker, fix_com: bool) -> ET.Element:
- timestep_element = ET.Element("timestep", units="femtosecond")
- timestep_element.text = str(walker.timestep)
-
- tau = ET.Element("tau", units="femtosecond")
- tau.text = "100"
- thermostat_pimd = ET.Element("thermostat", mode="pile_g")
- thermostat_pimd.append(tau)
- thermostat = ET.Element("thermostat", mode="langevin")
- thermostat.append(tau)
- if walker.nve:
- dynamics = ET.Element("dynamics", mode="nve")
- dynamics.append(timestep_element)
- elif walker.nvt:
- dynamics = ET.Element("dynamics", mode="nvt")
- dynamics.append(timestep_element)
- if walker.pimd:
- dynamics.append(thermostat_pimd)
- else:
- dynamics.append(thermostat)
- elif walker.npt:
+ if (ensemble := walker.ensemble) == Ensemble.NVE:
+ mode = "nve"
+ elif ensemble == Ensemble.NVT:
+ mode = "nvt"
+ elif ensemble in (Ensemble.NPT, Ensemble.NVST):
mode = "npt"
- dynamics = ET.Element("dynamics", mode=mode)
- dynamics.append(timestep_element)
- if walker.pimd:
- dynamics.append(thermostat_pimd)
- else:
- dynamics.append(thermostat)
- mode = "flexible"
- barostat = ET.Element("barostat", mode=mode)
- tau = ET.Element("tau", units="femtosecond")
- tau.text = "200"
- barostat.append(tau)
- barostat.append(thermostat) # never use thermostat_pimd here!
- dynamics.append(barostat)
else:
raise ValueError("invalid walker {}".format(walker))
+ dynamics = ET.Element("dynamics", mode=mode)
+ timestep_element = ET.Element("timestep", units="femtosecond")
+ timestep_element.text = str(walker.timestep)
+ dynamics.append(timestep_element)
+
+ # thermostat
+ if ensemble != Ensemble.NVE:
+ xml_thermo = """
+
+ {tau_thermo}
+
+ """.format(
+ mode=("pile_g" if walker.pimd else "langevin"),
+ tau_thermo=walker.tau_thermostat,
+ )
+ dynamics.append(ET.fromstring(xml_thermo))
+
+ # barostat - never use pile_g here!
+ if ensemble in (Ensemble.NPT, Ensemble.NVST):
+ xml_baro = """
+
+ {tau_baro}
+ {vol_constraint}
+
+ {tau_thermo}
+
+
+ """.format(
+ tau_baro=walker.tau_barostat,
+ tau_thermo=walker.tau_thermostat,
+ vol_constraint=(ensemble == Ensemble.NVST),
+ )
+ dynamics.append(ET.fromstring(xml_baro))
motion = ET.Element("motion", mode="dynamics")
motion.append(dynamics)
@@ -196,8 +173,9 @@ def setup_motion(walker: Walker, fix_com: bool) -> ET.Element:
return motion
-@typeguard.typechecked
-def setup_ensemble(components: list[HamiltonianComponent], weights_header: tuple[str, ...]) -> ET.Element:
+def setup_ensemble(
+ components: list[HamiltonianComponent], weights_header: tuple[str, ...]
+) -> ET.Element:
ensemble = ET.Element("ensemble")
# set TEMP in any case to avoid i-PI from throwing weird errors
@@ -236,32 +214,29 @@ def setup_ensemble(components: list[HamiltonianComponent], weights_header: tuple
return ensemble
-@typeguard.typechecked
-def setup_forces(hamiltonian_components: list[HamiltonianComponent]) -> ET.Element:
+def setup_forces(components: list[HamiltonianComponent]) -> ET.Element:
forces = ET.Element("forces")
- for comp in hamiltonian_components:
- if comp.shared: # only add components shared across all walkers
+ for comp in components:
+ if comp.shared: # only add components shared across all walkers
force = ET.Element("force", forcefield=comp.name, weight=comp.name.upper())
forces.append(force)
return forces
-@typeguard.typechecked
def setup_ffplumed(nplumed: int) -> list[ET.Element]:
ffplumed = []
for i in range(nplumed):
input_file = ET.Element("file", mode="xyz", cell_units="angstrom")
input_file.text = "start_0.xyz" # always present
- plumeddat = ET.Element("plumeddat")
- plumeddat.text = "metad_input{}.txt".format(i)
+ plumed_dat = ET.Element("plumed_dat")
+ plumed_dat.text = "metad_input{}.txt".format(i)
ff = ET.Element("ffplumed", name="metad{}".format(i), pbc="False")
ff.append(input_file)
- ff.append(plumeddat)
+ ff.append(plumed_dat)
ffplumed.append(ff)
return ffplumed
-@typeguard.typechecked
def setup_system_template(
walkers: list[Walker],
ensemble_table: EnsembleTable,
@@ -271,12 +246,14 @@ def setup_system_template(
) -> ET.Element:
system_template = ET.Element("system_template")
labels = ET.Element("labels")
- labels.text = create_xml_list(['INDEX'] + [w.upper() for w in ensemble_table.keys])
+ labels.text = create_xml_list(["INDEX"] + [w.upper() for w in ensemble_table.keys])
system_template.append(labels)
for i in range(len(ensemble_table)):
instance = ET.Element("instance")
- instance.text = create_xml_list([f'{i}'] + [str(w) for w in ensemble_table.get_index(i)])
+ instance.text = create_xml_list(
+ [f"{i}"] + [str(w) for w in ensemble_table.get_index(i)]
+ )
system_template.append(instance)
initialize = ET.Element("initialize", nbeads=str(walkers[0].nbeads))
@@ -284,10 +261,13 @@ def setup_system_template(
start.text = " start_INDEX.xyz "
initialize.append(start)
velocities = ET.Element("velocities", mode="thermal", units="kelvin")
- velocities.text = " TEMP " if "TEMP" in ensemble_table.keys else " 300 " # valid template parameter
+ velocities.text = (
+ " TEMP " if "TEMP" in ensemble_table.keys else " 300 "
+ ) # valid template parameter
initialize.append(velocities)
if walkers[0].masses is not None:
import ase.units
+
AMU_TO_AU = ase.units._amu / ase.units._me
masses = ET.Element("masses", mode="manual")
masses.text = create_xml_list([str(i) for i in walkers[0].masses * AMU_TO_AU])
@@ -305,7 +285,6 @@ def setup_system_template(
return system_template
-@typeguard.typechecked
def setup_output(
components: list[HamiltonianComponent],
observables: Optional[list[str]],
@@ -319,13 +298,16 @@ def setup_output(
n_forces = sum([comp.shared for comp in components], 0)
full_list = (
- DEFAULT_OBSERVABLES + potential_component_names(n_forces) + observables
+ DEFAULT_OBSERVABLES
+ + [potential_component_name(i) for i in range(n_forces)]
+ + observables
)
- if len(components) != n_forces: # store bias force components
- full_list.append('ensemble_bias{electronvolt}')
+ if len(components) != n_forces: # store bias force components
+ full_list.append("ensemble_bias{electronvolt}")
observables = list(set(full_list))
if step is None:
+ # TODO: this logic should be elsewhere
step = checkpoint_step
output = ET.Element("output", prefix="output")
@@ -356,7 +338,6 @@ def setup_output(
return output, observables
-@typeguard.typechecked
def setup_smotion(
coupling: Optional[Coupling], plumed_list: list[AppFuture]
) -> ET.Element:
@@ -381,76 +362,111 @@ def setup_smotion(
return smotion
-def make_start_command(command: str, input_xml: File, start_xyz: File, nwalkers: int = 1) -> str:
+def make_server_command(
+ command: str,
+ input_xml: File,
+ start_xyz: File,
+ output_xyz: File,
+ output_props: list[File],
+ output_trajs: list[File],
+) -> str:
""""""
- return f'{command} --nwalkers={nwalkers} --input_xml={input_xml.filepath} --start_xyz={start_xyz.filepath} &'
+ args = [
+ f"{command}",
+ f"--input_xml={input_xml.filepath}",
+ f"--start_xyz={start_xyz.filepath}",
+ f"--output_xyz={output_xyz.filepath}",
+ ]
+ if output_props:
+ args.append(
+ "--output_props=" + ",".join([file.filepath for file in output_props])
+ )
+ if output_trajs:
+ args.append(
+ "--output_traj=" + ",".join([file.filepath for file in output_trajs])
+ )
+ return " ".join(args + ["&"])
-def make_client_command(command: str, address: str, hamiltonian: File,
- start: File, arg: str, max_force: float = None) -> str:
+def make_driver_commands(
+ driver_kwargs: list[dict], file_xyz: File, files_hamiltonian: list[File]
+) -> list[str]:
""""""
- return '{c} --address={a} --path_hamiltonian={p} --start={s} {m} {arg} &'.format(
- c=command, a=address.lower(), p=hamiltonian.filepath, s=start.filepath, arg=arg,
- m=(f'--max_force={max_force}' if max_force is not None else ''),
- )
+ assert len(driver_kwargs) >= len(files_hamiltonian)
+ default = f'i-pi-driver-py -u -S "" -m custom -P {PATH_DRIVER} -a {{address}} -o {{options}} &'
+
+ commands = []
+ for kwargs in driver_kwargs:
+ # drivers know which hamiltonian to load by 'idx'
+ address, idx = kwargs.pop("address"), kwargs.pop("idx")
+ options = [
+ str(file_xyz),
+ str(files_hamiltonian[idx]),
+ *[f"{k}={v}" for k, v in kwargs.items() if v is not None],
+ ]
+ commands.append(default.format(address=address, options=",".join(options)))
+ return commands
+
+
+def make_wait_for_sockets_command(
+ addresses: Iterable[str], timeout: int = 10, interval: float = 0.1
+) -> str:
+ """"""
+ exist = " && ".join([f"[ -e {a} ]" for a in addresses])
+ return f"t=0; until {exist} || (( t++ >= {timeout} )); do sleep {interval}; done"
def _execute_ipi(
nwalkers: int,
- hamiltonian_names: list[str],
- client_args: list[str],
+ driver_kwargs: list[dict],
keep_trajectory: bool,
- max_force: Optional[float],
coupling_command: Optional[str],
command_server: str,
- command_client: str,
*plumed_list: str,
env_vars: dict = {},
- stdout: str = "",
- stderr: str = "",
+ stdout: str = parsl.AUTO_LOGNAME,
+ stderr: str = parsl.AUTO_LOGNAME,
inputs: list = [],
outputs: list = [],
parsl_resource_specification: Optional[dict] = {},
) -> str:
- write_command = '\n'.join([
- f'echo "{plumed_str}" > metad_input{i}.txt' for i, plumed_str in enumerate(plumed_list)
- ])
- env_command = 'export ' + ' '.join([f"{name}={value}" for name, value in env_vars.items()])
- command_start = make_start_command(command_server, inputs[0], inputs[1], nwalkers)
- commands_client = []
- for i, name in enumerate(hamiltonian_names):
- args = client_args[i]
- for arg in args:
- commands_client += make_client_command(command_client, name, inputs[2 + i], inputs[1], arg, max_force),
-
- command_end = f'{command_server} --cleanup --output_xyz={outputs[0].filepath}'
- commands_copy = []
- for i in range(nwalkers):
- 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 coupling_command is not None:
- commands_copy += coupling_command,
- command_copy = '; '.join(commands_copy)
+ """"""
+ file_xml, file_xyz_in, *files_in = inputs
+ file_xyz_out, files_props = outputs[0], outputs[1 : 1 + nwalkers]
+ files_traj = outputs[1 + nwalkers : 1 + 2 * nwalkers] if keep_trajectory else []
+ if coupling_command:
+ files_in = files_in[:-1] # exclude REX swapfile
+
+ write_command_args = [
+ f'echo "{plumed_str}" > metad_input{i}.txt'
+ for i, plumed_str in enumerate(plumed_list)
+ ]
+ command_start = make_server_command(
+ command_server, file_xml, file_xyz_in, file_xyz_out, files_props, files_traj
+ )
+ command_wait = make_wait_for_sockets_command(
+ set(d["address"] for d in driver_kwargs)
+ )
+ commands_driver = make_driver_commands(driver_kwargs, file_xyz_in, files_in)
command_list = [
TMP_COMMAND,
CD_COMMAND,
- write_command,
- env_command,
+ "\n".join(write_command_args),
+ export_env_command(env_vars),
command_start,
- *commands_client,
+ command_wait,
+ *commands_driver,
"wait",
- command_end,
- command_copy,
]
+ if coupling_command:
+ command_list.append(coupling_command)
return "\n".join(command_list)
execute_ipi = bash_app(_execute_ipi, executors=["ModelEvaluation"])
-@typeguard.typechecked
def _sample(
walkers: list[Walker],
steps: int,
@@ -466,12 +482,12 @@ def _sample(
verbosity: str = "medium",
) -> list[SimulationOutput]:
assert len(walkers) > 0
- hamiltonian_components, ensemble_table, plumed_list = template(walkers)
- coupling = walkers[0].coupling
-
if motion_defaults is not None:
raise NotImplementedError
+ # generate i-Pi input XML
+ hamiltonian_components, ensemble_table, plumed_list = template(walkers)
+ coupling = walkers[0].coupling
motion = setup_motion(walkers[0], fix_com)
ensemble = setup_ensemble(hamiltonian_components, ensemble_table.keys)
forces = setup_forces(hamiltonian_components)
@@ -510,7 +526,7 @@ def _sample(
verbosity=str(verbosity),
safe_stride=str(checkpoint_step),
)
- sockets = setup_sockets([comp.name for comp in hamiltonian_components])
+ sockets = setup_sockets(hamiltonian_components)
for socket in sockets:
simulation.append(socket)
ffplumed = setup_ffplumed(len(plumed_list))
@@ -530,24 +546,31 @@ def _sample(
prng.append(seed)
simulation.append(prng)
+ # app setup and IO
context = psiflow.context()
definition = context.definitions["ModelEvaluation"]
- input_future = save_xml(
- simulation,
- outputs=[context.new_file("input_", ".xml")],
- ).outputs[0]
+ input_file = context.new_file("input_", ".xml")
+ _save_xml(simulation, outputs=[input_file])
inputs = [
- input_future,
+ input_file,
Dataset([w.state for w in walkers]).extxyz,
]
- # remove any Harmonic instances because they are not implemented with sockets -- TODO: why?
+ # figure out i-Pi MD driver configuration
+ # how many drivers (force evaluators) with which arguments?
+ # remove any Harmonic instances because they are not implemented with sockets -- TODO: why?
max_nclients = int(sum([w.nbeads for w in walkers]))
- client_args = []
- for comp in hamiltonian_components:
+ driver_kwargs = []
+ for i, comp in enumerate(hamiltonian_components):
inputs.append(comp.hamiltonian.serialize_function())
- args = definition.get_client_args(comp.name, max_nclients, motion="dynamics")
- client_args.append(args)
+ kwargs = {"idx": i, "address": comp.address, "max_force": max_force}
+ if isinstance(comp.hamiltonian, MACEHamiltonian):
+ kwargs["dtype"] = "float32" # TODO: should this be configurable?
+ for instance_kwargs in definition.get_driver_devices(max_nclients):
+ driver_kwargs.append(kwargs | instance_kwargs)
+ else:
+ driver_kwargs.append(kwargs)
+
outputs = [context.new_file("data_", ".xyz")]
outputs += [context.new_file("simulation_", ".txt") for _ in walkers]
if keep_trajectory:
@@ -567,48 +590,42 @@ def _sample(
else:
coupling_copy_command = None
- command_server = definition.server_command()
- command_client = definition.client_command()
- resources = definition.wq_resources(max_nclients)
+ # TODO: an app to check for valid input? (e.g., PBC + barostat)
result = execute_ipi(
len(walkers),
- [comp.name for comp in hamiltonian_components],
- client_args,
+ driver_kwargs,
keep_trajectory,
- max_force,
coupling_copy_command,
- command_server,
- command_client,
- *plumed_list,
+ definition.server_command(),
+ *plumed_list, # futures
env_vars=dict(definition.env_vars),
- stdout=parsl.AUTO_LOGNAME,
- stderr=parsl.AUTO_LOGNAME,
inputs=inputs,
outputs=outputs,
- parsl_resource_specification=resources,
+ parsl_resource_specification=definition.wq_resources(max_nclients),
)
- final_states = Dataset(None, result.outputs[0])
- simulation_outputs = [SimulationOutput(observables) for _ in range(len(walkers))]
- for i, simulation_output in enumerate(simulation_outputs):
- state = final_states[i]
- if walkers[i].order_parameter is not None:
- state = walkers[i].order_parameter.evaluate(state)
- simulation_output.parse(result, state)
- simulation_output.parse_data(
+ # process MD output
+ simulation_outputs = []
+ final_states = read_frames(inputs=result.outputs[:1])
+ for idx, walker in enumerate(walkers):
+ # TODO: order_parameter out of commission
+ # if walker.order_parameter is not None:
+ # state = walker.order_parameter.evaluate(state)
+ if walker.metadynamics is not None:
+ walker.metadynamics.wait_for(result)
+ output = SimulationOutput.from_md(
+ walker,
+ final_states[idx],
+ observables,
+ hamiltonian_components, # order is important
start,
- result.outputs[i + 1],
- hamiltonians=[comp.hamiltonian for comp in hamiltonian_components],
+ result.outputs[idx + 1],
+ result.outputs[len(walkers) + 1 + idx] if keep_trajectory else None,
+ result.stdout,
+ result.stderr,
)
- if keep_trajectory:
- j = len(walkers) + 1 + i
- trajectory = Dataset(None, result.outputs[j])
- if start > 0:
- trajectory = trajectory[start:]
- simulation_output.trajectory = trajectory
- if walkers[i].metadynamics is not None:
- walkers[i].metadynamics.wait_for(result)
- simulation_output.update_walker(walkers[i])
+ output.update_walker()
+ simulation_outputs.append(output)
if coupling is not None:
coupling.update(result)
@@ -616,7 +633,6 @@ def _sample(
return simulation_outputs
-@typeguard.typechecked
def sample(
walkers: list[Walker],
steps: int,
diff --git a/psiflow/sampling/server.py b/psiflow/sampling/server.py
index aa42a81b..2fe3c71c 100644
--- a/psiflow/sampling/server.py
+++ b/psiflow/sampling/server.py
@@ -1,304 +1,59 @@
import argparse
-import ast
-import glob
-import os
-import signal
+import traceback
import xml.etree.ElementTree as ET
-from copy import deepcopy
+import subprocess
+import time
from pathlib import Path
+import ase
import numpy as np
from ase.data import atomic_numbers
from ase.geometry import Cell
from ase.io import read, write
from ase.units import Bohr, Ha, _e, _hbar
+from ipi.engine.simulation import Simulation
+from ipi.utils.softexit import softexit
from psiflow.geometry import Geometry
-from psiflow.sampling.utils import TimeoutException, timeout_handler
+from psiflow.sampling.utils import create_xml_list
+INPUT_XML = "input.xml"
NONPERIODIC_CELL = 1000 * np.eye(3)
-def remdsort(inputfile, prefix="SRT_"):
- from ipi.engine.outputs import CheckpointOutput, PropertyOutput, TrajectoryOutput
- from ipi.engine.properties import getkey
- from ipi.inputs.simulation import InputSimulation
- from ipi.utils.io.inputs import io_xml
- from ipi.utils.messages import verbosity
-
- verbosity.level = "low"
- # opens & parses the input file
- ifile = open(inputfile, "r")
- xmlrestart = io_xml.xml_parse_file(ifile) # Parses the file.
- ifile.close()
-
- # ugly hack to remove ffplumed objects to avoid messing up with plumed output files
- newfields = [f for f in xmlrestart.fields[0][1].fields if f[0] != "ffplumed"]
- xmlrestart.fields[0][1].fields = newfields
-
- isimul = InputSimulation()
- isimul.parse(xmlrestart.fields[0][1])
-
- simul = isimul.fetch()
- swapfile = ""
- if simul.smotion is None or (
- simul.smotion.mode != "remd" and simul.smotion.mode != "multi"
- ):
- raise ValueError("Simulation does not look like a parallel tempering one.")
- else:
- if simul.smotion.mode == "remd":
- swapfile = simul.smotion.swapfile
- else:
- for sm in simul.smotion.mlist:
- if sm.mode == "remd":
- if swapfile != "":
- raise ValueError(
- "I'm not equipped to deal with multiple REMD outputs, sorry"
- )
- swapfile = sm.swapfile
- if swapfile == "":
- raise ValueError(
- "Could not determine the REMD swapfile name. \
- Sorry, you'll have to look carefully at your inputs."
- )
-
- # reconstructs the list of the property and trajectory files that have been output
- # and that should be re-ordered
- lprop = [] # list of property files
- ltraj = [] # list of trajectory files
- nsys = len(simul.syslist)
- for o in simul.outtemplate:
- o = deepcopy(o) # avoids overwriting the actual filename
- if simul.outtemplate.prefix != "":
- o.filename = simul.outtemplate.prefix + "." + o.filename
- if (
- type(o) is CheckpointOutput
- ): # properties and trajectories are output per system
- pass
- elif type(o) is PropertyOutput:
- nprop = []
- isys = 0
- for s in simul.syslist: # create multiple copies
- if s.prefix != "":
- filename = s.prefix + "_" + o.filename
- else:
- filename = o.filename
- ofilename = prefix + filename
- nprop.append(
- {
- "filename": filename,
- "ofilename": ofilename,
- "stride": o.stride,
- "ifile": open(filename, "r"),
- "ofile": open(ofilename, "w"),
- }
- )
- isys += 1
- lprop.append(nprop)
- elif (
- type(o) is TrajectoryOutput
- ): # trajectories are more complex, as some have per-bead output
- if getkey(o.what) in [
- "positions",
- "velocities",
- "forces",
- "extras",
- ]: # multiple beads
- nbeads = simul.syslist[0].beads.nbeads
- for b in range(nbeads):
- ntraj = []
- isys = 0
- # zero-padded bead number
- padb = (
- "%0" + str(int(1 + np.floor(np.log(nbeads) / np.log(10)))) + "d"
- ) % (b)
- for s in simul.syslist:
- if s.prefix != "":
- filename = s.prefix + "_" + o.filename
- else:
- filename = o.filename
- ofilename = prefix + filename
- if o.ibead < 0 or o.ibead == b:
- if getkey(o.what) == "extras":
- filename = filename + "_" + padb
- ofilename = ofilename + "_" + padb
- # Sets format of extras as None
- ntraj.append(
- {
- "filename": filename,
- "format": None,
- "ofilename": ofilename,
- "stride": o.stride,
- "ifile": open(filename, "r"),
- "ofile": open(ofilename, "w"),
- }
- )
- else:
- # FIX
- if o.format == "ase":
- extension = "extxyz"
- else:
- extension = o.format
- filename = filename + "_" + padb + "." + extension
- ofilename = ofilename + "_" + padb + "." + extension
- print(filename, ofilename)
- ntraj.append(
- {
- "filename": filename,
- "format": o.format,
- "ofilename": ofilename,
- "stride": o.stride,
- "ifile": open(filename, "r"),
- "ofile": open(ofilename, "w"),
- }
- )
- isys += 1
- if ntraj != []:
- ltraj.append(ntraj)
-
- else:
- ntraj = []
- isys = 0
- for s in simul.syslist: # create multiple copies
- if s.prefix != "":
- filename = s.prefix + "_" + o.filename
- else:
- filename = o.filename
- filename = filename + "." + o.format
- ofilename = prefix + filename
- ntraj.append(
- {
- "filename": filename,
- "format": o.format,
- "ofilename": ofilename,
- "stride": o.stride,
- "ifile": open(filename, "r"),
- "ofile": open(ofilename, "w"),
- }
- )
-
- isys += 1
- ltraj.append(ntraj)
-
- ptfile = open(simul.outtemplate.prefix + "." + swapfile, "r")
-
- # now reads files one frame at a time,
- # and re-direct output to the appropriate location
-
- line = ptfile.readline().split()
- irep = list(range(nsys)) # Could this be harmful?
- step = 0
- while True:
- # reads one line from index file
- try:
- for prop in lprop:
- for isys in range(nsys):
- sprop = prop[isys]
- if step % sprop["stride"] == 0: # property transfer
- iline = sprop["ifile"].readline()
- if len(iline) == 0:
- raise EOFError # useful if line is blank
- while iline[0] == "#": # fast forward if line is a comment
- prop[irep[isys]]["ofile"].write(iline)
- iline = sprop["ifile"].readline()
- prop[irep[isys]]["ofile"].write(iline)
- for traj in ltraj:
- for isys in range(nsys):
- straj = traj[isys]
- if step % straj["stride"] == 0: # property transfer
- # reads one frame from the input file
- ibuffer = []
- if straj["format"] is None:
- ibuffer.append(straj["ifile"].readline())
- ibuffer.append(straj["ifile"].readline())
- traj[irep[isys]]["ofile"].write("".join(ibuffer))
- if straj["format"] in ["xyz", "ase"]:
- iline = straj["ifile"].readline()
- nat = int(iline)
- ibuffer.append(iline)
- ibuffer.append(straj["ifile"].readline())
- for _i in range(nat):
- ibuffer.append(straj["ifile"].readline())
- traj[irep[isys]]["ofile"].write("".join(ibuffer))
- elif straj["format"] == "pdb":
- iline = straj["ifile"].readline()
- while iline.strip() != "" and iline.strip() != "END":
- ibuffer.append(iline)
- iline = straj["ifile"].readline()
- ibuffer.append(iline)
- traj[irep[isys]]["ofile"].write("".join(ibuffer))
- except EOFError:
- break
-
- if len(line) > 0 and step == int(line[0]):
- irep = [int(i) for i in line[1:]]
- line = ptfile.readline()
- line = line.split()
-
- step += 1
-
-
-def parse_checkpoint(
- checkpoint: ET.ElementTree,
-) -> tuple[list[Geometry], np.ndarray]:
- systems = {s.attrib["prefix"]: s for s in checkpoint.iter(tag="system")}
-
- states = []
- walker_index = 0
- while True:
- system = systems.get("walker-{}".format(walker_index), None)
- if system is not None:
- beads = list(system.iter(tag="beads"))[0]
- natoms = int(beads.attrib["natoms"])
- nbeads = int(beads.attrib["nbeads"])
-
- text = "".join(list(beads.iter(tag="q"))[0].text.split())
- positions = np.array(ast.literal_eval(text))
- positions = positions.reshape(nbeads, natoms, 3) * Bohr
-
- text = "".join(list(beads.iter(tag="names"))[0].text.split())
- text = text.replace(
- ",",
- '","',
- )
- text = text.replace(
- "[",
- '["',
- )
- text = text.replace(
- "]",
- '"]',
- )
- symbols = ast.literal_eval(text)
- numbers = [atomic_numbers[s] for s in symbols]
-
- text = "".join(list(system.iter(tag="cell"))[0].text.split())
- box = (
- np.array(ast.literal_eval(text)).reshape(3, 3).T * Bohr
- ) # transpose for convention
-
- # get current internal system time
- ensemble = list(system.iter(tag="ensemble"))[0]
- conversion = (_hbar / (_e * Ha)) * 1e12
- try:
- time = float(list(ensemble.iter(tag="time"))[0].text) * conversion
- except IndexError:
- time = 0.0
-
- geometry = Geometry.from_data(
- numbers=np.array(numbers),
- positions=np.mean(positions, axis=0),
- cell=box,
- )
- geometry.order["time"] = time
- states.append(geometry)
- walker_index += 1
- else:
- break
- return states
-
-
-def insert_addresses(input_xml):
+def parse_checkpoint(file_xml: str | Path) -> list[Geometry]:
+ """"""
+ from ipi.inputs.system import InputSystem
+ from ipi.utils.io.inputs.io_xml import xml_parse_file
+
+ # see Simulation.load_from_xml
+ xml_sim = xml_parse_file(open(file_xml)).fields[0][1]
+ xml_systems = {v.attribs["prefix"]: v for k, v in xml_sim.fields if k == "system"}
+
+ geometries = []
+ for name, xml_sys in xml_systems.items():
+ sys = InputSystem()
+ sys.parse(xml_sys)
+
+ # get structural properties
+ beads = sys.beads.fetch()
+ natoms, nbeads = beads.natoms, beads.nbeads # average over beads
+ positions = np.asarray(beads.q).reshape(nbeads, natoms, 3).mean(axis=0)
+ numbers = np.array([atomic_numbers[s] for s in beads.names])
+ cell = np.asarray(sys.cell.fetch().h).T # transpose to undo i-Pi convention
+ geometry = Geometry.from_data(numbers, positions * Bohr, cell * Bohr)
+
+ # get current internal system time
+ ensemble_data = sys.ensemble.fetch()
+ time = ensemble_data.time * (_hbar / (_e * Ha)) * 1e12
+ geometry.order["time"], geometry.order["name"] = time, name
+ geometries.append(geometry)
+
+ return geometries
+
+
+def insert_addresses(input_xml: ET.Element) -> None:
+ """Prepend working directory to socket names"""
for child in input_xml:
if child.tag == "ffsocket":
for child_ in child:
@@ -308,165 +63,144 @@ def insert_addresses(input_xml):
address.text = str(Path.cwd() / address.text.strip())
-def insert_data_start(input_xml, data_start):
- for child in input_xml:
- if child.tag == "system":
- initialize = ET.Element("initialize", nbeads="1")
- initialize.text = "start_INDEX.xyz"
-
-
-def anisotropic_barostat_h0(input_xml, data_start):
- path = "system_template/template/system/motion/dynamics/barostat"
- barostat = input_xml.find(path)
- if barostat is not None and (barostat.attrib["mode"] == "anisotropic"):
- h0 = ET.SubElement(barostat, "h0", shape="(3, 3)", units="angstrom")
- cell = np.array(data_start[0].cell)
- cell = cell.flatten(order="F")
- h0.text = " [ {} ] ".format(" , ".join([str(a) for a in cell]))
-
- path = "system_template/template/system/ensemble"
- ensemble = input_xml.find(path)
- assert ensemble is not None
- pressure = ensemble.find("pressure")
- if pressure is not None:
- ensemble.remove(pressure)
- stress = ET.SubElement(ensemble, "stress", units="megapascal")
- stress.text = " [ PRESSURE, 0, 0, 0, PRESSURE, 0, 0, 0, PRESSURE ] "
-
-
-def start(args):
- from ipi.engine.simulation import Simulation
- from ipi.utils.softexit import softexit
-
- data_start = read(args.start_xyz, index=":")
- assert len(data_start) == args.nwalkers
- for i in range(args.nwalkers):
- atoms = data_start[i]
- if not sum(atoms.pbc): # set fake large cell for i-PI
- atoms.pbc = True
- atoms.cell = Cell(NONPERIODIC_CELL)
+
+
+def wait_for_clients(input_xml, timeout: int = 60) -> None:
+ """Make sure clients have initialised successfully"""
+
+ # find sockets opened by server
+ sockets = []
+ xml_str = ET.tostring(input_xml, encoding="unicode")
+ for line in xml_str.splitlines():
+ if 'address' in line:
+ sockets.append(line.split(">")[1].split("<")[0])
+
+ for _ in range(timeout):
+ # check socket status
+ out = subprocess.check_output("ss -xpl | grep psiflow", shell=True).decode()
+ connections = {}
+ for line in out.strip().split("\n"):
+ data = line.split()
+ connections[data[4]] = (data[1], int(data[2]))
+
+ # does every socket have a pending handshake
+ pending = [connections.get(s, [None, 0])[1] for s in sockets]
+ if all([i > 0 for i in pending]):
+ return
+ time.sleep(1)
+
+ msg = f"Timed out waiting for clients to initialise. Sockets: {connections}"
+ raise ConnectionError(msg)
+
+
+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)
+ if not any(at.pbc): # set fake large cell for i-PI
+ at.pbc = True
+ at.cell = Cell(NONPERIODIC_CELL)
write("start_{}.xyz".format(i), data_start[i])
- with open(args.input_xml, "r") as f:
+ with open(input_xml, "r") as f:
input_xml = ET.fromstring(f.read())
- insert_data_start(input_xml, data_start)
insert_addresses(input_xml)
- anisotropic_barostat_h0(input_xml, data_start)
- with open("input.xml", "wb") as f:
+ with open(INPUT_XML, "wb") as f:
f.write(ET.tostring(input_xml, encoding="utf-8"))
- simulation = Simulation.load_from_xml("input.xml", sockets_prefix="")
- try:
- simulation.run()
- softexit.trigger(status="success", message=" @ SIMULATION: Exiting cleanly.")
- except TimeoutException:
- print("simulation timed out -- killing gracefully")
+ simulation = Simulation.load_from_xml(open(INPUT_XML), sockets_prefix="")
+ wait_for_clients(input_xml)
+ simulation.run()
+ return
-def cleanup(args):
+def cleanup(output_xyz: str, output_props: str, output_trajs: str) -> None:
from psiflow.data.utils import _write_frames
- with open("input.xml", "r") as f:
+ print("Starting cleanup")
+ with open(INPUT_XML, "r") as f:
content = f.read()
if "vibrations" in content:
# do stuff
- pass
- else:
- checkpoint = ET.parse("output.checkpoint")
- states = parse_checkpoint(checkpoint)
- for state in states:
- if np.allclose(state.cell, NONPERIODIC_CELL):
- state.cell[:] = 0.0
- _write_frames(*states, outputs=[args.output_xyz])
- if "remd" in content:
- remdsort("input.xml")
- for filepath in glob.glob("SRT_*"):
- # does not use shutil because it is not instantaneous
- source = filepath
- target = filepath.replace("SRT_", "")
- os.remove(target)
- assert not Path(target).exists()
- os.rename(source, target)
- assert Path(target).exists()
- i = 0
- while i < len(states):
- # try all formattings of bead index (i-PI inconsistency)
- paths = [
- Path("walker-{}_output.trajectory_0.extxyz".format(i)),
- Path("walker-{}_output.trajectory_00.extxyz".format(i)),
- Path("walker-{}_output.trajectory_000.extxyz".format(i)),
- Path("walker-{}_output.trajectory_0000.extxyz".format(i)),
- ]
- exists = [p.exists() for p in paths]
- if not sum(exists):
- break
- else:
- assert sum(exists) == 1
- path = paths[exists.index(True)]
- traj = read(path, index=":")
- path.unlink()
- print(states, [s.periodic for s in states])
- for atoms in traj:
- if not states[i].periodic: # load and replace cell
- atoms.pbc = False
- atoms.cell = None
- atoms.info.pop("ipi_comment", None)
- write(paths[0], traj) # always the same path
- i += 1
+ return
+
+ # collect final checkpoint geometries
+ states = parse_checkpoint("output.checkpoint")
+ for state in states:
+ if np.allclose(state.cell, NONPERIODIC_CELL):
+ state.cell[:] = 0.0
+ _write_frames(*states, outputs=[output_xyz])
+ print("Moved checkpoint geometries")
+
+ prefix = ""
+ if "remd" in content:
+ # unshuffle simulation output according to ensemble
+ prefix = "SRT_"
+ out = subprocess.run(
+ "i-pi-remdsort input.xml", shell=True, capture_output=True, text=True
+ )
+ 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)
+ for idx, file in enumerate(output_props):
+ file_src = next(Path.cwd().glob(f"{prefix}*{idx}*.properties"))
+ Path(file).write_text(file_src.read_text())
+ print("Moved simulation observables")
+
+ if len(output_trajs) == 0:
+ return
+
+ # move trajectories
+ assert len(states) == len(output_trajs)
+ for idx, file in enumerate(output_trajs):
+ file_src = next(Path.cwd().glob(f"{prefix}*{idx}*.trajectory*.extxyz"))
+ atoms = read(file_src, ":")
+ periodic = states[idx].periodic
+ for at in atoms:
+ if not periodic: # load and replace cell
+ at.pbc, at.cell = False, None
+ at.info.pop("ipi_comment", None)
+ write(file, atoms)
+ print("Moved simulation trajectories")
+
+ return
def main():
- signal.signal(signal.SIGTERM, timeout_handler)
parser = argparse.ArgumentParser()
- parser.add_argument(
- "--nwalkers",
- type=int,
- default=None,
- )
- parser.add_argument(
- "--input_xml",
- type=str,
- default=None,
- )
- parser.add_argument(
- "--start_xyz",
- type=str,
- default=None,
- )
- parser.add_argument(
- "--output_xyz",
- type=str,
- default=None,
- )
- parser.add_argument(
- "--cleanup",
- action="store_true",
- default=False,
- )
- parser.add_argument(
- "--start",
- type=int,
- default=False,
- )
+ parser.add_argument("--input_xml", required=True)
+ parser.add_argument("--start_xyz", required=True)
+ parser.add_argument("--output_xyz", required=True)
+ parser.add_argument("--output_props", default="")
+ parser.add_argument("--output_trajs", default="")
args = parser.parse_args()
- if not args.cleanup:
- start(args)
- else:
- # try:
- cleanup(args)
- # except BaseException as e: # noqa: B036
- # print(e)
- # print("i-PI cleanup failed!")
- # print("files in directory:")
- # for filepath in Path.cwd().glob("*"):
- # print(filepath)
- # print("")
-
- # names = [p.name for p in Path.cwd().glob("*")]
- # if "output.checkpoint" in names:
- # with open("output.checkpoint", "r") as f:
- # print(f.read())
- # else:
- # print("no output.checkpoint found")
+ time_start = time.time()
+ try:
+ run(args.start_xyz, args.input_xml)
+ softexit.trigger(status="success", message="@PSIFLOW: We are done here.")
+ except ConnectionError:
+ # TODO: in this case, no output files are generated..
+ traceback.print_exc()
+ softexit.trigger(status="bad", message="@PSIFLOW: Clients failed to connect.")
+ except np.linalg.LinAlgError:
+ # some NaN / INF value appeared
+ traceback.print_exc()
+ softexit.trigger(status="bad", message="@PSIFLOW: Simulation went boom.")
+ except SystemExit:
+ # i-Pi intercepts SIG_INT and SIG_TERM by default,
+ # merges all threads and calls sys.exit() before we can clean up
+ pass
+ finally:
+ cleanup(args.output_xyz, args.output_props, args.output_trajs)
+ time_stop = time.time()
+ print("|- Done -|")
+ print(f"Total simulation time: {time_stop - time_start:.0f} seconds")
diff --git a/psiflow/sampling/utils.py b/psiflow/sampling/utils.py
index ddc601d1..72ffb363 100644
--- a/psiflow/sampling/utils.py
+++ b/psiflow/sampling/utils.py
@@ -1,18 +1,13 @@
+import os
+import time
from typing import Any, Optional
import numpy as np
-import typeguard
from ase.data import chemical_symbols
-from psiflow.functions import Function
-
# do not use psiflow apps; parsl config is not loaded in this process!
from psiflow.geometry import Geometry
-# only import stuff which does not issue useless warnings; otherwise
-# python -c 'import .client; print(client.__file__)' is going to be polluted
-# with those import-related warnings
-
class ForceMagnitudeException(Exception):
pass
@@ -26,7 +21,11 @@ def timeout_handler(signum, frame):
raise TimeoutException
-@typeguard.typechecked
+def create_xml_list(items: list[str]) -> str:
+ """Pure helper"""
+ return " [ {} ] ".format(", ".join(items))
+
+
def check_forces(
forces: np.ndarray,
geometry: Any,
@@ -36,79 +35,41 @@ def check_forces(
geometry = Geometry.from_atoms(geometry)
exceeded = np.linalg.norm(forces, axis=1) > max_force
- if np.sum(exceeded):
- indices = np.arange(len(geometry))[exceeded]
- numbers = geometry.per_atom.numbers[exceeded]
- symbols = [chemical_symbols[n] for n in numbers]
- raise ForceMagnitudeException(
- "\nforce exceeded {} eV/A for atoms {}"
- " with chemical elements {}\n".format(
- max_force,
- indices,
- symbols,
- )
+ if not np.sum(exceeded):
+ return
+ indices = np.arange(len(geometry))[exceeded]
+ numbers = geometry.per_atom.numbers[exceeded]
+ symbols = [chemical_symbols[n] for n in numbers]
+ raise ForceMagnitudeException(
+ "\nforce exceeded {} eV/A for atoms {}"
+ " with chemical elements {}\n".format(
+ max_force,
+ indices,
+ symbols,
)
- else:
- pass
-
-
-class FunctionDriver:
-
- def __init__(
- self,
- template: Geometry,
- function: Function,
- max_force: Optional[float],
- verbose: bool = True, # used by i-PI internally?
- error_msg="",
- ):
- self.verbose = verbose
- self.template = template
- self.function = function
- self.max_force = max_force
-
- def check_arguments(self):
- pass
-
- def __call__(self, cell, pos):
- from ipi.utils.units import unit_to_internal, unit_to_user
-
- pos = unit_to_user("length", "angstrom", pos)
- cell = unit_to_user("length", "angstrom", cell.T)
-
- self.template.per_atom.positions[:] = pos
- if self.template.periodic:
- self.template.cell[:] = cell
-
- outputs = self.function(self.template)
- energy = outputs["energy"]
- forces = outputs["forces"]
- stress = outputs["stress"]
-
- # check for max_force
- if self.max_force is not None:
- check_forces(forces, self.template, self.max_force)
-
- # converts to internal quantities
- pot_ipi = np.asarray(
- unit_to_internal("energy", "electronvolt", energy), np.float64
- )
- force_ipi = np.asarray(unit_to_internal("force", "ev/ang", forces), np.float64)
- vir_calc = -stress * self.template.volume
- vir_ipi = np.array(
- unit_to_internal("energy", "electronvolt", vir_calc.T), dtype=np.float64
- )
- extras = ""
-
- return pot_ipi, force_ipi, vir_ipi, extras
-
-
-def initialise_driver(driver) -> None:
- """This function does not do anything yet. Stay tuned."""
- pass
-
-
-def check_output(data: dict) -> None:
- """This function does not do anything yet. Stay tuned."""
- pass
-
+ )
+
+
+def initialise_driver(driver: 'Psiflow_driver') -> None:
+ """"""
+ function = driver.function
+ name = function.__class__.__name__
+ affinity = os.sched_getaffinity(os.getpid())
+ t0 = time.time()
+ for _ in range(10):
+ function(driver.geometry) # torch warm-up before simulation
+ t1 = time.time()
+ msg = [
+ "- Psiflow -",
+ f"Initialising driver for {name} with options {driver.kwargs}",
+ f"CPU affinity [PID {os.getpid()}]: {affinity}",
+ f"Time for 10 evaluations: {t1 - t0:.3f}",
+ "- - - - - -",
+ ]
+ print("\n".join(msg))
+
+
+def check_output(driver: 'Psiflow_driver', data: dict) -> None:
+ """"""
+ if max_force := driver.kwargs.get("max_force"):
+ check_forces(data["forces"], driver.geometry, max_force)
diff --git a/psiflow/sampling/walker.py b/psiflow/sampling/walker.py
index c98cf5ca..0b2cdec6 100644
--- a/psiflow/sampling/walker.py
+++ b/psiflow/sampling/walker.py
@@ -1,10 +1,10 @@
-from __future__ import annotations # necessary for type-guarding class methods
-
import xml.etree.ElementTree as ET
-from typing import Optional, Union
+from typing import Optional, Union, Callable
+from enum import Enum
+from dataclasses import dataclass, field
+from copy import deepcopy
import numpy as np
-import typeguard
from parsl.app.app import python_app
from parsl.app.futures import DataFuture
from parsl.data_provider.files import File
@@ -13,136 +13,141 @@
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, combine_hamiltonians
from psiflow.order_parameters import OrderParameter
from psiflow.sampling.metadynamics import Metadynamics
from psiflow.utils.apps import copy_app_future
-@typeguard.typechecked
class Coupling:
pass
-@typeguard.typechecked
-def _conditioned_reset(
- condition: bool,
+class Ensemble(Enum):
+ NVE = "NVE"
+ NVT = "NVT"
+ NPT = "NPT"
+ NVST = "(N, V, sigma=0, T)"
+
+
+def _conditional_reset(
state: Geometry,
start: Geometry,
+ flag: bool | None,
+ condition: Callable | None,
+ *args,
) -> Geometry:
- from copy import deepcopy # copy necessary!
-
- if condition:
- return deepcopy(start)
- else:
- return state
-
-
-conditioned_reset = python_app(_conditioned_reset, executors=["default_threads"])
+ """Reset geometry based on a flag or boolean condition function."""
+ if condition is not None:
+ flag = condition(args)
+
+ # copy necessary!
+ return deepcopy(start) if flag else state
+
+
+conditional_reset = python_app(_conditional_reset, executors=["default_threads"])
+
+
+def get_ensemble_kwargs(walker: "Walker") -> dict:
+ """Extract all walker properties that define the MD sampling ensemble."""
+ return dict(
+ hamiltonian=walker.hamiltonian,
+ timestep=walker.timestep,
+ temperature=walker.temperature,
+ pressure=walker.pressure,
+ tau_thermostat=walker.tau_thermostat,
+ tau_barostat=walker.tau_barostat,
+ volume_constrained=walker.volume_constrained,
+ masses=walker.masses,
+ nbeads=walker.nbeads,
+ metadynamics=walker.metadynamics,
+ )
-@typeguard.typechecked
@psiflow.serializable
+@dataclass
class Walker:
start: Union[Geometry, AppFuture]
- hamiltonian: Optional[Hamiltonian]
- state: Union[Geometry, AppFuture, None]
- temperature: Optional[float]
- pressure: Optional[float]
- masses: Union[np.ndarray, float, None]
- nbeads: int
- timestep: float
- coupling: Optional[Coupling]
- metadynamics: Optional[Metadynamics]
- order_parameter: Optional[OrderParameter]
-
- def __init__(
- self,
- start: Union[Geometry, AppFuture],
- hamiltonian: Optional[Hamiltonian] = None,
- state: Union[Geometry, AppFuture, None] = None,
- temperature: Optional[float] = 300,
- pressure: Optional[float] = None,
- masses: Union[np.ndarray, float, None] = None,
- nbeads: int = 1,
- timestep: float = 0.5,
- metadynamics: Optional[Metadynamics] = None,
- order_parameter: Optional[OrderParameter] = None,
- ):
- self.start = start
- if hamiltonian is None:
- hamiltonian = Zero()
- self.hamiltonian = hamiltonian
- if state is None:
- state = copy_app_future(self.start)
- self.state = state
-
- if order_parameter is not None:
- self.start = order_parameter.evaluate(self.start)
-
- self.temperature = temperature
- self.pressure = pressure
-
- if isinstance(masses, (float, int)):
- masses *= self.start.atomic_masses
- self.masses = masses
- if self.masses is not None:
- assert len(self.masses) == len(self.start), "Masses do not match number of atoms"
-
- self.nbeads = nbeads
- self.timestep = timestep
-
- self.metadynamics = metadynamics
- self.order_parameter = order_parameter
+ hamiltonian: Hamiltonian = Zero()
+ timestep: float = 0.5
+ temperature: Optional[float] = 300
+ pressure: Optional[float] = None
+ tau_thermostat: float = 100
+ tau_barostat: float = 300
+ volume_constrained: bool = False
+ masses: Union[np.ndarray, float, None] = None
+ nbeads: int = 1
+ metadynamics: Optional[Metadynamics] = None
+ order_parameter: Optional[OrderParameter] = None
+
+ state: Union[Geometry, AppFuture] = field(init=False)
+ coupling: Optional[Coupling] = field(init=False)
+
+ def __post_init__(self):
+ self.state = copy_app_future(self.start)
self.coupling = None
- def reset(self, condition: Union[AppFuture, bool] = True):
- self.state = conditioned_reset(condition, self.state, self.start)
+ if self.temperature is None: # NVE
+ assert self.pressure is None and not self.volume_constrained
+ if self.volume_constrained:
+ self.pressure = 0 # TODO: warning?
+
+ if isinstance(self.start, Geometry) and not self.start.periodic:
+ # we cannot check this for futures
+ assert self.pressure is None, "Pressure requires PBC"
+
+ if self.order_parameter is not None:
+ # TODO: order_parameter out of commission
+ self.start = self.order_parameter.evaluate(self.start)
+
+ if (m := self.masses) is None:
+ pass # do nothing
+ elif isinstance(m, (float, int)):
+ self.masses = self.start.atomic_masses * m
+ elif isinstance(m, np.ndarray) and len(m) != len(self.start):
+ raise ValueError("Supplied masses do not match number of atoms")
+
+ def reset(self):
+ self.state = conditional_reset(self.state, self.start, True, None)
+
+ def conditional_reset(
+ self, flag: bool | None = None, condition: Callable | None = None, *args
+ ):
+ assert (flag is None) != (condition is None) # xor
+ self.state = conditional_reset(self.state, self.start, flag, condition, *args)
def is_reset(self) -> AppFuture:
return check_equality(self.start, self.state)
- def multiply(self, nreplicas: int) -> list[Walker]:
+ def multiply(self, nreplicas: int) -> list["Walker"]:
if self.coupling is not None:
raise ValueError("Cannot multiply walkers after they are coupled")
walkers = []
+ kwargs = get_ensemble_kwargs(self)
for _i in range(nreplicas):
- if self.metadynamics is not None:
- metadynamics = self.metadynamics.copy()
- else:
- metadynamics = None
- walker = Walker(
- start=self.start,
- hamiltonian=self.hamiltonian,
- state=self.state,
- temperature=self.temperature,
- pressure=self.pressure,
- masses=self.masses,
- nbeads=self.nbeads,
- timestep=self.timestep,
- metadynamics=metadynamics,
- ) # no coupling
+ if kwargs["metadynamics"] is not None:
+ kwargs["metadynamics"] = kwargs["metadynamics"].copy()
+ walker = Walker(start=self.start, **kwargs) # no coupling
+ walker.state = self.state
walkers.append(walker)
return walkers
@property
- def pimd(self):
+ def pimd(self) -> bool:
return self.nbeads != 1
@property
- def nve(self):
- return (self.temperature is None) and (self.pressure is None)
-
- @property
- def nvt(self):
- return (self.temperature is not None) and (self.pressure is None)
-
- @property
- def npt(self):
- return (self.temperature is not None) and (self.pressure is not None)
+ def ensemble(self) -> Ensemble:
+ if self.temperature is None:
+ return Ensemble.NVE
+ elif self.pressure is None:
+ return Ensemble.NVT
+ elif self.volume_constrained:
+ return Ensemble.NVST
+ else:
+ return Ensemble.NPT
-@typeguard.typechecked
def partition(walkers: list[Walker]) -> list[list[int]]:
indices = []
for i, walker in enumerate(walkers):
@@ -158,50 +163,36 @@ def partition(walkers: list[Walker]) -> list[list[int]]:
return indices
-# typeguarding incompatible with * expansion
def _get_minimum_energy_states(
- coefficients: np.ndarray,
- *energies: np.ndarray,
+ coefficients: np.ndarray, *energies: np.ndarray
) -> list[int]:
- import numpy
-
assert len(coefficients.shape) == 2
assert len(energies) == coefficients.shape[1]
- energies = numpy.array(energies)
-
indices = []
- for c in coefficients:
- energy = numpy.sum(c.reshape(-1, 1) * energies, axis=0)
- indices.append(int(numpy.argmin(energy)))
+ energies = np.array(energies)
+ for c in coefficients: # every c corresponds to a Walker hamiltonian
+ energy = np.sum(c.reshape(-1, 1) * energies, axis=0)
+ indices.append(int(np.argmin(energy)))
return indices
get_minimum_energy_states = python_app(
- _get_minimum_energy_states,
- executors=["default_threads"],
+ _get_minimum_energy_states, executors=["default_threads"]
)
-@typeguard.typechecked
def quench(walkers: list[Walker], dataset: Dataset) -> None:
- all_hamiltonians = sum([w.hamiltonian for w in walkers], start=Zero())
- energies = [h.compute(dataset, "energy") for h in all_hamiltonians.hamiltonians]
-
- coefficients = []
- for walker in walkers:
- c = all_hamiltonians.get_coefficients(1.0 * walker.hamiltonian)
- assert c is not None
- coefficients.append(c)
- coefficients = np.array(coefficients)
-
- indices = get_minimum_energy_states(
- coefficients,
- *energies,
- )
- data = dataset[indices]
- for i, walker in enumerate(walkers):
- walker.start = data[i]
+ """Assign the lowest energy geometry in dataset to every walker"""
+ hamiltonians = combine_hamiltonians([w.hamiltonian for w in walkers])
+ energies = [h.compute(dataset, "energy") for h in hamiltonians.hamiltonians]
+ coefficients = [
+ hamiltonians.get_coefficients(walker.hamiltonian * 1.0) for walker in walkers
+ ]
+ indices = get_minimum_energy_states(np.array(coefficients), *energies)
+ geometries = dataset.geometries()
+ for walker, idx in zip(walkers, indices):
+ walker.start = geometries[idx]
walker.reset()
@@ -213,16 +204,15 @@ def _random_indices(nindices: int, nstates: int) -> list[int]:
random_indices = python_app(_random_indices, executors=["default_threads"])
-@typeguard.typechecked
def randomize(walkers: list[Walker], dataset: Dataset) -> None:
+ """Randomly assign geometries from dataset to walkers"""
indices = random_indices(len(walkers), dataset.length())
- data = dataset[indices]
- for i, walker in enumerate(walkers):
- walker.start = data[i]
+ geometries = dataset.geometries()
+ for walker, idx in zip(walkers, indices):
+ walker.start = geometries[idx]
walker.reset()
-@typeguard.typechecked
def validate_coupling(walkers: list[Walker]):
couplings = []
counts = []
@@ -240,7 +230,6 @@ def validate_coupling(walkers: list[Walker]):
assert coupling.nwalkers == counts[i]
-@typeguard.typechecked
@psiflow.serializable
class ReplicaExchange(Coupling):
trial_frequency: int
@@ -259,7 +248,7 @@ def __init__(
self.swapfile = psiflow.context().new_file("swap_", ".txt")
self.nwalkers = nwalkers
- def __eq__(self, other: Optional[ReplicaExchange]) -> bool:
+ def __eq__(self, other: Optional["ReplicaExchange"]) -> bool:
if other is None:
return False
trial = self.trial_frequency == other.trial_frequency
diff --git a/psiflow/serialization.py b/psiflow/serialization.py
index cf863259..c37eea69 100644
--- a/psiflow/serialization.py
+++ b/psiflow/serialization.py
@@ -4,6 +4,7 @@
import json
from pathlib import Path
from typing import ClassVar, Optional, Union, get_args, get_origin, get_type_hints
+from dataclasses import InitVar
import typeguard
from parsl.app.app import python_app
@@ -46,7 +47,7 @@ def setter(self, value: type_hint) -> None:
return setter
-def update_init(init_func): # TODO: why not have a Mixin class instead?
+def update_init(init_func): # TODO: why not have a Mixin class instead?
def wrapper(self, *args, **kwargs):
self._geoms = {}
self._files = {}
@@ -61,6 +62,7 @@ def serializable(cls):
"""decorator to make class serializable"""
class_dict = dict(cls.__dict__)
for name, type_hint in get_type_hints(cls).items():
+ kind = None
if get_origin(type_hint) in [Union, Optional, list, tuple]:
args = get_args(type_hint)
if (File in args) or (DataFuture in args):
@@ -75,9 +77,18 @@ def serializable(cls):
if issubclass(arg, Serializable): # weird
kind = "serial"
else:
- if get_origin(type_hint) is ClassVar:
+ # TODO: temporary hotfixes all around
+ origin = get_origin(type_hint)
+ if origin is ClassVar:
continue # do nothing for classvars
- if not inspect.isclass(type_hint):
+ elif origin == dict:
+ kind = "attrs"
+ elif isinstance(type_hint, str) and type_hint.startswith("dataclasses"):
+ continue
+ elif isinstance(type_hint, InitVar):
+ continue
+
+ if kind is None and not inspect.isclass(type_hint):
raise ValueError(
"{} is formally not a class ({})".format(type_hint, name)
)
@@ -269,7 +280,7 @@ def deserialize(data_str: str, custom_cls: Optional[list] = None):
from psiflow.metrics import Metrics
from psiflow.models import MACE
from psiflow.order_parameters import OrderParameter
- from psiflow.reference import CP2K, GPAW, ORCA
+ from psiflow.reference import CP2K, GPAW, ORCA, ReferenceDummy
from psiflow.sampling import Metadynamics, ReplicaExchange, SimulationOutput, Walker
SERIALIZABLES = {}
@@ -281,6 +292,7 @@ def deserialize(data_str: str, custom_cls: Optional[list] = None):
CP2K,
GPAW,
ORCA,
+ ReferenceDummy,
Zero,
MACEHamiltonian,
EinsteinCrystal,
diff --git a/psiflow/utils/__init__.py b/psiflow/utils/__init__.py
index 915cdc82..e103e43e 100644
--- a/psiflow/utils/__init__.py
+++ b/psiflow/utils/__init__.py
@@ -2,3 +2,7 @@
CD_COMMAND = 'cd $tmpdir; echo "tmpdir: $PWD"'
+def export_env_command(env_vars: dict) -> str:
+ return "export " + " ".join(
+ [f"{name}={value}" for name, value in env_vars.items()]
+ )
diff --git a/psiflow/utils/io.py b/psiflow/utils/io.py
index 80dc475a..c5e93855 100644
--- a/psiflow/utils/io.py
+++ b/psiflow/utils/io.py
@@ -1,3 +1,4 @@
+import json
import xml.etree.ElementTree as ET
from typing import Any
@@ -105,10 +106,6 @@ def _dump_json(
outputs: list = [],
**kwargs,
) -> None:
- import json
-
- import numpy as np
-
def convert_to_list(array):
if not type(array) is np.ndarray:
if type(array) is np.floating:
@@ -119,6 +116,7 @@ def convert_to_list(array):
as_list.append(convert_to_list(item))
return as_list
+ assert len(outputs) == 1
for name in list(kwargs.keys()):
value = kwargs[name]
if type(value) is np.ndarray:
@@ -131,3 +129,26 @@ def convert_to_list(array):
dump_json = python_app(_dump_json, executors=["default_threads"])
+
+
+def _save_npz(
+ data: dict[str, np.ndarray], outputs: list[File], **kwargs: np.ndarray
+) -> None:
+ """"""
+ assert len(outputs) == 1
+ np.savez(outputs[0].filepath, **(data | kwargs))
+
+
+save_npz = python_app(_save_npz, executors=["default_threads"])
+
+
+def _load_npz(file_in: File) -> dict[str, Any]:
+ """"""
+ data: dict[str, np.ndarray] = dict(np.load(file_in))
+ for k, v in data.items():
+ if v.size == 1:
+ data[k] = v.item() # convert to native python type
+ return data
+
+
+load_npz = python_app(_load_npz, executors=["default_threads"])
diff --git a/psiflow/utils/parse.py b/psiflow/utils/parse.py
index f60a77a5..99fbcd72 100644
--- a/psiflow/utils/parse.py
+++ b/psiflow/utils/parse.py
@@ -56,3 +56,8 @@ def get_task_logs(task_id: int) -> tuple[Path, Path]:
stdout = next(path.rglob(f"task_{task_id}_*.stdout"))
stderr = next(path.rglob(f"task_{task_id}_*.stderr"))
return stdout, stderr
+
+
+def get_task_name_id(logfile: str) -> tuple[str, str]:
+ _, task_id, task_name = Path(logfile).stem.split("_", maxsplit=2)
+ return task_name, task_id
diff --git a/pyproject.toml b/pyproject.toml
index 4d7092f1..1f20c233 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
[project]
name = "psiflow"
-version = "4.0.0"
+version = "4.0.1"
description = "Library for developing interatomic potentials"
readme = "README.md"
requires-python = ">=3.10"
@@ -17,6 +17,7 @@ dependencies = [
"prettytable",
"psutil",
"cp2k-input-tools @ git+https://github.com/cp2k/cp2k-input-tools.git@3b9929735dcb3c8c0620a548b1fe20efecbad077", # need 2024.1
+ "ipi @ git+https://github.com/i-pi/i-pi.git@v3.1.10",
"pytimeparse",
]
@@ -25,7 +26,6 @@ dependencies = [
psiflow-client = "psiflow.sampling.client:main"
psiflow-server = "psiflow.sampling.server:main"
psiflow-mace-train = "psiflow.models.mace_utils:main"
-psiflow-ase-opt = "psiflow.sampling._ase:main"
[project.optional-dependencies]
diff --git a/tests/conftest.py b/tests/conftest.py
index 5ed9039b..e09e86e0 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -133,8 +133,8 @@ def dataset_h2(context):
@pytest.fixture
-def checkpoint():
- checkpoint_str = """
+def checkpoint() -> str:
+ return """
"""
- return ET.ElementTree(element=ET.fromstring(checkpoint_str))
diff --git a/tests/test_free_energy.py b/tests/test_free_energy.py
index fa986520..e78ea599 100644
--- a/tests/test_free_energy.py
+++ b/tests/test_free_energy.py
@@ -18,7 +18,7 @@ def test_integration_simple(dataset):
geometry = optimize(
dataset[3],
einstein,
- mode='fix_cell',
+ mode="fix_cell",
f_max=1e-4,
)
hessian = compute_harmonic(
@@ -124,17 +124,12 @@ def test_dihydrogen(dataset_h2):
optimized = optimize(
geometry,
hamiltonian,
- mode='fix_cell',
+ mode="fix_cell",
f_max=1e-4,
).result()
assert optimized.energy is not None
assert np.linalg.norm(optimized.per_atom.forces) < 1e-2
- hessian = compute_harmonic(
- optimized,
- hamiltonian,
- asr="crystal",
- pos_shift=0.001,
- )
+ hessian = compute_harmonic(optimized, hamiltonian, asr="crystal", pos_shift=0.001)
frequencies = compute_frequencies(hessian, geometry).result()
# check that highest frequency in inv cm corresponds to 3500 - 4000
frequencies_invcm = (frequencies * second) / (_c * 1e2) # in invcm
diff --git a/tests/test_models.py b/tests/test_models.py
index 5fc0ae72..04275e4c 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -46,22 +46,13 @@ def test_mace_init(mace_config, dataset):
assert not np.any(np.allclose(energies, 0.0))
energy_Cu = 3
energy_H = 7
- atomic_energies = {
- "Cu": energy_Cu,
- "H": energy_H,
- }
- hamiltonian = MACEHamiltonian(
- hamiltonian.external,
- atomic_energies=atomic_energies,
- )
+ atomic_energies = {"Cu": energy_Cu, "H": energy_H}
+ hamiltonian = MACEHamiltonian(hamiltonian.external, atomic_energies=atomic_energies)
assert hamiltonian != model.create_hamiltonian() # atomic energies
- evaluated = dataset.evaluate(hamiltonian)
+ evaluated = dataset.evaluate(hamiltonian).subtract_offset(Cu=energy_Cu, H=energy_H).geometries().result()
for i in range(nstates):
- assert np.allclose(
- energies[i],
- evaluated.subtract_offset(Cu=energy_Cu, H=energy_H)[i].result().energy,
- )
+ assert np.allclose(energies[i], evaluated[i].energy)
energies = hamiltonian.compute(dataset, "energy").result()
second = psiflow.deserialize(psiflow.serialize(hamiltonian).result())
diff --git a/tests/test_sampling.py b/tests/test_sampling.py
index fc707747..7d955c6a 100644
--- a/tests/test_sampling.py
+++ b/tests/test_sampling.py
@@ -8,8 +8,15 @@
from psiflow.geometry import check_equality
from psiflow.hamiltonians import EinsteinCrystal, PlumedHamiltonian
from psiflow.models import MACE
-from psiflow.sampling.optimize import optimize as optimize_ipi, optimize_dataset as optimize_dataset_ipi
-from psiflow.sampling.ase import optimize as optimize_ase, optimize_dataset as optimize_dataset_ase
+from psiflow.sampling.optimize import (
+ optimize as optimize_ipi,
+ optimize_dataset as optimize_dataset_ipi,
+)
+from psiflow.sampling.ase import (
+ optimize as optimize_ase,
+ optimize_dataset as optimize_dataset_ase,
+ file_ase as file_script_ase,
+)
from psiflow.sampling.metadynamics import Metadynamics
from psiflow.sampling.sampling import sample, template, EnsembleTable
from psiflow.sampling.server import parse_checkpoint
@@ -19,7 +26,9 @@
quench,
randomize,
replica_exchange,
+ Ensemble,
)
+from psiflow.sampling.output import Status
def test_walkers(dataset):
@@ -36,25 +45,22 @@ def test_walkers(dataset):
einstein = EinsteinCrystal(dataset[0], force_constant=0.1)
einstein_ = EinsteinCrystal(dataset[0], force_constant=0.2)
walker = Walker(dataset[0], einstein, temperature=300, metadynamics=mtd0)
- assert walker.nvt
- assert not walker.npt
+ assert walker.ensemble == Ensemble.NVT
assert not walker.pimd
- walkers = [walker]
- walkers.append(Walker(dataset[0], 0.5 * einstein_, nbeads=4, metadynamics=mtd1))
- walkers.append(Walker(dataset[0], einstein + plumed, nbeads=4))
- walkers.append(
- Walker(dataset[0], einstein, pressure=0, temperature=300, metadynamics=mtd1)
- )
- walkers.append(
- Walker(dataset[0], einstein_, pressure=100, temperature=600, metadynamics=mtd1)
- )
- walkers.append(Walker(dataset[0], einstein, temperature=600, metadynamics=mtd1))
+ walkers = [
+ walker,
+ Walker(dataset[0], 0.5 * einstein_, nbeads=4, metadynamics=mtd1),
+ Walker(dataset[0], einstein + plumed, nbeads=4),
+ Walker(dataset[0], einstein, pressure=0, temperature=300, metadynamics=mtd1),
+ Walker(dataset[0], einstein_, pressure=100, temperature=600, metadynamics=mtd1),
+ Walker(dataset[0], einstein, temperature=600, metadynamics=mtd1),
+ ]
- # nvt
+ # NVT
_walkers = [walkers[0], walkers[-1]]
hamiltonian_components, ensemble_table, plumed_list = template(_walkers)
- assert _walkers[0].nvt
+ assert _walkers[0].ensemble == Ensemble.NVT
assert len(hamiltonian_components) == 1
keys_ref = ("TEMP", "EinsteinCrystal0", "METAD0", "METAD1")
weights_ref = np.array([[300, 1, 1, 0], [600, 1, 0, 1]])
@@ -65,7 +71,7 @@ def test_walkers(dataset):
_walkers[0].metadynamics = None
_walkers[1].metadynamics = None
hamiltonian_components, ensemble_table, plumed_list = template(_walkers)
- assert _walkers[0].nvt
+ assert _walkers[0].ensemble == Ensemble.NVT
assert len(hamiltonian_components) == 1
keys_ref = ("TEMP", "EinsteinCrystal0")
weights_ref = np.array([[300, 1], [600, 1]])
@@ -77,7 +83,13 @@ def test_walkers(dataset):
hamiltonian_components, ensemble_table, plumed_list = template(_walkers)
assert _walkers[0].pimd
assert len(hamiltonian_components) == 3
- keys_ref = ("TEMP", "EinsteinCrystal0", "EinsteinCrystal1", "PlumedHamiltonian0", "METAD0")
+ keys_ref = (
+ "TEMP",
+ "EinsteinCrystal0",
+ "EinsteinCrystal1",
+ "PlumedHamiltonian0",
+ "METAD0",
+ )
weights_ref = np.array([[300, 0.0, 0.5, 0.0, 1.0], [300, 1.0, 0.0, 1.0, 0.0]])
assert ensemble_table == EnsembleTable(keys_ref, weights_ref)
@@ -87,15 +99,26 @@ def test_walkers(dataset):
_ = template(_walkers)
_walkers[0].metadynamics = Metadynamics("METAD: FILE=bla")
hamiltonian_components, ensemble_table, plumed_list = template(_walkers)
- assert _walkers[0].npt
+ assert _walkers[0].ensemble == Ensemble.NPT
assert len(hamiltonian_components) == 2
- keys_ref = ("TEMP", "PRESSURE", "EinsteinCrystal0", "EinsteinCrystal1", "METAD0", "METAD1")
- weights_ref = np.array([[300, 0, 0.0, 1.0, 1.0, 0.0], [600, 100, 1.0, 0.0, 0.0, 1.0]])
+ keys_ref = (
+ "TEMP",
+ "PRESSURE",
+ "EinsteinCrystal0",
+ "EinsteinCrystal1",
+ "METAD0",
+ "METAD1",
+ )
+ weights_ref = np.array(
+ [[300, 0, 0.0, 1.0, 1.0, 0.0], [600, 100, 1.0, 0.0, 0.0, 1.0]]
+ )
assert ensemble_table == EnsembleTable(keys_ref, weights_ref)
def test_parse_checkpoint(checkpoint):
- states = parse_checkpoint(checkpoint)
+ file = psiflow.context().new_file("input_", ".xml")
+ Path(file.filepath).write_text(checkpoint)
+ states = parse_checkpoint(file.filepath)
assert "time" in states[0].order
assert np.allclose(
states[0].cell,
@@ -118,44 +141,32 @@ def test_sample(dataset, mace_config):
METAD ARG=CV PACE=5 SIGMA=0.05 HEIGHT=5
"""
metadynamics = Metadynamics(plumed_str)
- walker0 = Walker(
- start=dataset[0],
- temperature=300,
- pressure=None,
- metadynamics=metadynamics,
- hamiltonian=0.9 * plumed + einstein,
- )
- walker1 = Walker(
- start=dataset[0],
- temperature=600,
- pressure=None,
- hamiltonian=einstein,
- )
- simulation_outputs = sample(
- [walker0, walker1],
- steps=100,
- step=10,
- )
- assert simulation_outputs[0].trajectory.length().result() == 11
- pos0 = simulation_outputs[0].trajectory[-1].result().per_atom.positions
- pos1 = walker0.state.result().per_atom.positions
- assert np.allclose(
- pos0,
- pos1,
- )
- assert np.allclose(
- simulation_outputs[0].trajectory[-1].result().cell,
- walker0.start.result().cell,
- )
- e = simulation_outputs[0]["potential{electronvolt}"].result()
- assert len(e) == 11
-
- o = sample([walker0], steps=20, step=2, start=10, keep_trajectory=False)[0]
- assert o.trajectory is None
- assert len(o["potential{electronvolt}"].result()) == 6
- o = sample([walker0], steps=20, step=2, start=10, keep_trajectory=True)[0]
- assert o.trajectory.length().result() == 6
- assert len(o["potential{electronvolt}"].result()) == 6
+ walkers = [
+ Walker(
+ start=dataset[0],
+ temperature=300,
+ metadynamics=metadynamics,
+ hamiltonian=0.9 * plumed + einstein,
+ ),
+ Walker(start=dataset[0], temperature=600, hamiltonian=einstein),
+ Walker(start=dataset[0], temperature=600, hamiltonian=einstein),
+ Walker(start=dataset[0], temperature=600, hamiltonian=einstein),
+ ]
+ outputs = sample(walkers[:2], steps=50, step=5)
+ [o1] = sample(walkers[2:3], steps=20, step=2, start=10, keep_trajectory=False)
+ [o2] = sample(walkers[3:], steps=20, step=2, start=10, keep_trajectory=True)
+
+ assert o1.trajectory is None
+ assert len(o1["potential{electronvolt}"].result()) == 6
+ assert o2.trajectory.length().result() == 6
+ assert len(o2["potential{electronvolt}"].result()) == 6
+
+ assert outputs[0].trajectory.length().result() == 11
+ pos0 = outputs[0].trajectory[-1].result().per_atom.positions
+ pos1 = walkers[0].state.result().per_atom.positions
+ assert np.allclose(pos0, pos1)
+ assert check_equality(walkers[0].state, outputs[0].state).result()
+ assert len(outputs[0]["potential{electronvolt}"].result()) == 11
# check whether metadynamics file has correct dependency
with open(metadynamics.external.result().filepath, "r") as f:
@@ -163,70 +174,39 @@ def test_sample(dataset, mace_config):
nhills = len(content.split("\n"))
assert nhills > 3
- assert len(simulation_outputs) == 2
+ assert len(outputs) == 2
energies = [
- simulation_outputs[0]["potential{electronvolt}"].result(),
- simulation_outputs[1]["potential{electronvolt}"].result(),
+ outputs[0]["potential{electronvolt}"].result(),
+ outputs[1]["potential{electronvolt}"].result(),
]
energies_ = [
- (0.9 * plumed + einstein).compute(simulation_outputs[0].trajectory, "energy"),
- einstein.compute(simulation_outputs[1].trajectory, "energy"),
+ (0.9 * plumed + einstein).compute(outputs[0].trajectory, "energy"),
+ einstein.compute(outputs[1].trajectory, "energy"),
]
assert len(energies[0]) == len(energies_[0].result())
- assert np.allclose(
- energies[0],
- energies_[0].result(),
- )
- time = simulation_outputs[0]["time{picosecond}"].result()
- assert np.allclose(
- time,
- np.arange(11) * 5e-4 * 10,
- )
+ assert np.allclose(energies[0], energies_[0].result())
+ time = outputs[0]["time{picosecond}"].result()
+ assert np.allclose(time, np.arange(0, 51, 5) * 5e-4)
# check pot components
assert np.allclose(
- simulation_outputs[1].get_energy(walker1.hamiltonian).result(),
- simulation_outputs[1]["potential{electronvolt}"].result(),
+ outputs[1].get_energy(walkers[1].hamiltonian).result(),
+ outputs[1]["potential{electronvolt}"].result(),
)
assert np.allclose(
- simulation_outputs[0].get_energy(walker0.hamiltonian).result(),
- simulation_outputs[0]["potential{electronvolt}"].result(),
- )
-
- manual_total = simulation_outputs[0].get_energy(plumed).result() * 0.9
- manual_total += simulation_outputs[0].get_energy(einstein).result()
- assert np.allclose(
- manual_total,
- simulation_outputs[0]["potential{electronvolt}"].result(),
- )
-
- # check PIMD output
- walker = Walker(
- start=dataset[2],
- temperature=200,
- pressure=None,
- hamiltonian=einstein,
- nbeads=11,
+ outputs[0].get_energy(walkers[0].hamiltonian).result(),
+ outputs[0]["potential{electronvolt}"].result(),
)
- output = sample([walker], steps=10, step=5)[0]
- for state in output.trajectory.geometries().result():
- assert len(state) == len(dataset[2].result())
- assert output.trajectory.length().result() == 3
- assert output.temperature is not None
- assert np.abs(output.temperature.result() - 200 < 200)
+ manual_total = outputs[0].get_energy(plumed).result() * 0.9
+ manual_total += outputs[0].get_energy(einstein).result()
+ assert np.allclose(manual_total, outputs[0]["potential{electronvolt}"].result())
- simulation_outputs = sample(
- [walker0, walker1], steps=100, step=10, observables=["ensemble_bias"]
- )
- bias = np.stack(
- [
- simulation_outputs[0]["ensemble_bias"].result(),
- simulation_outputs[1]["ensemble_bias"].result(),
- ],
- axis=0,
- )
- assert np.allclose(bias[1, :], 0)
- assert np.all(bias[0, :] > 0)
+ observable = "ensemble_bias{electronvolt}"
+ outputs = sample(walkers[:2], steps=100, step=10, observables=[observable])
+ bias0 = outputs[0][observable].result()
+ bias1 = outputs[1][observable].result()
+ assert np.allclose(bias1, 0)
+ assert np.all(bias0 > 0)
# check that old hills are there too
with open(metadynamics.external.result().filepath, "r") as f:
@@ -235,121 +215,138 @@ def test_sample(dataset, mace_config):
assert new_nhills > 3
assert new_nhills > nhills
+ # check PIMD output
+ geom_start = dataset[2]
+ walker = Walker(start=geom_start, temperature=200, hamiltonian=einstein, nbeads=6)
+ [output] = sample([walker], steps=10, step=5)
+ for state in output.trajectory.geometries().result():
+ assert len(state) == len(geom_start.result())
+ assert output.trajectory.length().result() == 3
+ assert output.temperature is not None
+ assert np.abs(output.temperature.result() - 200 < 200) # what?
+
model = MACE(**mace_config)
model.initialize(dataset[:3])
hamiltonian = model.create_hamiltonian()
- walker = Walker(
- start=dataset[0],
- temperature=600,
- pressure=None,
- hamiltonian=hamiltonian,
- )
+ walker = Walker(start=dataset[0], temperature=600, hamiltonian=hamiltonian)
walker.state = dataset[1]
- simulation_output = sample(
+ [output] = sample(
[walker],
steps=10,
- observables=[
- "potential{electronvolt}",
- "kinetic_md{electronvolt}",
- "temperature{kelvin}",
- ],
+ observables=["kinetic_md{electronvolt}"],
checkpoint_step=1,
fix_com=True, # otherwise temperature won't match
- )[0]
+ )
assert np.allclose(
hamiltonian.compute(dataset, "energy").result()[0],
- simulation_output["potential{electronvolt}"].result()[0],
+ output["potential{electronvolt}"].result()[0],
atol=1e-3,
)
- assert np.allclose(
- simulation_output.time.result(),
- 10 * 0.5 / 1000,
- )
- T0 = simulation_output.temperature.result()
- T1 = simulation_output["temperature{kelvin}"].result()[-1]
- assert np.allclose(T0, T1)
-
-
-def test_npt(dataset):
- einstein = EinsteinCrystal(dataset[0], force_constant=1e-4)
- walker = Walker(dataset[0], einstein, temperature=600, pressure=0, nbeads=2)
- output = sample([walker], steps=30)[0]
- assert output.status.result() == 0
- assert output.trajectory is None
-
- # cell should have changed during NPT
- assert not np.allclose(
- walker.start.result().cell,
- output.state.result().cell,
- )
+ assert np.allclose(output.time.result(), 10 * 0.5 / 1000)
+ temp0 = output.temperature.result()
+ temp1 = output["temperature{kelvin}"].result()[-1]
+ assert np.allclose(temp0, temp1)
+
+ return
+
+
+def test_ensembles(dataset, dataset_h2):
+ einstein = EinsteinCrystal(dataset[0], force_constant=1e-1)
+ walker_nve = Walker(dataset[0], einstein, temperature=None)
+ walker_nvt = Walker(dataset[0], einstein, temperature=300)
+ walker_npt = Walker(dataset[0], einstein, temperature=300, pressure=0)
+ walker_nvst = Walker(dataset[0], einstein, temperature=300, volume_constrained=True)
+ einstein_h2 = EinsteinCrystal(dataset_h2[3], force_constant=1e-1)
+ walker_h2 = Walker(dataset_h2[0], einstein_h2)
+ walkers = [walker_nve, walker_nvt, walker_npt, walker_nvst, walker_h2]
+
+ outputs = sample(walkers, steps=20, observables=["kinetic_md{electronvolt}"])
+ for output in outputs:
+ assert output.status.result() == Status.DONE
+ assert output.trajectory is None
+ output_nve, output_nvt, output_npt, output_nvst, output_h2 = outputs
+
+ # NVE
+ energy_pot = output_nve["potential{electronvolt}"].result()
+ energy_kin = output_nve["kinetic_md{electronvolt}"].result()
+ energy = energy_pot + energy_kin
+ assert np.allclose(energy, energy[0])
+
+ # NVT
+ volume = output_nvt["volume{angstrom3}"].result()
+ assert np.allclose(volume, volume[0])
+
+ # NPT
+ volume = output_npt["volume{angstrom3}"].result()
+ cell_start = walker_npt.start.result().cell
+ cell_stop = output_npt.state.result().cell
+ assert not np.allclose(volume, volume[0])
+ assert not np.allclose(cell_start, cell_stop)
+
+ # NVST
+ volume = output_nvst["volume{angstrom3}"].result()
+ cell_start = walker_nvst.start.result().cell
+ cell_stop = output_nvst.state.result().cell
+ assert np.allclose(volume, volume[0])
+ assert not np.allclose(cell_start, cell_stop)
+
+ # nonperiodic
+ volume = output_h2["volume{angstrom3}"].result()
+ assert not output_h2.state.result().periodic
+ assert all(volume == 1e9) # arbitrary large cell
+
+ return
+
+
+def test_output_status(dataset):
+ """"""
+ geom = dataset[0].result()
+ einstein = EinsteinCrystal(geom, force_constant=0.1)
+ walker = Walker(start=geom, hamiltonian=einstein)
+ walker_boom = Walker(start=geom, hamiltonian=einstein, timestep=1e25, pressure=0)
+
+ # normal & exploded
+ outputs = sample([walker, walker_boom], steps=10)
+ assert outputs[0].status.result() == Status.DONE
+ assert outputs[1].status.result() == Status.EXPLODED
+
+ # max force
+ outputs = sample([walker], steps=100, step=1, max_force=1e-4)
+ assert outputs[0].status.result() == Status.FORCE_EXCEEDED
+ assert outputs[0].trajectory.length().result() == 1
+
+ # walltime
+ definition = psiflow.context().definitions["ModelEvaluation"]
+ definition.max_simulation_time = 5 / 60 # 5 seconds
+ outputs = sample([walker], steps=10000)
+ assert outputs[0].status.result() == Status.TIMEOUT
+ assert outputs[0].time.result() > 0
+
+ return
def test_reset(dataset):
einstein = EinsteinCrystal(dataset[0], force_constant=0.1)
- walker = Walker(
- start=dataset[0],
- temperature=300,
- pressure=None,
- hamiltonian=einstein,
- )
+ walker = Walker(start=dataset[0], temperature=300, hamiltonian=einstein)
walker.state = dataset[1]
assert not walker.is_reset().result()
assert not check_equality(walker.start, walker.state).result()
assert check_equality(walker.start, dataset[0]).result()
assert check_equality(walker.state, dataset[1]).result()
- walker.reset(False)
+ walker.conditional_reset(False)
assert not walker.is_reset().result()
walker.reset()
assert walker.is_reset().result()
- simulation_output = sample(
- [walker],
- steps=50,
- step=10,
- keep_trajectory=True,
- )[0]
- assert simulation_output.status.result() == 0
- assert not walker.is_reset().result()
- assert simulation_output.trajectory.length().result() == 6
- assert np.allclose(
- walker.state.result().per_atom.positions,
- simulation_output.trajectory[-1].result().per_atom.positions,
- )
-
- walker.hamiltonian = EinsteinCrystal(dataset[0], force_constant=1000)
- simulation_output = sample(
- [walker],
- steps=50,
- step=10,
- max_force=10,
- )[0]
- assert simulation_output.status.result() == 2
- assert walker.is_reset().result()
- assert not check_equality(walker.state, simulation_output.state).result()
- assert simulation_output.trajectory.length().result() == 1
-
- # check timeout
- simulation_output = sample(
- [walker],
- steps=5000000,
- step=100,
- )[0]
- assert simulation_output.status.result() == 1
+ [output] = sample([walker], steps=10)
assert not walker.is_reset().result()
- assert check_equality(walker.state, simulation_output.state).result()
- assert simulation_output.time.result() > 0
def test_quench(dataset):
dataset = dataset[:20]
einstein0 = EinsteinCrystal(dataset[3], force_constant=0.1)
einstein1 = EinsteinCrystal(dataset[11], force_constant=0.1)
- walkers = Walker(
- start=dataset[0],
- hamiltonian=einstein0,
- temperature=300,
- ).multiply(30)
-
+ walkers = Walker(start=dataset[0], hamiltonian=einstein0).multiply(10)
walkers[2].hamiltonian = einstein1
quench(walkers, dataset)
@@ -360,7 +357,7 @@ def test_quench(dataset):
def test_randomize(dataset):
- walkers = Walker(dataset[0]).multiply(300)
+ walkers = Walker(dataset[0]).multiply(50)
randomize(walkers, dataset)
length = dataset.length().result()
checks = []
@@ -371,33 +368,18 @@ def test_randomize(dataset):
assert not all(checks)
-def test_walker_nonperiodic(dataset_h2):
- einstein = EinsteinCrystal(dataset_h2[3], force_constant=1.0)
- walker = Walker(dataset_h2[0], einstein)
-
- output = sample([walker], steps=20, step=2)[0]
- assert not output.state.result().periodic
- for state in output.trajectory.geometries().result():
- assert not state.periodic
-
-
def test_rex(dataset):
einstein = EinsteinCrystal(dataset[0], force_constant=0.1)
- walker = Walker(
- dataset[0],
- hamiltonian=einstein,
- temperature=600,
- )
+ walker = Walker(dataset[0], hamiltonian=einstein, temperature=600)
walkers = walker.multiply(2)
- replica_exchange(walkers, trial_frequency=5)
+ replica_exchange(walkers, trial_frequency=2)
assert walkers[0].coupling.nwalkers == len(walkers)
assert len(partition(walkers)) == 1
assert len(partition(walkers)[0]) == 2
- outputs = sample(walkers, steps=50, step=10)
- assert outputs[0].trajectory.length().result() == 6
-
+ outputs = sample(walkers, steps=20, step=2)
swaps = np.loadtxt(walkers[0].coupling.swapfile.result().filepath)
+ assert outputs[0].trajectory.length().result() == 11
assert len(swaps) > 0 # at least some successful swaps
assert np.allclose(swaps[0, 1:], np.array([1, 0])) # 0, 1 --> 1, 0
@@ -421,30 +403,23 @@ def test_walker_serialization(dataset, tmp_path):
"""
metadynamics = Metadynamics(plumed_str)
walkers = Walker(
- dataset[0],
- hamiltonian=einstein,
- temperature=300,
- metadynamics=metadynamics,
+ dataset[0], hamiltonian=einstein, temperature=300, metadynamics=metadynamics
).multiply(3)
for i, walker in enumerate(walkers):
walker.hamiltonian *= 1 / (1 + i)
sample(walkers, steps=10, step=2)
- walkers[0].metadynamics.external.result()
- with open(walkers[0].metadynamics.external.filepath, "r") as f:
+ file = walkers[0].metadynamics.external.result()
+ with open(file.filepath, "r") as f:
assert len(f.read()) > 0
- data = []
- for obj in walkers:
- data.append(psiflow.serialize(obj, copy_to=tmp_path))
- for d in data:
- print(d.result())
-
+ data = [psiflow.serialize(obj, copy_to=tmp_path) for obj in walkers]
new_objects = [psiflow.deserialize(d.result()) for d in data]
psiflow.wait()
+ for d in data:
+ print(d.result())
walkers_ = new_objects[:3]
-
assert check_equality(walkers_[0].start, walkers[0].start).result()
assert check_equality(walkers_[0].state, walkers[0].state).result()
@@ -456,18 +431,21 @@ def test_walker_serialization(dataset, tmp_path):
def test_optimize_ipi(dataset):
einstein = EinsteinCrystal(dataset[2], force_constant=10)
- final = optimize_ipi(dataset[0], einstein, steps=1000000).result()
+ # i-PI optimizer's curvature guess fails in optimum --> don't start in dataset[2]
+ future, future_traj = optimize_ipi(
+ dataset[0], einstein, steps=1000000, keep_trajectory=True
+ )
+ future_dataset = optimize_dataset_ipi(dataset[3:5], einstein, steps=1000000)
+
+ final = future.result()
assert np.allclose(
final.per_atom.positions,
dataset[2].result().per_atom.positions,
atol=1e-4,
)
assert np.allclose(final.energy, 0.0) # einstein energy >= 0
-
- # i-PI optimizer's curvature guess fails in optimum --> don't start in dataset[2]
- optimized = optimize_dataset_ipi(dataset[3:5], einstein, steps=1000000)
- for g in optimized.geometries().result():
+ for g in future_dataset.geometries().result():
assert np.allclose(g.energy, 0.0)
@@ -475,16 +453,18 @@ def test_optimize_ase(dataset):
# TODO: test applied_pressure?
einstein = EinsteinCrystal(dataset[2], force_constant=10)
- final = optimize_ase(dataset[0], einstein, mode='fix_cell', f_max=1e-4).result()
+ future = optimize_ase(dataset[0], einstein, mode="fix_cell", f_max=1e-4)
+ optimized = optimize_dataset_ase(
+ dataset[3:5], einstein, mode="fix_cell", f_max=1e-4
+ )
+ final = future.result()
assert np.allclose(
final.per_atom.positions,
dataset[2].result().per_atom.positions,
atol=1e-4,
)
assert np.allclose(final.energy, 0.0) # einstein energy >= 0
-
- optimized = optimize_dataset_ase(dataset[3:5], einstein, mode='fix_cell', f_max=1e-4)
for g in optimized.geometries().result():
assert np.allclose(g.energy, 0.0)
@@ -506,17 +486,25 @@ def test_optimize_ase(dataset):
RESTRAINT ARG=cell.az AT=.1 KAPPA=1
"""
plumed_c = PlumedHamiltonian(plumed_input)
- final = optimize_ase(geom, plumed_c, mode='fix_shape', f_max=1e-8).result()
- ratio = geom.cell / final.cell
- assert np.allclose(final.cell[0], [3, 0, 0])
- assert np.allclose(ratio[ratio != 0], ratio[0, 0]) # check isotropic scaling (for nonzero components)
+ future1 = optimize_ase(geom, plumed_c, mode="fix_shape", f_max=1e-8)
+ future2 = optimize_ase(geom, plumed_c, mode="fix_volume", f_max=1e-8)
+ future3 = optimize_ase(geom, plumed_v + plumed_c, f_max=1e-8)
+ final1, final2, final3 = future1.result(), future2.result(), future3.result()
- final = optimize_ase(geom, plumed_c, mode='fix_volume', f_max=1e-8).result()
- assert np.allclose(final.cell[0], [3, .2, .1], atol=1e-4)
- assert np.allclose(geom.volume, final.volume)
+ ratio = geom.cell / final1.cell
+ assert np.allclose(final1.cell[0], [3, 0, 0])
+ assert np.allclose(
+ ratio[ratio != 0], ratio[0, 0]
+ ) # check isotropic scaling (for nonzero components)
- final = optimize_ase(geom, plumed_v + plumed_c, f_max=1e-8).result()
- assert np.allclose(final.cell[0], [3, .2, .1], atol=1e-4)
- assert np.allclose(final.volume, 75)
+ assert np.allclose(final2.cell[0], [3, 0.2, 0.1], atol=1e-4)
+ assert np.allclose(geom.volume, final2.volume)
+ assert np.allclose(final3.cell[0], [3, 0.2, 0.1], atol=1e-4)
+ assert np.allclose(final3.volume, 75)
+ # check 'custom' optimisation script
+ final = optimize_ase(
+ dataset[0], einstein, mode="fix_cell", f_max=1e-4, script=file_script_ase
+ ).result()
+ assert np.allclose(final.energy, 0.0)