Source code for directupsampling.fesurface

import pathlib
from typing import Callable

import numpy as np
import scipy.optimize

from directupsampling.eos import EOS, check_pressure_unit, read_eos
from directupsampling.fecontrib import FreeEnergyContrib
from directupsampling.io.io import read as _io_read
from directupsampling.io.io import write as _io_write




class FreeEnergySurface:
    """Class representing a free energy surface with contributions.

    Takes contributions as argument which should be a list of
    FreeEnergyContrib's.
    """

    directupsampling_objtype = "freeenergysurface"

    def __init__(
        self,
        contributions: list[FreeEnergyContrib | EOS | Callable] = [],
        bounds: list[list] | None = None,
    ):
        self.contributions = []
        self._bounds = bounds
        for contrib in contributions:
            self.add_contribution(contrib)



    def get_bounds_from_contributions(self):
        all_bounds = []
        for contrib in self.contributions:
            if isinstance(contrib, EOS):
                all_bounds.append([[0, np.inf], [0, np.inf]])
            else:
                all_bounds.append(contrib.bounds)
        bounds = [
            [max([_[0][0] for _ in all_bounds]), min([_[0][1] for _ in all_bounds])],
            [max([_[1][0] for _ in all_bounds]), min([_[1][1] for _ in all_bounds])],
        ]
        return bounds


    @property
    def bounds(self):
        if self._bounds is None:
            return self.get_bounds_from_contributions()
        else:
            return self._bounds

    @bounds.setter
    def bounds(self, bounds):
        if np.array(bounds).shape != (2, 2):
            raise ValueError("bounds must be 2 by 2")
        self._bounds = [list(_) for _ in bounds]

    def __repr__(self):
        return f"FreeEnergySurface({self.__dict__!r})"

    def __call__(self, volume, temperature=None):
        return self.get_free_energy(volume, temperature)

    def __iter__(self):
        yield from self.contributions

    def __len__(self):
        return len(self.contributions)

    def __add__(self, b):
        sum_fes = type(self)(self.contributions, self._bounds)
        try:
            sum_fes.add_contribution(b)
        except NotImplementedError:
            return NotImplemented
        return sum_fes

    def __radd__(self, a):
        return self.__add__(a)



    def add_contribution(self, contribution):
        if isinstance(contribution, FreeEnergySurface):
            for contrib in list(contribution):
                self.add_contribution(contrib)
        elif isinstance(contribution, FreeEnergyContrib):
            self.insert_contribution(-1, contribution)
        elif isinstance(contribution, EOS):
            self.insert_contribution(0, contribution)
        else:
            raise NotImplementedError




    def insert_contribution(self, index, contribution):
        if index < 0:
            index = len(self.contributions) + 1 + index
        if isinstance(contribution, EOS):
            exists_eos = False
            for c in self.contributions:
                if isinstance(c, EOS):
                    exists_eos = True
            if exists_eos:
                raise RuntimeError(
                    f"FreeEnergySurface already contains an EOS, "
                    f"cannot insert {contribution}"
                )
        self.contributions.insert(index, contribution)




    def get_eos(self) -> EOS:
        eos_list = [_ for _ in self.contributions if isinstance(_, EOS)]
        if len(eos_list) == 0:
            raise RuntimeError(f"{self} contains no EOS")
        return eos_list[0]




    def get_free_energy(self, volume, temperature=None):
        bshape = np.broadcast(volume, temperature).shape
        energy = np.zeros(bshape)
        for contrib in self.contributions:
            if isinstance(contrib, EOS):
                energy += contrib.get_energy(volume) * np.ones(bshape)
            else:
                energy += contrib.get_free_energy(volume, temperature)

        return energy




    def get_pressure(self, volume, temperature=None, unit="eV_per_Angstrom3"):
        """Calculates pressure by central differences.

        Parameters
        ----------
        volume : float
            The volume at which to calculate the pressure.
        temperature : float, optional
            The temperature at which to calculate the pressure. If the
            FreeEenrgySurface only contains an EOS contribution, the
            temperature is optional and ignored if given.
        unit : {"eV_per_Angstrom3", "GPa"}
            Specifies the unit of the returned pressure. The default is eV/Ang^3.

        Returns
        -------
        pressure : float
            Returns the calculated pressure.

        """

        unit_conversion = 1 / check_pressure_unit(unit)

        dv = 1e-4  # Ang3
        ef = self.get_free_energy(volume + dv, temperature)
        eb = self.get_free_energy(volume - dv, temperature)
        deriv = (ef - eb) / (2 * dv)
        pressure = -deriv * unit_conversion
        return pressure




    def get_volume(
        self,
        pressure: float,
        temperature: float,
        pressure_unit: str = "eV_per_Angstrom3",
    ) -> float:
        """Tries to find volume at p and T by optimization.

        Parameters
        ----------
        pressure, temperature: float
            The given conditions for which to optimize the volume.
        pressure_unit : {"eV_per_Angstrom3", "GPa"}
            The unit in which the input pressure is given. The default is eV/Ang^3.

        """
        try:
            self.get_eos()
        except RuntimeError:
            raise RuntimeError("No EOS found, volume optimization not allowed")

        unit_conversion = check_pressure_unit(pressure_unit)
        pressure *= unit_conversion
        temp = temperature

        def delta_p2(volume):
            return (self.get_pressure(volume, temp) - pressure) ** 2

        eos = self.get_eos()
        margin = np.array([0.95, 1.05])
        vmin, vmax = [eos.get_volume(pressure), eos.eq_volume] * margin
        opt_volume = optimize_global(delta_p2, [vmin, vmax])

        _check_pressure(self.get_pressure(opt_volume, temp), pressure, temp)

        return opt_volume




    def get_gibbs_energy(
        self,
        pressure: float,
        temperature: float | None = None,
        pressure_unit: str = "eV_per_Angstrom3",
    ):
        """Gives Gibbs energy at p and T by Legendre transform from FES.

        Parameters
        ----------
        pressure, temperature: float
            Conditions for which the Gibbs energy should be returned.
        pressure_unit: {"eV_per_Angstrom3", "GPa"}
            Unit in which the pressure is supplied. The default is eV/Ang^3

        """
        unit_conversion = check_pressure_unit(pressure_unit)
        pressure *= unit_conversion
        v = self.get_volume(pressure, temperature, pressure_unit="eV_per_Angstrom3")
        gibbs_energy = self(v, temperature) + pressure * v
        return gibbs_energy




    def todict(self):
        dct = self.__dict__.copy()
        dct["bounds"] = dct.pop("_bounds")
        return dct




    def write(self, file: str | pathlib.Path):
        """Writes the FreeEnergySurface instance to a JSON formatted file."""
        _io_write(self, file)




    @classmethod
    def read(cls, file: str | pathlib.Path):
        """Read a `.fes` file written with `free_energy_surface.write()`."""
        return _io_read(file)




    @classmethod
    def read_old_format(
        cls, files: list[str | pathlib.Path] | dict[str | pathlib.Path]
    ):
        """Reads and returns a FreeEnergySurface instance.

        Takes a list or dict of file paths for the respective parametrization.
        If a dict, then the keys should be contribution/component names, and the
        values file paths. If a list, then the names are determined from the
        file names (without preceding path).
        """
        if isinstance(files, list):
            contribs = [read_contrib_file(pathlib.Path(_)) for _ in files]
        elif isinstance(files, dict):
            contribs = [read_contrib_file(pathlib.Path(_)) for _ in files.values()]
            for c, n in zip(contribs, files.keys()):
                c.name = n
        else:
            raise TypeError(f"argument type must be list or dict, found {type(files)}")
        return cls(contribs)






def optimize_global(function, range):
    def finish_fun(fun, ini_guess, args):
        vrange = 0.02  # relative
        bbracket = ini_guess * np.array([1 - vrange, 1 + vrange])
        result = scipy.optimize.minimize_scalar(
            fun, bracket=bbracket, args=args, tol=1e-10
        )  # seems <100 Pa depending on BM
        return result

    result = scipy.optimize.brute(function, ranges=(range,), Ns=20, finish=finish_fun)
    return result  # [0]





def read_contrib_file(file: pathlib.Path):
    name = file.name
    if name.startswith("E"):
        contrib = read_eos(file)
    else:
        contrib = FreeEnergyContrib(name)
        contrib.read_old_format(file)

    return contrib



def _check_pressure(pressure, given_pressure, temperature):
    """Assert that the optimized volume gives the right pressure (+- 100 Pa)."""
    try:
        np.testing.assert_allclose(pressure, given_pressure, atol=1e-7, rtol=0)
    except AssertionError as e:
        raise RuntimeError(
            f"Volume optimization failed: could not find volume for p={given_pressure} "
            f"at T={temperature} K. Nearest found pressure: {pressure}"
        ) from e