# -*- coding: utf-8 -*-
# Copyright (c) 2023, PyRETIS Development Team.
# Distributed under the LGPLv2.1+ License. See LICENSE for more info.
"""Scientific core of the iterative ``infinit`` interface-placement driver.
Ported from the upstream ``inftools`` ``misc/infinit_helper.py``. The
``infinit`` tool repeatedly runs a short simulation, re-estimates the
crossing probability with WHAM, and re-places the interfaces so every
ensemble carries roughly the same local crossing probability -- converging
on a good interface set automatically.
This module ports the *pure, testable* heart of one such iteration:
* :func:`binless_pcross` -- build the binless WHAM crossing-probability
curve ``p(x)`` from the per-path weights returned by
:func:`pyretis.analysis.path_weights.get_path_weights`;
* :func:`estimate_updated_interfaces` -- given that curve and the current
interfaces, compute the next iteration's interfaces (and matching
shooting-move list), reproducing the upstream ``update_toml_interfaces``
logic: interface-cap filtering, the "not enough data" guard, geometric
placement via :func:`estimate_interface_positions`, rounding to the
lambda resolution, and the land-on-state-A/B fix-ups;
* :func:`interfaces_from_data` -- the convenience chain
data file -> weights -> Pcross -> updated interfaces.
The *outer* driver loop (running ``pyretisrun``, threading the
``[infinit]`` ``cstep`` state through ``restart.toml``, and accumulating
``combo`` data files across iterations) is orchestration and is **not**
included here; it can be built on top of these functions plus
:func:`pyretis.analysis.combine_data.combine_data`.
"""
import numpy as np
from pyretis.analysis.path_weights import get_path_weights
from pyretis.analysis.interface_placement import estimate_interface_positions
from pyretis.analysis.wham_analysis import get_path_data_matrix
__all__ = ['binless_pcross', 'estimate_updated_interfaces',
'interfaces_from_data']
[docs]
def binless_pcross(max_op, weight, first_interface):
"""Build the binless crossing-probability curve from path weights.
A path that reaches order parameter ``max_op`` with unbiased weight
``weight`` contributes to the crossing probability at every order
parameter up to ``max_op``. Sorting the paths by ``max_op`` and
accumulating their weights from the top down therefore gives the
(unnormalised) probability to cross each order-parameter value; dividing
by the total weight makes it a probability that starts at ``1`` at the
first interface.
Parameters
----------
max_op : array_like
The maximum order parameter of every positive-ensemble path.
weight : array_like
The matching WHAM-unbiased path weight (see
:func:`pyretis.analysis.path_weights.get_path_weights`).
first_interface : float
The first interface (state-A) order parameter; the curve is
anchored at ``(first_interface, 1.0)``.
Returns
-------
x : numpy.ndarray
The order-parameter values (no duplicates, increasing).
p : numpy.ndarray
The crossing probability at each ``x`` (decreasing, ``p[0] == 1``).
"""
max_op = np.asarray(max_op, dtype=float)
weight = np.asarray(weight, dtype=float)
res_x = [float(first_interface)]
res_y = [1.0]
if max_op.size == 0:
return np.array(res_x), np.array(res_y)
order = np.argsort(max_op)
mo_sorted = max_op[order]
w_sorted = weight[order]
sumw = float(np.sum(w_sorted))
if sumw <= 0.0:
return np.array(res_x), np.array(res_y)
for i, moi in enumerate(mo_sorted):
# A path lying exactly on an order-parameter value already recorded
# is folded into that point (no duplicate x).
if res_x[-1] == moi:
continue
res_y.append(float(np.sum(w_sorted[i:]) / sumw))
res_x.append(float(moi))
return np.array(res_x), np.array(res_y)
def _round_down_to_lamres(values, lamres):
"""Round each value *down* to the nearest multiple of ``lamres``.
Rounding down (never up) guarantees an estimated interface never lands
above the highest order parameter actually sampled by the highest path,
which would leave that ensemble without a valid crossing path.
"""
scaled = np.round(np.asarray(values, dtype=float) / lamres, decimals=10)
return np.round(np.floor(scaled) * lamres, decimals=10)
def _shift_off_endpoint(interior, endpoint, lamres, direction):
"""Move an interior interface off a state boundary it landed on.
If an interior interface coincides with ``endpoint`` (state A or B),
shift it by one ``lamres`` in ``direction`` (``+1`` away from A, ``-1``
away from B); if that slot is taken, drop the interface instead.
"""
if endpoint not in interior:
return interior
idx = interior.index(endpoint)
shifted = round(interior[idx] + direction * lamres, 10)
if shifted not in interior:
interior[idx] = shifted
else:
interior.pop(idx)
return interior
def _finalize_interfaces(estimated, x_last, state_a, state_b, lamres,
n_workers):
"""Round, dedupe and boundary-fix an estimated interface set.
Appends state B, keeps the second-to-last interface below the highest
sampled order parameter, rounds the interior down to the lambda grid
(restoring the unrounded set if that would drop below ``n_workers``
ensembles), and shifts any interface that landed on state A/B.
"""
intf = [float(v) for v in estimated] + [state_b]
# Only clamp a genuine interior interface below x_last. With no interior
# interface (intf == [state_a, state_b]), intf[-2] IS the fixed state-A
# interface and clamping it would silently move state A.
if len(intf) > 2 and intf[-2] + lamres >= x_last:
intf[-2] = float(x_last - lamres)
interior = list(np.unique(_round_down_to_lamres(intf[1:-1], lamres)))
interior = [float(v) for v in interior]
if len(interior) + 1 < n_workers:
interior = [float(v) for v in intf[1:-1]]
interior = _shift_off_endpoint(interior, round(state_a, 10), lamres, +1)
interior = _shift_off_endpoint(interior, round(state_b, 10), lamres, -1)
new_interfaces = [intf[0]] + interior + [intf[-1]]
# One shooting move per ensemble: [0-] and [0+] shoot ("sh"), every
# positive ensemble wire-fences ("wf"). Length matches the final
# interface list (the upstream used the pre-dedup length, which could
# mismatch after rounding removed an interface).
shooting_moves = ['sh', 'sh'] + ['wf'] * (len(new_interfaces) - 2)
return new_interfaces, shooting_moves
[docs]
def estimate_updated_interfaces(x, p, interfaces, *, pl_target=0.3,
num_ens=None, n_workers=1, lamres=0.001,
interface_cap=None):
"""Compute the next iteration's interfaces from a Pcross curve.
Reproduces the order-parameter mathematics of the upstream
``update_toml_interfaces`` (without the ``restart.toml`` / ``combo``
bookkeeping): filter at the interface cap, bail out if the crossing
probability could not be constructed, place interfaces at geometric
local-Pcross spacing, round down to the lambda resolution, and nudge
any interface that landed on a state boundary.
Parameters
----------
x, p : array_like
The binless crossing-probability curve (see :func:`binless_pcross`).
interfaces : list of float
The *current* interface positions (only the first and last,
state A and state B, are used here).
pl_target : float, optional
Target per-ensemble local crossing probability (default 0.3). The
effective target is ``max(pl_target, p[-1] ** (1 / (2 * n_workers)))``
so the ensemble count cannot blow up past about twice the worker
count.
num_ens : int, optional
If given, place exactly this many ensembles (overrides
``pl_target``).
n_workers : int, optional
The number of simulation workers; the per-ensemble target and a
lower bound on the ensemble count both depend on it.
lamres : float, optional
The lambda resolution; interfaces are rounded down to a multiple of
it (default 0.001).
interface_cap : float, optional
Do not place interfaces above this order parameter (default: the
last/state-B interface).
Returns
-------
result : dict or None
``None`` if the crossing probability could not be constructed (the
caller should keep the current interfaces). Otherwise a dict with
``interfaces`` (the new list), ``shooting_moves`` (``"sh","sh"``
then ``"wf"`` per positive ensemble) and ``pl_used`` (the actual
per-ensemble local crossing probability).
"""
x = np.asarray(x, dtype=float)
p = np.asarray(p, dtype=float)
state_a = float(interfaces[0])
state_b = float(interfaces[-1])
cap = state_b if interface_cap is None else float(interface_cap)
# Don't place interfaces above the cap / last interface.
above = np.where(x > cap)[0]
if len(above) > 0:
x = x[:above[0]]
p = p[:above[0]]
# With too little data the Pcross can contain leading/contiguous zeros;
# if it is all zero (or only the anchor is non-zero) keep the interfaces.
nonzero = np.where(p > 0)[0]
if len(nonzero) == 0 or nonzero[-1] == 0:
return None
p = p[:nonzero[-1] + 1]
x = x[:nonzero[-1] + 1]
ptot = p[-1]
if num_ens:
estimated, pl_used = estimate_interface_positions(
x, p, num_ens=num_ens)
else:
pl_eff = max(pl_target, ptot ** (1.0 / (2 * n_workers)))
estimated, pl_used = estimate_interface_positions(
x, p, pl_target=pl_eff)
new_interfaces, shooting_moves = _finalize_interfaces(
estimated, x[-1], state_a, state_b, lamres, n_workers)
return {
'interfaces': new_interfaces,
'shooting_moves': shooting_moves,
'pl_used': float(pl_used),
}
[docs]
def interfaces_from_data(run_dir, interfaces, nskip=0, **kwargs):
"""Estimate updated interfaces directly from a run's path data.
Convenience chain: read the data matrix (a literal ``infswap_data.txt``
if one exists in ``run_dir``, else reconstructed from native
per-ensemble output -- see
:func:`pyretis.analysis.wham_analysis.get_path_data_matrix`), compute
the WHAM-unbiased path weights, build the binless crossing
probability, and re-estimate the interfaces. Keyword arguments are
forwarded to :func:`estimate_updated_interfaces`.
Parameters
----------
run_dir : str
The run directory (see :func:`pyretis.analysis.wham_analysis.
get_path_data_matrix`).
interfaces : list of float
The current interface positions.
nskip : int, optional
Number of initial data rows to discard.
**kwargs
Forwarded to :func:`estimate_updated_interfaces`
(``pl_target``, ``num_ens``, ``n_workers``, ``lamres``,
``interface_cap``).
Returns
-------
result : dict or None
As returned by :func:`estimate_updated_interfaces` (``None`` if the
crossing probability could not be constructed).
"""
matrix = get_path_data_matrix(run_dir, nskip=nskip)
path_nr, max_op, weight = get_path_weights(matrix, interfaces)
if path_nr.size == 0:
return None
x, p = binless_pcross(max_op, weight, interfaces[0])
return estimate_updated_interfaces(x, p, interfaces, **kwargs)