Source code for mchammer.data_analysis

import numpy as np
import pandas as pd

from ase.data import chemical_symbols as ase_chemical_symbols
from scipy import stats




def analyze_data(data: np.ndarray, max_lag: int | None = None) -> dict:
    """
    Returns a statistical analysis of a data series.

    Parameters
    ----------
    data
        Data series to analyze.
    max_lag
        Maximum lag between two data points used for computing the
        autocorrelation function.
        By default the length of the data series minus one.

    Returns
    -------
    dict
        The mean, the standard deviation, the correlation length, 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.
    """
    summary = dict(mean=data.mean(),
                   std=data.std())
    acf = get_autocorrelation_function(data, max_lag)
    correlation_length = _estimate_correlation_length_from_acf(acf)
    if correlation_length is not None:
        error_estimate = _estimate_error(data, correlation_length, confidence=0.95)
        summary['correlation_length'] = correlation_length
        summary['error_estimate'] = error_estimate
    else:
        summary['correlation_length'] = np.nan
        summary['error_estimate'] = np.nan
    return summary





def get_autocorrelation_function(data: np.ndarray, max_lag: int | None = None) -> np.ndarray:
    """
    Returns the autocorrelation function of a data series.

    The autocorrelation function is computed using `pandas.Series.autocorr
    <https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.Series.autocorr.html>`_.

    Parameters
    ----------
    data
        Data series for which to compute the autocorrelation function.
    max_lag
        Maximum lag between two data points.
        By default the length of the data series minus one.

    Returns
    -------
    numpy.ndarray
        The autocorrelation for each lag from zero to :attr:`max_lag` minus one.

    Raises
    ------
    ValueError
        If :attr:`max_lag` is smaller than one or not smaller than the length
        of the data series.
    """
    if max_lag is None:
        max_lag = len(data) - 1
    if max_lag < 1 or max_lag >= len(data):
        raise ValueError('max_lag should be between 1 and len(data)-1.')
    series = pd.Series(data)
    acf = [series.autocorr(lag) for lag in range(0, max_lag)]
    return np.array(acf)





def get_correlation_length(data: np.ndarray) -> int | None:
    r"""
    Returns an estimate of the correlation length of a data series.

    The correlation length is the first lag at which the autocorrelation
    function drops below :math:`\exp(-2)`.
    If the autocorrelation function never drops below that value within the
    data series, the function returns ``None``.

    Parameters
    ----------
    data
        Data series for which to estimate the correlation length.
    """

    acf = get_autocorrelation_function(data)
    correlation_length = _estimate_correlation_length_from_acf(acf)
    if correlation_length is None:
        return None
    return correlation_length





def get_error_estimate(data: np.ndarray, confidence: float = 0.95) -> float | None:
    r"""
    Returns an estimate of the standard error of the mean of a data series at
    the given confidence level via

    .. math::

       \mathrm{error} = t_\mathrm{factor} * \mathrm{std}(\mathrm{data}) / \sqrt{N_s}

    where :math:`t_\mathrm{factor}` is the factor corresponding to the
    confidence level and :math:`N_s` is the number of independent measurements,
    which is the length of the data series divided by the correlation length.

    If the correlation length cannot be estimated because the autocorrelation
    function does not decay within the data series, the function returns
    ``None``.

    Parameters
    ----------
    data
        Data series for which to estimate the error.
    confidence
        Confidence level of the error estimate.
        By default 0.95.
    """
    correlation_length = get_correlation_length(data)
    if correlation_length is None:
        return None
    error_estimate = _estimate_error(data, correlation_length, confidence)
    return error_estimate



def _estimate_correlation_length_from_acf(acf: np.ndarray) -> int | None:
    r"""
    Returns the first lag at which the autocorrelation function drops below
    :math:`\exp(-2)`, or ``None`` if it never does.
    """
    for i, a in enumerate(acf):
        if a < np.exp(-2):
            return i
    return None  # np.nan


def _estimate_error(data: np.ndarray,
                    correlation_length: int,
                    confidence: float) -> float:
    """ Returns the error estimate of the mean given the correlation length. """
    t_factor: float = stats.t.ppf((1 + confidence) / 2, len(data) - 1)
    error: float = t_factor * np.std(data) / np.sqrt(len(data) / correlation_length)
    return error




def get_tracer_diffusivity(data_container, species: str, dimensions: int = 3,
                           start: float = 0.0) -> float:
    """Returns the tracer diffusion coefficient of a species.

    The coefficient is obtained from a linear fit of the mean squared
    displacement against time, divided by :math:`2 d`.
    The units follow from the stored data, meaning Ų/s when positions
    carry Å and time carries seconds.

    Parameters
    ----------
    data_container
        Data container of a kinetic Monte Carlo run.
    species
        Chemical symbol of the species.
    dimensions
        Dimension of the diffusion, :math:`d` above.
    start
        Time before which rows are discarded, to exclude the ballistic or
        transient regime from the fit.
    """
    times, msd = data_container.get_mean_squared_displacement(species)
    mask = times >= start
    if mask.sum() < 2:
        raise ValueError('Too few rows for a diffusivity fit')
    slope = np.polyfit(times[mask], msd[mask], 1)[0]
    return float(slope / (2 * dimensions))





def get_correlation_factor(data_container, species: str,
                           hop_distance: float) -> float:
    """Returns the correlation factor of a species.

    The correlation factor is the mean squared displacement per particle
    divided by the product of the number of hops per particle and the
    squared hop distance.

    Parameters
    ----------
    data_container
        Data container of a kinetic Monte Carlo run.
    species
        Chemical symbol of the species.
    hop_distance
        Distance of a single hop.
    """
    _, msd = data_container.get_mean_squared_displacement(species)
    hops = data_container.get_hop_counts().get(species)
    if hops is None or hops == 0:
        raise ValueError(f'No hops recorded for {species}')
    # Particle species are fixed for the lifetime of a run and equal the
    # initial occupations, which the container stores as its reference
    # structure.
    number = ase_chemical_symbols.index(species)
    structure = data_container.structure
    n_particles = int(np.sum(structure.get_atomic_numbers() == number))
    if n_particles == 0:
        raise ValueError(f'No particles of species {species}')
    hops_per_particle = hops / n_particles
    return float(msd[-1] / (hops_per_particle * hop_distance ** 2))