# -*- coding: utf-8 -*-
# Copyright (c) 2023, PyRETIS Development Team.
# Distributed under the LGPLv2.1+ License. See LICENSE for more info.
"""Combine several infinite-swapping runs into one data file.
Ported from the upstream ``inftools`` ``combine_results.py``. Several runs
-- possibly with *different* interface sets, e.g. successive ``infinit``
interface-placement steps, or independent runs to be pooled for better
statistics -- are merged into a single ``infswap_data.txt``-shaped data
file + config whose interfaces are the union of the inputs. Each run's
path data is read via :func:`pyretis.analysis.training_set.get_path_data`
(a literal ``infswap_data.txt`` if the run still has one, else
reconstructed from its native per-ensemble output -- the scheduler itself
no longer writes the file). Each run's ensemble columns are re-mapped onto
the union interfaces, and the rows are proportionally interleaved so a
downstream block-error estimate sees a representative mix from every run.
The merged output keeps the historical ``infswap_data.txt`` matrix shape
deliberately: it is a derived analysis artifact (an explicit merge a user
asks for), not the scheduler's own primary output, and downstream WHAM /
``interfaces_from_data`` consumers already read that shape.
"""
import tomllib
import numpy as np
from pyretis.analysis.training_set import get_path_data
__all__ = ['combine_data']
def _read_interfaces(toml_file):
"""Return the ``[simulation] interfaces`` list from a TOML file."""
with open(toml_file, 'rb') as infile:
return tomllib.load(infile)['simulation']['interfaces']
def _load_runs(tomls, run_dirs):
"""Read every run's interfaces + paths and the union interface set.
Returns ``(runs, union)`` where each run is ``{'intf': ..., 'paths':
...}`` and ``union`` is the sorted set of all interfaces. Raises if the
runs do not share the same first/last interface.
"""
runs, all_intfs = [], []
for toml_file, run_dir in zip(tomls, run_dirs):
intf = _read_interfaces(toml_file)
runs.append({'intf': intf, 'paths': get_path_data(run_dir)})
all_intfs += intf
firsts = {run['intf'][0] for run in runs}
lasts = {run['intf'][-1] for run in runs}
if len(firsts) != 1 or len(lasts) != 1:
raise ValueError('all runs must share the first/last interface.')
return runs, sorted(set(all_intfs))
def _remap_rows(run, n_skip, col_map, n_union):
"""Re-map one run's data rows onto the union interface columns.
Returns a list of ready-to-write ``infswap_data.txt`` lines whose
fraction/weight columns are laid out over the ``n_union`` union
ensembles (absent columns filled with ``"----"``).
"""
lines = []
for line_num, path in enumerate(run['paths']):
if line_num < n_skip:
continue
mapped = {col_map[c]: v for c, v in path['cols'].items()}
frac, weig = [], []
for col in range(n_union):
if col in mapped:
frac.append(mapped[col][0])
weig.append(mapped[col][1])
else:
frac.append('----')
weig.append('----')
row = (f"\t{path['pn']}\t{path['len']}\t{path['max_op']}\t"
+ '\t'.join(frac) + '\t' + '\t'.join(weig) + '\t\n')
lines.append(row)
return lines
def _interleave(datalines, out):
"""Proportionally interleave per-run rows into ``<out>.txt``.
Each run is emitted at a stride proportional to its share of the
total, so a downstream block-error estimate sees a representative mix
from every run throughout the file.
"""
total = sum(len(lines) for lines in datalines)
with open(f'{out}.txt', 'w', encoding='utf-8') as outfile:
if total == 0:
return
modulos = [max(1, int(np.round(total / len(lines)))) if lines else 1
for lines in datalines]
idx = 0
while sum(len(lines) for lines in datalines) > 0:
for mod, lines in zip(modulos, datalines):
if idx % mod == 0 and lines:
outfile.write(lines.pop(0))
idx += 1
[docs]
def combine_data(tomls, run_dirs, skip=(100,), out='combo'):
"""Combine several runs into ``<out>.txt`` + ``<out>.toml``.
Parameters
----------
tomls : list of str
The input TOML config files, one per run (interfaces are read
from each).
run_dirs : list of str
The matching run directories, in the same order -- a literal
``infswap_data.txt`` inside one is used if present, otherwise its
path data is reconstructed from native per-ensemble output (see
:func:`pyretis.analysis.training_set.get_path_data`).
skip : sequence of int, optional
Initial rows to skip per run -- either one value for all runs or
one per run.
out : str, optional
Output basename; writes ``<out>.txt`` and ``<out>.toml``.
Returns
-------
interfaces : list of float
The union interface set written to ``<out>.toml``.
"""
tomls = list(tomls)
run_dirs = list(run_dirs)
if len(set(tomls)) != len(tomls) or len(set(run_dirs)) != len(run_dirs):
raise ValueError('combine_data needs distinct toml/run directories.')
if len(tomls) != len(run_dirs):
raise ValueError('need one run directory per toml.')
skip = list(skip)
if len(skip) not in (1, len(tomls)):
raise ValueError('skip must be one value or one per run.')
if len(tomls) > 1 and len(skip) == 1:
skip = skip * len(tomls)
runs, union = _load_runs(tomls, run_dirs)
# The two zero-ensembles keep columns 0/1; the interior interfaces are
# re-mapped by interface value onto the union columns 2, 3, ...
col_of = {intf: i for i, intf in enumerate(union[1:-1], 2)}
with open(f'{out}.toml', 'w', encoding='utf-8') as outfile:
outfile.write('[simulation]\ninterfaces = [\n')
outfile.write(''.join(f' {v},\n' for v in union))
outfile.write(']\n')
n_union = len(union)
datalines = []
for run, n_skip in zip(runs, skip):
recol = [0, 1] + [col_of[i] for i in run['intf'][1:-1]]
col_map = dict(enumerate(recol))
datalines.append(_remap_rows(run, n_skip, col_map, n_union))
_interleave(datalines, out)
return union