Source code for pyretis.core.path_load

"""Load infinite-swapping paths from disk into :class:`.Path` objects.

These loaders were historically part of the inf-flavour ``Path`` module
(``pyretis.core._path_inf``). After the A3.1b Path collapse there is a
single :class:`pyretis.core.path.Path`, so the loaders live here, free of
the now-deleted ``_path_inf`` module. They build native ``Path`` objects
whose phasepoints are file-backed snapshot :class:`.System` objects (the
external-engine representation), exactly as the infinite-swapping
scheduler expects.
"""

from __future__ import annotations

import logging
import os
from typing import Any, Dict, List

from pyretis.core.path import Path
from pyretis.core._system_inf import System
from pyretis.inout._formatter_inf import (
    EnergyPathFile,
    OrderPathFile,
    PathExtFile,
)

logger = logging.getLogger(__name__)
logger.addHandler(logging.NullHandler())


def load_path(pdir: str) -> Path:
    """Load a path from the given directory."""
    trajtxt = os.path.join(pdir, "traj.txt")
    ordertxt = os.path.join(pdir, "order.txt")
    if not os.path.isfile(trajtxt):
        raise FileNotFoundError(trajtxt)
    if not os.path.isfile(ordertxt):
        raise FileNotFoundError(ordertxt)

    # load trajtxt
    with PathExtFile(trajtxt, "r") as trajfile:
        # Just get the first trajectory:
        traj = next(trajfile.load())

        # Update trajectory to use full path names:
        for i, snapshot in enumerate(traj["data"]):
            config = os.path.join(pdir, "accepted", snapshot[1])
            traj["data"][i][1] = config
            reverse = int(snapshot[3]) == -1
            idx = int(snapshot[2])
            traj["data"][i][2] = idx
            traj["data"][i][3] = reverse

        for config in set(frame[1] for frame in traj["data"]):
            if not os.path.isfile(config):
                raise FileNotFoundError(config)

    # load ordertxt
    with OrderPathFile(ordertxt, "r") as orderfile:
        orderdata = next(orderfile.load())["data"][:, 1:]

    path = Path()
    for snapshot, order in zip(traj["data"], orderdata):
        frame = System()
        frame.order = order
        frame.config = (snapshot[1], snapshot[2])
        frame.vel_rev = snapshot[3]
        path.phasepoints.append(frame)
    _load_energies_for_path(path, pdir)
    # TODO: CHECK PATH SOMEWHERE .acc, sta = _check_path(path, path_ensemble)
    return path


def _load_energies_for_path(path: Path, dirname: str) -> None:
    """Load energy data for a path.

    Args:
        path: The path we are to set up/fill.
        dirname: The path to the directory with the input files.
    """
    energy_file_name = os.path.join(dirname, "energy.txt")
    try:
        with EnergyPathFile(energy_file_name, "r") as energyfile:
            energy = next(energyfile.load())
            path.update_energies(
                energy["data"]["ekin"],
                energy["data"]["vpot"],
                energy["data"].get("etot", []),
                energy["data"].get("temp", []),
            )
    except FileNotFoundError:
        pass


def load_paths_from_disk(config: Dict[str, Any]) -> List[Path]:
    """Load paths from disk."""
    load_dir = config["simulation"]["load_dir"]
    paths = []
    for pnumber in config["current"]["active"]:
        new_path = load_path(os.path.join(load_dir, str(pnumber)))
        status = "re" if "restarted_from" in config["current"] else "ld"
        # TODO: important for shooting move if 'ld' is set. need a smart way
        # to remember if status is 'sh' or 'wf' etc. maybe in the toml file.
        new_path.generated = (status, float("nan"), 0, 0)
        new_path.maxlen = config["simulation"]["tis_set"]["maxlength"]
        paths.append(new_path)
        # assign pnumber
        paths[-1].path_number = pnumber
    return paths