# Copyright (c) 2026, PyRETIS Development Team.
# Distributed under the LGPLv2.1+ License. See LICENSE for more info.
"""The infinite-swapping scheduler's path-archive writer.
This module holds the one genuinely scheduler-specific piece of the old
``formatter_repex`` module: the path-data formatters and the path archive
used by the infinite-swapping coordinator. Everything else in
``formatter_repex`` duplicated the canonical PyRETIS file-IO/formatter
code under :py:mod:`pyretis.inout.formats`,
:py:class:`pyretis.inout.common.OutputBase` and
:py:class:`pyretis.inout.fileio.FileIO`; those duplicates have been
removed and the base classes are imported from their canonical location.
What is kept here is specific to the infinite-swapping data model and the
archive layout the scheduler writes, and is *not* a duplicate of the
native code:
OrderPathFormatter / EnergyPathFormatter / PathExtFormatter
Path-data formatters whose ``format`` reads the infinite-swapping
phase-point attributes (``phasepoint.order``,
``getattr(phasepoint, key)`` for the energy terms and
``phasepoint.config`` / ``phasepoint.vel_rev`` for the trajectory
references). The native path formatters in
:py:mod:`pyretis.inout.formats` instead read ``phasepoint.particles``
and reconstruct derived energy terms, so they cannot be reused
directly.
PathStorage
The infinite-swapping path archive. It writes ``order.txt``,
``energy.txt`` and ``traj.txt`` to ``<dir>/<path_number>/`` and moves
the trajectory files into ``<dir>/<path_number>/accepted/``. It is
called as ``output(step, {"path": ..., "dir": ...})`` and returns the
moved :py:class:`~pyretis.core.path.Path`. The native
:py:class:`pyretis.inout.archive.PathStorage` uses a different call
signature and an accepted/rejected archive layout, so it is not a
drop-in replacement.
"""
from __future__ import annotations
import logging
import os
import shutil
from typing import (
TYPE_CHECKING,
Any,
Dict,
List,
Optional,
Tuple,
TypedDict,
)
from pyretis.inout.common import OutputBase, make_dirs
from pyretis.inout.formats.energy import EnergyFormatter
from pyretis.inout.formats.formatter import OutputFormatter
from pyretis.inout.formats.order import OrderFormatter
logger = logging.getLogger(__name__) # pylint: disable=invalid-name
logger.addHandler(logging.NullHandler())
if TYPE_CHECKING: # pragma: no cover
from collections.abc import Iterable
from pyretis.core.path import Path as InfPath
[docs]
class FormattersEntry(TypedDict):
"""To store formatters and output files together."""
fmt: OutputFormatter # The formatter to use.
file: str # The file to write.
[docs]
class PathStorage(OutputBase):
"""A class for handling storage of external trajectories.
Attributes
----------
target
Determines the target for this output class. Here it will be a
file archive (i.e., a directory based collection of files).
formatters
This dict contains the formatters for writing path data, with
default filenames used for them.
out_dir_fmt
A format to use for creating directories within the archive.
This one is applied to the step number for the output.
"""
target = "file-archive"
formatters: Dict[str, FormattersEntry] = {
"order": {"fmt": OrderPathFormatter(), "file": "order.txt"},
"energy": {"fmt": EnergyPathFormatter(), "file": "energy.txt"},
"traj": {"fmt": PathExtFormatter(), "file": "traj.txt"},
}
out_dir_fmt = "{}"
[docs]
def __init__(self, keep_traj_fnames: Optional[list] = None):
"""Set up the storage.
Parameters
----------
keep_traj_fnames : list, optional
A list of file extensions matched against the source
directories of the trajectories; matching files are kept.
Notes
-----
No formatters are passed to the parent class. This is because
this class is less flexible and only intended to do one thing:
write path data for external trajectories.
"""
formatter = OutputFormatter("empty formatter", header=None)
super().__init__(formatter)
if keep_traj_fnames is None:
keep_traj_fnames = []
self.keep_traj_fnames = keep_traj_fnames
[docs]
def output_path_files(
self, step: int, data: List[Any], target_dir: str
) -> List[Tuple[str, str]]:
"""Write the output files for energy, path and order parameter.
Parameters
----------
step : int
The current simulation step.
data : list
A tuple containing:
- The path as an object like :py:class:`.Path`.
- A string containing the status of this path.
target_dir : str
The path to where we archive the files.
Returns
-------
files : list
The files created as a list of tuples. Each tuple contains:
- The full path to the file.
- A relative path to the file. The relative path is useful
for organizing internally in archives.
"""
path, status = data[0], data[1]
files = []
for key, val in self.formatters.items():
logger.debug("Storing: %s", key)
fmt = val["fmt"]
full_path = os.path.join(target_dir, val["file"])
relative_path = os.path.join(
self.out_dir_fmt.format(step), val["file"]
)
files.append((full_path, relative_path))
with open(full_path, mode="w", encoding="utf8") as output:
for line in fmt.format(step, (path, status)):
output.write(f"{line}\n")
return files
[docs]
@staticmethod
def _move_path(
path: InfPath,
target_dir: str,
keep_traj_fnames: list,
prefix: Optional[str] = None,
) -> InfPath:
"""Copy a path to a given target directory.
Parameters
----------
path : InfPath
The path to copy.
target_dir : str
The location where we are moving the path to.
keep_traj_fnames : list
A list of file extensions that are matched against the
source directories in which the trajectories are stored.
File extensions that match the pattern are also stored.
prefix : str, optional
A prefix for the file names of copied files.
Returns
-------
path_copy : InfPath
A copy of the input path.
"""
path_copy = path.copy()
new_pos, source = _generate_file_names(
path_copy, target_dir, prefix=prefix
)
# Keep any files whose extension matches the patterns in
# keep_traj_fnames:
if keep_traj_fnames:
for source_file in source.copy().keys():
source_dir, source_fname = os.path.split(source_file)
traj_name, _ = os.path.splitext(source_fname)
for ext in keep_traj_fnames:
new_fname = traj_name + ext
fpath = os.path.join(source_dir, new_fname)
if os.path.isfile(fpath):
source[fpath] = os.path.join(target_dir, new_fname)
# Update positions:
for pos, phasepoint in zip(new_pos, path_copy.phasepoints):
phasepoint.config = (pos[0], pos[1])
for src, dest in source.items():
if src != dest:
if os.path.exists(dest):
if os.path.isfile(dest):
logger.debug("Removing %s as it exists", dest)
os.remove(dest)
logger.debug("Copy %s -> %s", src, dest)
shutil.move(src, dest)
return path_copy
[docs]
def output(self, step: int, data: Any) -> InfPath:
"""Format the path data and store the path.
Parameters
----------
step : int
The current simulation step.
data : Any
A dictionary containing the path and the directory to
write to.
Returns
-------
path : InfPath
A copy of the path (moved to the new directory).
"""
path = data["path"]
home_dir = data["dir"]
# This is the path on the form: /path/to/000/traj/11
archive_path = os.path.join(
home_dir,
f"{path.path_number}",
)
# To organize things we create a subfolder for storing the
# files. This is on the form: /path/to/000/traj/11/accepted
traj_dir = os.path.join(archive_path, "accepted")
# Create the needed directories:
make_dirs(traj_dir)
# Write order, energy and traj files to the archive:
_ = self.output_path_files(step, [path, "ACC"], archive_path)
path = self._move_path(path, traj_dir, self.keep_traj_fnames)
return path
[docs]
def write(self, towrite: str, end: str = "\n") -> bool:
"""We do not need the write method for this object."""
logger.critical(
'%s does *not* support the "write" method!',
self.__class__.__name__,
)
return False
[docs]
def __str__(self) -> str:
"""Return basic info."""
return f"{self.__class__.__name__} - archive writer."
[docs]
def _generate_file_names(
path: InfPath, target_dir: str, prefix: Optional[str] = None
) -> Tuple[List[Tuple[str, int]], Dict[str, str]]:
"""Generate new file names for moving or copying paths.
Parameters
----------
path : InfPath
The path object we are going to store.
target_dir : str
The location where we are moving the path to.
prefix : str, optional
The prefix can be used to prefix the name of the files.
Returns
-------
out : tuple
A tuple containing:
- A list with new file names.
- A dict which defines the unique "source -> destination" for
the copy/move operations.
"""
source = {}
new_pos = []
for phasepoint in path.phasepoints:
pos_file, idx = phasepoint.config
if pos_file not in source:
localfile = os.path.basename(pos_file)
if prefix is not None:
localfile = f"{prefix}{localfile}"
dest = os.path.join(target_dir, localfile)
source[pos_file] = dest
dest = source[pos_file]
new_pos.append((dest, idx))
return new_pos, source