# -*- coding: utf-8 -*-
# Copyright (c) 2023, PyRETIS Development Team.
# Distributed under the LGPLv2.1+ License. See LICENSE for more info.
"""Iterative ``infinit`` interface-placement driver for PyRETIS.
Ported (and adapted to the PyRETIS run interface) from the upstream
``inftools`` ``infinit`` tool. ``infinit`` automatically places TIS/RETIS
interfaces by repeating:
1. run a short infinite-swapping simulation with the current interfaces;
2. re-estimate the crossing probability with WHAM and re-place the
interfaces so every ensemble carries roughly the same local crossing
probability (via
:func:`pyretis.analysis.interface_estimation.interfaces_from_data`);
3. re-select the most-decorrelated active paths from the run as the
initial paths for the next iteration's interface set;
4. repeat for the configured number of iterations.
The scientific core (binless WHAM crossing probability + geometric
interface placement) lives in
:mod:`pyretis.analysis.interface_estimation`; this module
is the orchestration around the PyRETIS infinite-swapping scheduler
(:func:`pyretis.simulation.scheduler.scheduler`).
Unlike the upstream driver -- which mutates one run file in place
and threads a ``cstep`` state machine through it -- each PyRETIS iteration
runs in its own ``<workdir>/step_<i>`` directory from a *fresh* config (no
``[current]`` section), loading a ``load/`` directory of re-selected paths
numbered ``0 .. N-1``. PyRETIS' :func:`~pyretis.simulation.setup.setup_config`
then starts that iteration cleanly. This keeps every iteration's I/O
isolated and reproducible.
The driver does **not** generate the very first set of initial paths
(``cstep == -1`` / zero-path generation in the upstream); supply a working
``load_dir`` for the starting interfaces, exactly as the infinite-swapping
examples do.
"""
import logging
import os
import shutil
import tomllib
import numpy as np
from pyretis.analysis.interface_estimation import interfaces_from_data
from pyretis.inout.settings import dump_toml
logger = logging.getLogger(__name__) # pylint: disable=invalid-name
logger.addHandler(logging.NullHandler())
__all__ = ['set_default_infinit', 'reselect_initial_paths', 'run_infinit']
[docs]
def set_default_infinit(config):
"""Validate and fill the ``[infinit]`` section of a config dict.
Parameters
----------
config : dict
The parsed TOML config; must contain an ``[infinit]`` table with at
least ``steps_per_iter`` (a list of per-iteration step counts).
Returns
-------
iset : dict
The validated ``[infinit]`` settings (defaults filled in place).
"""
if 'infinit' not in config:
raise ValueError("Config has no [infinit] section.")
iset = config['infinit']
pl_target = iset.get('pL', 0.3)
if not 0.0 < pl_target < 1.0:
raise ValueError("[infinit] pL must be strictly between 0 and 1.")
iset['pL'] = pl_target
if 'nskip' in iset:
raise ValueError("[infinit] 'nskip' is replaced by 'skip' "
"(a fraction between 0 and 1).")
skip = iset.get('skip', 0.1)
if not 0.0 <= skip < 1.0:
raise ValueError("[infinit] 'skip' must be a fraction in [0, 1).")
iset['skip'] = skip
lamres = iset.get('lamres', 0.001)
if lamres < 1e-5:
raise ValueError("[infinit] 'lamres' must be >= 1e-5.")
iset['lamres'] = lamres
if 'steps_per_iter' not in iset:
raise ValueError("[infinit] needs 'steps_per_iter' (a list).")
steps_per_iter = list(iset['steps_per_iter'])
workers = config.get('runner', {}).get('workers', 1)
if not all(steps >= workers for steps in steps_per_iter):
raise ValueError("Every 'steps_per_iter' entry must be >= the "
"number of workers.")
iset['steps_per_iter'] = steps_per_iter
iset.setdefault('cstep', 0)
return iset
def _read_toml(path):
"""Load a TOML file into a dict."""
with open(path, 'rb') as infile:
return tomllib.load(infile)
def _write_toml(config, path):
"""Write a config dict to a TOML file."""
with open(path, 'wb') as outfile:
dump_toml(config, outfile)
def _path_order(path_dir):
"""Return ``(op_second_frame, max_op)`` for a stored path, or ``None``.
``None`` is returned when the path directory is incomplete (missing
``order.txt`` / ``traj.txt`` / a referenced trajectory frame), so it is
skipped rather than trusted.
"""
order_file = os.path.join(path_dir, 'order.txt')
traj_file = os.path.join(path_dir, 'traj.txt')
if not (os.path.isfile(order_file) and os.path.isfile(traj_file)):
return None
frames = np.unique(np.loadtxt(traj_file, usecols=[1], dtype=str))
for frame in np.atleast_1d(frames):
if not os.path.isfile(os.path.join(path_dir, 'accepted', str(frame))):
return None
order = np.loadtxt(order_file, usecols=[0, 1])
order = np.atleast_2d(order)
return float(order[1, 1]), float(np.max(order[:, 1]))
[docs]
def reselect_initial_paths(src_load_dir, active_pnums, interfaces, dest_dir):
"""Re-select active paths for a new interface set into ``dest_dir``.
Ported from the upstream ``initial_path_from_iretis`` (the
restart-active branch). The ``[0-]`` path is kept; the remaining active
paths are sorted by their maximum order parameter and assigned to
increasing positive ensembles (a path is valid in the highest ensemble
whose interface it crosses). Any ensemble still empty is filled by
cloning a path from a higher ensemble. Selected paths are copied into
``dest_dir`` as ``0 .. len(interfaces)-1``.
Parameters
----------
src_load_dir : str
The directory holding the previous iteration's paths (one numbered
sub-directory per path).
active_pnums : list of int
The active path numbers from the previous run's ``output.toml``
(``[current].active``), one per old ensemble.
interfaces : list of float
The *new* interface positions.
dest_dir : str
The directory to create and fill with the re-selected paths.
Returns
-------
chosen : dict
Mapping ``ensemble index -> source path directory``.
"""
n_ens = len(interfaces)
# Gather every still-valid path surviving in the previous load dir (the
# active paths plus any others delete_old kept), so the re-selection has
# as many distinct candidates as possible.
candidates = {}
for name in sorted(os.listdir(src_load_dir)):
path_dir = os.path.join(src_load_dir, name)
if not os.path.isdir(path_dir):
continue
order = _path_order(path_dir)
if order is not None:
candidates[path_dir] = order
chosen = {}
# The [0-] ensemble (interfaces[0] does not change): prefer the previous
# active [0-], else any path that starts below interfaces[0].
minus_dir = os.path.join(src_load_dir, str(active_pnums[0]))
if minus_dir in candidates:
chosen[0] = minus_dir
else:
for path_dir, (op_one, _omax) in candidates.items():
if op_one < interfaces[0]:
chosen[0] = path_dir
break
# Positive ensembles: assign distinct paths, highest max order parameter
# first, each into the highest still-empty ensemble it is valid in (it
# crosses that ensemble's lower interface).
positives = sorted(
((path_dir, omax) for path_dir, (_op1, omax) in candidates.items()
if omax > interfaces[0]),
key=lambda item: item[1], reverse=True)
for path_dir, omax in positives:
valid_in = 0
for i, interface in enumerate(interfaces[:-1]):
if omax > interface:
valid_in = i + 1
ens = valid_in
while ens >= 1:
if ens not in chosen:
chosen[ens] = path_dir
break
ens -= 1
# Clone to fill any positive ensemble still empty: a path from the
# nearest filled higher ensemble crosses this ensemble's interface too;
# for the very top, the highest-max-op path does (the placement caps the
# top interface below it).
for i in range(1, n_ens):
if i in chosen:
continue
donor = None
for j in range(i + 1, n_ens):
if j in chosen:
donor = chosen[j]
break
if donor is None and positives:
donor = positives[0][0]
if donor is not None and candidates[donor][1] > interfaces[i - 1]:
chosen[i] = donor
missing = [i for i in range(n_ens) if i not in chosen]
if missing:
highest = max((o[1] for o in candidates.values()),
default=float('nan'))
raise ValueError(
f"No valid path for ensemble(s) {missing} with the new "
f"interfaces {interfaces} (highest sampled max order "
f"parameter {highest}).")
os.mkdir(dest_dir)
for ens in range(n_ens):
shutil.copytree(chosen[ens], os.path.join(dest_dir, str(ens)))
return chosen
def _run_scheduler(toml_path):
"""Run the infinite-swapping scheduler for one config file."""
# Imported lazily so importing this module does not pull the whole
# simulation stack (and to avoid a circular import via the analysis
# package).
# pylint: disable=import-outside-toplevel
from pyretis.simulation.scheduler import scheduler
from pyretis.simulation.setup import setup_config
config = setup_config(toml_path)
if config is not None:
scheduler(config)
def _prepare_step(step_dir, src_load, aux_files, step_config):
"""Create a step directory and stage its inputs.
Copies ``src_load`` into ``<step_dir>/load`` and the auxiliary files
(order-parameter module, ...) alongside it, and writes the fresh
per-iteration config. Returns the step's ``load`` directory.
"""
os.makedirs(step_dir, exist_ok=True)
step_load = os.path.join(step_dir, 'load')
if not os.path.isdir(step_load):
shutil.copytree(src_load, step_load)
for aux in aux_files:
dest = os.path.join(step_dir, os.path.basename(aux))
if os.path.isfile(aux) and not os.path.isfile(dest):
shutil.copy(aux, dest)
_write_toml(step_config, os.path.join(step_dir, 'output.toml'))
return step_load
def _step_config(config, interfaces, shooting_moves, n_steps):
"""Build a fresh per-iteration config (no ``[current]`` -> fresh run)."""
step_config = dict(config)
step_config['simulation'] = dict(config['simulation'])
step_config['simulation']['interfaces'] = interfaces
step_config['simulation']['shooting_moves'] = (
shooting_moves or ['sh', 'sh'] + ['wf'] * (len(interfaces) - 2))
step_config['simulation']['steps'] = int(n_steps)
step_config['simulation']['load_dir'] = 'load'
step_config.pop('current', None)
return step_config
def _run_in_dir(step_dir):
"""Run the scheduler with the working directory set to ``step_dir``."""
cwd = os.getcwd()
try:
os.chdir(step_dir)
_run_scheduler('output.toml')
finally:
os.chdir(cwd)
[docs]
def run_infinit(toml='infswap.toml', workdir='.', log_file='infinit.log'):
"""Run the iterative infinit interface-placement driver.
Parameters
----------
toml : str
The starting infinite-swapping config. Must contain a working
``[simulation] load_dir`` for its interfaces and an ``[infinit]``
section (see :func:`set_default_infinit`).
workdir : str
Directory in which the per-iteration ``step_<i>`` sub-directories
are created (default: the current directory).
log_file : str
File (under ``workdir``) for the per-iteration interface log.
Returns
-------
interfaces : list of float
The interface set after the final iteration.
"""
config = _read_toml(toml)
iset = set_default_infinit(config)
interfaces = list(config['simulation']['interfaces'])
shooting_moves = list(config['simulation'].get('shooting_moves', []))
workers = config.get('runner', {}).get('workers', 1)
lamres = iset['lamres']
skip = iset['skip']
pl_target = iset['pL']
interface_cap = config['simulation'].get('tis_set', {}).get(
'interface_cap')
toml_dir = os.path.dirname(os.path.abspath(toml))
src_load = os.path.abspath(
os.path.join(toml_dir, config['simulation'].get('load_dir', 'load')))
# Auxiliary files the workers need in their run directory (e.g. a custom
# order-parameter module), resolved relative to the starting config.
aux_files = [os.path.join(toml_dir, f)
for f in config.get('runner', {}).get('files', [])]
log_path = os.path.join(workdir, log_file)
for step, n_steps in enumerate(iset['steps_per_iter'][iset['cstep']:]):
step_dir = os.path.join(workdir, f'step_{step}')
step_load = _prepare_step(
step_dir, src_load, aux_files,
_step_config(config, interfaces, shooting_moves, n_steps))
_run_in_dir(step_dir)
# Re-estimate the interfaces from this iteration's data. step_dir
# is a run directory: interfaces_from_data auto-detects a literal
# infswap_data.txt if one exists there, else reconstructs the same
# matrix from the native per-ensemble output (the only output the
# scheduler writes now -- see get_path_data_matrix).
restart = _read_toml(os.path.join(step_dir, 'output.toml'))
result = interfaces_from_data(
step_dir, interfaces,
nskip=int(restart['current']['cstep'] * skip),
pl_target=pl_target, n_workers=workers, lamres=lamres,
interface_cap=interface_cap)
with open(log_path, 'a', encoding='utf-8') as logf:
if result is None:
logf.write(f"step {step}: Pcross not constructible, "
f"keeping interfaces {interfaces}\n")
src_load = step_load
continue
interfaces = result['interfaces']
shooting_moves = result['shooting_moves']
logf.write(f"step {step}: interfaces = {interfaces}\n")
# Re-select the active paths for the next iteration's interfaces.
next_load = os.path.join(step_dir, 'load_next')
reselect_initial_paths(step_load, restart['current']['active'],
interfaces, next_load)
src_load = next_load
return interfaces