Source code for mchammer.data_containers.data_container

""" Data container class. """

import numpy as np
from .base_data_container import BaseDataContainer
from ..data_analysis import analyze_data




class DataContainer(BaseDataContainer):
    """
    Data container for storing information concerned with
    Monte Carlo simulations performed with :program:`mchammer`.

    This is the data container the thermodynamic ensembles write.
    On top of the storage and access the base class provides, it analyzes
    scalar observables through :func:`get_average` and :func:`analyze_data`,
    which estimates the correlation length and the statistical error.
    The section :ref:`data_container` describes the container in more detail.

    Parameters
    ----------
    structure
        Reference atomic structure associated with the data container.
    ensemble_parameters
        Parameters associated with the underlying ensemble.
    metadata
        Metadata associated with the data container.

    Example
    -------
    The following snippet runs a short simulation in the canonical ensemble and
    analyzes the potential along the trajectory.
    The parameters of the cluster expansion are made up to keep the example
    self-contained::

        >>> from ase.build import bulk
        >>> from icet import ClusterExpansion, ClusterSpace
        >>> from mchammer.calculators import ClusterExpansionCalculator
        >>> from mchammer.ensembles import CanonicalEnsemble

        >>> # prepare cluster expansion
        >>> prim = bulk('Au')
        >>> cs = ClusterSpace(prim, cutoffs=[4.3], chemical_symbols=['Ag', 'Au'])
        >>> ce = ClusterExpansion(cs, [0, 0, 0.1, -0.02])

        >>> # prepare initial configuration
        >>> structure = prim.repeat(3)
        >>> for k in range(5):
        ...     structure[k].symbol = 'Ag'

        >>> # set up and run the simulation
        >>> calculator = ClusterExpansionCalculator(structure, ce)
        >>> mc = CanonicalEnsemble(structure=structure, calculator=calculator,
        ...                        temperature=600)
        >>> mc.run(len(structure) * 50)

        >>> # analyze the potential, skipping the first steps as equilibration
        >>> dc = mc.data_container
        >>> print(dc.data)
        >>> mean = dc.get_average('potential', start=len(structure) * 10)
        >>> summary = dc.analyze_data('potential', start=len(structure) * 10)
        >>> print(summary['error_estimate'])
    """



    def analyze_data(self, tag: str, start: int | None = None, max_lag: int | None = None) -> dict:
        """
        Returns a statistical analysis of a scalar observable.

        Parameters
        ----------
        tag
            Name of the observable to analyze.
        start
            Smallest trial step to include.
            By default all records are included.
        max_lag
            Maximum lag between two points in the data series when computing
            the autocorrelation function.
            By default the length of the data series minus one.

        Returns
        -------
        dict
            The mean, the standard deviation, the correlation length in trial
            steps, and the error estimate of the mean at 95 % confidence, under
            the keys ``mean``, ``std``, ``correlation_length``, and
            ``error_estimate``.
            The correlation length and the error estimate are ``nan`` when the
            autocorrelation function does not decay within the data series.

        Raises
        ------
        ValueError
            If the observable is not in the data container.
        ValueError
            If the observable is not scalar.
        ValueError
            If the observations are not evenly spaced in trial steps.
        """

        # get data for tag
        if tag in ['trajectory', 'occupations']:
            raise ValueError('{} is not scalar'.format(tag))
        steps, data = self.get('mctrial', tag, start=start)

        # check that steps are evenly spaced
        diff = np.diff(steps)
        step_length = diff[0]
        if not np.allclose(step_length, diff):
            raise ValueError('data records must be evenly spaced.')

        summary = analyze_data(data, max_lag=max_lag)
        summary['correlation_length'] *= step_length  # in mc-trials
        return summary




    def get_average(self, tag: str, start: int | None = None) -> float:
        """
        Returns the average of a scalar observable.

        Parameters
        ----------
        tag
            Name of the observable to average.
        start
            Smallest trial step to include.
            By default all records are included.

        Raises
        ------
        ValueError
            If the observable is not in the data container.
        ValueError
            If the observable is not scalar.
        """
        if tag in ['trajectory', 'occupations']:
            raise ValueError('{} is not scalar'.format(tag))
        data = self.get(tag, start=start)
        return np.mean(data)