from icet import ClusterSpace
from icet.core.sublattices import Sublattices
import numpy as np
from scipy.integrate import cumulative_trapezoid
from ase import Atoms
from ase.units import kB
from mchammer import DataContainer
import operator
from collections import Counter
from math import factorial
from functools import reduce
def _lambda_function_forward(n_steps: int, step: int) -> float:
    """
    Returns the current lambda value for a backward calculation.
    Parameters
    ----------
    n_steps
        Total number of steps..
    step
        Current step.
    """
    x = (step + 1) / (n_steps)
    lam = x**5*(70*x**4 - 315*x**3 + 540*x**2 - 420*x + 126)
    # Due to numerical precision lambda may be slightly outside the interval [0.0, 1.0]
    lam = np.clip(lam, 0.0, 1.0)
    return lam
def _lambda_function_backward(n_steps: int, step: int) -> float:
    """
    Returns the current lambda value for a backward calculation.
    Parameters
    ----------
    n_steps
        Total number of steps.
    step
        Current step.
    """
    return 1 - _lambda_function_forward(n_steps, step)
def _stirling(n: int) -> float:
    """ Stirling approximation of log(n!). """
    return n * np.log(n) - n + 0.5 * np.log(2*np.pi*n)
def _lognpermutations(array: list[int]) -> float:
    """If _npermutations becomes too large we resort to calculating the
    logarithm directly using Stirling's approximation, i.e., we
    calculate
    .. math::
        log(n!) - (log(a1!) + log(a2!) + log(a3!) + ... + log(ak!)
    where we denote in the code
    .. math::
        v  n = log(n!)
        den = (log(a1!) + log(a2!) + log(a3!) + ... + log(ak!)
    """
    n = _stirling(len(array))
    # Gets the number of each unique atom type
    mults = Counter(array).values()
    den = reduce(operator.add,
                 (_stirling(v) for v in mults),
                 1)
    return n - den
def _npermutations(array: list[int]) -> float:
    """
    Returns the total number of ways we can permutate the array which is given by
    .. math::
        n! / (a1! * a2! * a3! * ... * ak!)
    where we denote in the code
    .. math::
        den = (a1! * a2! * a3! * ... * ak!)
    """
    n = factorial(len(array))
    # Gets the number of each unique atom type
    a = Counter(array).values()
    den = reduce(operator.mul,
                 (factorial(v) for v in a),
                 1)
    return n / den
def _ideal_mixing_entropy(swap_sublattice_probabilities: list[int],
                          atoms_on_sublattices: np.ndarray,
                          boltzmann_constant: float) -> float:
    """
    Calculates the ideal mixing entropy for the supercell.
    Parameters
    ----------
    swap_sublattice_probabilites
        Sublattices that are active during the MC simulation.
    atoms_on_sublattices
        Sorted atomic numbers in sublattices lists.
    boltzmann_constant
        Value of Boltzmann constant in the units used in the MC simulation.
    """
    log_multiplicity = []
    sublattices = np.array(swap_sublattice_probabilities) > 0
    for aos in atoms_on_sublattices[sublattices]:
        try:
            log_multiplicity.append(np.log(_npermutations(aos)))
        except Exception:
            log_multiplicity.append(_lognpermutations(aos))
    return boltzmann_constant * np.sum(log_multiplicity)
def _get_atoms_on_sublattice(structure: Atoms,
                             sublattices: Sublattices) -> np.ndarray:
    """Sorts the atom symbols in the structure to lists involving
    specific sublattices, for example,
    [[5, 5, 3, 3, 3], [11, 11, 10, 10, 10 10, 10]]
    Parameters
    ----------
    structure
        Atomic structure.
    sublattices
        Sublattices for the supercell obtained from the cluster space.
    """
    _atoms_on_sublattices = []
    for sl in sublattices:
        atoms_on_sublattice = []
        for an in sl.atomic_numbers:
            natoms = np.where(structure.numbers == an)[0].size
            atoms_on_sublattice.extend([an] * natoms)
        _atoms_on_sublattices.append(atoms_on_sublattice)
    _atoms_on_sublattices = np.array(_atoms_on_sublattices,
                                     dtype=object)
    return _atoms_on_sublattices
[docs]
def get_free_energy_thermodynamic_integration(
        dc: DataContainer,
        cluster_space: ClusterSpace,
        forward: bool,
        max_temperature: float = np.inf,
        sublattice_probabilities: list[float] = None,
        boltzmann_constant: float = kB,
) -> (np.ndarray, np.ndarray):
    r"""Returns the free energy calculated via thermodynamic integration
    using the :class:`ThermodynamicIntegrationEnsemble
    <mchammer.ensembles.ThermodynamicIntegrationEnsemble>`.
    The temperature dependence of the free energy can be extracted
    from the thermodynamic integration as
    .. math::
        F(T) =  \frac{F_0(\lambda)}{\lambda} + \frac{T_0}{\lambda} S_\text{B}
    where :math:`S_\text{B}` is the Boltzmann entropy,
    .. math::
        T = \frac{T_0}{\lambda}
    and
    .. math::
        F_0(\lambda) = \int_0^\lambda \left\langle\frac{\mathrm{d}H(\lambda)}
                       {\mathrm{d}\lambda}\right\rangle_{H} \mathrm{d}\lambda
    Parameters
    ----------
    dc
        Data container from the thermodynamic integration simulation.
    cluster_space
        Cluster space used to construct the cluster expansion used for the simulation.
    forward
        Whether or not the thermodynamic integration was carried out forward or backward.
    max_temperature
        Largest temperature to extract from the thermodynamic integration.
    sublattice_probababilites
        Sublattice probabilties that were provided to the thermodynamic integration
        simulation.
    boltzmann_constant
        Boltzmann constant in the units used for the thermodynamic integration
        simulation.
    """
    lambdas, potentials = dc.get('lambda', 'potential')
    if not forward:
        # We want to integrate from high to low temperature even though
        # the simulation was from low to high.
        potentials = potentials[::-1]
        lambdas = lambdas[::-1]
    sublattices = cluster_space.get_sublattices(dc.structure)
    if sublattice_probabilities is None:
        sublattice_probabilities = [True] * len(sublattices)
    atoms_on_sublattices = _get_atoms_on_sublattice(dc.structure,
                                                    sublattices)
    ideal_mixing_entropy = _ideal_mixing_entropy(sublattice_probabilities,
                                                 atoms_on_sublattices,
                                                 boltzmann_constant)
    temperature = dc._ensemble_parameters['temperature']
    with np.errstate(divide='ignore'):
        # First value of lambdas will be zero (we ignore this warning)
        temperatures = (1 / lambdas) * temperature
    temp_inds = np.where(temperatures <= max_temperature)[0]
    free_energy_change = cumulative_trapezoid(potentials, x=lambdas, initial=0)
    with np.errstate(invalid='ignore', divide='ignore'):
        # First value of lambdas will be zero (we ignore this warning)
        free_energy = (free_energy_change / lambdas -
                       temperatures * ideal_mixing_entropy)
    free_energy[np.isnan(free_energy)] = 0
    return (temperatures[temp_inds], free_energy[temp_inds]) 
[docs]
def get_free_energy_temperature_integration(
        dc: DataContainer,
        cluster_space: ClusterSpace,
        forward: bool,
        temperature_reference: float,
        free_energy_reference: float = None,
        sublattice_probabilities: list[float] = None,
        max_temperature: float = np.inf,
        boltzmann_constant: float = kB,
) -> (np.ndarray, np.ndarray):
    r""" Returns the free energy calculated using temperature integration and
    the corresponding temperature.
    .. math::
        \frac{A(T_{2})}{T_{2}} = \frac{A(T_{1})}{T_{1}}
        - \int_{T_{1}}^{T_{2}}\frac{U(T)}{T^2}\mathrm{d}T
    Parameters
    ----------
    dc
        Data container from a canonical annealing simulation.
        The first (last for :attr:`forward=False`) temperature in the
        data container has to be at the same temperature as
        :attr:`temperature_reference`.
    cluster_space
        Cluster space used to construct the cluster expansion
        that was used in the simulations.
    forward
        If ``True`` the canonical annealing simulation was carried out
        from high to low temperature, otherwise the opposite is assumed.
    temperature_reference
        Temperature at which :attr:`free_energy_reference` was calculated
    free_energy_reference
        Reference free energy. If set to ``None`` it will be assumeed that the
        free energy at :attr:`temperature_reference` can be approximated by
        :math:`T S_B` where :math:`S_B` is the ideal mixing entropy.
    sublattice_probababilites
        Sublattice probabilties that were provided to the canonical annealing
        simulation.
    max_temperature
        Largest temperature to extract from the temperature integration.
    boltzmann_constant
        Boltzmann constant in the units used for the thermodynamic integration
        simulation.
    """
    temperatures, potentials = dc.get('temperature', 'potential')
    if not forward:
        # We want to integrate from high to low temperature even though
        # the simulation was from low to high.
        potentials = potentials[::-1]
        temperatures = temperatures[::-1]
    if not np.isclose(temperatures[0], temperature_reference):
        raise ValueError(
            'The first or the last temperature in the temperature list'
            ' has to equal the reference temperature.'
            f' reference_temperature: {temperature_reference}'
            f' first canonical temperature: {temperatures[0]}')
    temp_inds = np.where(temperatures <= max_temperature)[0]
    if free_energy_reference is None:
        sublattices = cluster_space.get_sublattices(dc.structure)
        if sublattice_probabilities is None:
            sublattice_probabilities = [True] * len(sublattices)
        atoms_on_sublattices = _get_atoms_on_sublattice(dc.structure,
                                                        sublattices)
        ideal_mixing_entropy = _ideal_mixing_entropy(sublattice_probabilities,
                                                     atoms_on_sublattices,
                                                     boltzmann_constant)
        free_energy_reference = -ideal_mixing_entropy * temperature_reference
    reference = free_energy_reference / temperature_reference
    integral = cumulative_trapezoid(potentials / temperatures**2,
                                    x=temperatures, initial=0)
    free_energy = (reference - integral) * temperatures
    return (temperatures[temp_inds], free_energy[temp_inds])