# -*- coding: utf-8 -*-
# Copyright (c) 2023, PyRETIS Development Team.
# Distributed under the LGPLv2.1+ License. See LICENSE for more info.
"""Unbiased per-path weights and the transmission coefficient.
This module ports two infinite-swapping analysis tools from the upstream
``inftools`` package (``tistools/path_weights.py`` and
``tistools/calc_transmission.py``):
* :func:`get_path_weights` -- the WHAM-unbiased equilibrium weight of every
positive-ensemble path, so that an observable can be averaged as
``<O> = sum_i w_i O_i`` (predictive-power analysis, binless crossing
probability, ...).
* :func:`transmission_coefficient` -- the conditional transmission
coefficient for a proposed transition-state value along a chosen
collective variable, defined as in Figure 8 of
`J. Chem. Theory Comput. (doi:10.1021/acs.jctc.5c01814)
<https://doi.org/10.1021/acs.jctc.5c01814>`_.
The path weights reuse the same WHAM unweighting already implemented in
:mod:`pyretis.analysis.wham_analysis`; the data are the
``infswap_data.txt``-shaped matrix returned by
:func:`pyretis.analysis.wham_analysis.get_path_data_matrix` (a literal
file if one still exists, otherwise reconstructed from the native
per-ensemble output the scheduler writes now).
"""
import numpy as np
__all__ = ['get_path_weights', 'transmission_coefficient']
[docs]
def get_path_weights(matrix, interfaces):
"""Compute the WHAM-unbiased weight of every positive-ensemble path.
The ``[0-]`` ensemble (the first ``Cxy``/HA-weight column) is excluded;
weights are computed for the paths sampled in the positive ensembles.
Parameters
----------
matrix : list of list of float
The ``infswap_data.txt`` matrix as returned by
:func:`pyretis.analysis.wham_analysis.read_data_matrix`. Each row is
``[path_nr, length, max_op, Cxy_0 ... Cxy_{n-1},
HA_0 ... HA_{n-1}]`` with ``"----"`` already mapped to ``0.0``.
interfaces : list of float
The interface positions (``n`` interfaces ``=> n`` ensemble
columns, ensemble 0 being ``[0-]``).
Returns
-------
path_nr : numpy.ndarray
The path index of every positive-ensemble path.
max_op : numpy.ndarray
The maximum order parameter of each such path.
weight : numpy.ndarray
The WHAM-unbiased equilibrium weight of each such path.
"""
nintf = len(interfaces)
data = np.asarray(matrix, dtype=float)
if data.size == 0:
return np.array([]), np.array([]), np.array([])
# Positive-ensemble paths are those NOT sampled in [0-]; in the data
# file their [0-] column (the first Cxy, index 3) is "----" -> 0.0.
plus = data[:, 3] == 0.0
data = data[plus]
if data.shape[0] == 0:
return np.array([]), np.array([]), np.array([])
path_nr = data[:, 0].astype(int)
max_op = data[:, 2].astype(float)
# Drop the [0-] ensemble (column index 3 / 3+nintf): the positive
# ensembles are columns 4 .. 3+nintf-1 (Cxy) and 4+nintf .. 3+2*nintf-1
# (high-acceptance weights).
path_f = data[:, 4:3 + nintf]
path_w = data[:, 4 + nintf:3 + 2 * nintf]
# Per-path WHAM unweighting Cxy / HA-weight, then rescale each ensemble
# column so its weights sum to the (fractional) number of samples.
with np.errstate(divide='ignore', invalid='ignore'):
weights = path_f / path_w
weights[np.isnan(weights)] = 0.0
col_sum = np.sum(weights, axis=0)
with np.errstate(divide='ignore', invalid='ignore'):
weights = weights / col_sum * np.sum(path_f, axis=0)
weights[np.isnan(weights)] = 0.0
wsum = np.sum(weights, axis=0)
# WHAM crossing-probability recursion over the positive interfaces.
ploc_wham = np.zeros(nintf)
ploc_wham[0] = 1.0
for i, intf_p1 in enumerate(interfaces[1:]):
crossed = max_op >= intf_p1
nj = wsum[:i + 1]
njl = np.sum(crossed[:, None] * weights[:, :i + 1], axis=0)
ploc_wham[i + 1] = np.sum(njl) / np.sum(nj / ploc_wham[:i + 1])
# The unbiased per-path weight A[j] = Q[K] * sum_e w[j, e].
q_factor = 1.0 / np.cumsum(wsum / ploc_wham[:-1])
weight = np.zeros(max_op.shape[0])
intf_arr = np.asarray(interfaces)
for j in range(max_op.shape[0]):
below = np.where(max_op[j] >= intf_arr)[0]
k_idx = min(below[-1], nintf - 2)
weight[j] = q_factor[k_idx] * np.sum(weights[j])
return path_nr, max_op, weight
[docs]
def transmission_coefficient(matrix, interfaces, order_files, ts, *, dim=1,
nskip=0):
"""Estimate the conditional transmission coefficient for a CV.
Counts, for every reactive-or-recrossing positive-ensemble path, how
many times the chosen collective variable ``dim`` crosses the proposed
transition-state value ``ts``, and forms the WHAM-weighted ratio of
reactive paths to total positive crossings.
Parameters
----------
matrix : list of list of float
The ``infswap_data.txt`` matrix (see :func:`get_path_weights`).
interfaces : list of float
The interface positions.
order_files : callable
``order_files(path_nr)`` must return the per-frame order-parameter
array of that path (columns ``[time, cv_0, cv_1, ...]``), or
``None`` if it is unavailable.
ts : float
The proposed transition-state value (along the original order
parameter; reactive paths are still defined by it).
dim : int, optional
The collective-variable column of the order array to count
crossings for (default 1, the first CV after the time column).
nskip : int, optional
Number of initial paths in ``matrix`` to discard.
Returns
-------
tcoeff : float
The conditional transmission coefficient.
"""
sub = matrix[nskip:]
path_nr, max_op, weight = get_path_weights(sub, interfaces)
if path_nr.size == 0:
return float('nan')
last_intf = interfaces[-1]
sel_w, sel_react, sel_cross = [], [], []
for nr, mop, wgt in zip(path_nr, max_op, weight):
# Only paths that reached the proposed transition state.
if mop < ts:
continue
order = order_files(int(nr))
if order is None:
continue
order = np.asarray(order)
crosses = _count_positive_crossings(order[:, dim], ts)
sel_w.append(wgt)
sel_react.append(1.0 if mop > last_intf else 0.0)
sel_cross.append(crosses)
sel_w = np.asarray(sel_w)
sel_react = np.asarray(sel_react)
sel_cross = np.asarray(sel_cross)
denom = np.sum(sel_w * sel_cross)
if denom == 0.0:
return float('nan')
return float(np.sum(sel_w * sel_react) / denom)
def _count_positive_crossings(values, ts):
"""Count upward crossings of ``ts`` along the value sequence."""
crossings = 0
for first, second in zip(values[:-1], values[1:]):
if second > ts > first:
crossings += 1
return crossings