Source code for mchammer.data_containers.kinetic_data_container

"""
This module provides the KineticDataContainer class.
"""

from typing import Any

from .base_data_container import BaseDataContainer




class KineticDataContainer(BaseDataContainer):
    """Data container for kinetic Monte Carlo simulations.

    In addition to the machinery of :class:`BaseDataContainer` this
    container persists the physical time and the particle tracking state,
    which cannot be reconstructed from stored occupations.
    A restart that lost either would silently corrupt every transport
    observable.

    The step counter :attr:`mctrial` counts hops, and the physical time in
    seconds is an ordinary column of the ensemble data.
    """

    def _update_last_state(self,
                           last_step: int,
                           occupations: list[int],
                           accepted_trials: int,
                           random_state: Any,
                           generator_state: Any,
                           time: float,
                           particle_ids: list[int],
                           displacements: list[list[float]],
                           hop_counts: dict,
                           clamp_count: int) -> None:
        """Updates the last state of the kinetic Monte Carlo simulation.

        Every field is required, since a restart reads all of them and a
        state that carries only some of them is of no use.

        Parameters
        ----------
        last_step
            Last step (hop counter).
        occupations
            Occupation vector observed during the last step.
        accepted_trials
            Number of accepted steps since the last write of ensemble
            data, from which the acceptance ratio of an interval is
            computed.
            Every kinetic Monte Carlo step is accepted, so the number
            equals the number of steps of the interval.
        random_state
            State of the generator the :mod:`random` module shares.
        generator_state
            State of the generator of the ensemble, which is what the
            sequence of a run is drawn from.
        time
            Physical time in seconds.
        particle_ids
            Particle identity at each site.
        displacements
            Accumulated unwrapped displacement of each particle.
        hop_counts
            Accumulated number of hops per chemical symbol.
        clamp_count
            Number of rate evaluations so far in which the kinetically
            resolved activation energy was raised to the clamp at half
            the magnitude of the energy difference.
        """
        super()._update_last_state(
            last_step=last_step,
            occupations=occupations,
            accepted_trials=accepted_trials,
            random_state=random_state)
        self._last_state['generator_state'] = generator_state
        self._last_state['time'] = time
        self._last_state['particle_ids'] = particle_ids
        self._last_state['displacements'] = displacements
        self._last_state['hop_counts'] = hop_counts
        self._last_state['clamp_count'] = clamp_count



    def get_mean_squared_displacement(self, species: str) -> tuple:
        """Returns time and mean squared displacement for a species.

        Parameters
        ----------
        species
            Chemical symbol of the species.
        """
        tag = f'msd_{species}'
        if tag not in self.data.columns:
            raise ValueError(f'No mean squared displacement data for {species}.'
                             f' Available columns: {list(self.data.columns)}')
        rows = self.data[['time', tag]].dropna()
        return rows['time'].to_numpy(), rows[tag].to_numpy()




    def get_hop_counts(self) -> dict:
        """Returns the accumulated number of hops per species."""
        counts = {}
        for column in self.data.columns:
            if column.startswith('hops_'):
                rows = self.data[column].dropna()
                if len(rows):
                    counts[column.removeprefix('hops_')] = int(rows.iloc[-1])
        return counts