Source code for mchammer.kinetics.local_cluster_expansion_calculator

"""
This module provides the LocalClusterExpansionCalculator class.
"""

import numpy as np

from ase import Atoms
from ase.data import atomic_numbers as ase_atomic_numbers

from icet.core.local_cluster_expansion import LocalClusterExpansion




class LocalClusterExpansionCalculator:
    """The per-supercell form of a local cluster expansion.

    The calculator locates every instance of the event prototype in the
    supercell at construction and builds a compiled evaluator from them,
    so that runtime evaluation reduces to a gather from a point-function
    table, a product over cluster sites, and a dot with the parameters.

    Every instance is one image of the prototype translated by a lattice
    vector, so the orbits are resolved once per image and an instance is a
    pair of an image and the cell it sits in.
    The setup and the memory of the calculator are therefore linear in the
    supercell with a small coefficient, and the term that grows with it is
    the table of supercell site indices the evaluator holds, one entry per
    cell and per lattice site the orbits reference.

    Parameters
    ----------
    structure
        Supercell for which to set up the calculator.
        Only the geometry of the structure enters the setup, occupations
        are supplied at evaluation time.
    local_cluster_expansion
        Local cluster expansion from which to build the calculator.
    """

    def __init__(self,
                 structure: Atoms,
                 local_cluster_expansion: LocalClusterExpansion) -> None:
        self._local_cluster_expansion = local_cluster_expansion
        self._structure = structure.copy()
        lcs = local_cluster_expansion.local_cluster_space
        parameters = np.array(local_cluster_expansion.parameters)
        self._migrating_number = ase_atomic_numbers[lcs.event.species]

        # The parameter vector splits into the local block, which the
        # compiled evaluator holds, and the augmentation block, which
        # multiplies global features that stay in Python.
        n_local = 1 + len(lcs.orbits)
        self._local_parameters = parameters[:n_local]
        self._augmentation_parameters = parameters[n_local:]
        self._evaluator, instances = lcs._build_evaluator(
            structure, self._local_parameters)
        self._instances = list(instances)
        self._augmentation = None
        if len(self._augmentation_parameters) > 0:
            self._augmentation = lcs._build_augmentation_evaluator(structure)

    @property
    def local_cluster_expansion(self) -> LocalClusterExpansion:
        """Local cluster expansion from which the calculator was set up."""
        return self._local_cluster_expansion

    @property
    def structure(self) -> Atoms:
        """Supercell for which the calculator was set up (copy)."""
        return self._structure.copy()

    @property
    def event_instances(self) -> list:
        """Undirected event instances as tuples of supercell site indices."""
        return list(self._instances)

    @property
    def migrating_number(self) -> int:
        """Atomic number of the migrating species."""
        return self._migrating_number



    def calculate_kra_many(self, instance_indices: np.ndarray,
                           occupations: np.ndarray) -> np.ndarray:
        """Returns the local properties of a batch of event instances.

        All instances are evaluated against the same occupation vector,
        one result per requested instance.
        Each value is the local part plus the augmentation contribution
        when the model carries one, consistent with
        :func:`LocalClusterExpansion.predict`.

        Parameters
        ----------
        instance_indices
            Indices of the instances in :attr:`event_instances`.
        occupations
            Occupation vector of the supercell by atomic number.
        """
        return self.calculate_kra_local_many(instance_indices, occupations) \
            + self.calculate_augmentation(occupations)




    def calculate_kra_local_many(self, instance_indices: np.ndarray,
                                 occupations: np.ndarray) -> np.ndarray:
        """Returns the local parts for a batch of event instances.

        The augmentation contribution is excluded, which is what a
        sampling loop needs when it accounts for the augmentation through
        a constant offset, see :func:`calculate_augmentation`.

        Parameters
        ----------
        instance_indices
            Indices of the instances in :attr:`event_instances`.
        occupations
            Occupation vector of the supercell by atomic number.
        """
        return self._evaluator.evaluate_many(instance_indices, occupations)




    def calculate_augmentation(self, occupations: np.ndarray) -> float:
        """Returns the augmentation contribution for one occupation vector.

        The contribution is zero for a model without augmentation.
        For a model with augmentation it depends only on quantities that
        are invariant during a kinetic Monte Carlo run whose hops conserve
        the occupation multiset of every sublattice group, which is what
        lets an ensemble evaluate it once and treat it as a constant
        offset per transition.

        Parameters
        ----------
        occupations
            Occupation vector of the supercell by atomic number.
        """
        if self._augmentation is None:
            return 0.0
        return float(np.dot(self._augmentation(np.asarray(occupations)),
                            self._augmentation_parameters))