"""Native per-ensemble output for the infinite-swap coordinator.
The coordinator internally shares each accepted path across ensembles
via a ``frac`` vector instead of keeping a distinct scalar-weight path
per ensemble. ``pyretisanalyse`` and the downstream crossing-probability
/ rate analysis expect the per-ensemble ``pathensemble.txt`` format --
this is the **only** output format the coordinator writes, regardless of
which TOML dialect (native or infretis-flavour) launched the run; the
former alternative inf-format ``infswap_data.txt`` is no longer written
(see ``MERGE_TODO.md``; WHAM-style analysis instead reconstructs the
same matrix on demand, :func:`reconstruct_path_data_matrix` below).
This module bridges the two by emitting, for every ensemble ``j`` the
accepted path contributed to in the current cycle, one native
``pathensemble.txt`` row -- reusing the canonical
:py:class:`pyretis.inout.formats.pathensemble.PathEnsembleFormatter` so
the byte layout matches the native flow exactly. Alongside each row,
one native ``order.txt`` and ``energy.txt`` block (the ``# Cycle:``
block format of
:py:class:`pyretis.inout.formats.order.OrderPathFormatter` /
:py:class:`pyretis.inout.formats.energy.EnergyPathFormatter`) is
appended in the same ensemble directory, because the native analysis
pairs those blocks with the ``pathensemble.txt`` rows one-to-one in row
order (matched on the cycle number).
The frac -> Weight mapping
--------------------------
The native crossing-probability analysis
(:py:func:`pyretis.analysis.path_analysis.analyse_path_ensemble`) reads
the ``Weight`` column as a *high-acceptance* weight: each row's
contribution to the order-parameter histogram is
``mc_weight * (1.0 / Weight)``
where ``mc_weight`` is the Monte-Carlo stationary count (1 for an
accepted row, incremented on each following rejected row). The
coordinator emits one **accepted** row per ensemble per cycle, so
``mc_weight == 1`` for every row and the contribution reduces to
``1.0 / Weight``. To make the native histogram reproduce the
infinite-swap frac-weighted histogram -- where a path contributes its
fractional occupancy ``frac[j]`` to ensemble ``j`` -- the coordinator
therefore writes
``Weight = 1.0 / frac[j]`` (the per-cycle frac increment)
so that ``1.0 / Weight == frac[j]``. The ``Step`` column carries the
global coordinator cycle. This is the empirically-validated realisation
of the ratified ``frac -> Weight`` cycle decision.
The ``ha_weight.txt`` sidecar
------------------------------
``Weight`` actually encodes ``ha_factor / frac_inc`` (see
:func:`write_native_pathensemble_data`): for wire-fencing/stone-skipping
moves ``ha_factor`` is the path's high-acceptance crossing-count factor
for this ensemble (1.0 for the indicator moves sh/wt and the minus
ensemble), and ``frac_inc`` is this cycle's fractional-occupancy
increment. The combined ``Weight`` column is exactly what the *native*
crossing-probability histogram needs, but it is **lossy** for WHAM-style
analysis (:py:mod:`pyretis.analysis.wham_analysis`,
:py:mod:`pyretis.analysis.path_weights`): those need ``frac_inc`` and
``ha_factor`` as two independently-meaningful numbers (e.g. each
ensemble's WHAM q-factor weight is the *raw*, un-unweighted sum of
``frac_inc``, not the sum of ``frac_inc / ha_factor``), and ``ha_factor``
cannot be recovered after the fact from the path's move code or order
parameter alone. Separately, the canonical ``PathEnsembleFormatter``
columns carry no path-identity column at all (the native serial loop
never needed one -- each ensemble owns its own independent path
numbering), so there is no way to tell which accepted path a row belongs
to either, which a reconstruction needs to group a path's rows across
cycles/ensembles.
So ``write_native_pathensemble_data`` also appends one ``(path number,
ha_factor)`` pair per ``pathensemble.txt`` row, in the same row order, to
``ha_weight.txt`` -- the two pieces of information that genuinely have no
other home in the native file set (both already live in memory at write
time -- the coordinator's per-cycle ``contributions`` loop variable and
the path's ``ha_weights`` vector), kept deliberately minimal (two numbers
per row) rather than re-introducing a second copy of the whole path-data
matrix.
"""
from __future__ import annotations
import logging
import math
import os
from typing import Any, Dict, List, Sequence
from pyretis.core.pathensemble import generate_ensemble_name
from pyretis.inout.formats.energy import EnergyPathFormatter
from pyretis.inout.formats.order import OrderPathFormatter
from pyretis.inout.formats.pathensemble import PathEnsembleFormatter
logger = logging.getLogger("main")
logger.addHandler(logging.NullHandler())
__all__ = [
"order_output_requested",
"energy_output_requested",
"init_native_pathensemble_files",
"write_native_pathensemble_data",
"map_move_to_native_mc",
"reconstruct_path_data_matrix",
]
# The native per-ensemble path-data file names.
PATHENSEMBLE_FILE_NAME = "pathensemble.txt"
ORDER_FILE_NAME = "order.txt"
ENERGY_FILE_NAME = "energy.txt"
HA_WEIGHT_FILE_NAME = "ha_weight.txt"
# ha_weight.txt is a flat, one-line-per-pathensemble.txt-row sidecar (unlike
# order.txt/energy.txt, which use per-cycle "# Cycle:" blocks), so it gets a
# single file-level header instead. It also carries the path number: the
# canonical PathEnsembleFormatter columns have no path-identity column at
# all (the native serial loop never needed one -- each ensemble owns its
# own independent path numbering), but path identity is exactly what is
# needed to group a path's rows across cycles/ensembles when reconstructing
# the infswap_data.txt matrix, and -- like ha_factor -- it is available
# in-memory at write time (the coordinator's global path_number) but has
# nowhere else to be persisted.
HA_WEIGHT_HEADER = "# Step PathNumber HA-weight"
HA_WEIGHT_FMT = "{:>10d} {:>10d} {:>22.15e}"
# The coordinator's shooting-move codes mapped onto the native ``Mc``
# column codes. Most codes are shared verbatim between the two flows
# (``sh`` / ``ss`` / ``wt`` / ``wf`` / ``tr``); only a few coordinator
# tags have a distinct native spelling. Unrecognised codes are passed
# through unchanged (truncated to the native two-character column).
_MOVE_TO_NATIVE_MC = {
"sh": "sh",
"ss": "ss",
"wt": "wt",
"wf": "wf",
"tr": "tr",
"ct": "00",
"ld": "ld",
"00": "00",
"s+": "s+",
"s-": "s-",
"ki": "ki",
}
[docs]
def map_move_to_native_mc(move: str) -> str:
"""Map a coordinator move tag to the native ``Mc`` column code.
Parameters
----------
move : string
The coordinator's move tag (``generated[0]``), e.g. ``"sh"`` or
``"tr"``.
Returns
-------
string
The native two-character ``Mc`` code.
"""
code = _MOVE_TO_NATIVE_MC.get(str(move), str(move))
return code[:2]
[docs]
def order_output_requested(config: Dict[str, Any]) -> bool:
"""Return True when per-ensemble ``order.txt`` output is requested.
The input shim carries the native ``[output] order-file`` interval
as ``order_file`` (default 1). The native analysis pairs **one**
``order.txt`` block with **each** ``pathensemble.txt`` row (see
:py:func:`pyretis.analysis.path_analysis.analyse_path_ensemble`),
and the coordinator writes a row every cycle a path contributes
frac -- so the interval reduces to an on (``> 0``) / off (``<= 0``)
switch here.
Parameters
----------
config : dict
The coordinator configuration dictionary.
Returns
-------
boolean
True when per-ensemble ``order.txt`` output is requested.
"""
return int(config.get("output", {}).get("order_file", 1)) > 0
[docs]
def energy_output_requested(config: Dict[str, Any]) -> bool:
"""Return True when per-ensemble ``energy.txt`` output is requested.
See :py:func:`order_output_requested`; this is the ``energy-file``
counterpart (carried as ``energy_file``, default 1).
Parameters
----------
config : dict
The coordinator configuration dictionary.
Returns
-------
boolean
True when per-ensemble ``energy.txt`` output is requested.
"""
return int(config.get("output", {}).get("energy_file", 1)) > 0
def _ensemble_dirs(state: Any) -> List[str]:
"""Return the per-ensemble native output directories.
The coordinator's ``frac`` vector has one entry per ensemble plus a
trailing sentinel column; the real ensembles are indices
``0 .. n - 2`` and map to native directories ``000``, ``001``, ...
(``000`` being the ``[0^-]`` minus ensemble).
Parameters
----------
state : object like :py:class:`pyretis.simulation.repex.InfSwapState`
The coordinator state.
Returns
-------
list of string
The per-ensemble directory paths (absolute).
"""
data_dir = state.config.get("output", {}).get("data_dir", "./")
dirs = []
if getattr(state, "single_tis", False):
# One [i^+] ensemble (state index 0); its native output directory
# is named by the ensemble_number (e.g. 001), not the internal
# index, so the native analysis/comparison finds it where it
# expects.
return [
os.path.join(
data_dir, generate_ensemble_name(state.single_tis_ens_num)
)
]
if getattr(state, "explore", False):
# Explore: N-1 positive ensembles, no [0^-]. The native output dirs
# are 001, 002, ... (the 1-based positive-ensemble numbers), matching
# the native explorer output and InfSwapState.initiate_ensembles.
return [
os.path.join(data_dir, generate_ensemble_name(j + 1))
for j in range(state.n)
]
if getattr(state, "pptis_no_minus", False):
# PPTIS without [0^-]: the native dirs are named by the coordinator's
# initiate_ensembles key (generate_ensemble_name(i)).
return [
os.path.join(data_dir, generate_ensemble_name(j))
for j in range(state.n)
]
for j in range(state.n - 1):
dirs.append(os.path.join(data_dir, generate_ensemble_name(j)))
return dirs
[docs]
def init_native_pathensemble_files(state: Any, overwrite: bool = True) -> None:
"""Create the per-ensemble directories and file headers.
Each ensemble directory gets a ``pathensemble.txt`` with the canonical
native header, plus empty ``order.txt`` / ``energy.txt`` files when
those are requested (the path-data files carry no file-level header in
the native flow -- each appended block starts with its own ``# Cycle:``
comment), and an ``ha_weight.txt`` sidecar (see the module docstring),
which is unconditional -- unlike order/energy it is small and is what
keeps WHAM-style analysis exact, so it is never opt-in.
Parameters
----------
state : object like :py:class:`pyretis.simulation.repex.InfSwapState`
The coordinator state.
overwrite : bool, optional
When True (a fresh start, the default) the files are (re)written
with a clean header. When False (a restart) only **missing** files
are created and existing ones are kept for appending. The restart
path must still (re)create missing files: a run that *always*
restarts (``method="restart"``, e.g. a native config seeded from an
infswap restart that never wrote native output) would otherwise
never get its ``pathensemble.txt`` and leave the analysis with
nothing to read.
"""
formatter = PathEnsembleFormatter()
header = formatter.header
extra_files = []
if order_output_requested(state.config):
extra_files.append(ORDER_FILE_NAME)
if energy_output_requested(state.config):
extra_files.append(ENERGY_FILE_NAME)
for ens_dir in _ensemble_dirs(state):
os.makedirs(ens_dir, exist_ok=True)
target = os.path.join(ens_dir, PATHENSEMBLE_FILE_NAME)
if overwrite or not os.path.isfile(target):
with open(target, "w", encoding="utf-8") as handle:
handle.write(header + "\n")
ha_target = os.path.join(ens_dir, HA_WEIGHT_FILE_NAME)
if overwrite or not os.path.isfile(ha_target):
with open(ha_target, "w", encoding="utf-8") as handle:
handle.write(HA_WEIGHT_HEADER + "\n")
for name in extra_files:
extra_target = os.path.join(ens_dir, name)
if overwrite or not os.path.isfile(extra_target):
with open(extra_target, "w", encoding="utf-8"):
pass
def _interface_flags(min_op: float, max_op: float, start_op: float,
end_op: float,
interfaces: Sequence[float]) -> List[str]:
"""Return the native ``l``, ``m``, ``r`` interface flags for a path.
Parameters
----------
min_op : float
The smallest order parameter on the path.
max_op : float
The largest order parameter on the path.
start_op : float
The order parameter of the first phase point.
end_op : float
The order parameter of the last phase point.
interfaces : sequence of float
The ensemble's ``(left, middle, right)`` interfaces.
Returns
-------
list of string
The ``[l, m, r]`` flags (each one of ``L`` / ``R`` / ``*``).
"""
left, right = min(interfaces), max(interfaces)
middle = interfaces[1]
def side(value: float) -> str:
if value <= left:
return "L"
if value >= right:
return "R"
return "*"
flag_l = side(start_op)
flag_r = side(end_op)
flag_m = "M" if min_op < middle <= max_op else "*"
return [flag_l, flag_m, flag_r]
def _format_native_row(formatter: PathEnsembleFormatter, step: int,
nacc: int, nshoot: int, row: Dict[str, Any],
interfaces: Sequence[float],
weight: float) -> str:
"""Format one native ``pathensemble.txt`` row via the formatter.
The canonical :py:class:`PathEnsembleFormatter` defines the byte
layout; this builds the per-path dict it expects and reuses its
``PATH_FMT`` so the format stays identical to the native flow.
Parameters
----------
formatter : object like :py:class:`PathEnsembleFormatter`
The canonical native formatter.
step : int
The global coordinator cycle (the native ``Step`` column).
nacc : int
The running accepted-path count for this ensemble.
nshoot : int
The running shoot count for this ensemble.
row : dict
The coordinator's per-path native-column data (see
:py:func:`build_native_row_data`).
interfaces : sequence of float
The ensemble's ``(left, middle, right)`` interfaces.
weight : float
The native ``Weight`` column value (``1.0 / frac[j]``).
Returns
-------
string
The formatted native row.
"""
flags = _interface_flags(
row["ordermin"][0], row["ordermax"][0],
row["start_op"], row["end_op"], interfaces,
)
generated = row["generated"]
# The shooting-point order parameter (``O-shoot``) is only meaningful
# for shooting moves; the native flow writes 0.0 for non-shooting
# moves (ld / 00 / tr / s+ / s-). The coordinator's ``generated[1]``
# can be NaN for an init/loaded path, which would (a) diverge from the
# native 0.0 and (b) break the reference comparison, since
# ``math.isclose(nan, nan)`` is False. Coerce a NaN O-shoot to 0.0.
o_shoot = generated[1]
if o_shoot is None or (isinstance(o_shoot, float) and math.isnan(o_shoot)):
o_shoot = 0.0
return formatter.PATH_FMT.format(
step,
nacc,
nshoot,
flags[0],
flags[1],
flags[2],
row["length"],
row["status"],
map_move_to_native_mc(generated[0]),
row["ordermin"][0],
row["ordermax"][0],
row["ordermin"][1],
row["ordermax"][1],
o_shoot,
generated[2],
generated[3],
weight,
)
def _block_comment(step: int, row: Dict[str, Any]) -> str:
"""Return the native ``# Cycle:`` comment heading a path-data block.
Mirrors the native path formatters
(:py:class:`pyretis.inout.formats.order.OrderPathFormatter` /
:py:class:`pyretis.inout.formats.energy.EnergyPathFormatter`), which
write ``# Cycle: {step}, status: {status}, move: {generated}``. The
move tuple is normalised to plain Python scalars so the rendering
matches the native flow regardless of any numpy scalar types in the
coordinator's ``generated`` tag.
Parameters
----------
step : int
The global coordinator cycle (the native ``Cycle`` number).
row : dict
The coordinator's per-path native-column data.
Returns
-------
string
The block comment line.
"""
generated = tuple(
value.item() if hasattr(value, "item") else value
for value in row["generated"]
)
return f"# Cycle: {step}, status: {row['status']}, move: {generated}"
def _format_order_block(formatter: OrderPathFormatter, step: int,
row: Dict[str, Any]) -> List[str]:
"""Format one native ``order.txt`` block for a path row.
Parameters
----------
formatter : object like :py:class:`OrderPathFormatter`
The canonical native order-path formatter (defines the header
and per-line layout).
step : int
The global coordinator cycle.
row : dict
The coordinator's per-path data; ``row["order_series"]`` holds
one order-parameter tuple per phase point.
Returns
-------
list of string
The block lines (comment, column header, one line per phase
point).
"""
lines = [_block_comment(step, row), formatter.header]
for i, order in enumerate(row["order_series"]):
lines.append(formatter.format_data(i, order))
return lines
def _format_energy_block(formatter: EnergyPathFormatter, step: int,
row: Dict[str, Any]) -> List[str]:
"""Format one native ``energy.txt`` block for a path row.
Parameters
----------
formatter : object like :py:class:`EnergyPathFormatter`
The canonical native energy-path formatter.
step : int
The global coordinator cycle.
row : dict
The coordinator's per-path data; ``row["energy_series"]`` holds
one ``{vpot, ekin, etot, temp}`` dict per phase point (missing
terms are ``None`` and render as ``nan``, the native
convention).
Returns
-------
list of string
The block lines (comment, column header, one line per phase
point).
"""
lines = [_block_comment(step, row), formatter.header]
for i, energy in enumerate(row["energy_series"]):
lines.append(formatter.apply_format(i, energy))
return lines
[docs]
def write_native_pathensemble_data(state: Any) -> None:
"""Append native rows for the cycle's frac contributions per ensemble.
Realises the ratified cycle decision: **each accepted path gets a row
in every ensemble ``j`` it has a non-zero per-cycle frac increment
in**. For ensemble ``j``, one row is appended to its
``pathensemble.txt`` for every live path that contributed frac to
``j`` this cycle. ``Weight`` is written as ``1.0 / frac_increment``
(so the native analysis, which reads ``Weight`` as a 1/HA-weight,
recovers ``frac_increment`` as the histogram weight -- see the module
docstring); ``Step`` is the global coordinator cycle.
When requested, every appended row is accompanied by one ``# Cycle:``
block in the ensemble's ``order.txt`` and ``energy.txt`` carrying the
path's per-phase-point order parameters and energies. The native
analysis (:py:func:`~pyretis.analysis.path_analysis.
analyse_path_ensemble`) consumes these in **lockstep** -- one block
per ``pathensemble.txt`` row, matched on the cycle number -- so the
blocks are emitted exactly one per row, in row order. Every row is
unconditionally also paired with one ``(path number, ha_factor)``
entry in ``ha_weight.txt`` (see the module docstring) -- the WHAM-side
consumers need ``frac_inc`` and ``ha_factor`` separately, not just
their ``Weight`` ratio, and path identity to group/match rows, neither
of which the canonical ``pathensemble.txt`` columns carry.
Parameters
----------
state : object like :py:class:`pyretis.simulation.repex.InfSwapState`
The coordinator state, after the cycle's frac increments have
been applied (``state.last_frac_increment`` populated, mapping
each ensemble index to a list of ``(path_number, frac)`` pairs).
"""
formatter = PathEnsembleFormatter()
order_formatter = OrderPathFormatter()
energy_formatter = EnergyPathFormatter()
write_order = order_output_requested(state.config)
write_energy = energy_output_requested(state.config)
step = state.cstep
ens_dirs = _ensemble_dirs(state)
increments = state.last_frac_increment
# The native analysis reads the histogram weight as 1/Weight, which must
# be the WHAM-unweighted occupancy Cxy / HA = frac_inc / ha_factor (the
# infretis unweighting step), where ha_factor is the path's cv entry for
# this ensemble: the wf/ss High-Acceptance crossing count, or 1 for the
# indicator moves (sh/wt) and the minus ensemble. So Weight =
# ha_factor / frac_inc -- which reduces to 1/frac_inc for the indicator
# moves, leaving sh/wt output unchanged. The HA weight vector lives on
# the path as the UNPADDED cv (one entry per ensemble in calc_cv_vector
# order); add_traj left-pads a +-path cv by ``_offset`` into the
# swap-matrix columns, so the native ensemble-dir index ``j`` (000 ==
# minus) maps to cv index ``j - offset`` (the minus columns carry HA 1).
offset = getattr(state, "_offset", 0)
for j, ens_dir in enumerate(ens_dirs):
contributions = increments.get(j, [])
if not contributions:
continue
interfaces = state.ensembles[j]["interfaces"]
lines = []
order_lines = []
energy_lines = []
ha_lines = []
for pn, frac_inc in contributions:
if frac_inc <= 0.0:
continue
row = state.native_rows.get(pn)
if row is None:
continue
ha_factor = 1.0
ha_vec = row.get("ha_weights")
k = j - offset
if ha_vec is not None and 0 <= k < len(ha_vec):
ha_factor = float(ha_vec[k])
# A zero HA weight carries no histogram mass (a wf/ss path
# that contributes nothing to this ensemble's crossing
# probability) -- skip the row.
if ha_factor <= 0.0:
continue
state.native_nacc[j] = state.native_nacc.get(j, 0) + 1
if map_move_to_native_mc(row["generated"][0]) in (
"sh", "ss", "wt", "wf"
):
state.native_nshoot[j] = state.native_nshoot.get(j, 0) + 1
# Weight = ha_factor / frac_inc (the WHAM Cxy/HA unweighting);
# ha_factor == 1 for sh/wt so this is 1/frac_inc there.
weight = ha_factor / float(frac_inc)
lines.append(_format_native_row(
formatter, step, state.native_nacc[j],
state.native_nshoot.get(j, 0), row, interfaces, weight,
))
# One (path number, ha_factor) entry per pathensemble.txt row,
# same order: frac_inc = ha_factor / Weight then recovers both
# numbers exactly (not just their ratio), and the path number
# is what lets a reconstruction group/match rows by path --
# see the module docstring.
ha_lines.append(HA_WEIGHT_FMT.format(step, pn, ha_factor))
if write_order:
order_lines.extend(
_format_order_block(order_formatter, step, row)
)
if write_energy:
energy_lines.extend(
_format_energy_block(energy_formatter, step, row)
)
if lines:
target = os.path.join(ens_dir, PATHENSEMBLE_FILE_NAME)
with open(target, "a", encoding="utf-8") as handle:
handle.write("\n".join(lines) + "\n")
if order_lines:
target = os.path.join(ens_dir, ORDER_FILE_NAME)
with open(target, "a", encoding="utf-8") as handle:
handle.write("\n".join(order_lines) + "\n")
if energy_lines:
target = os.path.join(ens_dir, ENERGY_FILE_NAME)
with open(target, "a", encoding="utf-8") as handle:
handle.write("\n".join(energy_lines) + "\n")
if ha_lines:
target = os.path.join(ens_dir, HA_WEIGHT_FILE_NAME)
with open(target, "a", encoding="utf-8") as handle:
handle.write("\n".join(ha_lines) + "\n")
def _read_active_paths(run_dir: str) -> set:
"""Return the currently-live (not-yet-archived) path numbers.
Read from the run file's ``[current] active`` list -- written by
every scheduler cycle (:py:meth:`pyretis.simulation.repex.
InfSwapState.write_toml`). The historical ``infswap_data.txt`` writer
only emitted a row for a path once it was *archived* (replaced in
every ensemble); a still-live path's total is not yet final, so this
filter keeps the reconstruction matching that convention exactly
instead of including in-progress totals.
``output.toml`` (the single run file: resolved config + running
state) is preferred; a run started before that consolidation
persisted the same ``[current]`` state in a separate ``restart.toml``
instead, so that name is tried as a fallback -- matching
:func:`pyretis.bin.pyretisrun.run_pyretis_path_sampling`'s own
prefer-then-fall-back handling of the same two names.
Parameters
----------
run_dir : string
The run directory (where the run file lives).
Returns
-------
set of int
The active path numbers, or an empty set if neither run file is
present (no run has happened / nothing to filter).
"""
try:
import tomllib as _toml
except ImportError:
import tomli as _toml
run_file = os.path.join(run_dir, "output.toml")
if not os.path.isfile(run_file):
run_file = os.path.join(run_dir, "restart.toml")
if not os.path.isfile(run_file):
return set()
with open(run_file, "rb") as handle:
restart = _toml.load(handle)
return set(restart.get("current", {}).get("active", []))
def _read_ha_weight_rows(ens_dir: str) -> List[tuple]:
"""Read ``ha_weight.txt`` into ``(path_number, ha_factor)`` tuples.
Parameters
----------
ens_dir : string
A native ensemble directory (e.g. ``.../000``).
Returns
-------
list of tuple
``(path_number, ha_factor)`` per row, in file order -- the same
order :func:`write_native_pathensemble_data` appended them in,
which is the same order as the ensemble's ``pathensemble.txt``
rows.
"""
rows = []
target = os.path.join(ens_dir, HA_WEIGHT_FILE_NAME)
with open(target, encoding="utf-8") as handle:
for line in handle:
stripped = line.strip()
if not stripped or stripped.startswith("#"):
continue
_step, path_number, ha_factor = stripped.split()
rows.append((int(path_number), float(ha_factor)))
return rows
[docs]
def reconstruct_path_data_matrix(
run_dir: str, nskip: int = 0
) -> List[List[float]]:
"""Rebuild the ``infswap_data.txt`` matrix from native per-ensemble output.
The native route never writes ``infswap_data.txt``; WHAM-style
analysis (:py:mod:`pyretis.analysis.wham_analysis`,
:py:mod:`pyretis.analysis.path_weights`,
:py:mod:`pyretis.analysis.training_set`) instead reconstructs the same
matrix from the ``pathensemble.txt`` + ``ha_weight.txt`` pair in every
numbered ensemble directory. For ensemble ``j`` and a captured row,
``frac_inc = ha_factor / Weight`` recovers the cycle's fractional
occupancy increment exactly (see the module docstring); a path's
``Cxy`` for that ensemble is the sum of ``frac_inc`` over every cycle
it was live there, and its ``HA`` for that ensemble is ``ha_factor``
itself (constant across those cycles by construction). ``Length`` and
``Max-O`` are path-intrinsic, so every row carrying a given path
number must agree -- a mismatch anywhere is a logic error in the
writer, not data to silently average over, hence the loud failures
below rather than a best-effort reconciliation.
Row order is **not** a byte-for-byte reproduction of the live
writer's append order (paths there are written when archived, in
roughly-but-not-exactly path-number order under infinite swapping).
Rows here are sorted by ascending path number instead. This is
immaterial to :func:`pyretis.analysis.path_weights.get_path_weights`
(and so to ``interfaces_from_data``/``infinit``, its caller) and to
:func:`pyretis.analysis.training_set.read_path_data` /
:mod:`pyretis.analysis.combine_data`'s per-path remapping -- all
operate on the row set, not its sequence. It DOES affect the
block-error / running-average machinery in
:func:`pyretis.analysis.wham_analysis.analyse_wham_output` (a
chronological proxy, not the original sequence); that estimator is
superseded by the native crossing-probability analysis once this
function is wired in (see ``MERGE_TODO.md``).
Two further, deliberate differences from the historical
``infswap_data.txt`` writer (verified empirically against it before
it was retired, not assumed):
* **Still-live paths are excluded**, via :func:`_read_active_paths`
(``restart.toml``'s ``[current] active`` list) -- matching the old
writer, which only emitted a row once a path was *archived*, never
for an in-progress (not-yet-final) total.
* **A path that never contributed a non-zero fraction anywhere has no
row at all**, unlike the old writer (which still wrote an
all-``"----"`` row for it on archival). Such a path leaves no trace
in any native per-cycle file (``frac_inc <= 0.0`` rows are never
written -- see :func:`write_native_pathensemble_data`), so there is
nothing to reconstruct it from. This is provably harmless: an
all-zero row cannot change any sum, ratio, or weight any consumer
computes from the matrix (it contributes exactly 0 everywhere it
would appear).
Parameters
----------
run_dir : string
The directory holding the numbered native ensemble directories
(``000``, ``001``, ...).
nskip : int, optional
Number of initial paths (by ascending path number) to discard.
Returns
-------
matrix : list of list of float
Exactly the shape :func:`pyretis.analysis.wham_analysis.
read_data_matrix` returns: one row per path, ``[path_nr, length,
max_op, Cxy_0 .. Cxy_{n-1}, HA_0 .. HA_{n-1}]``.
"""
ens_names = sorted(
name for name in os.listdir(run_dir)
if name.isdigit() and os.path.isdir(os.path.join(run_dir, name))
)
nintf = len(ens_names)
if nintf == 0:
return []
cxy: Dict[int, List[float]] = {}
ha: Dict[int, List[float]] = {}
length: Dict[int, int] = {}
max_op: Dict[int, float] = {}
for j, name in enumerate(ens_names):
ens_dir = os.path.join(run_dir, name)
path_target = os.path.join(ens_dir, PATHENSEMBLE_FILE_NAME)
path_rows = []
with open(path_target, encoding="utf-8") as handle:
for line in handle:
parsed = PathEnsembleFormatter.parse(line)
if parsed is not None:
path_rows.append(parsed)
ha_rows = _read_ha_weight_rows(ens_dir)
if len(path_rows) != len(ha_rows):
raise ValueError(
f'{path_target} has {len(path_rows)} rows but '
f'{os.path.join(ens_dir, HA_WEIGHT_FILE_NAME)} has '
f'{len(ha_rows)} -- the native writer keeps these in '
'lockstep, so a mismatch means the files are stale, '
'truncated, or from different runs.'
)
for parsed, (path_number, ha_factor) in zip(path_rows, ha_rows):
weight = parsed["weight"]
if weight <= 0.0:
raise ValueError(
f'Non-positive Weight ({weight}) for path '
f'{path_number} in {path_target}: cannot recover '
'frac_inc = ha_factor / Weight.'
)
frac_inc = ha_factor / weight
row_length = parsed["length"]
row_max_op = parsed["ordermax"][0]
if path_number in length:
if length[path_number] != row_length:
raise ValueError(
f'Path {path_number} has inconsistent Length '
f'across ensembles ({length[path_number]} vs '
f'{row_length}); Length is path-intrinsic and '
'must agree everywhere the path appears.'
)
if not math.isclose(
max_op[path_number], row_max_op, abs_tol=1e-12
):
raise ValueError(
f'Path {path_number} has inconsistent Max-O '
f'across ensembles ({max_op[path_number]} vs '
f'{row_max_op}).'
)
else:
length[path_number] = row_length
max_op[path_number] = row_max_op
cxy[path_number] = [0.0] * nintf
ha[path_number] = [0.0] * nintf
cxy[path_number][j] += frac_inc
if ha[path_number][j] != 0.0 and not math.isclose(
ha[path_number][j], ha_factor, rel_tol=1e-12
):
raise ValueError(
f'Path {path_number} has inconsistent HA-weight in '
f'ensemble {name} across cycles '
f'({ha[path_number][j]} vs {ha_factor}); ha_factor '
'must be constant for one path within one ensemble.'
)
ha[path_number][j] = ha_factor
active = _read_active_paths(run_dir)
matrix = [
[float(pn), float(length[pn]), float(max_op[pn])]
+ cxy[pn] + ha[pn]
for pn in sorted(cxy)
if pn not in active
]
del matrix[:nskip]
return matrix