diff --git a/psiflow/free_energy/phonons.py b/psiflow/free_energy/phonons.py index a73ddeb4..8b60223e 100644 --- a/psiflow/free_energy/phonons.py +++ b/psiflow/free_energy/phonons.py @@ -16,9 +16,7 @@ from psiflow.hamiltonians import Hamiltonian, MixtureHamiltonian from psiflow.sampling.sampling import ( setup_sockets, - label_forces, make_force_xml, - serialize_mixture, make_start_command, make_client_command ) @@ -139,7 +137,7 @@ def compute_harmonic( energy_shift: float = 0.00095, ) -> AppFuture: hamiltonian: MixtureHamiltonian = 1 * hamiltonian - names = label_forces(hamiltonian) + names = hamiltonian.get_named_components() sockets = setup_sockets(names) forces = make_force_xml(hamiltonian, names) @@ -175,7 +173,7 @@ def compute_harmonic( input_future, Dataset([state]).extxyz, ] - inputs += serialize_mixture(hamiltonian, dtype="float64") + inputs += hamiltonian.serialize(dtype="float64") client_args = [] for name in names: diff --git a/psiflow/hamiltonians.py b/psiflow/hamiltonians.py index 3faeb1b0..13b53603 100644 --- a/psiflow/hamiltonians.py +++ b/psiflow/hamiltonians.py @@ -3,7 +3,7 @@ import urllib from functools import partial from pathlib import Path -from typing import ClassVar, Optional, Union +from typing import ClassVar, Optional, Union, Callable import numpy as np import typeguard @@ -28,17 +28,21 @@ from psiflow.utils.io import dump_json +# TODO: executors not very consistent for some reason (default_threads <> default_htex) + @typeguard.typechecked @psiflow.serializable class Hamiltonian(Computable): outputs: ClassVar[tuple] = ("energy", "forces", "stress") batch_size = 1000 + app: ClassVar[Callable] # TODO: app is actually an instance variable, but serialization complains.. + function_name: ClassVar[str] def compute( self, arg: Union[Dataset, AppFuture[list], list, AppFuture, Geometry], *outputs: Optional[str], - batch_size: Optional[int] = -1, # if -1: take class default + batch_size: Optional[int] = -1, # if -1: take class default TODO: why not just None? ) -> Union[list[AppFuture], AppFuture]: if len(outputs) == 0: outputs = tuple(self.__class__.outputs) @@ -54,15 +58,14 @@ def compute( def __eq__(self, hamiltonian: Hamiltonian) -> bool: raise NotImplementedError - def __mul__(self, a: float) -> Hamiltonian: + def __mul__(self, a: float) -> MixtureHamiltonian: return MixtureHamiltonian([self], [a]) def __add__(self, hamiltonian: Hamiltonian) -> Hamiltonian: if type(hamiltonian) is Zero: return self - if type(hamiltonian) is MixtureHamiltonian: - return hamiltonian.__add__(self) - return 1.0 * self + 1.0 * hamiltonian + mixture = MixtureHamiltonian([self], [1.]) + return mixture.__add__(hamiltonian) def __sub__(self, hamiltonian: Hamiltonian) -> Hamiltonian: return self + (-1.0) * hamiltonian @@ -101,9 +104,10 @@ def __mul__(self, a: float) -> Hamiltonian: return Zero() def __add__(self, hamiltonian: Hamiltonian) -> Hamiltonian: - return hamiltonian + return hamiltonian # (Zero + Hamiltonian) is different from (Hamiltonian + Zero) + + __rmul__ = __mul__ # handle float * Zero - __rmul__ = __mul__ # handle float * Hamiltonian @typeguard.typechecked @@ -142,88 +146,90 @@ def compute( # override compute for efficient batching batch_size=batch_size, ) - def __eq__(self, hamiltonian: Hamiltonian) -> bool: + def __eq__(self, hamiltonian: Hamiltonian | MixtureHamiltonian) -> bool: if type(hamiltonian) is not MixtureHamiltonian: return False if len(self.coefficients) != len(hamiltonian.coefficients): return False - for i, h in enumerate(self.hamiltonians): - if h not in hamiltonian.hamiltonians: + for c, h in zip(self.coefficients, self.hamiltonians): + if (idx := hamiltonian.get_index(h)) is None: return False - coefficient = hamiltonian.coefficients[hamiltonian.hamiltonians.index(h)] - if self.coefficients[i] != coefficient: + if c != hamiltonian.coefficients[idx]: return False return True - def __mul__(self, a: float) -> Hamiltonian: + def __mul__(self, a: float) -> MixtureHamiltonian: return MixtureHamiltonian( self.hamiltonians, [c * a for c in self.coefficients], ) + __rmul__ = __mul__ # handle float * MixtureHamiltonian + def __len__(self) -> int: return len(self.coefficients) - def __add__(self, hamiltonian: Hamiltonian) -> Hamiltonian: + def __add__(self, hamiltonian: Hamiltonian | MixtureHamiltonian) -> MixtureHamiltonian: if type(hamiltonian) is Zero: return self if type(hamiltonian) is not MixtureHamiltonian: - coefficients = list(self.coefficients) - hamiltonians = list(self.hamiltonians) - try: - index = hamiltonians.index(hamiltonian) - coefficients[index] += 1.0 - except ValueError: - hamiltonians.append(hamiltonian) - coefficients.append(1.0) - else: - coefficients = list(hamiltonian.coefficients) - hamiltonians = list(hamiltonian.hamiltonians) - for c, h in zip(self.coefficients, self.hamiltonians): - try: - index = hamiltonians.index(h) - coefficients[index] += c - except ValueError: - hamiltonians.append(h) - coefficients.append(c) + hamiltonian = 1.0 * hamiltonian # turn into mixture + + coefficients = list(hamiltonian.coefficients) + hamiltonians = list(hamiltonian.hamiltonians) + for c, h in zip(self.coefficients, self.hamiltonians): + if (idx := hamiltonian.get_index(h)) is None: + hamiltonians.append(h) + coefficients.append(c) + else: + coefficients[idx] += c return MixtureHamiltonian(hamiltonians, coefficients) - def get_index(self, hamiltonian) -> Optional[int]: + def get_index(self, hamiltonian: Hamiltonian) -> Optional[int]: assert type(hamiltonian) is not MixtureHamiltonian - if hamiltonian not in self.hamiltonians: + try: + return self.hamiltonians.index(hamiltonian) + except ValueError: return None - return self.hamiltonians.index(hamiltonian) - def get_indices(self, mixture) -> 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 - for h in mixture.hamiltonians: - if h not in self.hamiltonians: - return None indices = [] for h in mixture.hamiltonians: indices.append(self.get_index(h)) + if any([idx is None for idx in indices]): + return None return tuple(indices) - def get_coefficient(self, hamiltonian) -> Optional[float]: + def get_coefficient(self, hamiltonian: Hamiltonian) -> Optional[float]: assert type(hamiltonian) is not MixtureHamiltonian - if hamiltonian not in self.hamiltonians: + if (idx := self.get_index(hamiltonian)) is None: return None - return self.coefficients[self.hamiltonians.index(hamiltonian)] + return self.coefficients[idx] - def get_coefficients(self, mixture) -> Optional[tuple[float, ...]]: + def get_coefficients(self, mixture: MixtureHamiltonian) -> Optional[tuple[float, ...]]: assert type(mixture) is MixtureHamiltonian - for h in mixture.hamiltonians: + for h in mixture.hamiltonians: # every mixture component must be a component of self if h not in self.hamiltonians: return None - coefficients = [] - for h in self.hamiltonians: - coefficient = mixture.get_coefficient(h) - if coefficient is None: - coefficient = 0.0 - coefficients.append(coefficient) + + # components of self that are not in mixture get coefficient 0 + coefficients = [(mixture.get_coefficient(h) or 0) for h in self.hamiltonians] return tuple(coefficients) - __rmul__ = __mul__ # handle float * Hamiltonian + def get_named_components(self) -> list[str]: + """Create unique string name for every hamiltonian component to be used in i-Pi sampling""" + names, counts = [], {} + for h in self.hamiltonians: + name = h.__class__.__name__ + counts.setdefault(name, 0) + names.append(f"{name}{counts[name]}") + counts[name] += 1 + return names + + def serialize(self, **kwargs) -> list[DataFuture]: + return [h.serialize_function(**kwargs) for h in self.hamiltonians] @typeguard.typechecked @@ -255,12 +261,12 @@ def parameters(self) -> dict: "volume": get_attribute(self.reference_geometry, "volume"), } - def __eq__(self, hamiltonian: Hamiltonian) -> bool: - if type(hamiltonian) is not EinsteinCrystal: - return False - if not np.allclose(self.force_constant, hamiltonian.force_constant): - return False - if self.reference_geometry != hamiltonian.reference_geometry: + def __eq__(self, hamiltonian: Hamiltonian | EinsteinCrystal) -> bool: + if ( + not isinstance(hamiltonian, EinsteinCrystal) or + not np.allclose(self.force_constant, hamiltonian.force_constant) or + self.reference_geometry != hamiltonian.reference_geometry + ): return False return True @@ -302,10 +308,11 @@ def parameters(self) -> dict: external = None return {"plumed_input": self.plumed_input, "external": external} - def __eq__(self, other: Hamiltonian) -> bool: - if type(other) is not type(self): - return False - if self.plumed_input != other.plumed_input: + def __eq__(self, other: Hamiltonian | PlumedHamiltonian) -> bool: + if ( + not isinstance(other, PlumedHamiltonian) or + self.plumed_input != other.plumed_input + ): return False return True @@ -322,7 +329,7 @@ def __init__( reference_geometry: Union[Geometry, AppFuture[Geometry]], hessian: Union[np.ndarray, AppFuture[np.ndarray]], ): - self.reference_geometry = reference_geometry + self.reference_geometry = reference_geometry # TODO: why not copy_app_future(geometry) like others? self.hessian = hessian self._create_apps() @@ -343,12 +350,13 @@ def parameters(self) -> dict: "hessian": self.hessian, } - def __eq__(self, hamiltonian: Hamiltonian) -> bool: + def __eq__(self, hamiltonian: Hamiltonian | Harmonic) -> bool: if type(hamiltonian) is not Harmonic: return False if hamiltonian.reference_geometry != self.reference_geometry: return False + # TODO: why this check? Is it not always an ndarray? # slightly different check for numpy arrays is_array0 = type(hamiltonian.hessian) is np.ndarray is_array1 = type(self.hessian) is np.ndarray @@ -409,7 +417,7 @@ def parameters(self) -> dict: "env_vars": evaluation.env_vars, } - def __eq__(self, hamiltonian) -> bool: + def __eq__(self, hamiltonian: Hamiltonian | MACEHamiltonian) -> bool: if type(hamiltonian) is not MACEHamiltonian: return False if self.external.filepath != hamiltonian.external.filepath: @@ -424,6 +432,8 @@ def __eq__(self, hamiltonian) -> bool: return False return True + # TODO: the methods below are outdated.. + @classmethod def mace_mp0(cls, size: str = "small") -> MACEHamiltonian: urls = dict( diff --git a/psiflow/sampling/ase.py b/psiflow/sampling/ase.py index 5bfca2b7..dfa63a6d 100644 --- a/psiflow/sampling/ase.py +++ b/psiflow/sampling/ase.py @@ -13,7 +13,6 @@ from psiflow.geometry import Geometry from psiflow.hamiltonians import Hamiltonian from psiflow.utils.io import dump_json -from psiflow.sampling.sampling import serialize_mixture, label_forces from psiflow.utils import TMP_COMMAND, CD_COMMAND from ._ase import ALLOWED_MODES @@ -79,8 +78,8 @@ def optimize( input_geometry = Dataset([state]).extxyz hamiltonian = 1.0 * hamiltonian # convert to mixture - names, coeffs = label_forces(hamiltonian), hamiltonian.coefficients - input_forces = serialize_mixture(hamiltonian, dtype="float64") # double precision for MLPs + names, coeffs = hamiltonian.get_named_components(), hamiltonian.coefficients + 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) ] diff --git a/psiflow/sampling/output.py b/psiflow/sampling/output.py index 309d716a..7f734586 100644 --- a/psiflow/sampling/output.py +++ b/psiflow/sampling/output.py @@ -25,8 +25,7 @@ @typeguard.typechecked def potential_component_names(n: int): - str_format = "pot_component_raw({})" - return [str_format.format(i) + "{electronvolt}" for i in range(n)] + return [f'pot_component_raw({i}){{electronvolt}}' for i in range(n)] @typeguard.typechecked diff --git a/psiflow/sampling/sampling.py b/psiflow/sampling/sampling.py index 71d28c68..4d69ffe1 100644 --- a/psiflow/sampling/sampling.py +++ b/psiflow/sampling/sampling.py @@ -2,11 +2,12 @@ import math import xml.etree.ElementTree as ET +from dataclasses import dataclass from typing import Optional, Union, Iterable -from collections import defaultdict import parsl import typeguard +import numpy as np from parsl.app.app import bash_app from parsl.data_provider.files import File from parsl.dataflow.futures import AppFuture, DataFuture @@ -24,44 +25,78 @@ from psiflow.utils import TMP_COMMAND, CD_COMMAND +@dataclass +class HamiltonianComponent: + name: str + hamiltonian: Hamiltonian + shared: bool + + +@dataclass +class EnsembleTable: + keys: tuple[str, ...] + weights: np.ndarray + + def get_index(self, idx: int) -> np.ndarray: + return self.weights[idx] + + def get_key(self, key: str) -> np.ndarray: + return self.weights[:, self.keys.index(key)] + + def __len__(self) -> int: + return self.weights.shape[0] + + 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 + ): + 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[dict[str, Hamiltonian], list[tuple], list[AppFuture]]: +) -> tuple[list[HamiltonianComponent], EnsembleTable, list[AppFuture]]: # multiply by 1.0 to ensure result is Mixture in case len(walkers) == 1 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) - - # TODO: take care of ordering in weights_table (names <> total_hamiltonian.hamiltonians) - weights_header = tuple(names) + # 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))) + 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_dict = {} + if walkers[0].nvt or walkers[0].npt: + weights_dict['TEMP'] = [w.temperature for w in walkers] if walkers[0].npt: - weights_header = ("TEMP", "PRESSURE") + weights_header - elif walkers[0].nvt: - weights_header = ("TEMP",) + weights_header - else: - pass - - weights_table = [weights_header] - for walker in walkers: - coefficients = total_hamiltonian.get_coefficients(1.0 * walker.hamiltonian) - if walker.npt: - ensemble = (walker.temperature, walker.pressure) - elif walker.nvt: - ensemble = (walker.temperature,) - else: - ensemble = () - weights_table.append(ensemble + tuple(coefficients)) + weights_dict['PRESSURE'] = [w.pressure for w in walkers] # inspect metadynamics attributes and allocate additional weights per MTD - walker_indices = [] metad_objects = [] for i, walker in enumerate(walkers): mtd = walker.metadynamics @@ -72,50 +107,25 @@ def template( "multiple walker metadynamics, use " "psiflow.sampling.multiple_walker_metadynamics" ) - walker_indices.append(i) metad_objects.append(mtd) + 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] - for i, row in enumerate(weights_table): - if i == 0: - to_append = ["METAD{}".format(i) for i in range(len(metad_objects))] - else: - to_append = [0] * len(metad_objects) - try: - index = walker_indices.index(i - 1) - to_append[index] = 1 - except ValueError: - pass - weights_table[i] = row + tuple(to_append) - - hamiltonians_map = {n: h for n, h in zip(names, hamiltonians)} - return hamiltonians_map, weights_table, plumed_list - - -def label_forces(hamiltonian: MixtureHamiltonian) -> list[str]: - """Create unique string name for every hamiltonian""" - # TODO: move into class? - assert isinstance(hamiltonian, MixtureHamiltonian) - names, counts = [], defaultdict(lambda: 0) - for h in hamiltonian.hamiltonians: - name = h.__class__.__name__ - names.append(f'{name}{counts[name]}') - counts[name] += 1 - return names + 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 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 serialize_mixture(hamiltonian: MixtureHamiltonian, **kwargs) -> list[DataFuture]: - # TODO: move into class? - return [h.serialize_function(**kwargs) for h in hamiltonian.hamiltonians] - - @typeguard.typechecked def setup_sockets( hamiltonian_labels: Iterable[str], @@ -187,52 +197,52 @@ def setup_motion(walker: Walker, fix_com: bool) -> ET.Element: @typeguard.typechecked -def setup_ensemble(weights_header: tuple[str, ...]) -> ET.Element: +def setup_ensemble(components: list[HamiltonianComponent], weights_header: tuple[str, ...]) -> ET.Element: ensemble = ET.Element("ensemble") - if "TEMP" in weights_header: - temperature = ET.Element("temperature", units="kelvin") - temperature.text = "TEMP" - ensemble.append(temperature) - else: # set TEMP in any case to avoid i-PI from throwing weird errors - temperature = ET.Element("temperature", units="kelvin") - temperature.text = " 300 " - ensemble.append(temperature) + + # set TEMP in any case to avoid i-PI from throwing weird errors + temperature = ET.Element("temperature", units="kelvin") + temperature.text = "TEMP" if "TEMP" in weights_header else " 300 " + ensemble.append(temperature) if "PRESSURE" in weights_header: pressure = ET.Element("pressure", units="megapascal") pressure.text = "PRESSURE" ensemble.append(pressure) bias = ET.Element("bias") - bias_weights = ET.Element("bias_weights") bias_weights_list = [] - count = 0 - while True: # metadynamics bias if present - name = "METAD{}".format(count) - if name in weights_header: + + # add hamiltonian components that are not shared between all walkers as bias + # TODO: do we always want to do this? + for comp in components: + if not comp.shared: + force = ET.Element("force", forcefield=comp.name) + bias.append(force) + bias_weights_list.append(comp.name.upper()) + + # add metadynamics bias if present + for idx in range(len(weights_header) - len(components)): + if (name := f"METAD{idx}") in weights_header: force = ET.Element("force", forcefield=name.lower()) bias.append(force) bias_weights_list.append(name) - count += 1 - else: - break - bias_weights.text = " [ " + ", ".join(bias_weights_list) + " ] " - if count > 0: + if bias_weights_list: ensemble.append(bias) + bias_weights = ET.Element("bias_weights") + bias_weights.text = create_xml_list(bias_weights_list) ensemble.append(bias_weights) + return ensemble @typeguard.typechecked -def setup_forces(weights_header: tuple[str, ...]) -> ET.Element: +def setup_forces(hamiltonian_components: list[HamiltonianComponent]) -> ET.Element: forces = ET.Element("forces") - for name in weights_header: - if name in ["TEMP", "PRESSURE"]: - continue - if name.startswith("METAD"): # added as bias, not as main force - continue - force = ET.Element("force", forcefield=name, weight=name.upper()) - forces.append(force) + for comp in hamiltonian_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 @@ -254,20 +264,21 @@ def setup_ffplumed(nplumed: int) -> list[ET.Element]: @typeguard.typechecked def setup_system_template( walkers: list[Walker], - weights_table: list[tuple], + ensemble_table: EnsembleTable, motion: ET.Element, ensemble: ET.Element, forces: ET.Element, ) -> ET.Element: system_template = ET.Element("system_template") labels = ET.Element("labels") - _labels = [w.upper() for w in weights_table[0]] - labels.text = "[ INDEX, {} ]".format(", ".join(_labels)) + labels.text = create_xml_list(['INDEX'] + [w.upper() for w in ensemble_table.keys]) system_template.append(labels) - for i, weights in enumerate(weights_table[1:]): + for i in range(len(ensemble_table)): + print(i) + print(ensemble_table.get_index(i)) instance = ET.Element("instance") - instance.text = "[ {}, {}]".format(i, ", ".join([str(w) for w in weights])) + 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)) @@ -275,16 +286,13 @@ def setup_system_template( start.text = " start_INDEX.xyz " initialize.append(start) velocities = ET.Element("velocities", mode="thermal", units="kelvin") - if "TEMP" in weights_table[0]: # valid template parameter - velocities.text = " TEMP " - else: - velocities.text = " 300 " + 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 = str(list(walkers[0].masses * AMU_TO_AU)).replace("[", "[ ").replace("]", " ]") # replace str(list()) with np.array_str()? + masses.text = create_xml_list([str(i) for i in walkers[0].masses * AMU_TO_AU]) initialize.append(masses) system = ET.Element("system", prefix="walker-INDEX") @@ -301,25 +309,28 @@ def setup_system_template( @typeguard.typechecked def setup_output( - nwalkers: int, - nhamiltonians: int, + components: list[HamiltonianComponent], observables: Optional[list[str]], step: Optional[int], keep_trajectory: bool, checkpoint_step: int, ) -> tuple[ET.Element, list]: - output = ET.Element("output", prefix="output") if observables is None: observables = [] + + n_forces = sum([comp.shared for comp in components], 0) full_list = ( - DEFAULT_OBSERVABLES + potential_component_names(nhamiltonians) + observables + DEFAULT_OBSERVABLES + potential_component_names(n_forces) + observables ) + if len(components) != n_forces: # store bias force components + full_list.append('ensemble_bias{electronvolt}') observables = list(set(full_list)) if step is None: step = checkpoint_step + output = ET.Element("output", prefix="output") checkpoint = ET.Element( "checkpoint", filename="checkpoint", @@ -342,12 +353,9 @@ def setup_output( filename="properties", stride=str(step), ) - properties.text = " [ " + ", ".join(observables) + " ] " + properties.text = create_xml_list(observables) output.append(properties) - - # TODO: check whether observables are valid - simulation_outputs = [SimulationOutput(observables) for i in range(nwalkers)] - return output, simulation_outputs + return output, observables @typeguard.typechecked @@ -460,18 +468,18 @@ def _sample( verbosity: str = "medium", ) -> list[SimulationOutput]: assert len(walkers) > 0 - hamiltonians_map, weights_table, plumed_list = template(walkers) + hamiltonian_components, ensemble_table, plumed_list = template(walkers) coupling = walkers[0].coupling if motion_defaults is not None: raise NotImplementedError motion = setup_motion(walkers[0], fix_com) - ensemble = setup_ensemble(weights_table[0]) - forces = setup_forces(weights_table[0]) + ensemble = setup_ensemble(hamiltonian_components, ensemble_table.keys) + forces = setup_forces(hamiltonian_components) system_template = setup_system_template( walkers, - weights_table, + ensemble_table, motion, ensemble, forces, @@ -491,9 +499,9 @@ def _sample( start = math.floor(start / step) # start is applied on subsampled quantities if step is None: keep_trajectory = False - output, simulation_outputs = setup_output( - len(walkers), - len(hamiltonians_map), # for potential components + # TODO: check whether observables are valid? + output, observables = setup_output( + hamiltonian_components, # for potential components observables, step, keep_trajectory, @@ -504,7 +512,7 @@ def _sample( verbosity=str(verbosity), safe_stride=str(checkpoint_step), ) - sockets = setup_sockets(hamiltonians_map.keys()) + sockets = setup_sockets([comp.name for comp in hamiltonian_components]) for socket in sockets: simulation.append(socket) ffplumed = setup_ffplumed(len(plumed_list)) @@ -535,19 +543,17 @@ def _sample( Dataset([w.state for w in walkers]).extxyz, ] - # remove any Harmonic instances because they are not implemented with sockets - hamiltonian_names = list(hamiltonians_map.keys()) - + # remove any Harmonic instances because they are not implemented with sockets -- TODO: why? max_nclients = int(sum([w.nbeads for w in walkers])) - inputs += [h.serialize_function() for h in hamiltonians_map.values()] client_args = [] - for name in hamiltonian_names: - args = definition.get_client_args(name, max_nclients, motion="dynamics") + for comp in hamiltonian_components: + inputs.append(comp.hamiltonian.serialize_function()) + args = definition.get_client_args(comp.name, max_nclients, motion="dynamics") client_args.append(args) outputs = [context.new_file("data_", ".xyz")] - outputs += [context.new_file("simulation_", ".txt") for w in walkers] + outputs += [context.new_file("simulation_", ".txt") for _ in walkers] if keep_trajectory: - outputs += [context.new_file("data_", ".xyz") for w in walkers] + outputs += [context.new_file("data_", ".xyz") for _ in walkers] assert len(outputs) == 2 * len(walkers) + 1 else: assert len(outputs) == len(walkers) + 1 @@ -568,7 +574,7 @@ def _sample( resources = definition.wq_resources(max_nclients) result = execute_ipi( len(walkers), - hamiltonian_names, + [comp.name for comp in hamiltonian_components], client_args, keep_trajectory, max_force, @@ -585,7 +591,7 @@ def _sample( ) 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: @@ -594,7 +600,7 @@ def _sample( simulation_output.parse_data( start, result.outputs[i + 1], - hamiltonians=list(hamiltonians_map.values()), + hamiltonians=[comp.hamiltonian for comp in hamiltonian_components], ) if keep_trajectory: j = len(walkers) + 1 + i diff --git a/psiflow/serialization.py b/psiflow/serialization.py index 7560caa1..882f153e 100644 --- a/psiflow/serialization.py +++ b/psiflow/serialization.py @@ -14,6 +14,10 @@ import psiflow from psiflow.geometry import Geometry + +# TODO: this is only used in the Learning class currently + + _DataFuture = Union[File, DataFuture] @@ -42,7 +46,7 @@ def setter(self, value: type_hint) -> None: return setter -def update_init(init_func): +def update_init(init_func): # TODO: why not have a Mixin class instead? def wrapper(self, *args, **kwargs): self._geoms = {} self._files = {} diff --git a/tests/test_sampling.py b/tests/test_sampling.py index e94fded1..fc707747 100644 --- a/tests/test_sampling.py +++ b/tests/test_sampling.py @@ -11,7 +11,7 @@ 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.metadynamics import Metadynamics -from psiflow.sampling.sampling import sample, template +from psiflow.sampling.sampling import sample, template, EnsembleTable from psiflow.sampling.server import parse_checkpoint from psiflow.sampling.walker import ( Walker, @@ -53,58 +53,45 @@ def test_walkers(dataset): # nvt _walkers = [walkers[0], walkers[-1]] - hamiltonians_map, weights_table, plumed_list = template(_walkers) + hamiltonian_components, ensemble_table, plumed_list = template(_walkers) assert _walkers[0].nvt - assert len(hamiltonians_map) == 1 - assert weights_table[0] == ("TEMP", "EinsteinCrystal0", "METAD0", "METAD1") + assert len(hamiltonian_components) == 1 + keys_ref = ("TEMP", "EinsteinCrystal0", "METAD0", "METAD1") + weights_ref = np.array([[300, 1, 1, 0], [600, 1, 0, 1]]) + assert ensemble_table == EnsembleTable(keys_ref, weights_ref) assert len(plumed_list) == 2 - assert weights_table[1] == (300, 1.0, 1.0, 0.0) - assert weights_table[2] == (600, 1.0, 0.0, 1.0) # remove _walkers[0].metadynamics = None _walkers[1].metadynamics = None - hamiltonians_map, weights_table, plumed_list = template(_walkers) + hamiltonian_components, ensemble_table, plumed_list = template(_walkers) assert _walkers[0].nvt - assert len(hamiltonians_map) == 1 - assert weights_table[0] == ("TEMP", "EinsteinCrystal0") + assert len(hamiltonian_components) == 1 + keys_ref = ("TEMP", "EinsteinCrystal0") + weights_ref = np.array([[300, 1], [600, 1]]) + assert ensemble_table == EnsembleTable(keys_ref, weights_ref) assert len(plumed_list) == 0 - assert weights_table[1] == (300, 1.0) - assert weights_table[2] == (600, 1.0) # pimd partition _walkers = [walkers[1], walkers[2]] - hamiltonians_map, weights_table, plumed_list = template(_walkers) + hamiltonian_components, ensemble_table, plumed_list = template(_walkers) assert _walkers[0].pimd - assert len(hamiltonians_map) == 3 - assert weights_table[0] == ( - "TEMP", - "EinsteinCrystal0", - "EinsteinCrystal1", - "PlumedHamiltonian0", - "METAD0", - ) - assert weights_table[1] == (300, 0.0, 0.5, 0.0, 1.0) - assert weights_table[2] == (300, 1.0, 0.0, 1.0, 0.0) + assert len(hamiltonian_components) == 3 + 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) # npt partition _walkers = [walkers[3], walkers[4]] with pytest.raises(AssertionError): # mtd objects were equal - hamiltonians_map, weights_table, plumed_list = template(_walkers) + _ = template(_walkers) _walkers[0].metadynamics = Metadynamics("METAD: FILE=bla") - hamiltonians_map, weights_table, plumed_list = template(_walkers) + hamiltonian_components, ensemble_table, plumed_list = template(_walkers) assert _walkers[0].npt - assert len(hamiltonians_map) == 2 - assert weights_table[0] == ( - "TEMP", - "PRESSURE", - "EinsteinCrystal0", - "EinsteinCrystal1", - "METAD0", - "METAD1", - ) - assert weights_table[1] == (300, 0, 0.0, 1.0, 1.0, 0.0) - assert weights_table[2] == (600, 100, 1.0, 0.0, 0.0, 1.0) + 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]]) + assert ensemble_table == EnsembleTable(keys_ref, weights_ref) def test_parse_checkpoint(checkpoint):