Source code for pyretis.inout.restart

# Copyright (c) 2026, PyRETIS Development Team.
# Distributed under the LGPLv2.1+ License. See LICENSE for more info.
"""This module defines how we write and read restart files.

Important methods defined here
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

read_restart_file (:py:func:`.read_restart_file`)
    A method for reading restart information from a file.

write_restart_file (:py:func:`.write_restart_file`)
    A method for writing the restart file.

write_ensemble_restart (:py:func:`.write_ensemble_restart`)
    A method for writing restart files for path ensembles.
"""
import logging
import copy
import os
import pickle  # nosec B403
from pyretis.inout.common import atomic_write
logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
logger.addHandler(logging.NullHandler())


__all__ = ['read_restart_file',
           'write_restart_file',
           'write_ensemble_restart',
           'system_restart_info',
           'load_system_restart']


def system_restart_info(system):
    """Collect restart info for a System (free-function form).

    Equivalent to the soon-to-go ``System.restart_info()`` method on
    the harness System. Returns the same dict shape; works on any
    System with the bridge surface (``units``, ``temperature``,
    ``post_setup``, ``order``, ``box``, ``particles``).

    Parameters
    ----------
    system : object like :class:`.System`
        The system whose state to serialise.

    Returns
    -------
    dict
        Restart-info dictionary.

    """
    info = {}
    for attr in ('units', 'temperature', 'post_setup', 'order'):
        info[attr] = getattr(system, attr, None)
    try:
        info['box'] = system.box.restart_info()
    except AttributeError:
        pass
    try:
        info['particles'] = system.particles.restart_info()
    except AttributeError:
        pass
    return info


def load_system_restart(system, info):
    """Restore System state from a restart-info dict.

    Equivalent to the soon-to-go ``System.load_restart_info(info)``
    method on the harness System.

    Parameters
    ----------
    system : object like :class:`.System`
        The system to populate.
    info : dict
        The restart info, as produced by :func:`system_restart_info`.

    """
    # Local imports to avoid pulling restart-only dependencies at
    # module import time.
    from pyretis.core.box import box_from_restart
    from pyretis.core.particles import particles_from_restart
    for attr in ('units', 'temperature', 'post_setup', 'order'):
        if attr in info:
            setattr(system, attr, info[attr])
    system.box = box_from_restart(info)
    system.particles = particles_from_restart(info)


def write_restart_file(filename, simulation):
    """Write restart info for a simulation.

    Parameters
    ----------
    filename : string
        The file we are going to write to.
    simulation : object like :py:class:`.Simulation`
        A simulation object we will get information from.

    """
    info = simulation.restart_info()

    # Write atomically: the file is swapped into place with os.replace,
    # so an abrupt interruption can never leave a half-written restart
    # (it is always either the complete old or the complete new file).
    atomic_write(
        filename,
        lambda outfile: pickle.dump(info, outfile),
        binary=True,
    )


def write_ensemble_restart(ensemble, settings_ens):
    """Write a restart file for a path ensemble.

    Parameters
    ----------
    ensemble : dict
        it contains:

        * `path_ensemble` : object like :py:class:`.PathEnsemble`
          The path ensemble we are writing restart info for.
        * ` system` : object like :py:class:`.System`
          System is used here since we need access to the temperature
          and to the particle list.
        * `order_function` : object like :py:class:`.OrderParameter`
          The class used for calculating the order parameter(s).
        * `engine` : object like :py:class:`.EngineBase`
          The engine to use for propagating a path.

    settings_ens : dict
        A dictionary with the ensemble settings.

    """
    info = copy.deepcopy(settings_ens)
    if ensemble.get('path_ensemble') is not None:
        info['path_ensemble'] = ensemble['path_ensemble'].restart_info()
    if ensemble.get('system') is not None:
        info['system'] = ensemble['system'].restart_info()
    if ensemble.get('engine') is not None:
        info['engine'] = ensemble['engine'].restart_info()
    if ensemble.get('rgen') is not None:
        info['rgen'] = ensemble['rgen'].get_state()
    if ensemble.get('order_function') is not None:
        info['order_function'] = ensemble['order_function'].restart_info()

    filename = os.path.join(
        settings_ens['simulation']['exe_path'],
        ensemble['path_ensemble'].ensemble_name_simple,
        'ensemble.restart')
    # Write atomically: the file is swapped into place with os.replace,
    # so an abrupt interruption can never leave a half-written restart
    # (it is always either the complete old or the complete new file).
    atomic_write(
        filename,
        lambda outfile: pickle.dump(info, outfile),
        binary=True,
    )


def read_restart_file(filename):
    """Read restart info for a simulation.

    Parameters
    ----------
    filename : string
        The file we are going to read from.

    """
    with open(filename, 'rb') as infile:
        info = pickle.load(infile)  # nosec B301
    return info