Source code for directupsampling.sampling.settings

"""Settings for MD snapshot sampling."""

from __future__ import annotations

from collections.abc import Callable
from dataclasses import asdict, dataclass, field, replace
from typing import TYPE_CHECKING, Any

import ase.units
import numpy as np

from directupsampling.parallel import DummyMPI
from directupsampling.thermostats import LangevinGB

if TYPE_CHECKING:
    import pathlib
    from collections.abc import Callable

    from ase.md.md import MolecularDynamics




@dataclass
class Setting:


    def replace(self, **changes) -> Setting:
        return replace(self, **changes)




    def todict(self) -> dict:
        return asdict(self)






@dataclass
class ThermostatSetting(Setting):
    class_type: type[MolecularDynamics] = LangevinGB
    timestep: float = 1 * ase.units.fs
    friction: float = 0.01 / ase.units.fs
    fixcm: bool = True



    def get_kwargs(self) -> dict:
        dct = self.todict()
        del dct["class_type"]
        return dct






@dataclass
class MDLogSetting(Setting):
    interval: int = 1000
    logfile: str = "-"
    header: bool = True
    stress: bool = False
    peratom: bool = True
    mode: str = "a"
    comm: Any = DummyMPI()





@dataclass
class SamplingSetting(Setting):
    """Setting for MD snapshot sampling.

    Attributes
    ----------
        seed : int, optional
            Random seed for sampling. If None, a random seed is generated.
        rng : np.random.Generator, optional
            Random number generator for sampling. If None, a new generator is
            created using the seed.

        sampling_interval : int
            Interval for MD timesteps to take samples.
        sampling_offset : int
            Offset for MD timesteps for equilibration.

        sample_atoms : bool
            Whether to sample the atomic positions and forces (through a copy of
            the ASE `Atoms` instance). If False, only kinetic and potential
            energy are sampled.
        file : str or pathlib.Path, optional
            File path to save the sampled snapshots. If None, snapshots are not
            saved to a file. If specified, snapshots are buffered and written to
            the file every `buffer_for_writing` samples. The written snapshots are
            discarded from the container to keep memory usage low.
        buffer_for_writing : int
            Number of samples to buffer before writing to the file. Ignored if
            `file` is None.

        thermostat_setting : ThermostatSetting
            Setting for the ASE style thermostat used in the MD simulation.

        log_md : bool or int
            Whether to log the MD simulation using ASE's MDLogger. If True,
            logging is done every `mdlogger_setting.interval` steps. If False,
            no logging is done. If an integer, logging is done every that many
            steps.
        mdlogger_setting : MDLogSetting
            Setting for the MDLogger used to log the MD simulation.

        attach : Callable or list of Callable
            Function(s) to attach to the MD simulation, passed through to the
            `attach` method of ASE's MD class.

    Notes
    -----
        If `file` is specified, but the sampled snapshots are desired after the run,
        the file needs to be read with `SnapshotContainer.read`.

    """

    seed: int | None = None
    rng: np.random.Generator = None

    sampling_offset: int = 1000
    sampling_interval: int = 100

    sample_atoms: bool = True
    file: str | pathlib.Path | None = None
    buffer_for_writing: int = 10

    thermostat_setting: ThermostatSetting = field(default_factory=ThermostatSetting)

    log_md: bool | int = False
    mdlogger_setting: MDLogSetting = field(default_factory=MDLogSetting)

    attach: Callable | list[Callable] = field(default_factory=list)

    def __post_init__(self) -> None:
        """Post-initialization processing for SamplingSetting."""
        if isinstance(self.thermostat_setting, dict):
            self.thermostat_setting = ThermostatSetting(**self.thermostat_setting)
        if isinstance(self.mdlogger_setting, dict):
            self.mdlogger_setting = MDLogSetting(**self.mdlogger_setting)

        if self.seed is None:
            self.seed = np.random.SeedSequence().entropy % 2**32

        if self.rng is None:
            self.rng = np.random.default_rng(self.seed)

        if not isinstance(self.attach, list):
            self.attach = [self.attach]



    @classmethod
    def from_any(
        cls,
        dct: dict | SamplingSetting | None = None,
        **kwargs: dict,
    ) -> SamplingSetting:
        """Create a SamplingSetting from a dict, an existing instance, or None.

        Parameters
        ----------
        dct : dict, SamplingSetting, or None
            Source settings. Nested dicts for `thermostat_setting` and
            `mdlogger_setting` are automatically cast to their dataclass types.
            If already a SamplingSetting, returned as-is.
        \*\*kwargs
            Additional keyword arguments merged into `dct` (ignored when `dct`
            is already a SamplingSetting instance).

        Returns
        -------
        SamplingSetting
             The created SamplingSetting instance.

        """
        if isinstance(dct, cls):
            return dct

        if dct is None:
            dct = {}

        dct = dct.copy()
        dct.update(kwargs)
        return cls(**dct)