Source code for mchammer.calculators.constituent_strain_calculator

from icet.tools import ConstituentStrain
from icet import ClusterExpansion
from mchammer.calculators import ClusterExpansionCalculator
import numpy as np




class ConstituentStrainCalculator(ClusterExpansionCalculator):
    """
    Calculator for a cluster expansion combined with constituent strain.

    The property this calculator evaluates is the sum of the cluster expansion
    energy and the constituent strain energy, both for the whole supercell.
    The constituent strain energy is evaluated by the :class:`ConstituentStrain
    <icet.tools.ConstituentStrain>` object, which describes the strain energy of
    a coherent structure in reciprocal space.
    The section on :ref:`constituent strain <constituent_strain_example>`
    describes the method and shows the calculator in a Monte Carlo simulation.

    The :class:`ConstituentStrain <icet.tools.ConstituentStrain>` object stores
    one structure factor per k-point and reads it as the value before a change
    instead of recomputing it.
    The calculator therefore rebuilds the structure factors in
    :func:`set_occupations` and advances them in :func:`accept_change`.
    Only changes of a single site can be evaluated, so the calculator works with
    ensembles that flip one site per trial step, such as
    :class:`SemiGrandCanonicalEnsemble
    <mchammer.ensembles.SemiGrandCanonicalEnsemble>` and
    :class:`VCSGCEnsemble <mchammer.ensembles.VCSGCEnsemble>`.

    Parameters
    ----------
    constituent_strain
        Object that defines the strain energy of the system.
        Its supercell is the structure the calculator is set up for, so it has
        to be the one the Monte Carlo simulation samples.
    cluster_expansion
        Cluster expansion from which to build the calculator.
    name
        Human-readable identifier for this calculator.
    scaling
        Scaling factor applied to the property value predicted by the
        cluster expansion.
        By default the number of sites in the supercell.
    """

    def __init__(self, constituent_strain: ConstituentStrain,
                 cluster_expansion: ClusterExpansion,
                 name: str = 'Constituent Strain Calculator',
                 scaling: float | None = None) -> None:
        self.constituent_strain = constituent_strain
        super().__init__(structure=constituent_strain.supercell,
                         cluster_expansion=cluster_expansion,
                         name=name,
                         scaling=scaling)



    def set_occupations(self, occupations: list[int]) -> None:
        """
        Sets the configuration this calculator describes, both for the
        underlying cluster expansion calculator and for the
        :class:`ConstituentStrain <icet.tools.ConstituentStrain>` object.

        Parameters
        ----------
        occupations
            The entire occupation vector by atomic number.
        """
        super().set_occupations(occupations)
        # The structure factor stored for each k-point describes one
        # configuration, and the change evaluation reads it as the value
        # before the change rather than recomputing it from the occupations
        # it is given. Moving to a configuration that is not reachable from
        # the stored one through accepted changes therefore has to rebuild it,
        # which a total evaluation does for every k-point.
        self.constituent_strain.get_constituent_strain(np.asarray(occupations))




    def calculate_total(self, *, occupations: np.ndarray) -> float:
        """
        Returns the total property value of a configuration, which is the sum
        of the cluster expansion energy and the constituent strain energy of the
        whole supercell.

        The configuration is given in full and is not adopted, so the
        configuration the calculator holds is left where it is.

        Parameters
        ----------
        occupations
            The entire occupation vector by atomic number.
        """
        e = super().calculate_total(occupations=occupations)
        # The structure factors are left where they are, since a total
        # evaluation scores the occupations it is given and must not move the
        # configuration the next change is evaluated against.
        e += len(occupations) * \
            self.constituent_strain.get_constituent_strain(
                occupations, update_structure_factors=False)
        return e




    def calculate_change(self, *, sites: list[int],
                         current_occupations: list[int],
                         new_site_occupations: list[int]) -> float:
        """
        Returns the change of the property caused by changing the occupation
        of the site in :attr:`sites`.

        Parameters
        ----------
        sites
            Index of the site whose occupation changes, as a list of length one.
        current_occupations
            The entire occupation vector by atomic number before the change.
        new_site_occupations
            Atomic number after the change, as a list of length one.

        Raises
        ------
        NotImplementedError
            If more than one site changes.
        """
        if len(new_site_occupations) > 1:
            raise NotImplementedError('Only single flips are currently allowed in '
                                      'conjunction with the constituent strain calculator.')
        e = super().calculate_change(sites=sites,
                                     current_occupations=current_occupations,
                                     new_site_occupations=new_site_occupations)
        de_cs = self.constituent_strain.get_constituent_strain_change(current_occupations,
                                                                      sites[0])
        e += len(current_occupations) * de_cs
        return e




    def accept_change(self, *, sites: list[int] | None = None,
                      species: list[int] | None = None) -> None:
        """
        Advances the configuration held by the underlying cluster expansion
        calculator and informs the :class:`ConstituentStrain
        <icet.tools.ConstituentStrain>` object that the most recent change was
        accepted, so that it stores the new structure factors.

        Parameters
        ----------
        sites
            Indices of the sites whose occupations changed.
        species
            New occupations by atomic number on those sites.
        """
        super().accept_change(sites=sites, species=species)
        # An empty change advances nothing, so the structure factors must not
        # advance either. Forwarding it would commit the candidate the last
        # change evaluation staged, which is the trial that was rejected.
        if len(sites) == 0:
            return
        self.constituent_strain.accept_change()