# -*- coding: utf-8 -*-
# Copyright (c) 2023, PyRETIS Development Team.
# Distributed under the LGPLv2.1+ License. See LICENSE for more info.
"""Collect a machine-learning training set from path-sampling output.
Ported from the upstream ``inftools`` tool
``tistools/collect_trainingset.py``. It selects a set of configurations
spread across the order-parameter range (one frame per chosen path, drawn
from the interface band that path belongs to) so they can be used to
train, e.g., a machine-learned committor / shooting-point selector.
The upstream tool drew its selections with ``numpy.random`` (not
reproducible) and read trajectories with ASE only. This port:
* threads an explicit :class:`pyretis.core.random_gen.RandomGenerator`
so a given seed reproduces the same selection (PyRETIS determinism);
* separates the reproducible *selection* (:func:`select_training_frames`,
no trajectory I/O) from the trajectory *export*
(:func:`collect_training_set`, which reads/writes frames with ASE and
is therefore limited to ASE-readable trajectory formats).
"""
import logging
import os
import numpy as np
logger = logging.getLogger(__name__) # pylint: disable=invalid-name
logger.addHandler(logging.NullHandler())
__all__ = [
'read_path_data', 'get_path_data', 'select_training_frames',
'collect_training_set',
]
[docs]
def read_path_data(filename):
"""Parse an ``infswap_data.txt`` file into per-path dictionaries.
Port of the upstream ``inftools`` ``data_reader``: each returned path
records the ensembles it participated in (those columns where both the
``Cxy`` and the high-acceptance weight are present).
Parameters
----------
filename : str
Path to the ``infswap_data.txt`` file.
Returns
-------
paths : list of dict
One dict per path with keys ``pn`` (path number, str), ``len``
(str), ``max_op`` (float) and ``cols`` (a dict mapping the ensemble
column index to its ``[Cxy, HA-weight]`` strings).
"""
paths = []
n_ens = 0
with open(filename, encoding='utf-8') as infile:
for line in infile:
if line.startswith('#'):
continue
items = line.rstrip().split()
if not items:
continue
if not n_ens:
n_ens = int(len(items[3:]) / 2)
path_nr, length, max_op = items[:3]
cxy = items[3:n_ens + 3]
haw = items[3 + n_ens:2 * n_ens + 3]
# Skip paths that carry no weights at all.
if set(cxy) == set(haw) == {'----'}:
continue
cols = {}
for col, (one_cxy, one_haw) in enumerate(zip(cxy, haw)):
if '----' in (one_cxy, one_haw):
continue
cols[col] = [one_cxy, one_haw]
paths.append({
'pn': path_nr,
'len': length,
'max_op': float(max_op),
'cols': cols,
})
return paths
[docs]
def get_path_data(run_dir, nskip=0):
"""Return :func:`read_path_data`'s shape for a run, from either source.
Like :func:`pyretis.analysis.wham_analysis.get_path_data_matrix`: a
literal ``infswap_data.txt``/``infretis_data.txt`` in ``run_dir``
takes precedence (read via :func:`read_path_data`); otherwise the
same per-path dicts are built from the native per-ensemble output
(:func:`pyretis.inout.pathensemble_output.reconstruct_path_data_matrix`,
converted column-for-column to the ``[Cxy, HA-weight]`` string-pair
shape :func:`read_path_data` returns, so existing consumers
(:func:`select_training_frames`, :mod:`pyretis.analysis.combine_data`)
need no further change).
Parameters
----------
run_dir : str
The run directory.
nskip : int, optional
Number of initial rows/paths to discard.
Returns
-------
paths : list of dict
See :func:`read_path_data`.
"""
from pyretis.analysis.wham_analysis import detect_infswap_output
data_file = detect_infswap_output(run_dir)
if data_file is not None:
paths = read_path_data(data_file)
return paths[nskip:] if nskip else paths
# Imported here, not at module load, to avoid pulling pyretis.inout
# into every pyretis.analysis import.
from pyretis.inout.pathensemble_output import (
reconstruct_path_data_matrix,
)
matrix = reconstruct_path_data_matrix(run_dir, nskip=nskip)
paths = []
for row in matrix:
nintf = (len(row) - 3) // 2
cxy = row[3:3 + nintf]
haw = row[3 + nintf:3 + 2 * nintf]
# A genuine contribution is always > 0 (see
# write_native_pathensemble_data); 0.0 is reconstruct_path_data_
# matrix's "this ensemble was never touched" sentinel, matching
# read_path_data's "----" skip exactly.
cols = {
col: [str(one_cxy), str(one_haw)]
for col, (one_cxy, one_haw) in enumerate(zip(cxy, haw))
if one_cxy != 0.0
}
if not cols:
continue
paths.append({
'pn': str(int(row[0])),
'len': str(int(row[1])),
'max_op': row[2],
'cols': cols,
})
return paths
def _frames_in_band(order_file, lower, upper):
"""Return ``(frame_index, order_value)`` pairs inside an interface band.
``lower is None`` selects frames with ``order < upper`` (the first
band); otherwise ``lower < order <= upper``.
"""
matches = []
with open(order_file, encoding='utf-8') as infile:
for line in infile:
if line.startswith('#'):
continue
cols = line.split()
if len(cols) < 2:
continue
order = float(cols[1])
if lower is None:
if order < upper:
matches.append((cols[0], order))
elif lower < order <= upper:
matches.append((cols[0], order))
return matches
[docs]
def select_training_frames(paths, interfaces, input_dir, n_frames, rgen):
"""Reproducibly select training frames spread over the interfaces.
Paths are grouped by the ensembles they visited; from each group an
even share of paths is drawn (without replacement, no path reused
across groups), and from each path one frame is drawn from that
group's interface band. All draws use ``rgen``, so a given generator
state reproduces the selection exactly.
Parameters
----------
paths : list of dict
As returned by :func:`read_path_data`.
interfaces : list of float
The interface positions.
input_dir : str
Directory holding the per-path ``<pn>/order.txt`` files.
n_frames : int
Target number of frames (split evenly across the groups).
rgen : object like :class:`.RandomGenerator`
The random generator used for every selection.
Returns
-------
selected : list of tuple
``(path_nr, frame_index, order_value, ensemble)`` for each chosen
frame, ordered by path number.
"""
grouped = {}
for path in paths:
for key in path['cols']:
grouped.setdefault(key, []).append(path)
if not grouped:
return []
n_select = int(np.ceil(n_frames / len(grouped)))
used = set()
chosen_by_group = {}
for key in sorted(grouped):
avail = [p for p in grouped[key] if p['pn'] not in used]
if not avail:
chosen_by_group[key] = []
continue
take = min(n_select, len(avail))
pick_idx = np.atleast_1d(
rgen.choice(len(avail), size=take, replace=False))
picks = [avail[int(i)] for i in pick_idx]
chosen_by_group[key] = picks
used.update(p['pn'] for p in picks)
selected = []
for key in sorted(chosen_by_group):
upper = interfaces[key]
lower = interfaces[key - 1] if key > 1 else None
for path in chosen_by_group[key]:
order_file = os.path.join(input_dir, str(path['pn']), 'order.txt')
if not os.path.isfile(order_file):
logger.warning('Missing order file: %s', order_file)
continue
matches = _frames_in_band(order_file, lower, upper)
if not matches:
continue
idx = rgen.random_integers(0, len(matches) - 1)
frame_index, order_value = matches[idx]
selected.append((path['pn'], frame_index, order_value, key))
selected.sort(key=lambda row: int(row[0]))
return selected
def _resolve_frame(input_dir, path_nr, frame_index):
"""Map an order-file frame index to ``(traj_file, frame_in_file)``."""
traj_file = os.path.join(input_dir, str(path_nr), 'traj.txt')
with open(traj_file, encoding='utf-8') as infile:
for line in infile:
if line.startswith('#'):
continue
cols = line.split()
if cols and cols[0] == frame_index:
return cols[1], cols[2]
return None, None
[docs]
def collect_training_set(paths, interfaces, input_dir, n_frames, rgen,
out='trainingset'):
"""Select training frames and export them to an ``.xyz`` + report.
This reads and writes trajectory frames with ASE, so only
ASE-readable trajectory formats (multi-frame ``.xyz``, ``.traj``,
NetCDF, ...) are supported. The reproducible selection step
(:func:`select_training_frames`) has no such restriction.
Parameters
----------
paths, interfaces, input_dir, n_frames, rgen
See :func:`select_training_frames`.
out : str, optional
Output basename; writes ``<out>.xyz`` and ``<out>_report.txt``.
Returns
-------
selected : list of tuple
The selected frames (see :func:`select_training_frames`).
"""
try:
from ase.io import read, write
except ImportError as exc: # pragma: no cover
raise ImportError(
'collect_training_set needs ASE to read/write trajectory '
'frames (only ASE-readable formats are supported); install '
'ase or use select_training_frames for the selection only.'
) from exc
selected = select_training_frames(paths, interfaces, input_dir,
n_frames, rgen)
frames = []
report = ['path, frame_index, trajectory_file, order_parameter, ensemble']
for path_nr, frame_index, order_value, ensemble in selected:
traj_name, frame_in_file = _resolve_frame(input_dir, path_nr,
frame_index)
if traj_name is None:
logger.warning('No trajectory frame for path %s frame %s',
path_nr, frame_index)
continue
traj_path = os.path.join(input_dir, str(path_nr), 'accepted',
traj_name)
frames.append(read(traj_path, index=int(frame_in_file)))
report.append(f'{path_nr}, {frame_in_file}, {traj_name}, '
f'{order_value}, {ensemble}')
write(f'{out}.xyz', frames)
with open(f'{out}_report.txt', 'w', encoding='utf-8') as outfile:
outfile.write('\n'.join(report) + '\n')
logger.info('Wrote %d training frames to %s.xyz', len(frames), out)
return selected