# -*- coding: utf-8 -*-
# Copyright (c) 2023, PyRETIS Development Team.
# Distributed under the LGPLv2.1+ License. See LICENSE for more info.
"""Multiresolution Wire Fencing (``mwf``) move for the infinite-swapping route.
Ported from the upstream infretis ``classes/multires_wf.py`` (branch
``multires_wf``). Multiresolution wire fencing generates the wire-fencing
sub-paths at TWO engine resolutions: the intermediate sub-paths are shot
with a small number of subcycles (cheap, "high resolution" in the number
of shooting attempts) and only the final sub-path of each set is shot with
the engine's full subcycle count, before the surviving sub-path is extended
to a full path and accepted with the usual sub-trajectory acceptance.
The high-level helpers (``wirefence_weight_and_pick`` / ``shoot`` /
``extender`` / ``subt_acceptance`` / ``check_kick``) are pyretis's own
infinite-swapping move helpers in :mod:`pyretis.core.moves` (the
faithful port of ``infretis.core.tis``), so this is a near-verbatim port;
the only adaptations are the lower-level kick calls in
:func:`_generate_segment_using_shoot` (pyretis's ``modify_velocities``
takes the random generator explicitly and ``check_kick`` takes the TIS
settings).
New ``[tis]`` keys: ``mwf_subcycle_small`` (and the per-ensemble
``mwf_subcycle_small_by_ensemble`` dict) and ``mwf_nsubpath``; the number
of sets reuses the wire-fencing ``n_jumps``.
"""
from __future__ import annotations
import logging
import os
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple
logger = logging.getLogger(__name__) # pylint: disable=invalid-name
logger.addHandler(logging.NullHandler())
[docs]
@dataclass
class _MWFSettings:
"""Resolved multiresolution-wire-fencing parameters for one ensemble."""
nsubpath: int
nsubset: int # reuses the wire-fencing ``n_jumps``
subcycle_small: int # the new high-resolution subcycle count
subcycle_large: int # the engine's full subcycle count
[docs]
@contextmanager
def _temporary_subcycles(engine, new_subcycles):
"""Temporarily switch the engine to ``new_subcycles`` subcycles."""
old = getattr(engine, 'subcycles', None)
try:
if new_subcycles is not None and old is not None:
engine.subcycles = int(max(1, new_subcycles))
yield
finally:
if old is not None:
engine.subcycles = old
[docs]
def _positive_int_or_none(value: Any) -> Optional[int]:
"""Return ``value`` if it is a positive int (not bool), else ``None``."""
if isinstance(value, bool):
return None
if not isinstance(value, int):
return None
return value if value > 0 else None
[docs]
def _read_mwf_settings(ens_set: Dict[str, Any], engine) -> _MWFSettings:
"""Resolve the multiresolution settings for the current ensemble."""
tis = ens_set['tis_set']
large = getattr(engine, 'subcycles', 10)
small_default = max(1, large // 5)
scalar_small = _positive_int_or_none(tis.get('mwf_subcycle_small'))
fallback_small = (scalar_small if scalar_small is not None
else small_default)
by_ensemble = tis.get('mwf_subcycle_small_by_ensemble')
override_small: Optional[int] = None
if isinstance(by_ensemble, dict):
ens_name = ens_set.get('ens_name')
if ens_name is not None and ens_name in by_ensemble:
override_small = _positive_int_or_none(by_ensemble[ens_name])
subcycle_small = (override_small if override_small is not None
else fallback_small)
return _MWFSettings(
nsubpath=int(tis.get('mwf_nsubpath', 3)),
nsubset=int(tis.get('n_jumps', 2)),
subcycle_small=int(subcycle_small),
subcycle_large=int(large),
)
[docs]
def _random_internal_indices(path, rgen, k: int) -> List[int]:
"""Up to ``k`` distinct interior indices of ``path``, shuffled."""
n = len(path.phasepoints)
if n <= 2:
return []
available = list(range(1, n - 1))
rgen.shuffle(available)
return available[:max(0, min(k, len(available)))]
[docs]
def _pick_shooting_idx(path, rgen, is_highres: bool, queue: List[int]):
"""Pick a shooting index according to resolution.
High resolution pops the head of ``queue`` (empty -> status ``"EQU"``,
a failed sub-path); low resolution draws a uniform interior index.
"""
if is_highres:
if not queue:
return 0, queue, 'EQU'
idx = queue.pop(0)
return int(idx), queue, 'OK'
n = len(path.phasepoints)
if n <= 2:
return 0, queue, 'OK'
return int(rgen.integers(1, n - 1)), queue, 'OK'
[docs]
def _generate_segment_using_shoot(si, idx, ens_set_sub, engine, start_cond):
"""Shoot a sub-path from ``si.phasepoints[idx]`` (MWF-chosen point).
The MWF caller owns the shooting-point policy, so we build the explicit
shooting point and pass it to :func:`pyretis.core.moves.shoot`,
performing the velocity kick (``modify_velocities`` + ``check_kick``)
here to preserve standard MC kick semantics.
"""
from pyretis.core.moves import shoot, check_kick
tis_set = ens_set_sub['tis_set']
rgen = ens_set_sub['rgen']
maxlen = tis_set['maxlength']
if idx <= 0 or idx >= len(si.phasepoints) - 1:
return False, si.empty_path(maxlen=maxlen), 'IDX'
shooting_point = si.phasepoints[idx].copy()
shooting_point.idx = idx
dek, _ = engine.modify_velocities(shooting_point, tis_set, rgen)
shooting_point.order = engine.calculate_order(shooting_point)
scratch = si.empty_path(maxlen=maxlen)
if not check_kick(shooting_point, ens_set_sub['interfaces'], scratch,
rgen, dek, tis_set):
return False, scratch, scratch.status
return shoot(ens_set_sub, si, engine, shooting_point, start_cond)
[docs]
def multires_wire_fencing(ens_set: Dict[str, Any], trial_path, engine,
start_cond: Tuple[str, ...] = ('L',)):
"""Perform a Multiresolution Wire Fencing (``mwf``) move.
Parameters
----------
ens_set : dict
Ensemble settings (``interfaces``, ``tis_set``, ``rgen``,
``ens_name``).
trial_path : object like :py:class:`pyretis.core.path.Path`
The current accepted path to move from.
engine : object like :py:class:`pyretis.engines.engine.EngineBase`
The MD engine used for propagating.
start_cond : tuple of str
The starting condition (``"L"`` / ``"R"``).
Returns
-------
out[0] : boolean
True if the path can be accepted.
out[1] : object like :py:class:`pyretis.core.path.Path`
The generated path.
out[2] : string
The status of the path.
"""
from pyretis.core.moves import (
wirefence_weight_and_pick,
extender,
subt_acceptance,
)
interfaces_full = ens_set['interfaces']
lambda_cap = ens_set['tis_set'].get('interface_cap', interfaces_full[2])
wf_int = [interfaces_full[1], interfaces_full[1], lambda_cap]
# Step 1: pick s0 between lambda_i and lambda_cap.
n_frames, s0 = wirefence_weight_and_pick(
trial_path, wf_int[0], wf_int[2], return_seg=True, ens_set=ens_set
)
if n_frames == 0:
logger.warning('MWF: no valid frames between i and cap; aborting')
return False, trial_path, 'NSG'
sub_ens = {
'interfaces': wf_int,
'rgen': ens_set['rgen'],
'ens_name': ens_set['ens_name'],
'start_cond': ('L', 'R'),
'tis_set': dict(ens_set['tis_set']),
}
sub_ens['tis_set']['allowmaxlength'] = True
sub_ens['tis_set']['maxlength'] = ens_set['tis_set']['maxlength']
mwf = _read_mwf_settings(ens_set, engine)
countset = 1
succ_sets = 0
total_succ_subpaths = 0
if len(s0.phasepoints) <= 2:
trial_path.status = 'NSG'
return False, trial_path, trial_path.status
si = s0.copy()
last_acc_si = si.copy()
while True:
i = 0
set_rejected = False
si = si.copy()
last_acc_in_set = si.copy()
subpaths_in_set = 0
# Fresh high-resolution shooting queue, strictly set-local.
shooting_queue = _random_internal_indices(
s0, sub_ens['rgen'], mwf.nsubpath - 1
)
last_acc_queue = list(shooting_queue)
while i < mwf.nsubpath:
is_highres = (i + 1) < mwf.nsubpath
idx, shooting_queue, pick_status = _pick_shooting_idx(
si, sub_ens['rgen'], is_highres, shooting_queue
)
if pick_status == 'EQU':
i += 1
gen_success = False
seg = si.empty_path(maxlen=sub_ens['tis_set']['maxlength'])
status = 'EQU'
elif (len(si.phasepoints) <= 2 or idx <= 0
or idx >= len(si.phasepoints) - 1):
gen_success = False
seg = si.empty_path(maxlen=sub_ens['tis_set']['maxlength'])
status = 'SHP'
else:
i += 1
if i == mwf.nsubpath:
subcycles = mwf.subcycle_large
else:
subcycles = mwf.subcycle_small
with _temporary_subcycles(engine, subcycles):
gen_success, seg, status = _generate_segment_using_shoot(
si, idx, sub_ens, engine, start_cond=('L', 'R')
)
subpaths_in_set += 1
# MWF-specific rejection: cap-cap path or a length-3 highres
# sub-path (a "no-move" shot) is not a usable link.
if gen_success:
start_side, end_side, _, _ = seg.check_interfaces(wf_int)
is_cap_cap = start_side == 'R' and end_side == 'R'
is_len3_highres = (subcycles == mwf.subcycle_small
and seg.length == 3)
if is_cap_cap or is_len3_highres:
gen_success = False
status = 'CCP' if is_cap_cap else 'L3H'
if not gen_success:
for rej_seg_trajname in seg.adress:
os.remove(rej_seg_trajname)
if i == 1 or i == mwf.nsubpath:
set_rejected = True
si = s0.copy()
break
si = last_acc_in_set.copy()
shooting_queue = list(last_acc_queue)
continue
# Accept the sub-path.
total_succ_subpaths += 1
si = seg.copy()
last_acc_in_set = seg.copy()
remaining_highres = (mwf.nsubpath - 1) - i
if remaining_highres > 0:
shooting_queue = _random_internal_indices(
seg, sub_ens['rgen'], remaining_highres
)
else:
shooting_queue = []
last_acc_queue = list(shooting_queue)
if i == mwf.nsubpath:
break
if not set_rejected:
succ_sets += 1
last_acc_si = si.copy()
s0 = si.copy()
if countset == mwf.nsubset:
if succ_sets == 0:
trial_path.status = 'NSG'
return False, trial_path, trial_path.status
ok, full_path, _ = extender(
last_acc_si, engine, ens_set, start_cond
)
if not ok:
return False, full_path, full_path.status
success, accepted = subt_acceptance(
full_path, ens_set, engine, start_cond
)
accepted.generated = ('mwf', 9001, succ_sets, accepted.length)
if not success:
return False, accepted, accepted.status
return True, accepted, accepted.status
countset += 1