Source code for pyretis.simulation.setup

"""Setup for the infinite-swapping scheduler.

Reads the TOML config, builds the replica-exchange (REPEX) state, and
hands the worker pool to
:py:class:`pyretis.simulation.async_runner.aiorunner`.
"""

import logging
import os
from typing import Optional, Tuple

import tomli

from pyretis.core.path_load import load_paths_from_disk
from pyretis.core.moves import run_md
from pyretis.core.system import System
from pyretis.engines.factory import create_engines
from pyretis.initiation.initiate_kick import run_parallel_kick_job
from pyretis.setup.createsimulation import prepare_system
from pyretis.simulation.async_runner import (
    aiorunner,
    future_list,
    get_worker_engines,
)
from pyretis.simulation.repex import InfSwapState, spawn_rng
from pyretis.inout.config_adapter import (
    _pad3,
    _write_external_path_files,
    _write_inf_path_files,
)
from pyretis.inout.pathensemble_output import init_native_pathensemble_files

logger = logging.getLogger("main")
logger.setLevel(logging.DEBUG)


[docs] class TOMLConfigError(Exception): """Raised when there is an error in the .toml configuration."""
# def __init__(self, message): # super().__init__(message)
[docs] def _build_ensemble_windows(config: dict) -> InfSwapState: """Build an ``InfSwapState`` with its ensemble windows initiated. Shared by :func:`setup_internal` (the steady-state scheduler state) and :func:`run_kick_phase` (a throwaway state used only to read off ``.ensembles``/``.rgen``/``.repptis`` for the parallel kick phase). The two calls build INDEPENDENT ``InfSwapState`` instances -- there is no shared mutable RNG state between the kick phase and the steady-state scheduler that follows it, which is acceptable under "init = validity not identity" (the kick phase does not need to be byte-reproducible against the steady-state run's own RNG stream). Parameters ---------- config : dict The configuration dictionary. Returns ------- InfSwapState A REPEX state with ``.ensembles`` already built (via ``initiate_ensembles``); no paths are loaded. """ # A single-ensemble TIS run has no [0^-] minus ensemble (it is one # [i^+] ensemble defined by its [left, middle, right] bounds), an # EXPLORE run has none either (N-1 independent positive ensembles), # and a PPTIS run without zero_left/flux likewise has no [0^-]. All # three are built without the minus offset; every other native route # keeps the [0^-]/[0^+] structure. no_minus = config["simulation"].get("single_tis", False) or \ config["simulation"].get("explore", False) or \ config["simulation"].get("pptis_no_minus", False) state = InfSwapState(config, minus=not no_minus) state.initiate_ensembles() return state
[docs] def setup_internal(config: dict) -> Tuple[dict, InfSwapState]: """Run the various setup functions. Parameters ---------- config : dict the configuration dictionary Returns ------- tuple A blank md_items dict An initialized REPEX state """ state = _build_ensemble_windows(config) # load paths from disk and add to repex paths = load_paths_from_disk(config) state.load_paths(paths) # Initialise the per-ensemble pathensemble.txt files (headers + # directories) -- the only output format the coordinator writes now. # A fresh start writes them clean; a restart keeps existing files # (appending) but still CREATES any that are missing -- a run that # ALWAYS restarts (method="restart", seeded from a restart that never # wrote native output) would otherwise leave the analysis with # nothing to read. init_native_pathensemble_files( state, overwrite="restarted_from" not in config["current"]) # create first md_items dict md_items = { "mc_moves": state.mc_moves, "interfaces": state.interfaces, "cap": state.cap, # Single-ensemble TIS: the worker weights the trial for the ONE # ensemble (a single compute_weight, not the multi-ensemble cv). "single_tis": state.single_tis, # Explore: each independent positive ensemble weights its trial for # its OWN window (per-ensemble compute_weight); an explore path is # accepted anywhere, so its state occupancy is a one-hot (add_traj). "explore": state.explore, } # setup the engine_occupation list _, engine_occ = create_engines(config) state.engine_occ = engine_occ return md_items, state
[docs] def _run_md_task(md_items: dict) -> dict: """Worker-side entry point for the runner. Runs in the worker process and injects that worker's process-local engine pool into :func:`pyretis.core.moves.run_md`, so the runner stays engine-agnostic and ``run_md`` takes its engines explicitly rather than reading a module global. """ return run_md(md_items, get_worker_engines())
[docs] def _run_kick_task(kick_item: dict) -> dict: """Worker-side entry point for one ensemble's parallel kick job. Mirrors :func:`_run_md_task`: runs in the worker process and injects that worker's process-local engine pool into :func:`pyretis.initiation.initiate_kick.run_parallel_kick_job`. """ return run_parallel_kick_job(kick_item, get_worker_engines())
[docs] def setup_runner(state: InfSwapState) -> Tuple[aiorunner, future_list]: """Set the task runner class up. Parameters ---------- state : InfSwapState A REPEX state from which to get the config dict """ # setup client with state.workers workers runner = aiorunner(state.config, state.config["runner"]["workers"]) # Attach the run_md task and start the runner's workers runner.set_task(_run_md_task) runner.start() # A managed list of futures futures = future_list() return runner, futures
[docs] def _build_bootstrap_system(config: dict, run_dir: str) -> System: """Build the shared starting point as a file-backed snapshot. Every kick job runs inside a scheduler worker, whose engines have already had ``setup_streaming``/process-local construction applied (:func:`pyretis.simulation.async_runner.worker_initializer`, ``prepare_streaming_engines``). For an internal engine, this means ``EngineBase._is_streaming`` -- which gates whether ``modify_velocities`` takes the streaming (raw-numpy-rgen-compatible) branch or the legacy native one (which needs a :py:class:`pyretis.core.random_gen.RandomGenerator`-wrapped rgen, the wrong type for the scheduler's rgen convention) -- is true only when the system handed to it carries NO in-memory particles (``has_particles=False``, a file-backed snapshot). External engines are file-backed by construction already. So the starting system for every kick job must be a snapshot, not the native in-memory :func:`pyretis.setup.createsimulation.prepare_system` output. This resolves the full native particle/box/forcefield setup once via ``prepare_system`` (handling every native particle-input convention: inline ``pos``, ``[particles.position] input_file``, per-engine construction), then writes that resolved state out as a one-frame bootstrap trajectory file and returns a fresh snapshot ``System`` pointing at it -- the same pattern :func:`pyretis.inout.config_adapter._write_inf_path_files` already uses to serialise an accepted in-memory path into the coordinator's on-disk, file-backed format. Parameters ---------- config : dict The translated scheduler config. run_dir : str The absolute directory the simulation runs from; the bootstrap file is written there. Returns ------- System A snapshot system (``has_particles`` False) with ``.config`` pointing at the written bootstrap file. """ native_system = prepare_system(config) names = list(config.get("particles", {}).get("name", ["X"])) bootstrap_file = os.path.join(run_dir, "kick_bootstrap.xyz") pos = native_system.pos vel = native_system.vel npart = pos.shape[0] with open(bootstrap_file, "w", encoding="utf-8") as xyz: xyz.write(f"{npart}\n") xyz.write("# Box: 0.0000 0.0000 0.0000\n") for j in range(npart): row_pos = _pad3(pos[j]) row_vel = _pad3(vel[j]) name = names[j] if j < len(names) else names[0] coords = "".join(f"{v:>16.9f}" for v in (*row_pos, *row_vel)) xyz.write(f"{name:<3s}{coords}\n") snapshot = System() snapshot.config = (bootstrap_file, 0) return snapshot
[docs] def _build_kick_items(state: InfSwapState, system0, config: dict, run_dir: str) -> list: """Build one kick job dict per ensemble window. Parameters ---------- state : InfSwapState A state with ``.ensembles``/``.rgen``/``.repptis`` already set (from :func:`_build_ensemble_windows`). system0 The shared starting :py:class:`.System`; a fresh copy is taken per ensemble. config : dict The translated scheduler config. run_dir : str The absolute directory the simulation runs from. Returns ------- list of dict One kick job per ensemble, see :func:`pyretis.initiation.initiate_kick.run_parallel_kick_job` for the dict's keys. """ ens_engs = config["simulation"]["ensemble_engines"] kick_items = [] for ens_num, ens in state.ensembles.items(): kick_items.append({ "ens_num": ens_num, # ens_num here is state.ensembles' own 0-indexed key (e.g. 0 # = [0^-]); ens_engs is indexed the same way. This differs # from prep_md_items' "ens_engs[ens_num + 1]" (repex.py), whose # ens_num comes from the PICKED dict, offset by -1 (ens_num=-1 # for [0^-]) -- "+ 1" there converts back to this same index. "eng_name": ens_engs[ens_num][0], "interfaces": ens["interfaces"], "start_cond": ens["start_cond"], # The body ensembles of a PPTIS/REPPTIS route must contain # only paths crossing the middle interface (state.repptis is # the already-translated, authoritative flag -- see # InfSwapState.initiate_ensembles, which gives every i != 0 # body ensemble a ["L", "R"] start_cond under that route). "must_cross_M": bool(state.repptis) and ens_num != 0, "tis_settings": ens["tis_set"], "system": system0.copy(), "run_dir": run_dir, "engine_rgen": spawn_rng(state.rgen), "kick_rgen": spawn_rng(state.rgen), "path_rgen": spawn_rng(state.rgen), }) return kick_items
[docs] def _dispatch_kick_items(config: dict, kick_items: list) -> list: """Run every kick job through a short-lived worker pool. Parameters ---------- config : dict The translated scheduler config (forwarded to the pool's worker initializer, see :class:`pyretis.simulation.async_runner.aiorunner`). kick_items : list of dict One kick job per ensemble, from :func:`_build_kick_items`. Returns ------- list of dict One ``run_parallel_kick_job`` result per kick item. """ runner = aiorunner(config, config["runner"]["workers"]) runner.set_task(_run_kick_task) runner.start() try: futures = future_list() for kick_item in kick_items: futures.add(runner.submit_work(kick_item)) results = [futures.as_completed().result() for _ in kick_items] runner.stop() except BaseException: runner.close() raise return results
[docs] def _write_kick_results(config: dict, run_dir: str, results: list) -> str: """Serialise every kick result into ``load_dir/<i>/``. Reuses the same on-disk contract :func:`pyretis.inout.config_adapter.generate_load_dir` already writes (and :func:`pyretis.core.path_load.load_paths_from_disk` already reads), via the existing :func:`pyretis.inout.config_adapter._write_external_path_files` / :func:`pyretis.inout.config_adapter._write_inf_path_files` writers. Parameters ---------- config : dict The translated scheduler config. run_dir : str The absolute directory the simulation runs from. results : list of dict One ``run_parallel_kick_job`` result per ensemble, from :func:`_dispatch_kick_items`. Returns ------- str The absolute path to the populated ``load_dir``. """ names = list(config.get("particles", {}).get("name", ["X"])) load_path = os.path.join(run_dir, config["simulation"]["load_dir"]) os.makedirs(load_path, exist_ok=True) for result in results: path = result["path"] pdir = os.path.join(load_path, str(result["ens_num"])) os.makedirs(pdir, exist_ok=True) # Mirror generate_load_dir's external-vs-in-memory dispatch: a # file-backed engine path has no in-memory particles (or no pos # on them), so it needs the trajectory-file archiver rather than # the xyz writer. first_pp = path.phasepoints[0] if path.phasepoints else None external = ( first_pp is None or not first_pp.has_particles or getattr(first_pp.particles, "pos", None) is None ) if external: _write_external_path_files(path, pdir) else: _write_inf_path_files(path, pdir, names) return load_path
[docs] def run_kick_phase(config: dict, run_dir: str) -> str: """Kick-initiate every ensemble in parallel, one job per ensemble. Builds a short-lived worker pool (the same ``aiorunner`` / ``worker_initializer`` / ``create_engines`` machinery :func:`setup_runner` uses for the steady-state scheduler loop), dispatches one kick job per ensemble, waits for all to complete, and serialises each accepted path into ``load_dir/<i>/`` via the existing :func:`pyretis.inout.config_adapter._write_external_path_files` / :func:`pyretis.inout.config_adapter._write_inf_path_files` writers -- the same on-disk contract :func:`pyretis.core.path_load.load_paths_from_disk` reads, so the steady-state scheduler that follows this call needs no changes. Only ``[initial-path] kick-from = "initial"`` is supported here: each ensemble's kick is initiated independently, from its own starting system, in parallel. ``kick-from = "previous"`` (each ensemble seeded from the closest phase point of the PREVIOUS ensemble's just-kicked path) is inherently sequential -- the caller is responsible for routing that case to the existing :func:`pyretis.inout.config_adapter.generate_load_dir` instead. Parameters ---------- config : dict The translated scheduler config (see :func:`pyretis.inout.config_adapter.native_to_infswap_config`), with ``apply_config_defaults`` already applied (so ``ensemble_engines``/``load_dir`` are populated). run_dir : str The absolute directory the simulation runs from. Returns ------- str The absolute path to the populated ``load_dir``. """ # InfSwapState requires config["current"]["size"] (and the # "restarted_from" rgen-resume check). At this point in the # pipeline -- before output.toml is ever written -- the caller's # config legitimately has no [current] yet (setup_config builds it # the first time the scheduler reads output.toml back). Build it on # a shallow COPY of config, not in place: dumping config to # output.toml below must stay WITHOUT a [current] (the documented # "write without current, let setup_config build it" contract). state_config = dict(config) state_config.setdefault("current", _build_fresh_current(config)) state = _build_ensemble_windows(state_config) system0 = _build_bootstrap_system(config, run_dir) kick_items = _build_kick_items(state, system0, config, run_dir) results = _dispatch_kick_items(config, kick_items) return _write_kick_results(config, run_dir, results)
[docs] def _build_fresh_current(config: dict) -> dict: """Build a fresh ``[current]`` state (cstep 0) for a new run. The ensemble count depends on the route: a single-ensemble TIS run is ONE [i^+] ensemble (its three interfaces are the [left, middle, right] bounds), an EXPLORE run has N-1 independent positive ensembles, and a PPTIS-without-zero_left run has no [0^-]; every other route keeps the one-ensemble-per-interface layout. Parameters ---------- config : dict The configuration dictionary. Returns ------- dict The fresh ``[current]`` sub-dict. """ interfaces = config["simulation"]["interfaces"] if config["simulation"].get("single_tis", False): # single-ensemble TIS: the three interfaces are ONE ensemble's # [left, middle, right] bounds, so there is exactly one ensemble # (not three). initiate_ensembles builds it from the triple. size = 1 elif config["simulation"].get("explore", False): size = len(interfaces) - 1 elif config["simulation"].get("pptis_no_minus", False): # pptis-without-zero_left: no [0^-]. The [0^+] is present # only when zero_ensemble=True (pptis default = False). has_zero_plus = not config["simulation"].get( "pptis_no_zero_plus", False) size = int(has_zero_plus) + max(len(interfaces) - 2, 0) else: size = len(interfaces) return { "traj_num": size, "cstep": 0, "active": list(range(size)), "locked": [], "size": size, "frac": {}, "wsubcycles": [0 for _ in range(config["runner"]["workers"])], "tsubcycles": 0, }
[docs] def setup_config( inp: str = "output.toml", re_inp: str = "output.toml" ) -> Optional[dict]: """Set dict from a TOML file up. The single run file is ``output.toml``: it carries the fully-resolved configuration (settings + applied defaults) and, once the scheduler has written it, the running ``[current]`` state. A fresh start and a resume therefore read the same file. The legacy ``infswap.toml`` / ``restart.toml`` names are still accepted (an existing on-disk run started before the consolidation keeps resuming) -- the caller just passes that filename as ``inp``. Parameters ---------- inp : str A string specifying the input file (def: output.toml). re_inp : str A string specifying the restart file (def: output.toml). When it differs from ``inp`` (a foreign run file is passed) and exists, it is rejected so a stray state file cannot shadow the run file. Returns ------- dict or None A dictionary containing the configuration parameters or None. """ # sets up the dict from a TOML file. # load input: if os.path.isfile(inp): with open(inp, mode="rb") as read: config = tomli.load(read) else: logger.info("%s file not found, exit.", inp) return None # When inp and re_inp name the SAME file (the output.toml default) # this guard is inert -- the consolidated file legitimately holds both # the config and the state. It only fires for the foreign-file case: # a separate state file found alongside a differently-named run file. if inp != re_inp and os.path.isfile(re_inp): msg = f"Restart file '{re_inp}' found, but its not the run file!" raise ValueError(msg) # A resume is discriminated solely by an ADVANCED [current] state: # a [current] whose cstep has progressed past zero. A freshly-written # output.toml carries a cstep=0 [current] (the scheduler writes it at # t0) -- that is NOT a resume, so it takes the fresh branch below. curr = config.get("current") if curr is not None and curr.get("cstep", 0) > 0: # --- resume: continue from the persisted [current] state --- # if cstep and steps are equal, we stop here. if curr.get("cstep") == curr.get("restarted_from", -1): return None # set 'restarted_from' curr["restarted_from"] = curr["cstep"] # check active paths: load_dir = config["simulation"].get("load_dir", "trajs") for act in curr["active"]: store_p = os.path.join(load_dir, str(act), "traj.txt") if not os.path.isfile(store_p): return None elif curr is None: # no 'current' in toml, start from step 0. config["current"] = _build_fresh_current(config) # The remaining case (a cstep=0 [current] already present, e.g. the # output.toml the scheduler writes at t0) keeps that [current] as-is: # it is neither a resume nor a fresh build. # Apply the configuration defaults so the resolved values are present # (and persisted into output.toml). Idempotent: re-applying it to an # output.toml that already carries the defaults is a no-op. apply_config_defaults(config) check_config(config) return config
[docs] def apply_config_defaults(config: dict) -> dict: """Fill in the resolved configuration defaults in place. Sets the per-ensemble engine list and the ``[simulation]`` / ``[simulation.tis_set]`` / ``[output]`` defaults, validates the output-deletion settings, and applies the quantis engine override. The function does NOT touch ``[current]`` (the running state), so it is safe to call both before a ``[current]`` exists (when dumping a fresh ``output.toml``) and from :func:`setup_config` after it has been built. Every assignment is a ``setdefault``/idempotent guard, so a repeated call is a no-op. Parameters ---------- config : dict The configuration dictionary, mutated in place. Returns ------- dict The same ``config`` dictionary, with defaults applied. """ # quantis or any other method requiring different engines in each ensemble has_ens_engs = config["simulation"].get("ensemble_engines", False) if not has_ens_engs: ens_engs = [] for _ in config["simulation"]["interfaces"]: ens_engs.append(["engine"]) config["simulation"]["ensemble_engines"] = ens_engs # set all keywords only once, so they appear in output.toml # and we can avoid the .get() in other parts if "seed" not in config["simulation"].keys(): config["simulation"]["seed"] = 0 # [simulation] defaults config["simulation"].setdefault("load_dir", "load") config["simulation"].setdefault("zeroswap", 0.5) config["simulation"].setdefault("pick_scheme", 0) # [simulation.tis_set] defaults config["simulation"]["tis_set"].setdefault("quantis", False) config["simulation"]["tis_set"].setdefault("lambda_minus_one", False) config["simulation"]["tis_set"].setdefault("accept_all", False) # we do not set default interface_cap here, it defaults to # interfaces[-1] in wf already. # [output] defaults config["output"].setdefault("keep_maxop_trajs", False) config["output"].setdefault("delete_old", False) config["output"].setdefault("delete_old_all", False) # validation for output settings keep_maxop_trajs = config["output"]["keep_maxop_trajs"] delete_old = config["output"]["delete_old"] delete_old_all = config["output"]["delete_old_all"] if not delete_old and keep_maxop_trajs: raise TOMLConfigError("keep_maxop_trajs=True requires delete_old=True") if delete_old_all and keep_maxop_trajs: msg = ( "delete_old_all=True will delete all trajectories. Set " "keep_maxop_trajs to False in the [output] section" ) raise TOMLConfigError(msg) # handle quantis configuration quantis = config["simulation"]["tis_set"]["quantis"] if quantis and not has_ens_engs: config["simulation"]["ensemble_engines"][0] = ["engine0"] return config
[docs] def check_config(config: dict) -> None: """Perform some checks on the settings from the .toml file. Parameters ---------- config : dict The configuration dictionary. """ intf = config["simulation"]["interfaces"] n_ens = len(config["simulation"]["interfaces"]) n_workers = config["runner"]["workers"] sh_moves = config["simulation"]["shooting_moves"] n_sh_moves = len(sh_moves) # Single-ensemble TIS: the three interfaces are ONE ensemble's # [left, middle, right] bounds (left may equal middle for the [0^+] # ensemble) with a single shooting move -- so the interface-as-N- # ensemble-positions checks (no duplicates, n_ens <= n_shooting_moves) # do not apply. The sorted (non-decreasing) check still holds. single_tis = config["simulation"].get("single_tis", False) pptis_no_minus = config["simulation"].get("pptis_no_minus", False) skip_intf_checks = single_tis or pptis_no_minus lambda_minus_one = config["simulation"]["tis_set"]["lambda_minus_one"] intf_cap = config["simulation"]["tis_set"].get("interface_cap", False) if lambda_minus_one is not False and lambda_minus_one >= intf[0]: raise TOMLConfigError( "lambda_minus_one interface must be less than the first interface!" ) if n_ens < 2: raise TOMLConfigError("Define at least 2 interfaces!") if n_workers > n_ens - 1: raise TOMLConfigError("Too many workers defined!") # The permeability mirror / target-swap moves persist a global # order-function mutation (the mirror flag / target index) on accept, # which is process-local per worker; with >1 worker a later pick of # the same ensemble may land on a worker with an un-mutated order # function -> inconsistent, non-reproducible sampling. The native # route already gates this (pyretisrun.scheduler_supported_features); # enforce it on the infinite-swapping route too. tis_set = config["simulation"]["tis_set"] if (tis_set.get("mirror_freq") or tis_set.get("target_freq")) \ and n_workers > 1: raise TOMLConfigError( "Permeability mirror / target-swap moves (mirror_freq / " "target_freq) require [runner] workers = 1." ) if sorted(intf) != intf: raise TOMLConfigError("Your interfaces are not sorted!") if not skip_intf_checks and len(set(intf)) != len(intf): raise TOMLConfigError("Your interfaces contain duplicate values!") if not skip_intf_checks and n_ens > n_sh_moves: raise TOMLConfigError( f"N_interfaces {n_ens} > N_shooting_moves {n_sh_moves}!" ) if intf_cap and intf_cap > intf[-1]: raise TOMLConfigError( f"Interface_cap {intf_cap} > interface[-1]={intf[-1]}" ) if intf_cap and intf_cap < intf[0]: raise TOMLConfigError( f"Interface_cap {intf_cap} < interface[-2]={intf[-2]}" ) # engine checks unique_engines = [] for engines in config["simulation"]["ensemble_engines"]: for engine in engines: if engine not in unique_engines: unique_engines.append(engine) for key1 in unique_engines: if key1 not in config.keys(): raise TOMLConfigError(f"Engine '{key1}' not defined!") # gromacs check for key1 in unique_engines: if config[key1]["class"] == "gromacs": eng1 = config[key1].copy() inp_path1 = eng1.pop("input_path") for key2 in unique_engines: eng2 = config[key2].copy() inp_path2 = eng2.pop("input_path") if eng1 != eng2 and inp_path1 == inp_path2: raise TOMLConfigError( "Found differing engine settings with identic" + "al 'input_path'. This would overwrite the" + " settings of one of the engines in" + " 'pyretis.mdp'!" ) # check wsubcycles and tsubcycles in case restarting from old version if "wsubcycles" not in config["current"]: list_of_zeros = [0 for _ in range(config["runner"]["workers"])] config["current"]["wsubcycles"] = list_of_zeros if "tsubcycles" not in config["current"]: config["current"]["tsubcycles"] = 0 # if increased number of workers wsub_num = len(config["current"]["wsubcycles"]) if wsub_num < config["runner"]["workers"]: extra = config["runner"]["workers"] - wsub_num config["current"]["wsubcycles"] += [0] * extra