diff --git a/src/reactionflow/_durable.py b/src/reactionflow/_durable.py index b76be4a..bcebeb7 100644 --- a/src/reactionflow/_durable.py +++ b/src/reactionflow/_durable.py @@ -2,19 +2,22 @@ from __future__ import annotations +import hashlib import os -from hashlib import sha256 +import re +import shutil +from collections.abc import Iterator +from contextlib import contextmanager from pathlib import Path +SAFE_NAME = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]*") + def file_digest(path: Path) -> str: """SHA-256 of a file's contents.""" - checksum = sha256() with path.open("rb") as handle: - for chunk in iter(lambda: handle.read(1024 * 1024), b""): - checksum.update(chunk) - return checksum.hexdigest() + return hashlib.file_digest(handle, "sha256").hexdigest() def sync_directory(path: Path) -> None: @@ -54,3 +57,16 @@ def publish(temporary: Path, final: Path) -> None: flush_to_disk(temporary) os.replace(temporary, final) sync_directory(final.parent) + + +@contextmanager +def publishing(temporary: Path, final: Path) -> Iterator[None]: + """Make a temporary sibling directory to fill, publish it, and remove it after an Exception.""" + + temporary.mkdir() + try: + yield + publish(temporary, final) + except Exception: + shutil.rmtree(temporary, ignore_errors=True) + raise diff --git a/src/reactionflow/adapters/_model_files.py b/src/reactionflow/adapters/_model_files.py index 4cdc677..f70acac 100644 --- a/src/reactionflow/adapters/_model_files.py +++ b/src/reactionflow/adapters/_model_files.py @@ -1,4 +1,4 @@ -"""Small, explicit helpers for optional model packages and cached weights.""" +"""Small, explicit helpers for adapter options, optional model packages and cached weights.""" from __future__ import annotations @@ -13,7 +13,17 @@ from urllib.request import urlopen from uuid import uuid4 -from .._durable import ensure_directory, file_digest, publish, sync_directory +from .._durable import ensure_directory, file_digest, publishing, sync_directory + + +def check_paths(options: Mapping[str, object], names: tuple[str, ...], message: str) -> None: + """Refuse each named option that is present but not an absolute path string.""" + + for name in names: + if name in options: + value = options[name] + if not isinstance(value, str) or not value or not Path(value).is_absolute(): + raise ValueError(message.format(name)) def require_package(name: str, expected: str, extra: str) -> None: @@ -31,13 +41,12 @@ def require_package(name: str, expected: str, extra: str) -> None: def model_cache(options: Mapping[str, object], backend: str) -> Path: raw = options.get("cache_dir", os.environ.get("REACTIONFLOW_MODEL_CACHE")) if raw is None: - root = Path.home() / ".cache" / "reactionflow" / "models" - else: - if not isinstance(raw, str) or not raw: - raise ValueError("cache_dir must be an absolute path string") - root = Path(raw).expanduser() - if not root.is_absolute(): - raise ValueError("cache_dir must be an absolute path") + return Path.home() / ".cache" / "reactionflow" / "models" / backend + if not isinstance(raw, str) or not raw: + raise ValueError("cache_dir must be an absolute path string") + root = Path(raw).expanduser() + if not root.is_absolute(): + raise ValueError("cache_dir must be an absolute path") return root / backend @@ -74,8 +83,7 @@ def cached_download(url: str, cache_dir: Path, *, download: bool) -> Path: if final.exists(): return _cached_file(final, url) temporary = cache_dir / f".{key}-{uuid4().hex}.tmp" - temporary.mkdir() - try: + with publishing(temporary, final): path = temporary / "model" with urlopen(url, timeout=60) as response, path.open("wb") as handle: expected_size = response.headers.get("Content-Length") @@ -90,8 +98,4 @@ def cached_download(url: str, cache_dir: Path, *, download: bool) -> Path: json.dumps({"url": url, "sha256": file_digest(path)}, sort_keys=True) + "\n", encoding="utf-8", ) - publish(temporary, final) - except Exception: - shutil.rmtree(temporary, ignore_errors=True) - raise return final / "model" diff --git a/src/reactionflow/adapters/_torch.py b/src/reactionflow/adapters/_torch.py index 34d4443..1832e17 100644 --- a/src/reactionflow/adapters/_torch.py +++ b/src/reactionflow/adapters/_torch.py @@ -9,6 +9,7 @@ from collections.abc import Iterator, Mapping from contextlib import contextmanager from importlib.metadata import version +from numbers import Integral from pathlib import Path from typing import Any @@ -44,31 +45,23 @@ def _load_torch(device: str) -> Any: def _torch_environment(torch: Any, device: str) -> dict[str, object]: + # Each setting's required value and its value now. settings = { - "deterministic_algorithms": torch.are_deterministic_algorithms_enabled(), - "deterministic_warn_only": torch.is_deterministic_algorithms_warn_only_enabled(), - "float32_matmul_precision": torch.get_float32_matmul_precision(), - "cuda_matmul_allow_tf32": torch.backends.cuda.matmul.allow_tf32, - "cudnn_allow_tf32": torch.backends.cudnn.allow_tf32, - "cudnn_deterministic": torch.backends.cudnn.deterministic, - "cudnn_benchmark": torch.backends.cudnn.benchmark, + "deterministic_algorithms": (True, torch.are_deterministic_algorithms_enabled()), + "deterministic_warn_only": (False, torch.is_deterministic_algorithms_warn_only_enabled()), + "float32_matmul_precision": ("highest", torch.get_float32_matmul_precision()), + "cuda_matmul_allow_tf32": (False, torch.backends.cuda.matmul.allow_tf32), + "cudnn_allow_tf32": (False, torch.backends.cudnn.allow_tf32), + "cudnn_deterministic": (True, torch.backends.cudnn.deterministic), + "cudnn_benchmark": (False, torch.backends.cudnn.benchmark), } - required = { - "deterministic_algorithms": True, - "deterministic_warn_only": False, - "float32_matmul_precision": "highest", - "cuda_matmul_allow_tf32": False, - "cudnn_allow_tf32": False, - "cudnn_deterministic": True, - "cudnn_benchmark": False, - } - changed = [name for name, value in required.items() if settings[name] != value] + changed = [name for name, (required, value) in settings.items() if value != required] if changed: raise RuntimeError(f"model backend changed required Torch inference settings: {changed}") result: dict[str, object] = { "torch_version": str(torch.__version__), "device": device, - **settings, + **{name: value for name, (_, value) in settings.items()}, "num_threads": torch.get_num_threads(), "num_interop_threads": torch.get_num_interop_threads(), "cublas_workspace_config": os.environ.get("CUBLAS_WORKSPACE_CONFIG"), @@ -104,6 +97,17 @@ def _preserve_rng(torch: Any) -> Iterator[None]: torch.set_default_dtype(dtype) +def require_charge_and_spin(atoms: ase.Atoms, label: str) -> None: + """Refuse a structure without an explicit integer charge and positive spin multiplicity.""" + + for name in ("charge", "spin"): + value = atoms.info.get(name) + if isinstance(value, bool) or not isinstance(value, Integral): + raise ValueError(f"{label} requires explicit integer atoms.info[{name!r}]") + if atoms.info["spin"] < 1: + raise ValueError(f"{label} spin must be a positive multiplicity") + + class TorchModelAdapter(ASELangevinBAOABAdapter): """Reuse the ASE runtime while binding every lease to the same local model.""" @@ -148,6 +152,7 @@ def _lease(self) -> Iterator[tuple[Calculator, ComponentState]]: try: if any(file_digest(path) != files[name]["sha256"] for name, path in self.files.items()): raise ValueError("model files changed while the calculator was being constructed") + here = Path(__file__) state = ComponentState( kind=self.kind, metadata={ @@ -158,11 +163,9 @@ def _lease(self) -> Iterator[tuple[Calculator, ComponentState]]: "ase_version": ase.__version__, "numpy_version": np.__version__, "adapter_source_sha256": file_digest(Path(inspect.getfile(type(self)))), - "torch_adapter_source_sha256": file_digest(Path(__file__)), - "ase_runtime_source_sha256": file_digest(Path(__file__).with_name("ase.py")), - "integrator_source_sha256": file_digest( - Path(__file__).parents[1] / "ase_npt.py" - ), + "torch_adapter_source_sha256": file_digest(here), + "ase_runtime_source_sha256": file_digest(here.with_name("ase.py")), + "integrator_source_sha256": file_digest(here.parents[1] / "ase_npt.py"), **_torch_environment(torch, self.device), **self._extra_metadata(calculator), }, diff --git a/src/reactionflow/adapters/aimnet2.py b/src/reactionflow/adapters/aimnet2.py index 3737e06..dac6463 100644 --- a/src/reactionflow/adapters/aimnet2.py +++ b/src/reactionflow/adapters/aimnet2.py @@ -13,9 +13,9 @@ from ase.calculators.calculator import Calculator, all_changes from .._durable import file_digest -from ..campaign import TrajectorySpec -from ._model_files import cached_download, model_cache, require_package -from ._torch import TorchModelAdapter +from ..campaign import TrajectorySpec, _reject_unknown +from ._model_files import cached_download, check_paths, model_cache, require_package +from ._torch import TorchModelAdapter, require_charge_and_spin AIMNET_VERSION = "0.2.0" _FAMILIES = { @@ -64,9 +64,7 @@ def catalog() -> dict[str, Any]: def _validate_options(options: Mapping[str, Any]) -> None: - unknown = set(options) - _OPTIONS - if unknown: - raise ValueError(f"unknown AIMNet2 adapter options: {sorted(unknown)}") + _reject_unknown(options, _OPTIONS, "AIMNet2 adapter options") if ("model" in options) == ("checkpoint" in options): raise ValueError("AIMNet2 requires exactly one of model or checkpoint") if "model" in options: @@ -78,11 +76,7 @@ def _validate_options(options: Mapping[str, Any]) -> None: raise ValueError("AIMNet2 model_index must be an integer from 0 to 3") elif "model_index" in options: raise ValueError("AIMNet2 model_index is only valid with a named model") - for name in ("checkpoint", "cache_dir"): - if name in options: - value = options[name] - if not isinstance(value, str) or not value or not Path(value).is_absolute(): - raise ValueError(f"AIMNet2 {name} must be an absolute path string") + check_paths(options, ("checkpoint", "cache_dir"), "AIMNet2 {} must be an absolute path string") if options.get("device", "cuda") not in ("cpu", "cuda"): raise ValueError("AIMNet2 device must be cpu or cuda") if options.get("coulomb_method", "auto") not in ("auto", "dsf", "ewald", "pme"): @@ -124,19 +118,12 @@ def prepare(options: Mapping[str, Any], *, download: bool = True) -> dict[str, P def _electronic_state(atoms: Atoms, predictor: Any) -> None: - for name in ("charge", "spin"): - value = atoms.info.get(name) - if isinstance(value, bool) or not isinstance(value, Integral): - raise ValueError(f"AIMNet2 requires explicit integer atoms.info[{name!r}]") + require_charge_and_spin(atoms, "AIMNet2") spin = atoms.info["spin"] - if spin < 1: - raise ValueError("AIMNet2 spin must be a positive multiplicity") - if "mult" in atoms.info and ( - isinstance(atoms.info["mult"], bool) - or not isinstance(atoms.info["mult"], Integral) - or atoms.info["mult"] != spin - ): - raise ValueError("AIMNet2 atoms.info['mult'] must agree with 'spin'") + if "mult" in atoms.info: + mult = atoms.info["mult"] + if isinstance(mult, bool) or not isinstance(mult, Integral) or mult != spin: + raise ValueError("AIMNet2 atoms.info['mult'] must agree with 'spin'") if not predictor.is_nse and spin != 1: raise ValueError( "this AIMNet2 model supports singlets only; use aimnet2-nse for open shells" @@ -186,10 +173,8 @@ def calculate(self, atoms=None, properties=None, system_changes=all_changes): chosen = ("dsf" if atoms.pbc.any() else "simple") if method == "auto" else method if chosen in ("ewald", "pme") and not atoms.pbc.all(): raise ValueError("AIMNet2 Ewald/PME requires a fully periodic structure") - if ( - chosen != "simple" - and (self.base_calc.metadata or {}).get("coulomb_mode") == "full_embedded" - ): + mode = (self.base_calc.metadata or {}).get("coulomb_mode") + if chosen != "simple" and mode == "full_embedded": raise ValueError( "this AIMNet2 checkpoint cannot change its embedded Coulomb method" ) diff --git a/src/reactionflow/adapters/ani1xnr.py b/src/reactionflow/adapters/ani1xnr.py index 5f1398a..617886d 100644 --- a/src/reactionflow/adapters/ani1xnr.py +++ b/src/reactionflow/adapters/ani1xnr.py @@ -10,7 +10,7 @@ from ase.calculators.calculator import Calculator from .._durable import file_digest -from ..campaign import TrajectorySpec +from ..campaign import TrajectorySpec, _reject_unknown from ._model_files import model_cache, require_package from ._torch import TorchModelAdapter @@ -51,9 +51,7 @@ def catalog() -> dict[str, Any]: def _validate(options: Mapping[str, Any]) -> None: - unknown = set(options) - _ALLOWED_OPTIONS - if unknown: - raise ValueError(f"unknown ANI-1xnr adapter options: {sorted(unknown)}") + _reject_unknown(options, _ALLOWED_OPTIONS, "ANI-1xnr adapter options") if options.get("device", "cuda") not in ("cpu", "cuda"): raise ValueError("ANI-1xnr device must be 'cpu' or 'cuda'") if options.get("dtype", "float32") not in ("float32", "float64"): diff --git a/src/reactionflow/adapters/ase.py b/src/reactionflow/adapters/ase.py index ebda752..bee7051 100644 --- a/src/reactionflow/adapters/ase.py +++ b/src/reactionflow/adapters/ase.py @@ -19,10 +19,9 @@ from .._durable import file_digest from ..ase_npt import restore_langevin_baoab, snapshot_langevin_baoab -from ..campaign import TrajectorySpec +from ..campaign import TrajectorySpec, _reject_unknown from ..restart import ComponentState, ExactRestartSnapshot -_CALCULATOR_KIND = "reactionflow.ase-calculator" _ALLOWED_OPTIONS = { "calculator_factory", "calculator_kwargs", @@ -141,9 +140,7 @@ class ASELangevinBAOABAdapter: """Shared exact NVT/NPT runtime for adapters that lease ASE calculators.""" def __init__(self, *, trajectory: TrajectorySpec) -> None: - unknown_conditions = set(trajectory.conditions) - _ALLOWED_CONDITIONS - if unknown_conditions: - raise ValueError(f"unknown ASE trajectory conditions: {sorted(unknown_conditions)}") + _reject_unknown(trajectory.conditions, _ALLOWED_CONDITIONS, "ASE trajectory conditions") self.trajectory = trajectory def _lease(self) -> AbstractContextManager[tuple[Calculator, ComponentState]]: @@ -160,9 +157,7 @@ def _dynamics(self, atoms: Atoms, calculator: Calculator) -> LangevinBAOAB: # SSNEB rotates each endpoint cell into lower-triangular form but cannot rotate the # field with it, so the cell must start in that form and keep its shape. _, rotation = atoms.cell.standard_form() - if not bool(conditions.get("hydrostatic", True)) or not np.allclose( - rotation, np.eye(3) - ): + if not conditions.get("hydrostatic", True) or not np.allclose(rotation, np.eye(3)): raise ValueError( "an external field under pressure needs hydrostatic NPT and a " "lower-triangular cell; rotate the structure and field together first" @@ -171,11 +166,9 @@ def _dynamics(self, atoms: Atoms, calculator: Calculator) -> LangevinBAOAB: atoms.calc = calculator if "momenta" not in atoms.arrays: MaxwellBoltzmannDistribution( - atoms, - temperature_K=self.trajectory.temperature_K, - rng=rng, + atoms, temperature_K=self.trajectory.temperature_K, rng=rng ) - if bool(conditions.get("zero_total_momentum", True)): + if conditions.get("zero_total_momentum", True): Stationary(atoms, preserve_temperature=True) pressure = self.trajectory.pressure_GPa return LangevinBAOAB( @@ -184,24 +177,15 @@ def _dynamics(self, atoms: Atoms, calculator: Calculator) -> LangevinBAOAB: temperature_K=self.trajectory.temperature_K, externalstress=None if pressure is None else -pressure * units.GPa, hydrostatic=bool(conditions.get("hydrostatic", True)), - T_tau=_positive_float( - conditions.get("thermostat_tau_fs", 100.0), - "thermostat_tau_fs", - ) + T_tau=_positive_float(conditions.get("thermostat_tau_fs", 100.0), "thermostat_tau_fs") * units.fs, P_tau=( None if pressure is None - else _positive_float( - conditions.get("barostat_tau_fs", 1000.0), - "barostat_tau_fs", - ) + else _positive_float(conditions.get("barostat_tau_fs", 1000.0), "barostat_tau_fs") * units.fs ), - P_mass=_optional_positive_float( - conditions.get("barostat_mass"), - "barostat_mass", - ), + P_mass=_optional_positive_float(conditions.get("barostat_mass"), "barostat_mass"), disable_cell_langevin=bool(conditions.get("disable_cell_langevin", False)), rng=rng, logfile=None, @@ -240,15 +224,9 @@ class ASECalculatorAdapter(ASELangevinBAOABAdapter): """Load one deterministic ASE calculator from a campaign configuration.""" def __init__(self, *, trajectory: TrajectorySpec, options: Mapping[str, Any]) -> None: - unknown_options = set(options) - _ALLOWED_OPTIONS - if unknown_options: - raise ValueError(f"unknown ASE calculator adapter options: {sorted(unknown_options)}") + _reject_unknown(options, _ALLOWED_OPTIONS, "ASE calculator adapter options") factory = options.get("calculator_factory") - if ( - not isinstance(factory, str) - or factory.count(":") != 1 - or any(not part for part in factory.split(":")) - ): + if not isinstance(factory, str) or factory.count(":") != 1 or not all(factory.split(":")): raise ValueError("calculator_factory must be 'module:callable'") kwargs = options.get("calculator_kwargs", {}) if not isinstance(kwargs, Mapping) or any(not isinstance(key, str) for key in kwargs): @@ -256,10 +234,8 @@ def __init__(self, *, trajectory: TrajectorySpec, options: Mapping[str, Any]) -> super().__init__(trajectory=trajectory) self.calculator_factory = factory self.calculator_kwargs = dict(kwargs) - self.model_files = tuple( - Path(item).expanduser().resolve() - for item in _string_sequence(options.get("model_files", []), "model_files") - ) + files = _string_sequence(options.get("model_files", []), "model_files") + self.model_files = tuple(Path(item).expanduser().resolve() for item in files) self.packages = _string_sequence(options.get("packages", []), "packages") self.torch_device = options.get("torch_device") if self.torch_device not in {None, "cpu", "cuda"}: @@ -267,23 +243,14 @@ def __init__(self, *, trajectory: TrajectorySpec, options: Mapping[str, Any]) -> self._model_hashes: dict[str, str] | None = None self._contract: ComponentState | None = None - def _factory(self) -> tuple[Any, Any]: - module_name, factory_name = self.calculator_factory.split(":", 1) - module = import_module(module_name) - factory = getattr(module, factory_name, None) - if not callable(factory): - raise TypeError(f"ASE calculator factory {self.calculator_factory!r} is not callable") - return module, factory - def _model_file_hashes(self) -> dict[str, str]: - if self._model_hashes is not None: - return self._model_hashes - hashes: dict[str, str] = {} - for path in self.model_files: - if not path.is_file(): - raise FileNotFoundError(f"MLIP model file is missing: {path}") - hashes[str(path)] = file_digest(path) - self._model_hashes = hashes + if self._model_hashes is None: + hashes: dict[str, str] = {} + for path in self.model_files: + if not path.is_file(): + raise FileNotFoundError(f"MLIP model file is missing: {path}") + hashes[str(path)] = file_digest(path) + self._model_hashes = hashes return self._model_hashes def _package_versions(self, module_name: str) -> dict[str, str]: @@ -313,7 +280,7 @@ def _calculator_state( None if source_path is None or not source_path.is_file() else file_digest(source_path) ) self._contract = ComponentState( - kind=_CALCULATOR_KIND, + kind="reactionflow.ase-calculator", metadata={ "factory": self.calculator_factory, "calculator_type": calculator_type, @@ -333,7 +300,11 @@ def _calculator_state( @contextmanager def _lease(self) -> Iterator[tuple[Calculator, ComponentState]]: self._model_file_hashes() - module, factory = self._factory() + module_name, factory_name = self.calculator_factory.split(":", 1) + module = import_module(module_name) + factory = getattr(module, factory_name, None) + if not callable(factory): + raise TypeError(f"ASE calculator factory {self.calculator_factory!r} is not callable") torch_environment: Mapping[str, object] = {} if self.torch_device is None: calculator = factory(**self.calculator_kwargs) diff --git a/src/reactionflow/adapters/mace.py b/src/reactionflow/adapters/mace.py index 018e342..be588e3 100644 --- a/src/reactionflow/adapters/mace.py +++ b/src/reactionflow/adapters/mace.py @@ -4,7 +4,6 @@ from collections.abc import Mapping from copy import deepcopy -from numbers import Integral from pathlib import Path from typing import Any @@ -13,9 +12,9 @@ from ase.calculators.calculator import Calculator, all_changes from ase.calculators.fd import calculate_numerical_stress -from ..campaign import TrajectorySpec -from ._model_files import cached_download, model_cache, require_package -from ._torch import TorchModelAdapter +from ..campaign import TrajectorySpec, _reject_unknown +from ._model_files import cached_download, check_paths, model_cache, require_package +from ._torch import TorchModelAdapter, require_charge_and_spin MACE_VERSION = "0.3.16" _POLAR_STRAIN = 1e-5 @@ -53,7 +52,11 @@ def catalog() -> dict[str, Any]: """Describe the pinned models without importing MACE or downloading weights.""" - models = [{"name": name, "family": "mp"} for name in _MP_MODELS] + details: dict[str, dict[str, Any]] = { + "mh-0": {"description": "MACE-MH-0; select an explicit checkpoint head"}, + "mh-1": {"description": "MACE-MH-1 multi-head foundation model", "heads": list(_MH1_HEADS)}, + } + models = [{"name": name, "family": "mp", **details.get(name, {})} for name in _MP_MODELS] models.extend( {"name": name, "family": "off", "description": "MACE-OFF23 organic molecules"} for name in ("small", "medium", "large") @@ -62,13 +65,6 @@ def catalog() -> dict[str, Any]: {"name": f"polar-1-{size}", "family": "polar", "description": f"MACE-POLAR-1 {label}"} for size, label in (("s", "small"), ("m", "medium"), ("l", "large")) ) - for model in models: - if model["name"] == "mh-1": - model.update( - description="MACE-MH-1 multi-head foundation model", heads=list(_MH1_HEADS) - ) - elif model["name"] == "mh-0": - model.update(description="MACE-MH-0; select an explicit checkpoint head") return { "backend": "mace", "factory": "reactionflow.adapters.mace:create_adapter", @@ -98,23 +94,15 @@ def catalog() -> dict[str, Any]: def _validate(options: Mapping[str, Any]) -> None: - unknown = set(options) - _ALLOWED_OPTIONS - if unknown: - raise ValueError(f"unknown MACE adapter options: {sorted(unknown)}") + _reject_unknown(options, _ALLOWED_OPTIONS, "MACE adapter options") if ("model" in options) == ("checkpoint" in options): raise ValueError("MACE requires exactly one of 'model' or 'checkpoint'") family = options.get("family", "mp") if family not in ("mp", "off", "polar"): raise ValueError("MACE family must be 'mp', 'off', or 'polar'") - if "model" in options: - model = options["model"] - if not isinstance(model, str) or not model: - raise ValueError("MACE model must be an explicit non-empty model name") - for name in ("checkpoint", "cache_dir"): - if name in options: - value = options[name] - if not isinstance(value, str) or not value or not Path(value).is_absolute(): - raise ValueError(f"MACE {name} must be an absolute path") + if "model" in options and (not isinstance(options["model"], str) or not options["model"]): + raise ValueError("MACE model must be an explicit non-empty model name") + check_paths(options, ("checkpoint", "cache_dir"), "MACE {} must be an absolute path") if options.get("device", "cuda") not in ("cpu", "cuda"): raise ValueError("MACE device must be 'cpu' or 'cuda'") if options.get("dtype", "float64") not in ("float32", "float64"): @@ -158,8 +146,7 @@ def prepare(options: Mapping[str, Any], *, download: bool = True) -> dict[str, P raise ValueError( f"unknown MACE {family} model {model!r}; choose from {sorted(registry)}" ) - url = registry[model] - path = cached_download(url, model_cache(options, "mace"), download=download) + path = cached_download(registry[model], model_cache(options, "mace"), download=download) return {"checkpoint": path} @@ -174,12 +161,7 @@ def _require_polar() -> None: def _polar_inputs(atoms: Atoms) -> tuple[int, int, tuple[float, ...]]: - for name in ("charge", "spin"): - value = atoms.info.get(name) - if isinstance(value, bool) or not isinstance(value, Integral): - raise ValueError(f"MACE-POLAR requires explicit integer atoms.info[{name!r}]") - if name == "spin" and value < 1: - raise ValueError("MACE-POLAR spin must be a positive multiplicity") + require_charge_and_spin(atoms, "MACE-POLAR") field = np.asarray(atoms.info.get("external_field", [0.0, 0.0, 0.0]), dtype=float) if field.shape != (3,) or not np.isfinite(field).all(): raise ValueError("MACE-POLAR external_field must be a finite three-vector") @@ -261,13 +243,15 @@ def _new_calculator(self) -> Calculator: from mace.calculators import MACECalculator from mace.modules import PolarMACE + polar = self.options.get("family") == "polar" + dtype = self.options.get("dtype", "float64") calculator_class = MACECalculator kwargs = { "model_paths": [str(self.files["checkpoint"])], "device": self.device, - "default_dtype": self.options.get("dtype", "float64"), + "default_dtype": dtype, } - if self.options.get("family") == "polar": + if polar: _require_polar() kwargs["model_type"] = "PolarMACE" @@ -294,23 +278,17 @@ def calculate(self, atoms=None, properties=("energy",), system_changes=all_chang if "head" in self.options: kwargs["head"] = self.options["head"] calculator = calculator_class(**kwargs) - polar = self.options.get("family") == "polar" if any(isinstance(model, PolarMACE) != polar for model in calculator.models): raise ValueError( "MACE checkpoint architecture does not match family: PolarMACE checkpoints " - "require family='polar'; other checkpoints must " - "use family='mp' or 'off'" + "require family='polar'; other checkpoints must use family='mp' or 'off'" ) if "head" in self.options and calculator.head != self.options["head"]: raise ValueError( f"MACE did not select requested head {self.options['head']!r}; " f"resolved to {calculator.head!r}" ) - return ( - _NumericalPolarStressCalculator(calculator, dtype=self.options.get("dtype", "float64")) - if polar - else calculator - ) + return _NumericalPolarStressCalculator(calculator, dtype=dtype) if polar else calculator def _extra_metadata(self, calculator: Calculator) -> dict[str, Any]: metadata = {"resolved_head": calculator.head} diff --git a/src/reactionflow/adapters/uma.py b/src/reactionflow/adapters/uma.py index 960fbac..d284f18 100644 --- a/src/reactionflow/adapters/uma.py +++ b/src/reactionflow/adapters/uma.py @@ -6,16 +6,14 @@ import re from collections.abc import Mapping from importlib.metadata import distribution -from numbers import Integral from pathlib import Path from typing import Any -from ase import Atoms from ase.calculators.calculator import Calculator, all_changes -from ..campaign import TrajectorySpec -from ._model_files import model_cache, require_package -from ._torch import TorchModelAdapter +from ..campaign import TrajectorySpec, _reject_unknown +from ._model_files import check_paths, model_cache, require_package +from ._torch import TorchModelAdapter, require_charge_and_spin _VERSION = "2.23.0" _OPTIONS = { @@ -75,16 +73,8 @@ def catalog() -> dict[str, Any]: } -def _absolute_path(value: object, name: str) -> Path: - if not isinstance(value, str) or not value or not Path(value).is_absolute(): - raise ValueError(f"UMA {name} must be an absolute path string") - return Path(value) - - def _validate_options(options: Mapping[str, Any]) -> None: - unknown = set(options) - _OPTIONS - if unknown: - raise ValueError(f"unknown UMA adapter options: {sorted(unknown)}") + _reject_unknown(options, _OPTIONS, "UMA adapter options") if ("model" in options) == ("checkpoint" in options): raise ValueError("UMA requires exactly one of model or checkpoint") task = options.get("task") @@ -92,8 +82,7 @@ def _validate_options(options: Mapping[str, Any]) -> None: raise ValueError("UMA requires an explicit task") if options.get("device", "cuda") not in ("cpu", "cuda"): raise ValueError("UMA device must be cpu or cuda") - if "cache_dir" in options: - _absolute_path(options["cache_dir"], "cache_dir") + check_paths(options, ("cache_dir",), "UMA {} must be an absolute path string") if "model" in options: if not isinstance(options["model"], str) or not options["model"]: raise ValueError("UMA model must name an official UMA checkpoint") @@ -107,9 +96,8 @@ def _validate_options(options: Mapping[str, Any]) -> None: ) if {"atom_refs", "form_elem_refs"} & options.keys(): raise ValueError("UMA local reference paths are only valid with checkpoint") - if "revision" in options and ( - not isinstance(options["revision"], str) or not options["revision"] - ): + revision = options.get("revision") + if "revision" in options and (not isinstance(revision, str) or not revision): raise ValueError("UMA revision must be a non-empty Hugging Face revision") elif "revision" in options: raise ValueError("UMA revision is only valid with a named model") @@ -144,11 +132,9 @@ def prepare(options: Mapping[str, Any], *, download: bool = True) -> dict[str, P _validate_options(options) require_package("fairchem-core", _VERSION, "uma") if "checkpoint" in options: - files = { - name: _absolute_path(options[name], name) - for name in ("checkpoint", "atom_refs", "form_elem_refs") - if name in options - } + names = ("checkpoint", "atom_refs", "form_elem_refs") + check_paths(options, names, "UMA {} must be an absolute path string") + files = {name: Path(options[name]) for name in names if name in options} else: from huggingface_hub import hf_hub_download from huggingface_hub.errors import HfHubHTTPError, LocalEntryNotFoundError @@ -201,15 +187,6 @@ def fetch(file_spec: Mapping[str, Any]) -> Path: return {name: path.resolve() for name, path in files.items()} -def _validate_omol(atoms: Atoms) -> None: - for name in ("charge", "spin"): - value = atoms.info.get(name) - if isinstance(value, bool) or not isinstance(value, Integral): - raise ValueError(f"UMA omol requires explicit integer atoms.info[{name!r}]") - if name == "spin" and value < 1: - raise ValueError("UMA omol spin must be a positive multiplicity") - - def _deterministic_backbone(model: Any) -> None: """Make UMA's backbone give the same result for the same positions, on CPU and GPU. @@ -301,11 +278,11 @@ def _new_calculator(self) -> Calculator: class ExplicitOMOLCalculator(FAIRChemCalculator): def check_state(self, atoms, tol=1e-15): - _validate_omol(atoms) + require_charge_and_spin(atoms, "UMA omol") return super().check_state(atoms, tol=tol) def calculate(self, atoms=None, properties=("energy",), system_changes=all_changes): - _validate_omol(atoms if atoms is not None else self.atoms) + require_charge_and_spin(atoms if atoms is not None else self.atoms, "UMA omol") return super().calculate(atoms, properties, system_changes) return ExplicitOMOLCalculator(predictor, task_name=task) diff --git a/src/reactionflow/ase_npt.py b/src/reactionflow/ase_npt.py index 8ea0195..3970729 100644 --- a/src/reactionflow/ase_npt.py +++ b/src/reactionflow/ase_npt.py @@ -15,10 +15,6 @@ _STATE_KIND = "ase.langevin_baoab" -def _copy_attribute(dynamics: LangevinBAOAB, name: str) -> np.ndarray: - return np.asarray(getattr(dynamics, name)).copy() - - def snapshot_langevin_baoab(dynamics: LangevinBAOAB) -> ComponentState: """Capture every mutable quantity used by the next BAOAB step.""" @@ -44,12 +40,13 @@ def snapshot_langevin_baoab(dynamics: LangevinBAOAB) -> ComponentState: metadata["rng_bit_generator"] = type(dynamics.rng.bit_generator).__name__ metadata["rng_state"] = dynamics.rng.bit_generator.state - arrays = {"accel": _copy_attribute(dynamics, "accel")} + # ComponentState takes its own copy of each array. + arrays = {"accel": np.asarray(dynamics.accel)} if dynamics.externalstress is not None: arrays.update( - p_eps=_copy_attribute(dynamics, "p_eps"), - force_eps=_copy_attribute(dynamics, "force_eps"), - gamma_mod=_copy_attribute(dynamics, "gamma_mod"), + p_eps=np.asarray(dynamics.p_eps), + force_eps=np.asarray(dynamics.force_eps), + gamma_mod=np.asarray(dynamics.gamma_mod), ) return ComponentState(kind=_STATE_KIND, metadata=metadata, arrays=arrays) @@ -67,9 +64,7 @@ def _restore_rng(metadata: dict[str, Any]) -> np.random.Generator | None: def restore_langevin_baoab( - atoms: Atoms, - calculator: Calculator, - state: ComponentState, + atoms: Atoms, calculator: Calculator, state: ComponentState ) -> LangevinBAOAB: """Reconstruct BAOAB so its next step matches an uninterrupted run.""" diff --git a/src/reactionflow/campaign.py b/src/reactionflow/campaign.py index ec15796..d1843f4 100644 --- a/src/reactionflow/campaign.py +++ b/src/reactionflow/campaign.py @@ -4,18 +4,16 @@ import json import math -import re -from collections.abc import Mapping, Sequence -from dataclasses import dataclass, field +from collections.abc import Iterable, Mapping +from dataclasses import dataclass, field, fields from numbers import Integral, Real from pathlib import Path from typing import Any +from ._durable import SAFE_NAME from .pathway import _RETIRED_KEYS from .run import DEFAULT_MAX_HOURS_PER_BARRIER, ReactionRunConfig -_TRAJECTORY_ID = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]*\Z") - def _mapping(value: object, name: str) -> Mapping[str, Any]: if not isinstance(value, Mapping) or any(not isinstance(key, str) for key in value): @@ -23,10 +21,10 @@ def _mapping(value: object, name: str) -> Mapping[str, Any]: return value -def _positive_int(value: object, name: str) -> int: - if isinstance(value, bool) or not isinstance(value, Integral) or value < 1: - raise ValueError(f"{name} must be a positive integer") - return int(value) +def _reject_unknown(given: Iterable[str], allowed: Iterable[str], label: str) -> None: + unknown = set(given) - set(allowed) + if unknown: + raise ValueError(f"unknown {label}: {sorted(unknown)}") def _finite_float(value: object, name: str, *, minimum: float | None = None) -> float: @@ -57,15 +55,9 @@ class AdapterSpec: @classmethod def from_dict(cls, value: object) -> AdapterSpec: mapping = _mapping(value, "adapter") - unknown = set(mapping) - {"factory", "options"} - if unknown: - raise ValueError(f"unknown adapter keys: {sorted(unknown)}") + _reject_unknown(mapping, {"factory", "options"}, "adapter keys") factory = mapping.get("factory") - if ( - not isinstance(factory, str) - or factory.count(":") != 1 - or any(not part for part in factory.split(":")) - ): + if not isinstance(factory, str) or factory.count(":") != 1 or not all(factory.split(":")): raise ValueError("adapter.factory must be 'module:callable'") return cls( factory=factory, @@ -88,24 +80,9 @@ class TrajectorySpec: @classmethod def from_dict(cls, value: object) -> TrajectorySpec: mapping = _mapping(value, "trajectory") - allowed = { - "id", - "total_steps", - "timestep_fs", - "temperature_K", - "pressure_GPa", - "seed", - "conditions", - } - unknown = set(mapping) - allowed - if unknown: - raise ValueError(f"unknown trajectory keys: {sorted(unknown)}") + _reject_unknown(mapping, [item.name for item in fields(TrajectorySpec)], "trajectory keys") trajectory_id = mapping.get("id") - if ( - not isinstance(trajectory_id, str) - or not _TRAJECTORY_ID.fullmatch(trajectory_id) - or trajectory_id in {".", ".."} - ): + if not isinstance(trajectory_id, str) or not SAFE_NAME.fullmatch(trajectory_id): raise ValueError("trajectory.id must be a safe non-empty file name") seed = mapping.get("seed") if isinstance(seed, bool) or not isinstance(seed, Integral) or not 0 <= seed < 2**64: @@ -118,9 +95,12 @@ def from_dict(cls, value: object) -> TrajectorySpec: ) if timestep_fs == 0: raise ValueError("trajectory.timestep_fs must be positive") + steps = mapping.get("total_steps") + if isinstance(steps, bool) or not isinstance(steps, Integral) or steps < 1: + raise ValueError("trajectory.total_steps must be a positive integer") return cls( id=trajectory_id, - total_steps=_positive_int(mapping.get("total_steps"), "trajectory.total_steps"), + total_steps=int(steps), timestep_fs=timestep_fs, temperature_K=_finite_float( mapping.get("temperature_K"), @@ -138,18 +118,14 @@ def from_dict(cls, value: object) -> TrajectorySpec: def _run_config(value: object) -> ReactionRunConfig: overrides = _mapping(value, "reaction_run") merged = ReactionRunConfig().to_dict() - unknown = set(overrides) - set(merged) - if unknown: - raise ValueError(f"unknown reaction_run keys: {sorted(unknown)}") + _reject_unknown(overrides, merged, "reaction_run keys") for key, item in overrides.items(): if key in {"detector", "pathway"}: nested = _mapping(item, f"reaction_run.{key}") - known = set(merged[key]) # type: ignore[arg-type] + known = set(merged[key]) # type: ignore[call-overload] if key == "pathway": known |= _RETIRED_KEYS - unknown_nested = set(nested) - known - if unknown_nested: - raise ValueError(f"unknown reaction_run.{key} keys: {sorted(unknown_nested)}") + _reject_unknown(nested, known, f"reaction_run.{key} keys") merged[key] = {**merged[key], **nested} # type: ignore[dict-item] else: merged[key] = item @@ -179,7 +155,7 @@ def load(cls, path: str | Path) -> CampaignConfig: raise ValueError("unsupported campaign schema; use schema_version 2") if value.get("mode", "md") != "md": raise ValueError("MD campaign requires mode 'md'") - unknown = set(value) - { + known = ( "schema_version", "mode", "structure", @@ -190,9 +166,8 @@ def load(cls, path: str | Path) -> CampaignConfig: "refine_after_md", "adapter_profiles", "max_hours_per_barrier", - } - if unknown: - raise ValueError(f"unknown campaign keys: {sorted(unknown)}") + ) + _reject_unknown(value, known, "campaign keys") output_value = value.get("output_root") if not isinstance(output_value, str) or not output_value: raise ValueError("campaign.output_root must be a path string") @@ -206,29 +181,23 @@ def structure_path(raw: object, name: str) -> Path: return path # A trajectory's own structure overrides the campaign-wide one. - default_structure = ( - structure_path(value["structure"], "campaign.structure") - if "structure" in value - else None - ) + default_structure = None + if "structure" in value: + default_structure = structure_path(value["structure"], "campaign.structure") raw_profiles = _mapping(value.get("adapter_profiles"), "campaign.adapter_profiles") if not raw_profiles: raise ValueError("campaign.adapter_profiles must be a non-empty object") adapter_profiles: dict[str, AdapterSpec] = {} for name, raw_profile in raw_profiles.items(): - if not _TRAJECTORY_ID.fullmatch(name) or name in {".", ".."}: + if not SAFE_NAME.fullmatch(name): raise ValueError( f"adapter profile name must be a safe non-empty identifier: {name!r}" ) adapter_profiles[name] = AdapterSpec.from_dict(raw_profile) raw_trajectories = value.get("trajectories") - if ( - not isinstance(raw_trajectories, Sequence) - or isinstance(raw_trajectories, (str, bytes)) - or not raw_trajectories - ): + if not isinstance(raw_trajectories, list) or not raw_trajectories: raise ValueError("campaign.trajectories must be a non-empty array") trajectories: list[TrajectorySpec] = [] trajectory_profiles: list[str] = [] @@ -251,8 +220,7 @@ def structure_path(raw: object, name: str) -> Path: trajectories.append(TrajectorySpec.from_dict(trajectory)) trajectory_profiles.append(profile) trajectory_structures.append(structure) - identifiers = [trajectory.id for trajectory in trajectories] - if len(set(identifiers)) != len(identifiers): + if len({trajectory.id for trajectory in trajectories}) != len(trajectories): raise ValueError("trajectory IDs must be unique") max_hours = value.get("max_hours_per_barrier", DEFAULT_MAX_HOURS_PER_BARRIER) if max_hours is not None: diff --git a/src/reactionflow/candidates.py b/src/reactionflow/candidates.py index 7c66714..ad806d2 100644 --- a/src/reactionflow/candidates.py +++ b/src/reactionflow/candidates.py @@ -4,8 +4,7 @@ import json from collections.abc import Collection -from dataclasses import dataclass -from numbers import Integral +from dataclasses import dataclass, replace from pathlib import Path import networkx as nx @@ -13,7 +12,7 @@ from ase.io import read, write from ._durable import file_digest -from .detection import Bond, atom_ids, canonical_copy, transport_copy +from .detection import Bond, _is_count, _symbols_by_id, canonical_copy, transport_copy @dataclass(frozen=True, slots=True) @@ -42,19 +41,17 @@ class _ReactantOrigin: @dataclass(slots=True) class _PendingTopology: atom_ids: tuple[int, ...] - reactant: Atoms - reactant_bonds: frozenset[Bond] - reactant_frame: int + reactant: _ReactantOrigin bonds: frozenset[Bond] product: Atoms product_frame: int + origins: tuple[_ReactantOrigin, ...] count: int = 0 - origins: tuple[_ReactantOrigin, ...] = () def _within(bonds: frozenset[Bond], region: Collection[int]) -> frozenset[Bond]: ids = set(region) - return frozenset(bond for bond in bonds if set(bond) <= ids) + return frozenset(bond for bond in bonds if bond[0] in ids and bond[1] in ids) def _select_origin( @@ -90,72 +87,34 @@ def _bonds(values: Collection[Bond], valid_ids: set[int]) -> frozenset[Bond]: def _changed_regions( - reactant: frozenset[Bond], - product: frozenset[Bond], - stable: frozenset[Bond] | None = None, - radius: int | None = None, + reactant: frozenset[Bond], product: frozenset[Bond], stable: frozenset[Bond], radius: int | None ) -> tuple[tuple[int, ...], ...]: - stable = product if stable is None else stable changed_atoms = { atom_id for bond in (reactant ^ product) | (stable ^ product) for atom_id in bond } - neighbors: dict[int, set[int]] = {} - for first, second in reactant | product | stable: - neighbors.setdefault(first, set()).add(second) - neighbors.setdefault(second, set()).add(first) - - if radius is not None: - # Each changed atom and the atoms within `radius` bonds of it; overlapping ones merge. - merged: list[set[int]] = [] - for atom_id in sorted(changed_atoms): - region = frontier = {atom_id} - for _ in range(radius): - frontier = {n for current in frontier for n in neighbors.get(current, ())} - region - region = region | frontier - for other in [item for item in merged if item & region]: - merged.remove(other) - region |= other - merged.append(region) - return tuple(sorted(tuple(sorted(region)) for region in merged)) - - regions: list[tuple[int, ...]] = [] - unseen = set(changed_atoms) - while unseen: - start = min(unseen) - region = {start} - stack = [start] - while stack: - current = stack.pop() - for neighbor in neighbors.get(current, ()): - if neighbor not in region: - region.add(neighbor) - stack.append(neighbor) - unseen -= region - regions.append(tuple(sorted(region))) - return tuple(sorted(regions)) + if not changed_atoms: + return () + graph = nx.from_edgelist(reactant | product | stable) + # Each changed atom and the atoms within `radius` bonds of it, or its whole bonded component + # if `radius` is None; overlapping ones merge. + if radius is None: + merged = [c for c in nx.connected_components(graph) if not changed_atoms.isdisjoint(c)] + else: + regions = nx.utils.UnionFind() + for atom_id in changed_atoms: + distances = nx.single_source_shortest_path_length(graph, atom_id, cutoff=radius) + regions.union(*distances) + merged = regions.to_sets() + return tuple(sorted(tuple(sorted(region)) for region in merged)) def _reaction_graph( candidate: ReactionCandidate, *, reverse: bool = False, radius: int | None = None ) -> nx.Graph: - if radius is not None and ( - isinstance(radius, bool) or not isinstance(radius, Integral) or radius < 0 - ): + if radius is not None and not _is_count(radius, 0): raise ValueError("radius must be a non-negative integer or None") - reactant_symbols = dict( - zip( - atom_ids(candidate.reactant), - candidate.reactant.get_chemical_symbols(), - strict=True, - ) - ) - product_symbols = dict( - zip( - atom_ids(candidate.product), - candidate.product.get_chemical_symbols(), - strict=True, - ) - ) + reactant_symbols = _symbols_by_id(candidate.reactant) + product_symbols = _symbols_by_id(candidate.product) region = set(candidate.atom_ids) if ( not region <= reactant_symbols.keys() @@ -164,22 +123,17 @@ def _reaction_graph( ): raise ValueError("candidate endpoint identities do not match") - reactant = set(_bonds(candidate.reactant_bonds, region)) - product = set(_bonds(candidate.product_bonds, region)) + reactant = _bonds(candidate.reactant_bonds, region) + product = _bonds(candidate.product_bonds, region) + if reverse: + reactant, product = product, reactant graph = nx.Graph() graph.add_nodes_from( (atom_id, {"element": reactant_symbols[atom_id]}) for atom_id in candidate.atom_ids ) - for first, second in reactant | product: - if (first, second) in reactant and (first, second) in product: - change = "unchanged" - elif (first, second) in product: - change = "formed" - else: - change = "broken" - if reverse: - change = {"formed": "broken", "broken": "formed"}.get(change, change) - graph.add_edge(first, second, change=change) + graph.add_edges_from(reactant & product, change="unchanged") + graph.add_edges_from(product - reactant, change="formed") + graph.add_edges_from(reactant - product, change="broken") if radius is None: return graph # The local change: the atoms of the changed bonds and every atom within `radius` bonds. @@ -191,6 +145,10 @@ def _reaction_graph( return graph.subgraph(nodes).copy() +def _directions(candidate: ReactionCandidate, radius: int | None) -> list[nx.Graph]: + return [_reaction_graph(candidate, reverse=reverse, radius=radius) for reverse in (False, True)] + + def _isomorphic(first: nx.Graph, second: nx.Graph) -> bool: return nx.is_isomorphic( first, @@ -214,18 +172,13 @@ def same_reaction( """ graph = _reaction_graph(first, radius=radius) - return _isomorphic(graph, _reaction_graph(second, radius=radius)) or _isomorphic( - graph, _reaction_graph(second, reverse=True, radius=radius) - ) + return any(_isomorphic(graph, direction) for direction in _directions(second, radius)) def reaction_key(candidate: ReactionCandidate, *, radius: int | None = None) -> str: """Hash shared by forward/reverse-equivalent candidates; confirm matches with same_reaction.""" - return min( - _graph_hash(_reaction_graph(candidate, reverse=reverse, radius=radius)) - for reverse in (False, True) - ) + return min(map(_graph_hash, _directions(candidate, radius))) class ReactionClasses: @@ -241,10 +194,7 @@ def __init__(self, radius: int | None = None) -> None: def find(self, candidate: ReactionCandidate) -> str | None: """Return the name of the class equivalent to this candidate, if one is known.""" - directions = [ - _reaction_graph(candidate, reverse=reverse, radius=self.radius) - for reverse in (False, True) - ] + directions = _directions(candidate, self.radius) for name, graph in self._graphs.get(min(map(_graph_hash, directions)), ()): if any(_isomorphic(graph, direction) for direction in directions): return name @@ -265,17 +215,9 @@ class ReactionTracker: """ def __init__(self, *, stability_frames: int = 3, region_radius: int | None = 2) -> None: - if ( - isinstance(stability_frames, bool) - or not isinstance(stability_frames, Integral) - or stability_frames < 1 - ): + if not _is_count(stability_frames, 1): raise ValueError("stability_frames must be a positive integer") - if region_radius is not None and ( - isinstance(region_radius, bool) - or not isinstance(region_radius, Integral) - or region_radius < 0 - ): + if region_radius is not None and not _is_count(region_radius, 0): raise ValueError("region_radius must be a non-negative integer or None") self.stability_frames = int(stability_frames) self.region_radius = None if region_radius is None else int(region_radius) @@ -295,16 +237,18 @@ def last_frame(self) -> int | None: def _candidate( self, pending: _PendingTopology, *, observed_frame: int, resolved: bool ) -> ReactionCandidate | None: - reactant_bonds = _within(pending.reactant_bonds, pending.atom_ids) + """The region's change, or None for a region that ended where it began.""" + + reactant_bonds = _within(pending.reactant.bonds, pending.atom_ids) if reactant_bonds == pending.bonds: return None return ReactionCandidate( - reactant=pending.reactant.copy(), + reactant=pending.reactant.atoms.copy(), product=pending.product.copy(), atom_ids=pending.atom_ids, reactant_bonds=reactant_bonds, product_bonds=pending.bonds, - reactant_frame=pending.reactant_frame, + reactant_frame=pending.reactant.frame, product_frame=pending.product_frame, observed_frame=observed_frame, resolved=resolved, @@ -328,12 +272,12 @@ def process( if self._last_frame is not None and frame <= self._last_frame: raise ValueError("frame numbers must increase") - ids = atom_ids(atoms) - symbols = dict(zip(ids, atoms.get_chemical_symbols(), strict=True)) + symbols = _symbols_by_id(atoms) if self._symbols is not None and symbols != self._symbols: raise ValueError("atom identities changed between frames") - bonds = _bonds(stable_bonds, set(ids)) - proposal = bonds if pending_bonds is None else _bonds(pending_bonds, set(ids)) + ids = set(symbols) + bonds = _bonds(stable_bonds, ids) + proposal = bonds if pending_bonds is None else _bonds(pending_bonds, ids) if self._accepted_bonds is None and pending_bonds is not None: raise ValueError("start the tracker before the detector has pending changes") snapshot = atoms.copy() @@ -367,11 +311,7 @@ def process( for origin in item.origins: support = tuple(sorted(set(origin.atom_ids).intersection(region))) if support: - origins.append( - _ReactantOrigin( - support, origin.atoms, origin.bonds, origin.frame - ) - ) + origins.append(replace(origin, atom_ids=support)) covered.update(support) added = tuple(sorted(set(region) - covered)) if added: @@ -380,28 +320,24 @@ def process( added, self._accepted, self._last_bonds, self._accepted_frame ) ) - reactant = _select_origin(origins, region, self._accepted_bonds, local_product) pending = _PendingTopology( atom_ids=region, - reactant=reactant.atoms, - reactant_bonds=reactant.bonds, - reactant_frame=reactant.frame, + reactant=_select_origin( + origins, region, self._accepted_bonds, local_product + ), bonds=local_product, product=snapshot, product_frame=frame, origins=tuple(origins), ) - if _within(bonds ^ proposal, region): - pending.count = 0 - else: - pending.count += 1 + pending.count = 0 if _within(bonds ^ proposal, region) else pending.count + 1 if pending.count < self.stability_frames: remaining.append(pending) continue # If a region joins an independently accepted event, the oldest snapshot can # precede that event. Keep this overlapping history as unresolved, rather than # claiming a reactant topology that does not match the retained geometry. - resolved = _within(pending.reactant_bonds, region) == _within( + resolved = _within(pending.reactant.bonds, region) == _within( self._accepted_bonds, region ) candidate = self._candidate(pending, observed_frame=frame, resolved=resolved) @@ -494,10 +430,7 @@ def write_checkpoint(self, path: str | Path) -> Path: @classmethod def read_checkpoint( - cls, - path: str | Path, - *, - stability_frames: int | None = None, + cls, path: str | Path, *, stability_frames: int | None = None ) -> ReactionTracker: """Restore a tracker written by ``write_checkpoint``.""" @@ -538,13 +471,10 @@ def read_checkpoint( ) for origin in item["origins"] ) - reactant = _select_origin(origins, region, tracker._accepted_bonds, bonds) tracker._pending.append( _PendingTopology( atom_ids=region, - reactant=reactant.atoms, - reactant_bonds=reactant.bonds, - reactant_frame=reactant.frame, + reactant=_select_origin(origins, region, tracker._accepted_bonds, bonds), bonds=bonds, product=snapshots[item["product"]], product_frame=int(item["product_frame"]), @@ -561,9 +491,7 @@ def _checkpoint_bonds(values: Collection[Bond] | None) -> list[list[int]] | None return [list(bond) for bond in sorted(values)] -def _restore_bonds( - values: Collection[Collection[int]] | None, -) -> frozenset[Bond] | None: +def _restore_bonds(values: Collection[Collection[int]] | None) -> frozenset[Bond] | None: if values is None: return None return frozenset( diff --git a/src/reactionflow/cli.py b/src/reactionflow/cli.py index 70f1ce3..35f47e7 100644 --- a/src/reactionflow/cli.py +++ b/src/reactionflow/cli.py @@ -21,6 +21,7 @@ from ._durable import ensure_directory, file_digest, publish, sync_directory from ._version import __version__ from .campaign import CampaignConfig +from .detection import _is_count from .mlip import load_mlip_adapter from .model_setup import install_and_prepare, model_catalog, prepare_adapter, prepare_campaign from .pathway import check_saddle @@ -34,10 +35,7 @@ def _trajectory_contract(campaign: CampaignConfig, index: int) -> dict[str, obje trajectory = campaign.trajectory(index) adapter = campaign.adapter_for(index) scientific_config = { - "adapter": { - "factory": adapter.factory, - "options": dict(adapter.options), - }, + "adapter": {"factory": adapter.factory, "options": dict(adapter.options)}, "reaction_run": campaign.reaction_run.to_dict(), "trajectory": asdict(trajectory), } @@ -55,11 +53,15 @@ def _trajectory_contract(campaign: CampaignConfig, index: int) -> dict[str, obje } +def _publish_json(value: dict[str, object], temporary: Path, path: Path) -> None: + temporary.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n", encoding="utf-8") + publish(temporary, path) + + def _bind_trajectory_contract(root: Path, contract: dict[str, object]) -> None: path = root / "trajectory-contract.json" if path.is_file(): - stored = json.loads(path.read_text(encoding="utf-8")) - if stored == contract: + if json.loads(path.read_text(encoding="utf-8")) == contract: sync_directory(root) return if (root / "state.json").exists(): @@ -71,11 +73,7 @@ def _bind_trajectory_contract(root: Path, contract: dict[str, object]) -> None: ensure_directory(root) temporary = root / f".trajectory-contract-{uuid4().hex}.tmp" try: - temporary.write_text( - json.dumps(contract, indent=2, sort_keys=True) + "\n", - encoding="utf-8", - ) - publish(temporary, path) + _publish_json(contract, temporary, path) except Exception: temporary.unlink(missing_ok=True) raise @@ -98,11 +96,7 @@ def resolve_task_index( ) -> int: """Map one Slurm process—or one explicit local index—to one trajectory.""" - if ( - isinstance(campaign_size, bool) - or not isinstance(campaign_size, Integral) - or campaign_size < 1 - ): + if not _is_count(campaign_size, 1): raise ValueError("campaign size must be a positive integer") variables = os.environ if environment is None else environment slurm_tasks = _environment_integer(variables, "SLURM_NTASKS") @@ -159,8 +153,7 @@ def run_selected_trajectory( if download: prepare_adapter(adapter_spec) root = campaign.output_root / trajectory.id - state_path = root / "state.json" - if state_path.is_file(): + if (root / "state.json").is_file(): _bind_trajectory_contract(root, _trajectory_contract(campaign, index)) adapter = load_mlip_adapter(adapter_spec, trajectory) run = ReactionRun.open(root) @@ -214,8 +207,7 @@ def check_band( path = band.with_name(f"{band.stem}.check.json") temporary = path.with_name(f".{path.name}-{uuid4().hex}.tmp") try: - temporary.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8") - publish(temporary, path) + _publish_json(result, temporary, path) finally: temporary.unlink(missing_ok=True) return path, result @@ -345,6 +337,34 @@ def _parser() -> argparse.ArgumentParser: return parser +def _run_and_record(campaign: CampaignConfig, arguments: argparse.Namespace) -> None: + index = resolve_task_index(len(campaign.trajectories), arguments.index) + trajectory_id = campaign.trajectory(index).id + payload: dict[str, object] = { + "trajectory_id": trajectory_id, + "trajectory_index": index, + "hostname": socket.gethostname(), + } + # Every Slurm task shares one log, so each worker also records its own error by trajectory. + error_path = campaign.output_root / trajectory_id / "last-error.json" + with suppress(OSError): + error_path.unlink() + try: + summary = run_selected_trajectory(campaign, index=index, download=arguments.download) + except Exception as error: + payload["error"] = f"{type(error).__name__}: {error}" + print(json.dumps(payload, sort_keys=True)) + temporary = error_path.with_name(f".last-error-{uuid4().hex}.tmp") + with suppress(OSError): + try: + error_path.parent.mkdir(parents=True, exist_ok=True) + _publish_json(payload, temporary, error_path) + finally: + temporary.unlink(missing_ok=True) + raise + print(json.dumps({**payload, **asdict(summary)}, sort_keys=True)) + + def main(argv: Sequence[str] | None = None) -> int: parser = _parser() arguments = parser.parse_args(argv) @@ -388,23 +408,19 @@ def main(argv: Sequence[str] | None = None) -> int: return install_and_prepare(campaign, index=arguments.index) models = prepare_campaign(campaign, index=arguments.index) print(json.dumps({"campaign": str(campaign.source), "models": models}, sort_keys=True)) - return 0 - if arguments.command == "status": + elif arguments.command == "status": if arguments.radius < 0: parser.error("--radius must not be negative") report = campaign_status(campaign, radius=arguments.radius) - print( - json.dumps(report, indent=2, sort_keys=True) - if arguments.json - else format_status(report) - ) - return 0 - if arguments.command == "scan": + if arguments.json: + print(json.dumps(report, indent=2, sort_keys=True)) + else: + print(format_status(report)) + elif arguments.command == "scan": if arguments.frame_ps is not None: parser.error("--frame-ps is for a trajectory file; a run records its own times") - indices = range(len(campaign.trajectories)) - if arguments.index is not None: - indices = [resolve_task_index(len(indices), arguments.index, environment={})] + count = len(campaign.trajectories) + indices = range(count) if arguments.index is None else (arguments.index,) for index in indices: trajectory = campaign.trajectory(index) root = campaign.output_root / trajectory.id @@ -422,8 +438,7 @@ def main(argv: Sequence[str] | None = None) -> int: ) out = root / "scan" if arguments.out is None else arguments.out / trajectory.id _print_scan(result.write(out), result, trajectory_id=trajectory.id) - return 0 - if arguments.command == "validate": + elif arguments.command == "validate": payload = { "campaign": str(campaign.source), "trajectories": len(campaign.trajectories), @@ -435,9 +450,7 @@ def main(argv: Sequence[str] | None = None) -> int: "adapter_profile_counts": dict(Counter(campaign.trajectory_profiles)), } print(json.dumps(payload, sort_keys=True)) - return 0 - - if arguments.command == "check": + elif arguments.command == "check": if arguments.index is None and len(campaign.trajectories) > 1: parser.error("--index picks the trajectory when a campaign has several") # Slurm tasks pick trajectories for run only, so a check inside any job reads --index. @@ -450,37 +463,8 @@ def main(argv: Sequence[str] | None = None) -> int: path, result = check_band(campaign, arguments.band, index=index, pressure_GPa=pressure) summary = {key: result[key] for key in ("status", "barrier_eV", "message")} print(json.dumps({"check": str(path), **summary}, sort_keys=True)) - return 0 - - index = resolve_task_index(len(campaign.trajectories), arguments.index) - trajectory_id = campaign.trajectory(index).id - payload: dict[str, object] = { - "trajectory_id": trajectory_id, - "trajectory_index": index, - "hostname": socket.gethostname(), - } - # Every Slurm task shares one log, so each worker also records its own error by trajectory. - error_path = campaign.output_root / trajectory_id / "last-error.json" - with suppress(OSError): - error_path.unlink() - try: - summary = run_selected_trajectory(campaign, index=index, download=arguments.download) - except Exception as error: - payload["error"] = f"{type(error).__name__}: {error}" - print(json.dumps(payload, sort_keys=True)) - temporary = error_path.with_name(f".last-error-{uuid4().hex}.tmp") - with suppress(OSError): - try: - error_path.parent.mkdir(parents=True, exist_ok=True) - temporary.write_text( - json.dumps(payload, indent=2, sort_keys=True) + "\n", - encoding="utf-8", - ) - publish(temporary, error_path) - finally: - temporary.unlink(missing_ok=True) - raise - print(json.dumps({**payload, **asdict(summary)}, sort_keys=True)) + else: + _run_and_record(campaign, arguments) return 0 diff --git a/src/reactionflow/detection.py b/src/reactionflow/detection.py index 41a365d..09dc07a 100644 --- a/src/reactionflow/detection.py +++ b/src/reactionflow/detection.py @@ -4,6 +4,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass, field +from functools import cache from numbers import Integral from typing import Any, Literal @@ -27,6 +28,10 @@ def _element_pair(value: object) -> ElementPair: return tuple(sorted((first, second))) # type: ignore[return-value] +def _is_count(value: object, minimum: int) -> bool: + return isinstance(value, Integral) and not isinstance(value, bool) and not value < minimum + + def _bond(first: int, second: int) -> Bond: if first == second: raise ValueError("a bond must connect two distinct atom IDs") @@ -43,14 +48,23 @@ def atom_ids(atoms: Atoms) -> tuple[int, ...]: raise ValueError("stable atom IDs are required") if getattr(raw, "dtype", None) is not None and not np.issubdtype(raw.dtype, np.integer): raise ValueError("atom IDs must be integers") - if any(isinstance(value, (bool, np.bool_)) or not isinstance(value, Integral) for value in raw): - raise ValueError("atom IDs must be integers") - ids = tuple(int(value) for value in raw) + if type(raw) is np.ndarray and raw.ndim == 1 and raw.dtype.kind in "iu": + ids = tuple(raw.tolist()) # an integer array holds only integers + else: + if any( + isinstance(value, (bool, np.bool_)) or not isinstance(value, Integral) for value in raw + ): + raise ValueError("atom IDs must be integers") + ids = tuple(int(value) for value in raw) if len(ids) != len(atoms) or len(set(ids)) != len(ids): raise ValueError("each atom needs one unique stable ID") return ids +def _symbols_by_id(atoms: Atoms) -> dict[int, str]: + return dict(zip(atom_ids(atoms), atoms.get_chemical_symbols(), strict=True)) + + def assign_atom_ids(atoms: Atoms) -> Atoms: """Assign missing IDs once and store them in the canonical ASE array.""" @@ -109,11 +123,7 @@ def __post_init__(self) -> None: breaking = float(self.break_scale) if not 0 < form < breaking: raise ValueError("scales must satisfy 0 < form_scale < break_scale") - if ( - isinstance(self.persistence_frames, bool) - or not isinstance(self.persistence_frames, Integral) - or self.persistence_frames < 1 - ): + if not _is_count(self.persistence_frames, 1): raise ValueError("persistence_frames must be positive") thresholds = { _element_pair(pair): _threshold(threshold) @@ -126,7 +136,7 @@ def __post_init__(self) -> None: def threshold_for(self, first: str, second: str) -> tuple[float, float]: pair = _element_pair((first, second)) if pair in self.pair_thresholds: - return self.pair_thresholds[pair] # type: ignore[index,return-value] + return self.pair_thresholds[pair] # type: ignore[return-value] radius = covalent_radii[atomic_numbers[first]] + covalent_radii[atomic_numbers[second]] return self.form_scale * radius, self.break_scale * radius @@ -171,7 +181,7 @@ def bond_distances(atoms: Atoms, config: BondDetectorConfig) -> dict[Bond, float def classify_bonds(atoms: Atoms, config: BondDetectorConfig) -> tuple[set[Bond], set[Bond]]: """Return the pairs within their formation distance and those inside the hysteresis gap.""" - symbols = dict(zip(atom_ids(atoms), atoms.get_chemical_symbols(), strict=True)) + symbols = _symbols_by_id(atoms) bonded: set[Bond] = set() ambiguous: set[Bond] = set() for bond, distance in bond_distances(atoms, config).items(): @@ -253,12 +263,9 @@ def pending_changes(self) -> frozenset[tuple[Bond, EventType]]: def last_frame(self) -> int | None: return self._last_frame - def _threshold(self, bond: Bond) -> tuple[float, float]: + def _elements(self, bond: Bond) -> ElementPair: assert self._numbers is not None - return self.config.threshold_for( - chemical_symbols[self._numbers[bond[0]]], - chemical_symbols[self._numbers[bond[1]]], - ) + return chemical_symbols[self._numbers[bond[0]]], chemical_symbols[self._numbers[bond[1]]] def _distances(self, atoms: Atoms, ids: tuple[int, ...]) -> dict[Bond, float]: distances = bond_distances(atoms, self.config) @@ -266,11 +273,7 @@ def _distances(self, atoms: Atoms, ids: tuple[int, ...]) -> dict[Bond, float]: indices = {atom_id: index for index, atom_id in enumerate(ids)} for bond in self._stable - distances.keys(): distances[bond] = float( - atoms.get_distance( - indices[bond[0]], - indices[bond[1]], - mic=bool(atoms.pbc.any()), - ) + atoms.get_distance(indices[bond[0]], indices[bond[1]], mic=bool(atoms.pbc.any())) ) return distances @@ -281,15 +284,18 @@ def process(self, atoms: Atoms, *, frame: int) -> list[BondEvent]: raise ValueError("frame numbers must increase") ids = atom_ids(atoms) numbers = dict(zip(ids, map(int, atoms.numbers), strict=True)) - if self._numbers is not None and numbers != self._numbers: - raise ValueError("atom identities changed between frames") if self._numbers is None: self._numbers = numbers + elif numbers != self._numbers: + raise ValueError("atom identities changed between frames") distances = self._distances(atoms, ids) + thresholds = cache(lambda first, second: self.config.threshold_for(first, second)) if self._last_frame is None: self._stable = { - bond for bond, distance in distances.items() if distance <= self._threshold(bond)[0] + bond + for bond, distance in distances.items() + if distance <= thresholds(*self._elements(bond))[0] } self._last_frame = frame return [] @@ -297,7 +303,8 @@ def process(self, atoms: Atoms, *, frame: int) -> list[BondEvent]: active: set[tuple[Bond, EventType]] = set() events: list[BondEvent] = [] for bond, distance in sorted(distances.items()): - form_distance, break_distance = self._threshold(bond) + pair = self._elements(bond) + form_distance, break_distance = thresholds(*pair) event_type: EventType | None = None if bond in self._stable and distance >= break_distance: event_type = "broken" @@ -321,10 +328,7 @@ def process(self, atoms: Atoms, *, frame: int) -> list[BondEvent]: BondEvent( event_type=event_type, atom_ids=bond, - symbols=( - chemical_symbols[numbers[bond[0]]], - chemical_symbols[numbers[bond[1]]], - ), + symbols=pair, distance=distance, form_distance=form_distance, break_distance=break_distance, @@ -363,10 +367,7 @@ def export_state(self) -> dict[str, Any]: @classmethod def from_state( - cls, - state: Mapping[str, Any], - *, - config: BondDetectorConfig | None = None, + cls, state: Mapping[str, Any], *, config: BondDetectorConfig | None = None ) -> BondChangeDetector: """Restore a detector from a version-1 checkpoint.""" @@ -387,8 +388,7 @@ def from_state( bond = _bond(*map(int, item["atom_ids"])) event_type = item["event_type"] detector._pending[(bond, event_type)] = _PendingChange( - first_seen_frame=int(item["first_seen_frame"]), - count=int(item["count"]), + first_seen_frame=int(item["first_seen_frame"]), count=int(item["count"]) ) detector._last_frame = state.get("last_frame") return detector diff --git a/src/reactionflow/model_setup.py b/src/reactionflow/model_setup.py index 22f9f2f..0f511ea 100644 --- a/src/reactionflow/model_setup.py +++ b/src/reactionflow/model_setup.py @@ -6,6 +6,7 @@ import subprocess import sys from collections.abc import Iterable +from contextlib import suppress from importlib import import_module from importlib.metadata import PackageNotFoundError, version from pathlib import Path @@ -32,17 +33,10 @@ def model_catalog(backend: str | None = None) -> list[dict]: } if backend is not None: backend = backend.replace("-", "_") - if backend is not None and backend not in modules: - raise ValueError(f"unknown backend {backend!r}; choose from {sorted(modules)}") - selected = modules.values() if backend is None else (modules[backend],) - return [ - { - "backend": module.rsplit(".", 1)[1], - "factory": f"{module}:create_adapter", - **import_module(module).catalog(), - } - for module in selected - ] + if backend not in modules: + raise ValueError(f"unknown backend {backend!r}; choose from {sorted(modules)}") + modules = {backend: modules[backend]} + return [import_module(module).catalog() for module in modules.values()] def _selected_adapters(campaign: CampaignConfig, index: int | None) -> tuple[AdapterSpec, ...]: @@ -72,23 +66,14 @@ def prepare_campaign(campaign: CampaignConfig, *, index: int | None = None) -> l results = [] for spec in _selected_adapters(campaign, index): - if spec.factory not in _PACKAGES: + if spec.factory in _PACKAGES: + files = {name: str(path) for name, path in prepare_adapter(spec).items()} + results.append({"factory": spec.factory, "status": "prepared", "files": files}) + else: + message = "This adapter manages its own dependencies and model files." results.append( - { - "factory": spec.factory, - "status": "external_setup", - "message": "This adapter manages its own dependencies and model files.", - } + {"factory": spec.factory, "status": "external_setup", "message": message} ) - continue - files = prepare_adapter(spec) - results.append( - { - "factory": spec.factory, - "status": "prepared", - "files": {name: str(path) for name, path in files.items()}, - } - ) return results @@ -96,13 +81,11 @@ def backend_conflict(factories: Iterable[str]) -> str | None: """Return why these built-in backends cannot share one Python environment, or None.""" packages = {_PACKAGES[factory] for factory in factories if factory in _PACKAGES} - mace = "mace-torch==0.3.16" in packages - uma = "fairchem-core==2.23.0" in packages - ani = "torchani==2.8.4" in packages - if mace and uma: - return "MACE requires e3nn==0.4.4, but UMA requires e3nn>=0.5" - if ani and uma: - return "ANI-1xnr pins Torch 2.11, but UMA's FAIR-Chem version requires Torch 2.13" + if "fairchem-core==2.23.0" in packages: + if "mace-torch==0.3.16" in packages: + return "MACE requires e3nn==0.4.4, but UMA requires e3nn>=0.5" + if "torchani==2.8.4" in packages: + return "ANI-1xnr pins Torch 2.11, but UMA's FAIR-Chem version requires Torch 2.13" return None @@ -123,28 +106,19 @@ def install_and_prepare( packages = list( dict.fromkeys(_PACKAGES[spec.factory] for spec in selected if spec.factory in _PACKAGES) ) - mace = "mace-torch==0.3.16" in packages - uma = "fairchem-core==2.23.0" in packages ani = "torchani==2.8.4" in packages conflicts = [] - if ani: - conflicts.append(("fairchem-core", "2.23.0")) - if uma: - conflicts.extend( - [("mace-torch", "0.3.16"), ("mace-torch", "0.3.15"), ("torchani", "2.8.4")] - ) - if mace: - conflicts.append(("fairchem-core", "2.23.0")) + if "fairchem-core==2.23.0" in packages: # backend_conflict has ruled out MACE and ANI-1xnr + conflicts = [("mace-torch", "0.3.16"), ("mace-torch", "0.3.15"), ("torchani", "2.8.4")] + elif ani or "mace-torch==0.3.16" in packages: + conflicts = [("fairchem-core", "2.23.0")] for name, pinned in conflicts: - try: - installed_version = version(name) - except PackageNotFoundError: - continue - if installed_version == pinned: - raise ValueError( - f"selected backends conflict with installed {name}=={pinned}; " - "use a fresh Python environment to preserve the existing model setup" - ) + with suppress(PackageNotFoundError): + if version(name) == pinned: + raise ValueError( + f"selected backends conflict with installed {name}=={pinned}; " + "use a fresh Python environment to preserve the existing model setup" + ) if any( spec.factory == "reactionflow.adapters.mace:create_adapter" and spec.options.get("family") == "polar" diff --git a/src/reactionflow/pathway.py b/src/reactionflow/pathway.py index 35029ee..cd01eca 100644 --- a/src/reactionflow/pathway.py +++ b/src/reactionflow/pathway.py @@ -18,7 +18,7 @@ from ase.optimize import FIRE from .candidates import ReactionCandidate -from .detection import Bond, BondDetectorConfig, atom_ids, canonical_copy, classify_bonds +from .detection import Bond, BondDetectorConfig, _is_count, atom_ids, canonical_copy, classify_bonds from .ssneb import SSNEB, cell_filter CalculatorProvider = Callable[[str], AbstractContextManager[Calculator]] @@ -95,8 +95,7 @@ def __post_init__(self) -> None: or self.curvature_cutoff_eV_A2 < 0 ): raise ValueError("cell-change limits and the curvature cutoff must be non-negative") - lookback = self.reactant_lookback - if isinstance(lookback, bool) or not isinstance(lookback, Integral) or lookback < 0: + if not _is_count(self.reactant_lookback, 0): raise ValueError("reactant_lookback must be a non-negative whole number") if not isinstance(self.split_steps, bool): raise ValueError("split_steps must be true or false") @@ -227,21 +226,19 @@ def _standard_cell(atoms: Atoms) -> None: def _prepare_frames( - frames: Sequence[Atoms], - config: PathwayConfig, - *, - variable_cell: bool = False, -) -> list[Atoms]: + frames: Mapping[int, Atoms], config: PathwayConfig, *, variable_cell: bool +) -> dict[int, Atoms]: """Constraint-free copies, each in the first frame's atom order and placed next to it.""" - reactant = canonical_copy(frames[0]) + frame_numbers = list(frames) + reactant = canonical_copy(frames[frame_numbers[0]]) reactant.set_constraint() first_ids = atom_ids(reactant) if variable_cell: _standard_cell(reactant) - prepared = [reactant] - for frame in frames[1:]: - product = canonical_copy(frame) + prepared = {frame_numbers[0]: reactant} + for frame_number in frame_numbers[1:]: + product = canonical_copy(frames[frame_number]) product.set_constraint() second_ids = atom_ids(product) if set(first_ids) != set(second_ids): @@ -252,7 +249,6 @@ def _prepare_frames( raise ValueError("reactant and product elements differ") if reactant.pbc.tolist() != product.pbc.tolist(): raise _EndpointUnresolved("reactant and product periodicity differs") - cells_differ = not np.array_equal(reactant.cell, product.cell) if variable_cell: _standard_cell(product) reference_positions = product.get_scaled_positions(wrap=False) @ reactant.cell @@ -262,7 +258,7 @@ def _prepare_frames( product.set_scaled_positions( np.linalg.solve(reactant.cell.T, (reactant.positions + displacement).T).T ) - elif cells_differ: + elif not np.array_equal(reactant.cell, product.cell): if not reactant.pbc.all() or reactant.cell.rank < 3 or product.cell.rank < 3: raise _EndpointUnresolved( "changed cells require fully periodic nonsingular endpoints" @@ -284,7 +280,7 @@ def _prepare_frames( if reactant.pbc.any(): displacement = find_mic(displacement, cell=reactant.cell, pbc=reactant.pbc)[0] product.positions = reactant.positions + displacement - prepared.append(product) + prepared[frame_number] = product return prepared @@ -292,9 +288,8 @@ def _local(candidate: ReactionCandidate) -> Callable[[Collection[Bond]], set[Bon """Pairs inside the reacting region, plus any pair involving an atom of a changed bond.""" region = set(candidate.atom_ids) - changed_atoms = { - atom_id for bond in candidate.reactant_bonds ^ candidate.product_bonds for atom_id in bond - } + changed = candidate.reactant_bonds ^ candidate.product_bonds + changed_atoms = {atom_id for bond in changed for atom_id in bond} return lambda bonds: { bond for bond in bonds if region.issuperset(bond) or changed_atoms.intersection(bond) } @@ -328,10 +323,7 @@ def _made( def _matches_outside( - relaxed: Atoms, - frame: Atoms, - changed: Collection[Bond], - detector_config: BondDetectorConfig, + relaxed: Atoms, frame: Atoms, changed: Collection[Bond], detector_config: BondDetectorConfig ) -> bool: """Whether a relaxed structure has the frame's bonds everywhere but the changed ones. @@ -374,8 +366,12 @@ def _endpoint_status( return None -def _snapshots(images: list[Atoms]) -> tuple[Atoms, ...]: - return tuple(image.copy() for image in images) +def _ended(status: str, images: list[Atoms], message: object) -> PathwayOutcome: + """An outcome with no barrier and calculator-free snapshots of the images it stopped at.""" + + return PathwayOutcome( + status, images=tuple(image.copy() for image in images), message=str(message) + ) def _hessian( @@ -397,10 +393,7 @@ def _hessian( def _validate_frequencies( - images: Sequence[Atoms], - profile: Sequence[float], - calculator: Calculator, - config: PathwayConfig, + images: Sequence[Atoms], profile: Sequence[float], calculator: Calculator, config: PathwayConfig ) -> FrequencyValidation: transition_state_index: int | None = None active_atom_ids: tuple[int, ...] = () @@ -431,6 +424,9 @@ def _validate_frequencies( weighted = np.linalg.eigvalsh(hessian * np.outer(weights, weights)) # hbar sqrt(eV / Ų / amu) in cm⁻¹ scale = units._hbar * 1e10 / np.sqrt(units._e * units._amu) / units.invcm + frequencies = np.sign(weighted) * scale * np.sqrt(np.abs(weighted)) + lowest = vectors[:, 0] + primary = tuple(tuple(map(float, row)) for row in lowest.reshape(-1, 3)) if modes else () along = tangent[active].reshape(-1) length = float(np.linalg.norm(along)) return FrequencyValidation( @@ -443,15 +439,9 @@ def _validate_frequencies( curvature_cutoff_eV_A2=config.curvature_cutoff_eV_A2, active_max_force_eV_A=active_max_force, eigenvalues_eV_A2=tuple(map(float, values)), - frequencies_cm1=tuple( - map(float, np.sign(weighted) * scale * np.sqrt(np.abs(weighted))) - ), - tangent_overlap=float(abs(vectors[:, 0] @ along) / length) if length else None, - primary_mode=( - tuple(tuple(map(float, row)) for row in vectors[:, 0].reshape(-1, 3)) - if modes - else () - ), + frequencies_cm1=tuple(map(float, frequencies)), + tangent_overlap=float(abs(lowest @ along) / length) if length else None, + primary_mode=primary, ) except Exception as exc: return FrequencyValidation( @@ -466,10 +456,7 @@ def _validate_frequencies( def _minimize( - atoms: Atoms, - calculator: Calculator, - config: PathwayConfig, - pressure_GPa: float | None, + atoms: Atoms, calculator: Calculator, config: PathwayConfig, pressure_GPa: float | None ) -> tuple[bool, int]: """Relax like an endpoint (cell filter at pressure for NPT); return (converged, steps).""" @@ -485,16 +472,19 @@ def _minimize( atoms.calc = None -def _relax( - atoms: Atoms, - *, - stage: str, - calculator_provider: CalculatorProvider, - config: PathwayConfig, - pressure_GPa: float | None = None, -) -> bool: - with calculator_provider(stage) as calculator: - return _minimize(atoms, calculator, config, pressure_GPa)[0] +@dataclass(frozen=True, slots=True) +class _Refinement: + """What every relaxation and band of one ``refine_pathway`` call shares.""" + + calculator_provider: CalculatorProvider + options: PathwayConfig + detector_config: BondDetectorConfig + pressure_GPa: float | None + + +def _relax(atoms: Atoms, stage: str, setup: _Refinement) -> bool: + with setup.calculator_provider(stage) as calculator: + return _minimize(atoms, calculator, setup.options, setup.pressure_GPa)[0] def _rmsd(atoms: Atoms, reference: Atoms) -> float: @@ -613,15 +603,14 @@ def check_saddle( energies.append(float(image.get_potential_energy())) finally: image.calc = None - volumes = () if pressure_GPa is None else tuple(float(image.get_volume()) for image in band) - enthalpies = ( - () - if pressure_GPa is None - else tuple( + volumes: tuple[float, ...] = () + enthalpies: tuple[float, ...] = () + if pressure_GPa is not None: + volumes = tuple(float(image.get_volume()) for image in band) + enthalpies = tuple( energy + pressure_GPa * units.GPa * volume for energy, volume in zip(energies, volumes, strict=True) ) - ) profile = energies if pressure_GPa is None else enthalpies common: dict[str, Any] = { "energies": tuple(energies), @@ -687,14 +676,15 @@ def refine_pathway( isinstance(pressure_GPa, bool) or not np.isfinite(pressure_GPa) ): raise ValueError("pressure_GPa must be a finite number or None") - outcome = _refine_pathway( - candidate, - calculator_provider=calculator_provider, - config=config, - detector_config=detector_config, - pressure_GPa=pressure_GPa, - frames=frames or {}, - ) + frames = frames or {} + if not candidate.resolved: + outcome = PathwayOutcome("unresolved", message="candidate topology was not resolved") + elif candidate.reactant_bonds == candidate.product_bonds: + outcome = PathwayOutcome("unresolved", message="candidate has no changed bonds") + else: + options = config or PathwayConfig() + setup = _Refinement(calculator_provider, options, detector_config, pressure_GPa) + outcome = _refine_pathway(candidate, setup, frames) method = "neb" if pressure_GPa is None else "ssneb" steps = tuple( replace(step, outcome=replace(step.outcome, method=method, pressure_GPa=pressure_GPa)) @@ -704,20 +694,9 @@ def refine_pathway( def _refine_pathway( - candidate: ReactionCandidate, - *, - calculator_provider: CalculatorProvider, - config: PathwayConfig | None, - detector_config: BondDetectorConfig, - pressure_GPa: float | None, - frames: Mapping[int, Atoms], + candidate: ReactionCandidate, setup: _Refinement, frames: Mapping[int, Atoms] ) -> PathwayOutcome: - - if not candidate.resolved: - return PathwayOutcome("unresolved", message="candidate topology was not resolved") - if candidate.reactant_bonds == candidate.product_bonds: - return PathwayOutcome("unresolved", message="candidate has no changed bonds") - options = config or PathwayConfig() + options, detector_config = setup.options, setup.detector_config changed = candidate.reactant_bonds ^ candidate.product_bonds first, last = candidate.reactant_frame, candidate.product_frame saved = {**frames, first: candidate.reactant, last: candidate.product} @@ -732,32 +711,17 @@ def _refine_pathway( ), first, ) - numbers = [frame for frame in sorted(saved) if start <= frame <= last] + window = {frame: saved[frame] for frame in sorted(saved) if start <= frame <= last} images: list[Atoms] = [] try: - variable_cell = pressure_GPa is not None - prepared = dict( - zip( - numbers, - _prepare_frames( - [saved[frame] for frame in numbers], options, variable_cell=variable_cell - ), - strict=True, - ) - ) + prepared = _prepare_frames(window, options, variable_cell=setup.pressure_GPa is not None) if not set(candidate.atom_ids) <= set(atom_ids(prepared[first])): raise ValueError("candidate atom IDs are missing from its endpoints") images = [prepared[start], prepared[last]] def relaxed(frame: int, stage: str) -> tuple[Atoms, bool]: atoms = prepared[frame].copy() - converged = _relax( - atoms, - stage=stage, - calculator_provider=calculator_provider, - config=options, - pressure_GPa=pressure_GPa, - ) + converged = _relax(atoms, stage, setup) return atoms, converged reactant, converged = relaxed(start, "relax_reactant") @@ -772,7 +736,7 @@ def relaxed(frame: int, stage: str) -> tuple[Atoms, bool]: # that relaxes into it. ends: list[tuple[frozenset[Bond], int]] = [] if options.split_steps and len(changed) > 1 and state is not None and state != changed: - for frame in numbers: + for frame in prepared: if not start < frame < last: continue atoms, done = relaxed(frame, "relax_frame") @@ -783,56 +747,34 @@ def relaxed(frame: int, stage: str) -> tuple[Atoms, bool]: ends.append((made, frame)) ends.append((changed, last)) except BarrierTimeLimit as exc: - return PathwayOutcome("time_limit", images=_snapshots(images), message=str(exc)) + return _ended("time_limit", images, exc) except _EndpointUnresolved as exc: - return PathwayOutcome("unresolved", images=_snapshots(images), message=str(exc)) + return _ended("unresolved", images, exc) except Exception as exc: - return PathwayOutcome("failed", images=_snapshots(images), message=str(exc)) + return _ended("failed", images, exc) + initial = candidate.reactant_bonds steps: list[PathwayStep] = [] before: frozenset[Bond] = frozenset() if options.split_steps and state is not None and 0 < len(state) < len(changed): # The reactant frame relaxes into an intermediate, so the steps from there still count. before = state - steps.append( - PathwayStep( - PathwayOutcome( - "unresolved", - images=_snapshots([reactant]), - message="the reactant frame relaxes partway along the event", - ), - candidate.reactant_bonds, - candidate.reactant_bonds ^ state, - ) + partway = _ended( + "unresolved", [reactant], "the reactant frame relaxes partway along the event" ) + steps.append(PathwayStep(partway, initial, initial ^ state)) source: int = start end: Atoms | None = reactant for made, frame in ends: - step = replace( - candidate, - reactant_bonds=candidate.reactant_bonds ^ before, - product_bonds=candidate.reactant_bonds ^ made, - ) + step = replace(candidate, reactant_bonds=initial ^ before, product_bonds=initial ^ made) if steps and steps[-1].outcome.status == "time_limit": outcome = replace(steps[-1].outcome, images=()) elif end is None: outcome = PathwayOutcome("unresolved", message="the step before ended elsewhere") else: - outcome, end = _refine_step( - step, - end, - converged, - prepared[frame], - calculator_provider=calculator_provider, - options=options, - detector_config=detector_config, - pressure_GPa=pressure_GPa, - ) - steps.append( - PathwayStep( - replace(outcome, frames=(source, frame)), step.reactant_bonds, step.product_bonds - ) - ) + outcome, end = _refine_step(step, end, converged, prepared[frame], setup) + outcome = replace(outcome, frames=(source, frame)) + steps.append(PathwayStep(outcome, step.reactant_bonds, step.product_bonds)) before, source, converged = made, frame, True if len(steps) == 1: return steps[0].outcome @@ -849,11 +791,7 @@ def _refine_step( reactant: Atoms, reactant_converged: bool, frame: Atoms, - *, - calculator_provider: CalculatorProvider, - options: PathwayConfig, - detector_config: BondDetectorConfig, - pressure_GPa: float | None, + setup: _Refinement, ) -> tuple[PathwayOutcome, Atoms | None]: """Refine the candidate from its relaxed reactant to a product built from the frame. @@ -861,23 +799,19 @@ def _refine_step( next step starts. """ + options = setup.options images = [reactant] end: Atoms | None = None try: - reacting = { - atom_id - for bond in candidate.reactant_bonds ^ candidate.product_bonds - for atom_id in bond - } + changed = candidate.reactant_bonds ^ candidate.product_bonds + reacting = {atom_id for bond in changed for atom_id in bond} index = {atom_id: position for position, atom_id in enumerate(atom_ids(reactant))} - moved = sorted( - { - index[atom_id] - for bond in candidate.reactant_bonds | candidate.product_bonds - if reacting.intersection(bond) - for atom_id in bond - } - ) + moved = [ + index[atom_id] + for bond in candidate.reactant_bonds | candidate.product_bonds + if reacting.intersection(bond) + for atom_id in bond + ] # The product starts from the relaxed reactant, cell included, with only the atoms in # changed bonds and the atoms bonded to them where they sat in the product frame. Every # other atom starts at the same minimum in both endpoints, so the band spends its images @@ -892,41 +826,26 @@ def _refine_step( product.set_cell(reactant.cell) product.set_scaled_positions(start) images = [reactant, product] - product_converged = _relax( - product, - stage="relax_product", - calculator_provider=calculator_provider, - config=options, - pressure_GPa=pressure_GPa, - ) + product_converged = _relax(product, "relax_product", setup) if not reactant_converged or not product_converged: - return PathwayOutcome( - "relaxation_failed", - images=_snapshots(images), - message="endpoint relaxation did not converge", - ), None + return _ended("relaxation_failed", images, "endpoint relaxation did not converge"), None - endpoint_status = _endpoint_status(candidate, reactant, product, detector_config) + endpoint_status = _endpoint_status(candidate, reactant, product, setup.detector_config) if endpoint_status is not None: status, message = endpoint_status - return PathwayOutcome(status, images=_snapshots(images), message=message), None + return _ended(status, images, message), None end = product images = [reactant] images.extend(reactant.copy() for _ in range(options.images - 2)) images.append(product) - if pressure_GPa is None: - band = NEB( - images, - climb=False, - allow_shared_calculator=True, - method="improvedtangent", - ) + if setup.pressure_GPa is None: + band = NEB(images, climb=False, allow_shared_calculator=True, method="improvedtangent") band.interpolate(method="idpp", mic=True) else: - band = SSNEB(images, pressure_GPa=pressure_GPa) + band = SSNEB(images, pressure_GPa=setup.pressure_GPa) band.interpolate() - with calculator_provider("neb") as calculator: + with setup.calculator_provider("neb") as calculator: for image in images: image.calc = calculator try: @@ -936,47 +855,33 @@ def _refine_step( endpoint.calc = SinglePointCalculator( endpoint, energy=endpoint.get_potential_energy() ) - neb_converged = bool( - FIRE(band, logfile=None).run( - fmax=options.neb_fmax, - steps=options.neb_steps, - ) + neb_converged = FIRE(band, logfile=None).run( + fmax=options.neb_fmax, steps=options.neb_steps ) if not neb_converged: - return PathwayOutcome( - "neb_failed", - images=_snapshots(images), - message="initial NEB did not converge", - ), end + return _ended("neb_failed", images, "initial NEB did not converge"), end band.climb = True - ci_neb_converged = bool( - FIRE(band, logfile=None).run( - fmax=options.neb_fmax, - steps=options.ci_neb_steps, - ) + ci_neb_converged = FIRE(band, logfile=None).run( + fmax=options.neb_fmax, steps=options.ci_neb_steps ) if not ci_neb_converged: - return PathwayOutcome( - "ci_neb_failed", - images=_snapshots(images), - message="climbing-image NEB did not converge", + return _ended( + "ci_neb_failed", images, "climbing-image NEB did not converge" ), end return check_saddle( images, calculator, - detector_config=detector_config, + detector_config=setup.detector_config, config=options, - pressure_GPa=pressure_GPa, + pressure_GPa=setup.pressure_GPa, ), end finally: for image in images: image.calc = None except BarrierTimeLimit as exc: - return PathwayOutcome("time_limit", images=_snapshots(images), message=str(exc)), None - except _EndpointUnresolved as exc: - return PathwayOutcome("unresolved", images=_snapshots(images), message=str(exc)), end + return _ended("time_limit", images, exc), None except Exception as exc: - return PathwayOutcome("failed", images=_snapshots(images), message=str(exc)), end + return _ended("failed", images, exc), end __all__ = [ diff --git a/src/reactionflow/restart.py b/src/reactionflow/restart.py index 62e0f46..eb2d89e 100644 --- a/src/reactionflow/restart.py +++ b/src/reactionflow/restart.py @@ -123,7 +123,6 @@ def write(self, path: str | Path) -> Path: final.mkdir(parents=True) atoms_path = final / "atoms.traj" arrays_path = final / "arrays.npz" - state_path = final / "restart.json" write(atoms_path, self.atoms, format="traj") archive: dict[str, np.ndarray] = {} @@ -136,10 +135,7 @@ def write(self, path: str | Path) -> Path: archive[archive_name] = array atom_arrays[name] = archive_name components: dict[str, dict[str, Any]] = {} - for label, component in ( - ("dynamics", self.dynamics), - ("calculator", self.calculator), - ): + for label, component in (("dynamics", self.dynamics), ("calculator", self.calculator)): names: dict[str, str] = {} for index, (name, array) in enumerate(sorted(component.arrays.items())): archive_name = f"{label}_{index:04d}" @@ -162,9 +158,8 @@ def write(self, path: str | Path) -> Path: "arrays.npz": file_digest(arrays_path), }, } - state_path.write_text( - json.dumps(manifest, indent=2, sort_keys=True, allow_nan=False) + "\n", - encoding="utf-8", + (final / "restart.json").write_text( + json.dumps(manifest, indent=2, sort_keys=True, allow_nan=False) + "\n", encoding="utf-8" ) return final @@ -193,11 +188,9 @@ def read(cls, path: str | Path) -> ExactRestartSnapshot: atom_arrays = manifest.get("atom_arrays", {}) if not isinstance(atom_arrays, dict): raise ValueError("exact-restart manifest has invalid atom arrays") + # The manifest is JSON, so every name is a string. if any( - not isinstance(name, str) - or not name - or not isinstance(archive_name, str) - or not archive_name + not name or not archive_name or not isinstance(archive_name, str) for name, archive_name in atom_arrays.items() ): raise ValueError("exact-restart manifest has an invalid atom array") @@ -206,11 +199,9 @@ def read(cls, path: str | Path) -> ExactRestartSnapshot: atoms.set_array(name, np.asarray(archive[archive_name]).copy()) for label in ("dynamics", "calculator"): value = components[label] - referenced_arrays.update(value.get("arrays", {}).values()) - arrays = { - name: np.asarray(archive[archive_name]).copy() - for name, archive_name in value.get("arrays", {}).items() - } + names = value.get("arrays", {}) + referenced_arrays.update(names.values()) + arrays = {name: archive[archive_name] for name, archive_name in names.items()} restored[label] = ComponentState( kind=value["kind"], version=value["version"], @@ -220,11 +211,7 @@ def read(cls, path: str | Path) -> ExactRestartSnapshot: ) if referenced_arrays != set(archive.files): raise ValueError("exact-restart array archive does not match its manifest") - return cls( - atoms=atoms, - dynamics=restored["dynamics"], - calculator=restored["calculator"], - ) + return cls(atoms=atoms, dynamics=restored["dynamics"], calculator=restored["calculator"]) __all__ = ["ComponentState", "ExactRestartSnapshot"] diff --git a/src/reactionflow/run.py b/src/reactionflow/run.py index 2ecd969..1a36d02 100644 --- a/src/reactionflow/run.py +++ b/src/reactionflow/run.py @@ -18,7 +18,7 @@ from ase import Atoms from ase.io.trajectory import Trajectory -from ._durable import ensure_directory, flush_to_disk, publish, sync_directory +from ._durable import ensure_directory, flush_to_disk, publish, publishing, sync_directory from ._version import __version__ from .candidates import ReactionCandidate, ReactionTracker from .detection import BondChangeDetector, BondDetectorConfig, canonical_copy, transport_copy @@ -31,7 +31,7 @@ ) from .restart import ExactRestartSnapshot from .runtime import ExactDynamicsRuntime, ExactRuntimeProvider -from .store import OccurrenceRecord, OccurrenceStore +from .store import OccurrenceStore _PHASES = {"new", "running", "refining", "completed", "failed"} @@ -51,10 +51,8 @@ class ReactionRunConfig: pathway: PathwayConfig = field(default_factory=PathwayConfig) def __post_init__(self) -> None: - for name, value in ( - ("observation_interval", self.observation_interval), - ("candidate_stability_frames", self.candidate_stability_frames), - ): + for name in ("observation_interval", "candidate_stability_frames"): + value = getattr(self, name) if isinstance(value, bool) or not isinstance(value, Integral) or value < 1: raise ValueError(f"{name} must be a positive integer") radius = self.candidate_region_radius @@ -182,12 +180,7 @@ def __init__(self, root: Path, config: ReactionRunConfig) -> None: self._refine_after_md = False @classmethod - def create( - cls, - root: str | Path, - *, - config: ReactionRunConfig | None = None, - ) -> ReactionRun: + def create(cls, root: str | Path, *, config: ReactionRunConfig | None = None) -> ReactionRun: """Create a new durable run directory.""" path = Path(root).resolve() @@ -217,8 +210,7 @@ def open(cls, root: str | Path) -> ReactionRun: run._barrier_seconds = float(value.get("barrier_seconds", 0.0)) if value.get("detector_state") is not None: run._detector = BondChangeDetector.from_state( - value["detector_state"], - config=run.config.detector, + value["detector_state"], config=run.config.detector ) if value.get("active_checkpoint") is not None: run._active_checkpoint = run.root / value["active_checkpoint"] @@ -294,16 +286,11 @@ def _checkpoint(self, runtime: ExactDynamicsRuntime) -> None: ) final = self.runtime_checkpoints / name temporary = self.runtime_checkpoints / f".{name}.tmp" - temporary.mkdir() - try: + # The previous checkpoint is removed once state.json names this one, so this one must + # already be on disk if the node fails. + with publishing(temporary, final): snapshot.write(temporary / "exact-restart") self._tracker.write_checkpoint(temporary / "tracker") - # The previous checkpoint is removed once state.json names this one, so this one must - # already be on disk if the node fails. - publish(temporary, final) - except Exception: - shutil.rmtree(temporary, ignore_errors=True) - raise self._active_checkpoint = final self._snapshot = snapshot self._write_state() @@ -355,27 +342,22 @@ def _start(self, atoms: Atoms) -> None: self._initial = initial self._phase = "running" - def _queue_pathway(self, record: OccurrenceRecord, candidate: ReactionCandidate) -> None: - # Every resolved occurrence is refined, also one that repeats an earlier reaction, - # because the reactions around it can change its barrier. - if ( - candidate.resolved - and record.occurrence_id not in self._pending - and not (self.pathways / record.occurrence_id).is_dir() - ): - self._pending.append(record.occurrence_id) - def _register(self, candidates: tuple[ReactionCandidate, ...], *, label: str) -> None: # Registration is idempotent, so an exact replay of an observation that was interrupted # before its checkpoint queues the same pathways again. for index, candidate in enumerate(candidates): occurrence_id = f"segment-{self._generation:04d}-frame-{label}-{index:04d}" - record, _ = self.occurrences.register( - occurrence_id, - candidate, - detector_config=self.config.detector, + self.occurrences.register( + occurrence_id, candidate, detector_config=self.config.detector ) - self._queue_pathway(record, candidate) + # Every resolved occurrence is refined, also one that repeats an earlier reaction, + # because the reactions around it can change its barrier. + if ( + candidate.resolved + and occurrence_id not in self._pending + and not (self.pathways / occurrence_id).is_dir() + ): + self._pending.append(occurrence_id) def _observe(self, atoms: Atoms, global_step: int, total_steps: int) -> None: """Advance the monitor one frame and register every completed occurrence.""" @@ -455,11 +437,7 @@ def _record_environment_change(self, resumed: ExactRestartSnapshot) -> None: """Log a restore under other software or hardware than the checkpoint recorded.""" assert self._snapshot is not None - before = { - **self._snapshot.calculator.metadata, - **_dynamics_versions(self._snapshot.dynamics.metadata), - } - after = {**resumed.calculator.metadata, **_dynamics_versions(resumed.dynamics.metadata)} + before, after = _environment(self._snapshot), _environment(resumed) changed = { key: [before.get(key), after.get(key)] for key in sorted(before.keys() | after.keys()) @@ -500,8 +478,7 @@ def _saved_frames(self, first: int, last: int) -> dict[int, Atoms]: def _publish_outcome(self, occurrence_id: str, outcome: PathwayOutcome) -> None: final = self.pathways / occurrence_id temporary = self.pathways / f".{occurrence_id}-{uuid4().hex}.tmp" - temporary.mkdir() - try: + with publishing(temporary, final): identity = { "schema_version": 1, "occurrence_id": occurrence_id, @@ -517,10 +494,6 @@ def _publish_outcome(self, occurrence_id: str, outcome: PathwayOutcome) -> None: folder = temporary / f"step-{number}" folder.mkdir() _write_result(folder, step.outcome, {**identity, "step": number, **bonds}) - publish(temporary, final) - except Exception: - shutil.rmtree(temporary, ignore_errors=True) - raise def _time_limited( self, provider: CalculatorProvider, max_hours: float | None @@ -562,9 +535,7 @@ def checked(*args: Any, **kwargs: Any) -> Any: return lease def _refine_pending( - self, - calculator_provider: CalculatorProvider, - pressure_GPa: float | None, + self, calculator_provider: CalculatorProvider, pressure_GPa: float | None ) -> None: """Refine each queued occurrence serially and publish its immutable result. @@ -699,8 +670,10 @@ def run( return self.summary() -def _dynamics_versions(metadata: Mapping[str, Any]) -> dict[str, Any]: - return {f"dynamics_{key}": metadata.get(key) for key in ("ase_version", "numpy_version")} +def _environment(snapshot: ExactRestartSnapshot) -> dict[str, Any]: + versions = ("ase_version", "numpy_version") + dynamics = {f"dynamics_{key}": snapshot.dynamics.metadata.get(key) for key in versions} + return {**snapshot.calculator.metadata, **dynamics} __all__ = ["ReactionRun", "ReactionRunConfig", "RunSummary"] diff --git a/src/reactionflow/runtime.py b/src/reactionflow/runtime.py index c3ba671..3c70f16 100644 --- a/src/reactionflow/runtime.py +++ b/src/reactionflow/runtime.py @@ -30,8 +30,7 @@ class ExactRuntimeProvider(Protocol): def start(self, atoms: Atoms) -> AbstractContextManager[ExactDynamicsRuntime]: ... def restore( - self, - snapshot: ExactRestartSnapshot, + self, snapshot: ExactRestartSnapshot ) -> AbstractContextManager[ExactDynamicsRuntime]: ... diff --git a/src/reactionflow/scan.py b/src/reactionflow/scan.py index 38d8669..3710b6d 100644 --- a/src/reactionflow/scan.py +++ b/src/reactionflow/scan.py @@ -10,13 +10,14 @@ from pathlib import Path from typing import Any +import networkx as nx from ase import Atoms from ase.formula import Formula from ase.io import iread from ase.io.trajectory import Trajectory from .candidates import ReactionCandidate, ReactionClasses, ReactionTracker -from .detection import Bond, BondChangeDetector, BondDetectorConfig, atom_ids, canonical_copy +from .detection import Bond, BondChangeDetector, BondDetectorConfig, _symbols_by_id, canonical_copy from .status import _describe type Frame = tuple[int, float | None, Atoms] @@ -123,29 +124,19 @@ def recorded_occurrences(root: str | Path) -> dict[OccurrenceKey, str]: continue metadata = json.loads(path.read_text(encoding="utf-8")) data = metadata["candidate"] - changes = {tuple(bond) for bond in data["reactant_bonds"]} ^ { - tuple(bond) for bond in data["product_bonds"] - } + changes = set(map(tuple, data["reactant_bonds"])) ^ set(map(tuple, data["product_bonds"])) frames = (data["reactant_frame"], data["product_frame"], data["observed_frame"]) recorded[(frames, frozenset(changes))] = metadata["occurrence_id"] return recorded -def _components(atoms: Iterable[int], bonds: Iterable[Bond]) -> dict[int, int]: - """Map each atom ID to the smallest atom ID of its bonded component.""" - - parent = {atom_id: atom_id for atom_id in atoms} - - def root(atom_id: int) -> int: - while parent[atom_id] != atom_id: - parent[atom_id] = parent[parent[atom_id]] - atom_id = parent[atom_id] - return atom_id +def _components(atoms: Iterable[int], bonds: Iterable[Bond]) -> nx.utils.UnionFind: + """Partition the atom IDs into bonded components.""" + components = nx.utils.UnionFind(atoms) for first, second in bonds: - first, second = root(first), root(second) - parent[max(first, second)] = min(first, second) - return {atom_id: root(atom_id) for atom_id in parent} + components.union(first, second) + return components def _changes(candidate: ReactionCandidate, symbols: Mapping[int, str]) -> str: @@ -184,7 +175,7 @@ def scan_frames( labels: list[str] = [] times: dict[int, float | None] = {} symbols: dict[int, str] = {} - molecule: dict[int, int] = {} + molecule = nx.utils.UnionFind() bond_changes: list[dict[str, Any]] = [] candidates: list[dict[str, Any]] = [] species: list[dict[str, Any]] = [] @@ -223,7 +214,7 @@ def add(candidate: ReactionCandidate, *, at_end: bool) -> None: times[frame] = time_ps events = detector.process(atoms, frame=frame) if not symbols: - symbols = dict(zip(atom_ids(atoms), atoms.get_chemical_symbols(), strict=True)) + symbols = _symbols_by_id(atoms) molecule = _components(symbols, detector.stable_bonds) # A crossing that leaves the pending list without being confirmed reverted in time. pending = detector.pending_changes @@ -253,15 +244,12 @@ def add(candidate: ReactionCandidate, *, at_end: bool) -> None: ): add(candidate, at_end=False) if frame % species_every == 0: - members: dict[int, list[int]] = {} - for atom_id, root in _components(symbols, detector.stable_bonds).items(): - members.setdefault(root, []).append(atom_id) counts = Counter( ( Formula.from_list([symbols[atom_id] for atom_id in group]).format("hill"), len({molecule[atom_id] for atom_id in group}), ) - for group in members.values() + for group in _components(symbols, detector.stable_bonds).to_sets() ) species.extend( { diff --git a/src/reactionflow/search/cli.py b/src/reactionflow/search/cli.py index f487ee0..9cf4b20 100644 --- a/src/reactionflow/search/cli.py +++ b/src/reactionflow/search/cli.py @@ -17,7 +17,7 @@ from ase import Atoms from ase.io import read -from ..campaign import AdapterSpec, TrajectorySpec +from ..campaign import AdapterSpec, TrajectorySpec, _reject_unknown from ..detection import canonical_copy from ..mlip import load_mlip_adapter from ..model_setup import install_and_prepare, prepare_adapter, prepare_campaign @@ -61,7 +61,7 @@ def load(cls, path: str | Path) -> SearchCampaign: raise ValueError("unsupported campaign schema; use schema_version 2") if data.get("mode") != "eon": raise ValueError("search campaign requires mode 'eon'") - unknown = data.keys() - { + known = { "schema_version", "mode", "structure", @@ -71,8 +71,7 @@ def load(cls, path: str | Path) -> SearchCampaign: "akmc", "eon", } - if unknown: - raise ValueError(f"unknown EON campaign keys: {sorted(unknown)}") + _reject_unknown(data, known, "EON campaign keys") def resolve(value: Any, label: str) -> Path: if not isinstance(value, str) or not value: @@ -99,16 +98,15 @@ def trajectories(self) -> tuple[TrajectorySpec]: Adapters are built for MD. EON uses only their calculator, so the MD fields are unused. """ - return ( - TrajectorySpec( - id="eon", - total_steps=1, - timestep_fs=1.0, - temperature_K=self.akmc.temperature_K, - pressure_GPa=None, - seed=self.akmc.seed, - ), + spec = TrajectorySpec( + id="eon", + total_steps=1, + timestep_fs=1.0, + temperature_K=self.akmc.temperature_K, + pressure_GPa=None, + seed=self.akmc.seed, ) + return (spec,) def adapter_for(self, index: int) -> AdapterSpec: if index != 0: @@ -221,9 +219,8 @@ def _format_status(report: dict[str, Any]) -> str: for row in report["step_rows"]: kind = "" if row["chemical"] is not None: - kind = f"{'chemical' if row['chemical'] else 'into unexplored conformer'}, " + ( - f"basin of {row['basin_states']}; " - ) + name = "chemical" if row["chemical"] else "into unexplored conformer" + kind = f"{name}, basin of {row['basin_states']}; " lines.append( f"step {row['step']}: {row['from']} -> {row['to']} via {row['process']} " f"({kind}{row['barrier_eV']:.4g} eV; {row['bonds']}), dt {row['dt_s']:.4g} s, " @@ -260,21 +257,17 @@ def handle_command(arguments: argparse.Namespace) -> int: command = arguments.command if command == "prepare": if arguments.install: - return install_and_prepare( - campaign, - index=arguments.index, - extra_packages=(f"pyeonclient[ase]=={EON_VERSION}",), - ) + extra = (f"pyeonclient[ase]=={EON_VERSION}",) + return install_and_prepare(campaign, index=arguments.index, extra_packages=extra) models = prepare_campaign(campaign, index=arguments.index) print(json.dumps({"campaign": str(campaign.source), "models": models}, sort_keys=True)) return 0 if command == "status": report = search_status(campaign) - print( - json.dumps(report, indent=2, sort_keys=True) - if arguments.json - else _format_status(report) - ) + if arguments.json: + print(json.dumps(report, indent=2, sort_keys=True)) + else: + print(_format_status(report)) return 0 if command == "run": from ..cli import resolve_task_index, visible_gpu @@ -287,24 +280,20 @@ def handle_command(arguments: argparse.Namespace) -> int: atoms = campaign.starting_atoms() warnings = [PINNED] if free_rotation(atoms) == "pinned" else [] if command == "validate": - print( - json.dumps( - { - "campaign": str(campaign.source), - "mode": "eon", - "atoms": len(atoms), - "output_root": str(campaign.output_root), - "tasks": 1, - "kmc_steps": campaign.akmc.steps, - "temperature_K": campaign.akmc.temperature_K, - "adapter_factory": campaign.adapter.factory, - "require_gpu": campaign.require_gpu, - "rigid_rotation": campaign.eon.rotation_for(atoms), - "warnings": warnings, - }, - sort_keys=True, - ) - ) + report = { + "campaign": str(campaign.source), + "mode": "eon", + "atoms": len(atoms), + "output_root": str(campaign.output_root), + "tasks": 1, + "kmc_steps": campaign.akmc.steps, + "temperature_K": campaign.akmc.temperature_K, + "adapter_factory": campaign.adapter.factory, + "require_gpu": campaign.require_gpu, + "rigid_rotation": campaign.eon.rotation_for(atoms), + "warnings": warnings, + } + print(json.dumps(report, sort_keys=True)) return 0 adapter = load_mlip_adapter(campaign.adapter, campaign.trajectories[0]) preflight = getattr(adapter, "preflight", None) @@ -322,16 +311,13 @@ def handle_command(arguments: argparse.Namespace) -> int: # One model load gives the calculator and the same identity an MD restart records. with lease() as (calculator, identity), _eon_output(log): backend = EONBackend(calculator, campaign.eon, campaign.akmc.temperature_K) + contract = { + "adapter": asdict(campaign.adapter), + "calculator": asdict(identity), + "eon": backend.contract, + } run_exploration( - campaign.output_root, - atoms, - backend=backend, - contract={ - "adapter": asdict(campaign.adapter), - "calculator": asdict(identity), - "eon": backend.contract, - }, - config=campaign.akmc, + campaign.output_root, atoms, backend=backend, contract=contract, config=campaign.akmc ) print(json.dumps(search_status(campaign), sort_keys=True)) return 0 diff --git a/src/reactionflow/search/eon.py b/src/reactionflow/search/eon.py index bac78b2..ae9d84f 100644 --- a/src/reactionflow/search/eon.py +++ b/src/reactionflow/search/eon.py @@ -39,7 +39,6 @@ _AFTER_SADDLE = {NOT_CONNECTED, 7, 9, 10} # EON's own words for the endings a guided search reports itself. _ENDINGS = { - 0: "Success", 3: "Barrier too high", 5: "Too many iterations", NOT_CONNECTED: "Saddle is not connected to initial state", @@ -74,12 +73,8 @@ def _guide_entry(value: object) -> str | list[int]: def _guide(value: object) -> dict[str, Any]: - if not isinstance(value, Mapping) or not set(value) <= { - "form", - "break", - "within_A", - "together", - }: + allowed = {"form", "break", "within_A", "together"} + if not isinstance(value, Mapping) or not set(value) <= allowed: raise ValueError("eon.guide takes 'form', 'break', 'within_A', and 'together'") together = value.get("together", False) if type(together) is not bool: @@ -87,7 +82,7 @@ def _guide(value: object) -> dict[str, Any]: guide: dict[str, Any] = {} for kind in ("form", "break"): pairs = value.get(kind, []) - if isinstance(pairs, str) or not isinstance(pairs, (list, tuple)): + if not isinstance(pairs, (list, tuple)): raise ValueError( f"eon.guide.{kind} must be a list of element pairs such as 'C-N' or atom ID pairs" ) @@ -146,7 +141,7 @@ def __post_init__(self) -> None: if self.prefactor is not None: object.__setattr__(self, "prefactor", _number(self.prefactor, "prefactor")) centers = self.displace_centers - if isinstance(centers, str) or not isinstance(centers, (list, tuple)): + if not isinstance(centers, (list, tuple)): raise ValueError("eon.displace_centers must be a list of elements or atom indices") for center in centers: if isinstance(center, bool) or not isinstance(center, (str, Integral)): @@ -257,10 +252,8 @@ def guided_push( symbols = atoms.get_chemical_symbols() ids = atom_ids(atoms) index = {atom_id: number for number, atom_id in enumerate(ids)} - bonded = { - tuple(sorted(index[i] for i in bond)) - for bond in classify_bonds(atoms, BondDetectorConfig())[0] - } + bonds = classify_bonds(atoms, BondDetectorConfig())[0] + bonded = {tuple(sorted(index[i] for i in bond)) for bond in bonds} neighbors: list[set[int]] = [set() for _ in atoms] for i, j in bonded: neighbors[i].add(j) @@ -280,18 +273,13 @@ def elements(i: int, j: int) -> str: for i, j in sorted(bonded): pairs[i, j] = atoms.get_distance(i, j, mic=True, vector=True), -1.0 - def matches(entry: str | list[int], i: int, j: int) -> bool: - if isinstance(entry, str): - return elements(i, j) == entry - return sorted((ids[i], ids[j])) == entry - candidates = [ [ (i, j) for (i, j), (_, sign) in pairs.items() if (sign > 0) == (kind == "form") and not (fixed[i] and fixed[j]) - and matches(entry, i, j) + and entry in (elements(i, j), sorted((ids[i], ids[j]))) ] for kind in ("form", "break") for entry in guide[kind] @@ -315,13 +303,8 @@ def matches(entry: str | list[int], i: int, j: int) -> bool: unit = vector / np.linalg.norm(vector) mode[i] += sign * share[0] * unit mode[j] -= sign * share[1] * unit - records.append( - { - "push": "form" if sign > 0 else "break", - "atom_ids": [int(ids[i]), int(ids[j])], - "elements": elements(i, j), - } - ) + push = "form" if sign > 0 else "break" + records.append({"push": push, "atom_ids": [ids[i], ids[j]], "elements": elements(i, j)}) positions = atoms.positions + settings.displace_size_A * mode free = np.flatnonzero(~fixed) positions[free] += rng.normal(0.0, _JITTER_A / math.sqrt(3 * len(free)), (len(free), 3)) @@ -333,8 +316,6 @@ def _linked_sets(candidates: list[list[tuple[int, int]]]) -> list[tuple[tuple[in """Every set that takes one pair from each list, with no pair twice, whose pairs link up through shared atoms.""" - if not all(candidates): - return [] touching = [] for found in candidates: table: dict[int, list[tuple[int, int]]] = {} @@ -379,14 +360,10 @@ def calculate(self, atoms=None, properties=None, system_changes=all_changes): # EON passes a zero force-evaluation cell for isolated molecules while # retaining Matter's original cell for geometry/I/O. Restore the ASE # cell from the source; a periodic calculation must keep its real cell. - isolated_force_cell = not self.template.pbc.any() and not atoms.cell.array.any() - if not isolated_force_cell and not np.allclose( - atoms.cell, self.template.cell, rtol=0, atol=1e-12 - ): + isolated = not self.template.pbc.any() and not atoms.cell.array.any() + if not (isolated or np.allclose(atoms.cell, self.template.cell, rtol=0, atol=1e-12)): raise ValueError("EON changed the fixed cell") - physical = self.template.copy() - physical.info = copy.deepcopy(self.template.info) - physical.set_positions(atoms.positions, apply_constraint=False) + physical = _moved(self.template, atoms.positions) physical.calc = self.calculator energy = physical.get_potential_energy() # EON applies the fixed-coordinate mask itself. @@ -399,15 +376,12 @@ def calculate(self, atoms=None, properties=None, system_changes=all_changes): def _ini(sections: dict[str, dict[str, object]]) -> str: - def text(value: object) -> str: - if isinstance(value, bool): - return str(value).lower() - return repr(value) if isinstance(value, float) else str(value) - lines = [] for section, values in sections.items(): lines.append(f"[{section}]") - lines.extend(f"{key} = {text(value)}" for key, value in values.items()) + for key, value in values.items(): + text = str(value).lower() if isinstance(value, bool) else str(value) + lines.append(f"{key} = {text}") return "\n".join(lines) + "\n" @@ -494,8 +468,8 @@ def _parameters( parameters.write_movies = False return parameters - def _matter(self, atoms: Atoms, fixed: np.ndarray, parameters, *, calculator=None): - bridge = _PreservingCalculator(calculator or self.calculator, atoms) + def _matter(self, atoms: Atoms, fixed: np.ndarray, parameters): + bridge = _PreservingCalculator(self.calculator, atoms) potential = self._eon.potential_from_ase(bridge) geometry = atoms.copy() geometry.set_constraint() @@ -555,12 +529,9 @@ def search(self, atoms: Atoms, *, seed: int) -> ProcessResult: ) center, radius, size = push metadata.update(center=center, displace_magnitude_A=size) + rotation = self.settings.rotation_for(source) parameters = self._parameters( - seed=seed, - rotation=self.settings.rotation_for(source), - center=center, - radius=radius, - size=size, + seed=seed, rotation=rotation, center=center, radius=radius, size=size ) matter, potential = self._matter(source, fixed, parameters) # pyeonclient 0.4.2 binds the full job in _core without re-exporting it. @@ -571,9 +542,8 @@ def search(self, atoms: Atoms, *, seed: int) -> ProcessResult: results = _results(Path("results.dat")) message = results.get("termination_reason_text", "") reason = int(results["termination_reason"]) - metadata.update( - force_calls=int(results.get("total_force_calls", 0)), termination_reason=reason - ) + metadata["force_calls"] = int(results.get("total_force_calls", 0)) + metadata["termination_reason"] = reason if reason in _AFTER_SADDLE: return ProcessResult( "failed", @@ -620,51 +590,52 @@ def _guided(self, source: Atoms, fixed: np.ndarray, seed: int, metadata: dict) - ) metadata["termination_reason"] = status.value if status.value != 0: - return ProcessResult( - "failed", _ENDINGS.get(status.value, status.name), metadata=metadata - ) + message = _ENDINGS.get(status.value, status.name) + return ProcessResult("failed", message, metadata=metadata) # The saddle's lowest mode, converged from the direction the search climbed. EON wraps # atoms into the cell, so the climb is taken by the minimum image. dimer = self._eon.ImprovedDimer(saddle, parameters, potential) - climb, _ = find_mic( - np.asarray(saddle.positions) - source.positions, source.cell, source.pbc - ) + climb, _ = find_mic(saddle.positions - source.positions, source.cell, source.pbc) dimer.compute(saddle, np.ascontiguousarray(climb / np.linalg.norm(climb))) axis = np.asarray(dimer.eigenvector).reshape(-1, 3) axis /= np.linalg.norm(axis) minima, converged = [], True for sign in (1, -1): side, _ = self._matter(source, fixed, parameters) - side.positions = np.asarray(saddle.positions) + sign * _OFFSET_A * axis + side.positions = saddle.positions + sign * _OFFSET_A * axis relaxed, done = side.relax() minima.append((_moved(source, relaxed.positions), float(relaxed.potential_energy))) converged = converged and bool(done) - reached = { - "saddle": _moved(source, saddle.positions), - "saddle_energy": float(saddle.potential_energy), - } + saddle_atoms = _moved(source, saddle.positions) + saddle_energy = float(saddle.potential_energy) at_start = [self.same(source, energy, *minimum) for minimum in minima] # As in EON, a saddle joins the start when exactly one side is the start. code = 9 if not converged else 0 if sum(at_start) == 1 else NOT_CONNECTED if code == 0: product, product_energy = minima[1] if at_start[0] else minima[0] - prefactors = self.prefactors(source, reached["saddle"], product) + prefactors = self.prefactors(source, saddle_atoms, product) code = 7 if prefactors is None else 0 metadata["termination_reason"] = code if code: return ProcessResult( - "failed", _ENDINGS[code], metadata=metadata, minima=tuple(minima), **reached + "failed", + _ENDINGS[code], + saddle=saddle_atoms, + saddle_energy=saddle_energy, + metadata=metadata, + minima=tuple(minima), ) return ProcessResult( "good", - _ENDINGS[0], + "Success", + saddle=saddle_atoms, product=product, reactant_energy=energy, + saddle_energy=saddle_energy, product_energy=product_energy, prefactor=prefactors[0], reverse_prefactor=prefactors[1], metadata=metadata, - **reached, ) def prefactors(self, first: Atoms, saddle: Atoms, second: Atoms) -> tuple[float, float] | None: diff --git a/src/reactionflow/search/labels.py b/src/reactionflow/search/labels.py index 0ee2441..03f8236 100644 --- a/src/reactionflow/search/labels.py +++ b/src/reactionflow/search/labels.py @@ -6,19 +6,17 @@ from ase import Atoms -from ..detection import BondDetectorConfig, atom_ids, classify_bonds +from ..detection import BondDetectorConfig, _symbols_by_id, classify_bonds -def bond_changes( - reactant: Atoms, product: Atoms, config: BondDetectorConfig | None = None -) -> dict[str, list[dict]]: +def bond_changes(reactant: Atoms, product: Atoms) -> dict[str, list[dict]]: """A bond forms when it is inside the formation distance in the product and was outside the breaking distance in the reactant; breaking is the reverse. Atoms keep their stable IDs.""" - config = config or BondDetectorConfig() + config = BondDetectorConfig() before, before_gap = classify_bonds(reactant, config) after, after_gap = classify_bonds(product, config) - symbols = dict(zip(atom_ids(reactant), reactant.get_chemical_symbols(), strict=True)) + symbols = _symbols_by_id(reactant) def listed(bonds: set) -> list[dict]: return [ diff --git a/src/reactionflow/search/records.py b/src/reactionflow/search/records.py index d3f3556..4c675b0 100644 --- a/src/reactionflow/search/records.py +++ b/src/reactionflow/search/records.py @@ -3,7 +3,6 @@ from __future__ import annotations import json -import re from dataclasses import dataclass from hashlib import sha256 from pathlib import Path @@ -12,17 +11,10 @@ from ase import Atoms from ase.io.jsonio import decode, encode -from .._durable import ensure_directory, publish, sync_directory +from .._durable import SAFE_NAME, ensure_directory, publish, sync_directory from ..detection import assign_atom_ids, atom_ids _KINDS = {"states", "processes", "attempts", "steps"} -_NAME = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]*\Z") - - -def _name(value: str) -> str: - if not isinstance(value, str) or not _NAME.fullmatch(value): - raise ValueError("record IDs must be safe non-empty names") - return value def _json(value: object) -> str: @@ -70,18 +62,17 @@ def _path(self, kind: str, record_id: str | None = None) -> Path: if kind not in _KINDS: raise ValueError(f"unsupported search record kind: {kind!r}") directory = self.root / kind - return directory if record_id is None else directory / f"{_name(record_id)}.json" + if record_id is None: + return directory + if not isinstance(record_id, str) or not SAFE_NAME.fullmatch(record_id): + raise ValueError("record IDs must be safe non-empty names") + return directory / f"{record_id}.json" def ids(self, kind: str) -> tuple[str, ...]: """List published record IDs in lexical order, ignoring interrupted temporary files.""" - return tuple( - sorted( - path.stem - for path in self._path(kind).glob("*.json") - if path.is_file() and _NAME.fullmatch(path.stem) - ) - ) + names = [path.stem for path in self._path(kind).glob("*.json") if path.is_file()] + return tuple(sorted(name for name in names if SAFE_NAME.fullmatch(name))) def write( self, @@ -138,11 +129,8 @@ def read(self, kind: str, record_id: str) -> SearchRecord: payload = document["payload"] if sha256(_json(payload).encode()).hexdigest() != document["sha256"]: raise ValueError(f"search record failed integrity check: {kind}/{record_id}") - if (payload.get("schema_version"), payload.get("kind"), payload.get("record_id")) != ( - 1, - kind, - record_id, - ): + header = payload.get("schema_version"), payload.get("kind"), payload.get("record_id") + if header != (1, kind, record_id): raise ValueError("unsupported or mismatched search record") if not isinstance(payload["data"], dict): raise ValueError("record data must be a JSON object") diff --git a/src/reactionflow/search/runner.py b/src/reactionflow/search/runner.py index 2e3dc62..f729bac 100644 --- a/src/reactionflow/search/runner.py +++ b/src/reactionflow/search/runner.py @@ -76,15 +76,6 @@ def _processes(store: SearchStore, identifiers: list[str]) -> list[dict[str, Any return [{**store.read("processes", pid).data, "id": pid} for pid in identifiers] -def _in_window(processes: list[dict[str, Any]], config: AKMCConfig) -> list[dict[str, Any]]: - """Processes inside the thermal window above the lowest barrier.""" - - if not processes: - return [] - lowest = min(process["barrier_eV"] for process in processes) - return [p for p in processes if in_window(p["barrier_eV"], lowest, config)] - - def is_internal(process: dict[str, Any]) -> bool: """A process that forms and breaks no bond stays inside its chemical basin.""" @@ -92,13 +83,17 @@ def is_internal(process: dict[str, Any]) -> bool: def kept_processes(processes: list[dict[str, Any]], config: AKMCConfig) -> list[dict[str, Any]]: - """Processes in the rate table. With basins, the window compares chemical exits only, so a - fast conformational change does not push the chemistry out of the window.""" + """Processes in the rate table, inside the thermal window above the lowest barrier. With + basins, the window compares chemical exits only, so a fast conformational change does not push + the chemistry out of the window.""" if config.basins == "none": - return _in_window(processes, config) - internal = [p for p in processes if is_internal(p)] - return internal + _in_window([p for p in processes if not is_internal(p)], config) + internal, compared = [], processes + else: + internal = [p for p in processes if is_internal(p)] + compared = [p for p in processes if not is_internal(p)] + lowest = min((p["barrier_eV"] for p in compared), default=0.0) + return internal + [p for p in compared if in_window(p["barrier_eV"], lowest, config)] def _repeat(store: SearchStore, backend: Any, known: list[dict], saddle: Atoms, energy: float): @@ -123,22 +118,19 @@ def _reverse( if _repeat(store, backend, known, structures["saddle"], data["saddle_energy_eV"]): return None reactant = store.read("states", data["state"]) - store.write( - "processes", - reverse_id, - { - "state": data["product"], - "product": data["state"], - "barrier_eV": data["saddle_energy_eV"] - data["product_energy_eV"], - "saddle_energy_eV": data["saddle_energy_eV"], - "product_energy_eV": reactant.data["energy_eV"], - "prefactor_s": data["reverse_prefactor_s"], - "reverse_prefactor_s": data["prefactor_s"], - "bonds": reversed_changes(data["bonds"]), - "source": f"reverse of {forward_id}", - }, - structures={"saddle": structures["saddle"], "product": reactant.structures["atoms"]}, - ) + process = { + "state": data["product"], + "product": data["state"], + "barrier_eV": data["saddle_energy_eV"] - data["product_energy_eV"], + "saddle_energy_eV": data["saddle_energy_eV"], + "product_energy_eV": reactant.data["energy_eV"], + "prefactor_s": data["reverse_prefactor_s"], + "reverse_prefactor_s": data["prefactor_s"], + "bonds": reversed_changes(data["bonds"]), + "source": f"reverse of {forward_id}", + } + keep = {"saddle": structures["saddle"], "product": reactant.structures["atoms"]} + store.write("processes", reverse_id, process, structures=keep) return reverse_id @@ -188,9 +180,8 @@ def _attempt( return {**search, "status": "new", "process": forward_id, "reverse": reverse}, {} origin = store.read("states", origin_id) energy = origin.data["energy_eV"] - data: dict[str, Any] = dict(active) result = backend.search(origin.structures["atoms"].copy(), seed=active["seed"]) - data.update(backend=result.metadata, message=result.message) + data = {**active, "backend": result.metadata, "message": result.message} if result.status == "no_push": return {**data, "status": "no_push"}, {} reference, extra = energy, {} @@ -228,30 +219,27 @@ def _attempt( return {**data, "status": "outside_window"}, {} repeated = _repeat(store, backend, known, result.saddle, result.saddle_energy) if repeated is not None: - relevant = repeated["id"] in {process["id"] for process in kept_processes(known, config)} + relevant = repeated in kept_processes(known, config) return {**data, "status": "repeat", "process": repeated["id"], "relevant": relevant}, {} data["product"] = _state_id(store, backend, result.product, result.product_energy) if data["product"] == origin_id: # Equivalent atoms exchanged places: no state change, so no part of the kinetics. return {**data, "status": "same_state"}, {} data["status"] = "new" - store.write( - "processes", - forward_id, - { - "state": origin_id, - "product": data["product"], - "barrier_eV": barrier, - "saddle_energy_eV": result.saddle_energy, - "product_energy_eV": result.product_energy, - "prefactor_s": result.prefactor, - "reverse_prefactor_s": result.reverse_prefactor, - "bonds": bonds, - "source": active["id"], - "search": data, - }, - structures={"saddle": result.saddle, "product": result.product, **extra}, - ) + process = { + "state": origin_id, + "product": data["product"], + "barrier_eV": barrier, + "saddle_energy_eV": result.saddle_energy, + "product_energy_eV": result.product_energy, + "prefactor_s": result.prefactor, + "reverse_prefactor_s": result.reverse_prefactor, + "bonds": bonds, + "source": active["id"], + "search": data, + } + keep = {"saddle": result.saddle, "product": result.product, **extra} + store.write("processes", forward_id, process, structures=keep) reverse = _reverse(store, backend, state, forward_id, store.read("processes", forward_id)) return {**data, "process": forward_id, "reverse": reverse}, {} @@ -277,19 +265,17 @@ def _commit_attempt(state: dict[str, Any], result: dict[str, Any]) -> None: state["active"] = None -def _basin_step( +def _basin_exit( store: SearchStore, state: dict[str, Any], config: AKMCConfig -) -> dict[str, Any] | None: - """Leave the explored part of the current chemical basin, or return None if it is closed. +) -> tuple[dict[str, Any], float, dict[str, Any]] | None: + """Pick the exit from the explored part of the current chemical basin, its time, and what the + step records about the basin, or return None if the basin is closed. The basin holds every state reached from the current one through processes that form and break no bond and that has been searched to the confidence target. Its exits are chemical processes and bond-preserving processes into states not yet searched to confidence. """ - identifier = f"step-{state['steps']:06d}" - if identifier in store.ids("steps"): - return store.read("steps", identifier).data current = state["current"] explored = { name @@ -321,48 +307,44 @@ def _basin_step( order = sorted(probabilities) pick, dt = kmc_step([probabilities[e] / mean for e in order], [config.seed, 1, state["steps"]]) chosen = exits[order[pick]] - counters = state["states"][current] - data = { - "from": current, - "to": chosen["product"], - "process": chosen["id"], + details = { "exit_state": chosen["state"], "chemical": not is_internal(chosen), - "dt_s": dt, - "time_s": state["time_s"] + dt, "total_rate_s": 1.0 / mean, "mean_escape_time_s": mean, "escape_probability": probabilities[chosen["id"]], "basin_states": len(basin), "processes_in_window": len(order), - "confidence": confidence(counters["repeats"]), - "searches": counters["searches"], } - store.write("steps", identifier, data) - return data + return chosen, dt, details def _step(store: SearchStore, state: dict[str, Any], config: AKMCConfig) -> dict[str, Any] | None: - if config.basins == "bonds": - return _basin_step(store, state, config) identifier = f"step-{state['steps']:06d}" if identifier in store.ids("steps"): return store.read("steps", identifier).data current = state["current"] counters = state["states"][current] - window = _in_window(_processes(store, counters["processes"]), config) - rates = [rate(p["prefactor_s"], p["barrier_eV"], config.kT) for p in window] - pick, dt = kmc_step(rates, [config.seed, 1, state["steps"]]) + if config.basins == "bonds": + found = _basin_exit(store, state, config) + if found is None: + return None + chosen, dt, details = found + else: + window = kept_processes(_processes(store, counters["processes"]), config) + rates = [rate(p["prefactor_s"], p["barrier_eV"], config.kT) for p in window] + pick, dt = kmc_step(rates, [config.seed, 1, state["steps"]]) + chosen = window[pick] + details = {"total_rate_s": float(sum(rates)), "processes_in_window": len(window)} data = { "from": current, - "to": window[pick]["product"], - "process": window[pick]["id"], + "to": chosen["product"], + "process": chosen["id"], "dt_s": dt, "time_s": state["time_s"] + dt, - "total_rate_s": float(sum(rates)), - "processes_in_window": len(window), "confidence": confidence(counters["repeats"]), "searches": counters["searches"], + **details, } store.write("steps", identifier, data) return data @@ -428,15 +410,15 @@ def run_exploration( state["chemical_steps"] = 0 _save_state(root, state) counted = "chemical_steps" if config.basins == "bonds" else "steps" + stop_reason = "steps_completed" while state[counted] < config.steps: if state["active"] is None: counters = state["states"][state["current"]] if confidence(counters["repeats"]) >= config.confidence: step = _step(store, state, config) if step is None: - state["stop_reason"] = "no_chemical_exit" - _save_state(root, state) - return state + stop_reason = "no_chemical_exit" + break state.update(current=step["to"], time_s=step["time_s"]) state["steps"] += 1 if step.get("chemical"): @@ -460,10 +442,9 @@ def run_exploration( _commit_attempt(state, result) if result["status"] == "no_push": # Every later search from this state would select nothing as well. - state["stop_reason"] = "no_push" - _save_state(root, state) - return state + stop_reason = "no_push" + break _save_state(root, state) - state["stop_reason"] = "steps_completed" + state["stop_reason"] = stop_reason _save_state(root, state) return state diff --git a/src/reactionflow/ssneb.py b/src/reactionflow/ssneb.py index d2cb14a..9f99647 100644 --- a/src/reactionflow/ssneb.py +++ b/src/reactionflow/ssneb.py @@ -16,7 +16,22 @@ class _TriangularCellFilter(UnitCellFilter): - """Constrain the deformation gradient itself to six independent entries.""" + """Constrain the deformation gradient itself to six independent entries. + + Lower-triangular cells have upper-triangular deformation gradients. All six independent + components permit volume and shape changes, while removing rigid cell rotation from the path + coordinates. + """ + + def __init__(self, atoms: Atoms, pressure_GPa: float, reference_cell: Cell) -> None: + # J = sqrt(N) * (V/N)^(1/3), the usual SSNEB atomic/cell length scale. + scale = len(atoms) ** (1 / 6) * reference_cell.volume ** (1 / 3) + super().__init__( + atoms, + orig_cell=reference_cell, + cell_factor=scale, + scalar_pressure=pressure_GPa * units.GPa, + ) def set_positions(self, positions, **kwargs): positions = positions.copy() @@ -29,22 +44,7 @@ def get_forces(self, **kwargs): return forces -def cell_filter(atoms: Atoms, pressure_GPa: float, reference_cell: Cell) -> UnitCellFilter: - """Use one reference and length scale for all images, including spectators.""" - # J = sqrt(N) * (V/N)^(1/3), the usual SSNEB atomic/cell length scale. - scale = len(atoms) ** (1 / 6) * reference_cell.volume ** (1 / 3) - # Lower-triangular cells have upper-triangular deformation gradients. - # All six independent components permit volume and shape changes, while - # removing rigid cell rotation from the path coordinates. - return _TriangularCellFilter( - atoms, - orig_cell=reference_cell, - cell_factor=scale, - scalar_pressure=pressure_GPa * units.GPa, - ) - - -class _SolidStateImage: +class _SolidStateImage(_TriangularCellFilter): """Expose a cell filter as an NEB image in nonperiodic generalized space. Physical atoms retain their own periodic cells. Periodic wrapping is only @@ -54,30 +54,17 @@ class _SolidStateImage: pbc = np.zeros(3, dtype=bool) cell = Cell.new() - def __init__(self, filtered: UnitCellFilter) -> None: - self.filtered = filtered - - def __len__(self) -> int: - return len(self.filtered) - - @property - def calc(self): - return self.filtered.atoms.calc - def get_atomic_numbers(self): - return np.concatenate([self.filtered.atoms.numbers, np.zeros(3, dtype=int)]) - - def get_positions(self): - return self.filtered.get_positions() + return np.concatenate([self.atoms.numbers, np.zeros(3, dtype=int)]) - def set_positions(self, positions, **kwargs): - self.filtered.set_positions(positions, **kwargs) + def get_potential_energy(self): + # UnitCellFilter asks for the force-consistent energy unless told otherwise. + return super().get_potential_energy(force_consistent=False) - def get_forces(self, **kwargs): - return self.filtered.get_forces(**kwargs) - def get_potential_energy(self): - return self.filtered.get_potential_energy(force_consistent=False) +def cell_filter(atoms: Atoms, pressure_GPa: float, reference_cell: Cell) -> UnitCellFilter: + """Use one reference and length scale for all images, including spectators.""" + return _TriangularCellFilter(atoms, pressure_GPa, reference_cell) class SSNEB(NEB): @@ -87,7 +74,7 @@ def __init__(self, images: list[Atoms], *, pressure_GPa: float) -> None: self.physical_images = images reference = images[0].cell.copy() super().__init__( - [_SolidStateImage(cell_filter(image, pressure_GPa, reference)) for image in images], + [_SolidStateImage(image, pressure_GPa, reference) for image in images], climb=False, allow_shared_calculator=True, method="improvedtangent", diff --git a/src/reactionflow/status.py b/src/reactionflow/status.py index d5092fd..a144f27 100644 --- a/src/reactionflow/status.py +++ b/src/reactionflow/status.py @@ -15,7 +15,7 @@ from .campaign import CampaignConfig from .candidates import ReactionCandidate, ReactionClasses, _reaction_graph -from .detection import atom_ids +from .detection import _symbols_by_id from .pathway import CONVERGED_STATUSES from .store import _read_bundle @@ -42,9 +42,7 @@ def _occurrences(root: Path) -> list[sqlite3.Row]: def _describe(candidate: ReactionCandidate, radius: int) -> str: - symbols = dict( - zip(atom_ids(candidate.reactant), candidate.reactant.get_chemical_symbols(), strict=True) - ) + symbols = _symbols_by_id(candidate.reactant) changes = Counter( f"{kind} {'-'.join(sorted(symbols[atom_id] for atom_id in bond))}" for kind, bonds in ( @@ -186,23 +184,18 @@ def campaign_status(campaign: CampaignConfig, *, radius: int = 2) -> dict[str, A entry["connectivity"][connectivity["status"]] += 1 if result["status"] == "saddle_validated": entry["barriers_eV"].append(result["barrier_eV"]) - reactions = [ - { - **entry, - "pathways": dict(entry["pathways"]), - "frequency": dict(entry["frequency"]), - "connectivity": dict(entry["connectivity"]), - } - for entry in entries.values() - ] - reactions.sort( + for entry in entries.values(): + for key in ("pathways", "frequency", "connectivity"): + entry[key] = dict(entry[key]) + reactions = sorted( + entries.values(), key=lambda item: ( item["model"], item["pressure_GPa"] is not None, item["pressure_GPa"] or 0.0, -item["events"], item["reaction"], - ) + ), ) return { "campaign": str(campaign.source), diff --git a/src/reactionflow/store.py b/src/reactionflow/store.py index 8b4958d..4baabbf 100644 --- a/src/reactionflow/store.py +++ b/src/reactionflow/store.py @@ -3,7 +3,6 @@ from __future__ import annotations import json -import re import shutil import sqlite3 from contextlib import closing @@ -15,7 +14,7 @@ from ase import Atoms from ase.io import read, write -from ._durable import ensure_directory, publish, sync_directory +from ._durable import SAFE_NAME, ensure_directory, publishing, sync_directory from .candidates import ReactionCandidate, ReactionClasses from .detection import BondDetectorConfig, atom_ids, canonical_copy, transport_copy @@ -152,10 +151,7 @@ def _class_id(self, db: sqlite3.Connection, candidate: ReactionCandidate) -> str return class_id def _insert( - self, - db: sqlite3.Connection, - occurrence_id: str, - candidate: ReactionCandidate, + self, db: sqlite3.Connection, occurrence_id: str, candidate: ReactionCandidate ) -> tuple[OccurrenceRecord, bool]: existing = self._row(db, occurrence_id) if existing is not None: @@ -186,8 +182,9 @@ def _recover(self) -> None: with closing(self._connect()) as db: db.execute("BEGIN IMMEDIATE") registered = {row["occurrence_id"] for row in db.execute("SELECT * FROM occurrences")} - for occurrence_id in sorted(registered - bundles): - raise FileNotFoundError(self.candidates / occurrence_id) + missing = registered - bundles + if missing: + raise FileNotFoundError(self.candidates / min(missing)) # A complete bundle left before its row was committed is registered now. for occurrence_id in sorted(bundles - registered): metadata, candidate = _read_bundle(self.candidates / occurrence_id) @@ -197,10 +194,7 @@ def _recover(self) -> None: db.commit() def _publish( - self, - occurrence_id: str, - candidate: ReactionCandidate, - metadata: dict[str, object], + self, occurrence_id: str, candidate: ReactionCandidate, metadata: dict[str, object] ) -> ReactionCandidate: final = self.candidates / occurrence_id if final.exists(): @@ -213,18 +207,12 @@ def _publish( temporary = self.candidates / f".{occurrence_id}.tmp" if temporary.exists(): shutil.rmtree(temporary) - temporary.mkdir() - try: + with publishing(temporary, final): write(temporary / "reactant.traj", transport_copy(candidate.reactant)) write(temporary / "product.traj", transport_copy(candidate.product)) (temporary / "candidate.json").write_text( - json.dumps(metadata, indent=2, sort_keys=True) + "\n", - encoding="utf-8", + json.dumps(metadata, indent=2, sort_keys=True) + "\n", encoding="utf-8" ) - publish(temporary, final) - except Exception: - shutil.rmtree(temporary, ignore_errors=True) - raise return _read_bundle(final)[1] def register( @@ -236,7 +224,7 @@ def register( ) -> tuple[OccurrenceRecord, bool]: """Retain an occurrence and return its assignment plus whether it was inserted.""" - if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]*", occurrence_id): + if not SAFE_NAME.fullmatch(occurrence_id): raise ValueError("occurrence_id must be a filesystem-safe name") metadata = { "schema_version": 1, diff --git a/src/reactionflow/wizard.py b/src/reactionflow/wizard.py index 5168f7e..2f9a203 100644 --- a/src/reactionflow/wizard.py +++ b/src/reactionflow/wizard.py @@ -18,7 +18,7 @@ from collections import Counter from collections.abc import Callable from contextlib import suppress -from dataclasses import dataclass, field +from dataclasses import dataclass from pathlib import Path from typing import Any @@ -30,24 +30,15 @@ from .model_setup import backend_conflict, model_catalog from .run import ReactionRunConfig -# A catalog list field that a model requires one value from, and the option it sets. -_CHOICES = {"heads": "head", "tasks": "task", "modals": "modal"} _UNSAFE = re.compile(r"[^A-Za-z0-9._-]+") -class _Quit(Exception): - """The session ended before the file was written.""" - - def _ask(question: str, default: str | None = None) -> str: """Return one answer. Enter takes the default; a None default means an answer is required.""" shown = f" [{default}]" if default else "" while True: - try: - answer = input(f"{question}{shown}: ").strip() - except EOFError: - raise _Quit from None + answer = input(f"{question}{shown}: ").strip() if answer or default is not None: return answer or default @@ -160,10 +151,10 @@ def parse(text: str) -> list[int]: return picks if many else picks[0] -@dataclass(frozen=True) +@dataclass class _Structure: path: Path - atoms: Atoms = field(compare=False) + atoms: Atoms @property def periodic(self) -> bool: @@ -186,8 +177,6 @@ def _structures(text: str) -> list[_Structure]: except Exception as error: raise ValueError(f"cannot read {name}: {error}") from None found.setdefault(path, _Structure(path, atoms)) - if not found: - raise ValueError("enter at least one structure file") for structure in found.values(): atoms, cell = structure.atoms, "periodic" if structure.periodic else "not periodic" formula = atoms.get_chemical_formula() @@ -235,7 +224,7 @@ def _absolute_paths(text: str) -> list[str]: def _builtin(backend: dict[str, Any]) -> _Profile: - """Turn one catalog model, and any head, task, or modal it needs, into adapter options.""" + """Turn one catalog model, and any head or task it needs, into adapter options.""" name = backend["backend"] models = backend["models"] @@ -250,11 +239,10 @@ def _builtin(backend: dict[str, Any]) -> _Profile: options: dict[str, Any] = {} if name == "ani1xnr" else {"model": entry["name"]} if "family" in entry: options["family"] = entry["family"] - for choices, option in _CHOICES.items(): - values = entry.get(choices) or [] - if values: - described = backend.get(choices) if isinstance(backend.get(choices), dict) else {} - labels = [f"{value:<16} {described.get(value, '')}".rstrip() for value in values] + # A catalog list field that a model requires one value from, and the option it sets. + for choices, option in (("heads", "head"), ("tasks", "task")): + if values := entry.get(choices): + labels = [f"{value:<16} {backend[choices][value]}" for value in values] options[option] = values[_menu(f"{option.capitalize()} for {entry['name']}", labels)] if name == "mace" and entry["name"] == "mh-0": options["head"] = _ask("Head name in the MACE-MH-0 checkpoint") @@ -262,7 +250,7 @@ def _builtin(backend: dict[str, Any]) -> _Profile: parts = [] if name.replace("_", "") in model.lower().replace("-", "") else [name] if options.get("family", "mp") != "mp" and options["family"] not in model: parts.append(options["family"]) - parts += [model, *(options[key] for key in ("head", "task", "modal") if key in options)] + parts += [model, *(options[key] for key in ("head", "task") if key in options)] return _Profile(_safe("-".join(parts)), backend["factory"], options, name) @@ -293,11 +281,10 @@ def _profile_name(text: str) -> str: def _needs_charge_and_spin(profile: _Profile) -> bool: - options = profile.options return ( profile.backend == "aimnet2" - or (profile.backend == "uma" and options.get("task") == "omol") - or (profile.backend == "mace" and options.get("family") == "polar") + or (profile.backend == "uma" and profile.options.get("task") == "omol") + or (profile.backend == "mace" and profile.options.get("family") == "polar") ) @@ -305,9 +292,7 @@ def _models( catalog: list[dict], previous: list[_Profile] | None, *, single: bool = False ) -> list[_Profile]: labels = [ - f"{backend['backend']:<11} {len(backend['models'])} " - f"model{'s' if len(backend['models']) > 1 else ''}" - for backend in catalog + f"{backend['backend']:<11} {_count(len(backend['models']), 'model')}" for backend in catalog ] labels.append("another model with an ASE calculator") if single: @@ -379,31 +364,22 @@ def _group(number: int, catalog: list[dict], previous: _Group | None, used: list ) hydrostatic = True if any(pressure is not None for pressure in pressures): - hydrostatic = ( - _menu( - "Under pressure, the cell", - ["keeps its shape and scales uniformly", "may also change shape"], - default="1" if not previous or previous.hydrostatic else "2", - ) - == 0 + cell_choice = _menu( + "Under pressure, the cell", + ["keeps its shape and scales uniformly", "may also change shape"], + default="1" if not previous or previous.hydrostatic else "2", ) + hydrostatic = cell_choice == 0 combinations = len(structures) * len(profiles) * len(temperatures) * len(pressures) runs = _ask_until_valid( f"Runs of each of these {combinations} combinations", _positive(), str(previous.runs) if previous else "1", ) - factors = [ - (structures, "structure"), - (profiles, "model"), - (temperatures, "temperature"), - (pressures, "pressure"), - (range(runs), "run"), - ] print( - " " - + " x ".join(_count(len(items), word) for items, word in factors) - + f" = {_count(combinations * runs, 'trajectory')}" + f" {_count(len(structures), 'structure')} x {_count(len(profiles), 'model')}" + f" x {_count(len(temperatures), 'temperature')} x {_count(len(pressures), 'pressure')}" + f" x {_count(len(range(runs)), 'run')} = {_count(combinations * runs, 'trajectory')}" ) return _Group(structures, profiles, temperatures, pressures, hydrostatic, runs) @@ -427,25 +403,25 @@ def _campaign(path: Path, groups: list[_Group], settings: dict[str, Any]) -> dic for profile in group.profiles: profiles.setdefault(profile.key, profile) names: dict[str, str] = {} + adapter_profiles: dict[str, Any] = {} for key, profile in profiles.items(): name, number = profile.name, 2 - while name in names.values(): + while name in adapter_profiles: name, number = f"{profile.name}-{number}", number + 1 names[key] = name + options = dict(profile.options) + if profile.backend: + options["device"] = "cuda" if settings["gpu"] else "cpu" + adapter_profiles[name] = {"factory": profile.factory, "options": options} structures = list(dict.fromkeys(s.path for group in groups for s in group.structures)) labels = _labels(structures) several_structures, several_models = len(structures) > 1, len(profiles) > 1 - - def relative(structure: Path) -> str: - return os.path.relpath(structure, path.parent) - stems, trajectories = [], [] seed = secrets.randbelow(2**32) for group in groups: - combinations = itertools.product( + for structure, profile, temperature, pressure in itertools.product( group.structures, group.profiles, group.temperatures, group.pressures - ) - for structure, profile, temperature, pressure in combinations: + ): anisotropic = pressure is not None and not group.hydrostatic parts = [labels[structure.path]] if several_structures else [] if several_models: @@ -458,7 +434,7 @@ def relative(structure: Path) -> str: for _ in range(group.runs): trajectory = {"adapter_profile": names[profile.key]} if several_structures: - trajectory["structure"] = relative(structure.path) + trajectory["structure"] = os.path.relpath(structure.path, path.parent) trajectory |= { "total_steps": settings["steps"], "timestep_fs": settings["timestep"], @@ -478,15 +454,9 @@ def relative(structure: Path) -> str: counted[stem] += 1 trajectories[index] = {"id": f"{stem}-{counted[stem]:0{width}d}", **trajectories[index]} - adapter_profiles = {} - for key, profile in profiles.items(): - options = dict(profile.options) - if profile.backend: - options["device"] = "cuda" if settings["gpu"] else "cpu" - adapter_profiles[names[key]] = {"factory": profile.factory, "options": options} campaign: dict[str, Any] = {"schema_version": 2} if not several_structures: - campaign["structure"] = relative(structures[0]) + campaign["structure"] = os.path.relpath(structures[0], path.parent) return campaign | { "output_root": settings["output"], "require_gpu": settings["gpu"], @@ -496,7 +466,7 @@ def relative(structure: Path) -> str: } -def _write(path: Path, campaign: dict[str, Any], load: Callable = CampaignConfig.load) -> None: +def _write(path: Path, campaign: dict[str, Any], load: Callable) -> None: path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_name(f".{path.name}.{secrets.token_hex(4)}.tmp") temporary.write_text(json.dumps(campaign, indent=2) + "\n", encoding="utf-8") @@ -558,12 +528,15 @@ def _eon_session(path: Path, catalog: list[dict]) -> int: "structure, which has none. In an extxyz file, add charge=0 spin=1 (your values) " "to the comment line." ) + temperature = _ask_until_valid("Temperature in K", _positive(float), "300") + steps = _ask_until_valid("Kinetic Monte Carlo steps", _positive(), "100") + confidence = _ask_until_valid( + "Confidence to reach in each state before stepping", _fraction, "0.95" + ) akmc = { - "temperature_K": _ask_until_valid("Temperature in K", _positive(float), "300"), - "steps": _ask_until_valid("Kinetic Monte Carlo steps", _positive(), "100"), - "confidence": _ask_until_valid( - "Confidence to reach in each state before stepping", _fraction, "0.95" - ), + "temperature_K": temperature, + "steps": steps, + "confidence": confidence, "seed": secrets.randbelow(2**32), } eon: dict[str, Any] = {} @@ -655,11 +628,10 @@ def _session(path: Path) -> int: while True: used = [profile for group in groups for profile in group.profiles] group = _group(len(groups) + 1, catalog, groups[-1] if groups else None, used) - if not _yes("Keep this group?", True): - continue - groups.append(group) - if not _yes("Add another group, for example a different model or structure?", False): - break + if _yes("Keep this group?", True): + groups.append(group) + if not _yes("Add another group, for example a different model or structure?", False): + break print() steps = _ask_until_valid("MD steps per trajectory", _positive(), "100000") @@ -694,7 +666,7 @@ def _session(path: Path) -> int: if not _yes(f"Write {_shown(path)}?", True): print("Nothing was written.") return 1 - _write(path, campaign) + _write(path, campaign, CampaignConfig.load) shown = shlex.quote(_shown(path)) print( f"Wrote {shown}. Next:\n" @@ -707,7 +679,7 @@ def _session(path: Path) -> int: def _relative_directory(text: str) -> str: - if os.path.isabs(text) or not text: + if os.path.isabs(text): raise ValueError("give a directory relative to the campaign file") return text @@ -719,7 +691,7 @@ def create_campaign(path: Path) -> int: import readline # noqa: F401 (line editing and history in input()) try: return _session(Path(path).resolve()) - except (KeyboardInterrupt, _Quit): + except (KeyboardInterrupt, EOFError): print("\nStopped. Nothing was written.") return 1