Source code for pyretis.core.shooting_selection

# Copyright (c) 2026, PyRETIS Development Team.
# Distributed under the LGPLv2.1+ License. See LICENSE for more info.
"""Strategy objects for biased shooting-point selection.

This module defines a small, dependency-free strategy abstraction used by
the *biased* shooting move (``shooting_move = "bias"``). A selector assigns
a strictly positive weight ``w(x)`` to every *interior* frame of a path
(the two end points are excluded, exactly as the uniform ``rnd`` move
excludes them). The shooting point is then drawn with probability
``w_i / W`` where ``W = sum_j w_j`` is the selection partition function,
and detailed balance is restored by the Metropolis acceptance factor
``min(1, W_old / W_new)`` applied in :py:func:`pyretis.core.tis.shoot_bias`.

The weight ``w_i`` of the chosen frame cancels in the acceptance ratio
because the shooting point is the *same physical frame* shared by the old
and the new path; only the two partition functions ``W_old`` and ``W_new``
survive. See ``docs/decisions/biased_shooting_acceptance.md`` and the user
documentation (``docs/user/section/tis.rst``) for the full derivation.

The module is intentionally free of any machine-learning or heavy
dependency: it only uses :py:mod:`numpy`. A learned (committor) selector
can later be added as one more :class:`ShootingPointSelector` subclass
whose :py:meth:`~ShootingPointSelector.weights` queries a model.

Important
---------
* Every random draw must use the **per-ensemble** random generator that is
  passed in to :py:meth:`~ShootingPointSelector.select`, never a path's own
  ``rgen`` (which can be aliased across ensembles and would break
  continue-vs-restart equivalence). :py:meth:`select` draws exactly one
  uniform number, so the random stream advances deterministically.
* Weights must be strictly positive and finite. Non-positive or non-finite
  weights raise instead of being silently clamped -- a wrong selection
  weight silently biases the sampled rate, which is the worst possible
  failure mode for this code.
"""
import logging
import numpy as np
from pyretis.core.common import generic_factory

logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
logger.addHandler(logging.NullHandler())


__all__ = [
    'ShootingPointSelector',
    'UniformSelector',
    'GaussianOrderSelector',
    'shooting_selector_factory',
    'selection_bias_accept',
]


def draw_uniform(rgen):
    """Draw a single uniform float in ``[0, 1)``, generator-agnostic.

    The biased move is used from both the pyretis-native flow (whose
    :class:`.RandomGenerator`/:class:`.PCG64Generator` expose ``rand``,
    returning a length-1 array) and the infinite-swapping flow (which
    passes a raw :class:`numpy.random.Generator`, whose ``random`` returns
    a scalar). Mirror the ``_draw_int`` detection used by
    :py:meth:`pyretis.core.path.Path.get_shooting_point` so the same code
    draws the byte-identical value through either generator.

    Parameters
    ----------
    rgen : object
        A pyretis-native random generator or a ``numpy.random.Generator``.

    Returns
    -------
    out : float
        A pseudo-random float in ``[0, 1)``.

    """
    if hasattr(rgen, 'rand'):
        return float(rgen.rand()[0])
    return float(rgen.random())


[docs] def selection_bias_accept(rgen, sum_weight_old, sum_weight_new): """Apply the selection-bias Metropolis test ``min(1, W_o/W_n)``. This is the sole acceptance correction for the aimless biased shooting move (the dynamics and the shooting-point weight cancel; see :py:func:`pyretis.core.tis.shoot_bias`). It is a tiny, side-effect-free function -- shared by the native and infinite-swapping flows -- so the exact arithmetic, in particular the orientation ``W_old / W_new`` and **not** its inverse, can be unit-tested without an engine. Parameters ---------- rgen : object The per-ensemble random generator; one uniform number is drawn (via :py:func:`draw_uniform`, so either generator flavour works). sum_weight_old : float ``W_o``, the selection partition function of the current path. sum_weight_new : float ``W_n``, the selection partition function of the trial path. Must be strictly positive. Returns ------- out : bool True if the trial path is accepted by the selection-bias test. """ if not sum_weight_new > 0.0: raise ValueError( 'Selection partition function of the trial path must be ' f'strictly positive, got W_new = {sum_weight_new}.' ) accept_ratio = min(1.0, sum_weight_old / sum_weight_new) return bool(draw_uniform(rgen) < accept_ratio)
def _interior_order_values(path, index=0): """Return the order-parameter values of the interior frames. Parameters ---------- path : object like :py:class:`.PathBase` The path to read the order parameter from. The first and last phase points are excluded (they are never shooting points). index : int, optional Which component of the per-frame order-parameter list to use. ``0`` is the progress coordinate; larger indices select extra collective variables. Returns ------- out : numpy.ndarray A 1-D array of length ``path.length - 2`` with the requested order-parameter component for every interior frame. """ interior = path.phasepoints[1:-1] values = np.array([point.order[index] for point in interior], dtype=np.float64) return values
[docs] class ShootingPointSelector: """Strategy for choosing a shooting point and its acceptance weight. A selector defines a strictly positive per-frame weight ``w(x)`` over a path's *interior* frames (the end points are excluded). The shooting point is drawn with probability ``w_i / W`` (``W = sum_j w_j``) and the biased shooting move is made detailed-balance correct by the Metropolis factor ``min(1, W_old / W_new)``. Subclasses only need to implement :py:meth:`weights`; the selection and the partition-function bookkeeping are provided here so every selector shares the exact same, reproducible draw. """
[docs] def weights(self, path): """Return the interior-frame selection weights. Parameters ---------- path : object like :py:class:`.PathBase` The path to weight. Only the interior frames ``path.phasepoints[1:-1]`` are considered. Returns ------- out : numpy.ndarray A 1-D array of strictly positive, finite weights, one per interior frame (length ``path.length - 2``). """ raise NotImplementedError
[docs] def _validated_weights(self, path): """Return the interior weights after checking the contract. Raises ------ ValueError If the path has no interior frame, or if any weight is non-finite or not strictly positive. Failing loud here keeps a broken weight from silently biasing the sampled rate. """ if path.length < 3: raise ValueError( 'Cannot select a shooting point: a path needs at least ' '3 frames to have an interior, got length ' f'{path.length}.' ) weights = np.asarray(self.weights(path), dtype=np.float64) expected = path.length - 2 if weights.shape != (expected,): raise ValueError( f'{type(self).__name__}.weights returned shape ' f'{weights.shape}, expected ({expected},) (one weight per ' 'interior frame).' ) if not np.all(np.isfinite(weights)): raise ValueError( f'{type(self).__name__}.weights returned non-finite ' 'values; selection weights must be finite.' ) if not np.all(weights > 0.0): raise ValueError( f'{type(self).__name__}.weights returned non-positive ' 'values; selection weights must be strictly positive. ' '(For exp-type weights this usually means the bias is too ' 'sharp and some frames underflow to 0 -- reduce alpha.)' ) return weights
[docs] def sum_weight(self, path): """Return the selection partition function ``W`` of a path. Parameters ---------- path : object like :py:class:`.PathBase` The path whose interior weights are summed. Returns ------- out : float ``W = sum_j w_j`` over the interior frames. Strictly positive. """ return float(self._validated_weights(path).sum())
[docs] def select(self, path, rgen): """Draw an interior frame with probability proportional to ``w_i``. Exactly one uniform random number is drawn from ``rgen`` and mapped to a frame through the inverse cumulative weight distribution, so the draw is deterministic given the random stream and independent of the concrete generator implementation. Parameters ---------- path : object like :py:class:`.PathBase` The path to select a shooting point from. rgen : object like :py:class:`.RandomGenerator` The **per-ensemble** random generator. A path's own ``rgen`` must not be used (see the module docstring). Returns ------- out[0] : int The index of the selected frame in ``path.phasepoints``. It lies in ``[1, path.length - 2]`` (an interior frame). out[1] : float The weight ``w_i`` of the selected frame. out[2] : float The selection partition function ``W`` of ``path``. """ weights = self._validated_weights(path) total = float(weights.sum()) # Inverse-CDF sampling from a single uniform draw (generator-agnostic: # native ``rand`` or inf ``random``). ``searchsorted`` on the # cumulative weights turns it into a weighted choice with no extra # draws. uniform = draw_uniform(rgen) cumulative = np.cumsum(weights) target = uniform * total interior_idx = int(np.searchsorted(cumulative, target, side='right')) # Guard the (measure-zero) edge where ``uniform`` rounds to exactly 1. interior_idx = min(interior_idx, weights.size - 1) idx = interior_idx + 1 # Shift past the excluded first frame. logger.debug('Biased shooting point at index %d (weight %g, W %g)', idx, weights[interior_idx], total) return idx, float(weights[interior_idx]), total
[docs] class UniformSelector(ShootingPointSelector): """Selector with constant weights (every interior frame equally likely). With ``w == 1`` the partition function is ``W = L - 2`` (the number of interior frames) and the biased-move acceptance reduces to the standard flexible-length two-way-shooting factor ``min(1, (L_o - 2)/(L_n - 2))``. This selector is the correctness anchor of the biased move: running ``shooting_move = "bias"`` with :class:`UniformSelector` must reproduce the statistics of the native uniform ``rnd`` move. It is **not** wired as the default shooting move -- the uniform ``rnd`` path is left untouched. """
[docs] def weights(self, path): """Return an array of ones, one per interior frame.""" return np.ones(path.length - 2, dtype=np.float64)
[docs] class GaussianOrderSelector(ShootingPointSelector): r"""Gaussian bias toward a target order-parameter value. The interior-frame weights are .. math:: w(x) = \\exp\\!\\big(-\\alpha\\,(\\lambda(x) - \\lambda_0)^2\\big), where :math:`\\lambda(x)` is the ``index``-th order-parameter component of a frame. This concentrates the shooting points near :math:`\\lambda_0` (for example a barrier-top / transition-state value), which is where re-shooting is most informative. It is a cheap, analytic, machine-learning-free bias and the reference selector for the statistical-equivalence test of the biased move. Parameters ---------- alpha : float Inverse-width of the Gaussian; must be strictly positive. Larger values concentrate the selection more tightly around ``lambda0``. Very large values can make distant frames underflow to zero weight, which :py:meth:`select` reports as an error -- reduce ``alpha`` then. lambda0 : float The order-parameter value the selection is biased toward. index : int, optional Which order-parameter component to read (default ``0``, the progress coordinate). """
[docs] def __init__(self, alpha, lambda0, index=0): """Initialise the Gaussian selector and validate its parameters.""" alpha = float(alpha) if alpha <= 0.0 or not np.isfinite(alpha): raise ValueError( f'GaussianOrderSelector needs a finite alpha > 0, got ' f'{alpha}.' ) if index < 0: raise ValueError( f'GaussianOrderSelector needs index >= 0, got {index}.' ) self.alpha = alpha self.lambda0 = float(lambda0) self.index = int(index)
[docs] def weights(self, path): """Return ``exp(-alpha (lambda - lambda0)**2)`` per interior frame.""" order = _interior_order_values(path, index=self.index) return np.exp(-self.alpha * (order - self.lambda0)**2)
[docs] def __str__(self): """Return a short human-readable description.""" return (f'GaussianOrderSelector(alpha={self.alpha}, ' f'lambda0={self.lambda0}, index={self.index})')
[docs] def shooting_selector_factory(settings): """Create a built-in shooting-point selector from settings. Mirrors :py:func:`pyretis.orderparameter.order_factory`: the ``class`` key selects one of the built-in selectors and the remaining keys are passed to its constructor. User-defined selectors are loaded from a ``module`` instead (handled by :py:func:`pyretis.setup.common.create_shooting_selector`). Parameters ---------- settings : dict The ``[shooting-selector]`` settings, including the ``class`` key and any constructor arguments (e.g. ``alpha``, ``lambda0``). Returns ------- out : object like :py:class:`ShootingPointSelector` The created selector, or ``None`` if the class is unknown. """ factory_map = { 'uniform': {'cls': UniformSelector}, 'uniformselector': {'cls': UniformSelector}, 'gaussianorderselector': {'cls': GaussianOrderSelector}, 'gaussian': {'cls': GaussianOrderSelector}, } return generic_factory(settings, factory_map, name='shooting-selector')