eos.py

# coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.


"""
This module implements various equation of states.

Note: Most of the code were initially adapted from ASE and deltafactor by
@gmatteo but has since undergone major refactoring.
"""

from copy import deepcopy
from abc import ABCMeta, abstractmethod
import logging
import warnings

import numpy as np
from scipy.optimize import leastsq, minimize

from pymatgen.core.units import FloatWithUnit
from pymatgen.util.plotting import pretty_plot, add_fig_kwargs, get_ax_fig_plt

__author__ = "Kiran Mathew, gmatteo"
__credits__ = "Cormac Toher"

logger = logging.getLogger(__file__)


class EOSBase(metaclass=ABCMeta):
    """
    Abstract class that must be subcalssed by all equation of state
    implementations.
    """

    def __init__(self, volumes, energies):
        """
         Args:
             volumes (list/numpy.array): volumes in Ang^3
             energies (list/numpy.array): energy in eV
        """
        self.volumes = np.array(volumes)
        self.energies = np.array(energies)
        # minimum energy(e0), buk modulus(b0),
        # derivative of bulk modulus wrt pressure(b1), minimum volume(v0)
        self._params = None
        # the eos function parameters. It is the same as _params except for
        # equation of states that uses polynomial fits(deltafactor and
        # numerical_eos)
        self.eos_params = None

    def _initial_guess(self):
        """
        Quadratic fit to get an initial guess for the parameters.

        Returns:
            tuple: (e0, b0, b1, v0)
        """
        a, b, c = np.polyfit(self.volumes, self.energies, 2)
        self.eos_params = [a, b, c]

        v0 = -b / (2 * a)
        e0 = a * (v0 ** 2) + b * v0 + c
        b0 = 2 * a * v0
        b1 = 4  # b1 is usually a small number like 4

        vmin, vmax = min(self.volumes), max(self.volumes)

        if not vmin < v0 and v0 < vmax:
            raise EOSError('The minimum volume of a fitted parabola is '
                           'not in the input volumes\n.')

        return e0, b0, b1, v0

    def fit(self):
        """
        Do the fitting. Does least square fitting. If you want to use custom
        fitting, must override this.
        """
        # the objective function that will be minimized in the least square
        # fitting
        self._params = self._initial_guess()
        self.eos_params, ierr = leastsq(lambda pars, x, y: y - self._func(x, pars),
                                        self._params, args=(self.volumes, self.energies))
        # e0, b0, b1, v0
        self._params = self.eos_params
        if ierr not in [1, 2, 3, 4]:
            raise EOSError("Optimal parameters not found")

    @abstractmethod
    def _func(self, volume, params):
        """
        The equation of state function. This must be implemented by all classes
        that derive from this abstract class.

        Args:
            volume (float/numpy.array)
             params (list/tuple): values for the parameters other than the
                volume used by the eos.
        """
        pass

    def func(self, volume):
        """
        The equation of state function with the paramters other than volume set
        to the ones obtained from fitting.

        Args:
             volume (list/numpy.array)

        Returns:
            numpy.array
        """
        return self._func(np.array(volume), self.eos_params)

    def __call__(self, volume):
        return self.func(volume)

    @property
    def e0(self):
        """
        Returns the min energy.
        """
        return self._params[0]

    @property
    def b0(self):
        """
        Returns the bulk modulus.
        Note: the units for the bulk modulus: unit of energy/unit of volume^3.
        """
        return self._params[1]

    @property
    def b0_GPa(self):
        """
        Returns the bulk modulus in GPa.
        Note: This assumes that the energy and volumes are in eV and Ang^3
            respectively
        """
        return FloatWithUnit(self.b0, "eV ang^-3").to("GPa")

    @property
    def b1(self):
        """
        Returns the derivative of bulk modulus wrt pressure(dimensionless)
        """
        return self._params[2]

    @property
    def v0(self):
        """
        Returns the minimum or the reference volume in Ang^3.
        """
        return self._params[3]

    @property
    def results(self):
        """
        Returns a summary dict.

        Returns:
            dict
        """
        return dict(e0=self.e0, b0=self.b0, b1=self.b1, v0=self.v0)

    def plot(self, width=8, height=None, plt=None, dpi=None, **kwargs):
        """
        Plot the equation of state.

        Args:
            width (float): Width of plot in inches. Defaults to 8in.
            height (float): Height of plot in inches. Defaults to width *
                golden ratio.
            plt (matplotlib.pyplot): If plt is supplied, changes will be made
                to an existing plot. Otherwise, a new plot will be created.
            dpi:
            kwargs (dict): additional args fed to pyplot.plot.
                supported keys: style, color, text, label

        Returns:
            Matplotlib plot object.
        """
        plt = pretty_plot(width=width, height=height, plt=plt, dpi=dpi)

        color = kwargs.get("color", "r")
        label = kwargs.get("label", "{} fit".format(self.__class__.__name__))
        lines = ["Equation of State: %s" % self.__class__.__name__,
                 "Minimum energy = %1.2f eV" % self.e0,
                 "Minimum or reference volume = %1.2f Ang^3" % self.v0,
                 "Bulk modulus = %1.2f eV/Ang^3 = %1.2f GPa" %
                 (self.b0, self.b0_GPa),
                 "Derivative of bulk modulus wrt pressure = %1.2f" % self.b1]
        text = "\n".join(lines)
        text = kwargs.get("text", text)

        # Plot input data.
        plt.plot(self.volumes, self.energies, linestyle="None", marker="o",
                 color=color)

        # Plot eos fit.
        vmin, vmax = min(self.volumes), max(self.volumes)
        vmin, vmax = (vmin - 0.01 * abs(vmin), vmax + 0.01 * abs(vmax))
        vfit = np.linspace(vmin, vmax, 100)

        plt.plot(vfit, self.func(vfit), linestyle="dashed", color=color,
                 label=label)

        plt.grid(True)
        plt.xlabel("Volume $\\AA^3$")
        plt.ylabel("Energy (eV)")
        plt.legend(loc="best", shadow=True)
        # Add text with fit parameters.
        plt.text(0.4, 0.5, text, transform=plt.gca().transAxes)

        return plt

    @add_fig_kwargs
    def plot_ax(self, ax=None, fontsize=12, **kwargs):
        """
        Plot the equation of state on axis `ax`

        Args:
            ax: matplotlib :class:`Axes` or None if a new figure should be created.
            fontsize: Legend fontsize.
            color (str): plot color.
            label (str): Plot label
            text (str): Legend text (options)

        Returns:
            Matplotlib figure object.
        """
        ax, fig, plt = get_ax_fig_plt(ax=ax)

        color = kwargs.get("color", "r")
        label = kwargs.get("label", "{} fit".format(self.__class__.__name__))
        lines = ["Equation of State: %s" % self.__class__.__name__,
                 "Minimum energy = %1.2f eV" % self.e0,
                 "Minimum or reference volume = %1.2f Ang^3" % self.v0,
                 "Bulk modulus = %1.2f eV/Ang^3 = %1.2f GPa" %
                 (self.b0, self.b0_GPa),
                 "Derivative of bulk modulus wrt pressure = %1.2f" % self.b1]
        text = "\n".join(lines)
        text = kwargs.get("text", text)

        # Plot input data.
        ax.plot(self.volumes, self.energies, linestyle="None", marker="o", color=color)

        # Plot eos fit.
        vmin, vmax = min(self.volumes), max(self.volumes)
        vmin, vmax = (vmin - 0.01 * abs(vmin), vmax + 0.01 * abs(vmax))
        vfit = np.linspace(vmin, vmax, 100)

        ax.plot(vfit, self.func(vfit), linestyle="dashed", color=color, label=label)

        ax.grid(True)
        ax.set_xlabel("Volume $\\AA^3$")
        ax.set_ylabel("Energy (eV)")
        ax.legend(loc="best", shadow=True)
        # Add text with fit parameters.
        ax.text(0.5, 0.5, text, fontsize=fontsize, horizontalalignment='center',
                verticalalignment='center', transform=ax.transAxes)

        return fig


class Murnaghan(EOSBase):

    def _func(self, volume, params):
        """
        From PRB 28,5480 (1983)
        """
        e0, b0, b1, v0 = tuple(params)
        return (e0 +
                b0 * volume / b1 * (((v0 / volume) ** b1) / (b1 - 1.0) + 1.0) -
                v0 * b0 / (b1 - 1.0))


class Birch(EOSBase):

    def _func(self, volume, params):
        """
        From Intermetallic compounds: Principles and Practice, Vol. I:
        Principles Chapter 9 pages 195-210 by M. Mehl. B. Klein,
        D. Papaconstantopoulos.
        case where n=0
        """
        e0, b0, b1, v0 = tuple(params)
        return (e0
                + 9.0 / 8.0 * b0 * v0 * ((v0 / volume) ** (2.0 / 3.0) - 1.0) ** 2
                + 9.0 / 16.0 * b0 * v0 * (b1 - 4.) *
                ((v0 / volume) ** (2.0 / 3.0) - 1.0) ** 3)


class BirchMurnaghan(EOSBase):

    def _func(self, volume, params):
        """
        BirchMurnaghan equation from PRB 70, 224107
        """
        e0, b0, b1, v0 = tuple(params)
        eta = (v0 / volume) ** (1. / 3.)
        return (e0 +
                9. * b0 * v0 / 16. * (eta ** 2 - 1) ** 2 *
                (6 + b1 * (eta ** 2 - 1.) - 4. * eta ** 2))


class PourierTarantola(EOSBase):

    def _func(self, volume, params):
        """
        Pourier-Tarantola equation from PRB 70, 224107
        """
        e0, b0, b1, v0 = tuple(params)
        eta = (volume / v0) ** (1. / 3.)
        squiggle = -3. * np.log(eta)
        return e0 + b0 * v0 * squiggle ** 2 / 6. * (3. + squiggle * (b1 - 2))


class Vinet(EOSBase):

    def _func(self, volume, params):
        """
        Vinet equation from PRB 70, 224107
        """
        e0, b0, b1, v0 = tuple(params)
        eta = (volume / v0) ** (1. / 3.)
        return (e0 + 2. * b0 * v0 / (b1 - 1.) ** 2
                * (2. - (5. + 3. * b1 * (eta - 1.) - 3. * eta)
                   * np.exp(-3. * (b1 - 1.) * (eta - 1.) / 2.)))


class PolynomialEOS(EOSBase):
    """
    Derives from EOSBase. Polynomial based equations of states must subclass
    this.
    """

    def _func(self, volume, params):
        return np.poly1d(list(params))(volume)

    def fit(self, order):
        """
        Do polynomial fitting and set the parameters. Uses numpy polyfit.

        Args:
             order (int): order of the fit polynomial
        """
        self.eos_params = np.polyfit(self.volumes, self.energies, order)
        self._set_params()

    def _set_params(self):
        """
        Use the fit polynomial to compute the parameter e0, b0, b1 and v0
        and set to the _params attribute.
        """
        fit_poly = np.poly1d(self.eos_params)
        # the volume at min energy, used as the intial guess for the
        # optimization wrt volume.
        v_e_min = self.volumes[np.argmin(self.energies)]
        # evaluate e0, v0, b0 and b1
        min_wrt_v = minimize(fit_poly, v_e_min)
        e0, v0 = min_wrt_v.fun, min_wrt_v.x[0]
        pderiv2 = np.polyder(fit_poly, 2)
        pderiv3 = np.polyder(fit_poly, 3)
        b0 = v0 * np.poly1d(pderiv2)(v0)
        db0dv = np.poly1d(pderiv2)(v0) + v0 * np.poly1d(pderiv3)(v0)
        # db/dp
        b1 = - v0 * db0dv / b0
        self._params = [e0, b0, b1, v0]


class DeltaFactor(PolynomialEOS):

    def _func(self, volume, params):
        x = volume ** (-2. / 3.)
        return np.poly1d(list(params))(x)

    def fit(self, order=3):
        """
        Overriden since this eos works with volume**(2/3) instead of volume.
        """
        x = self.volumes ** (-2. / 3.)
        self.eos_params = np.polyfit(x, self.energies, order)
        self._set_params()

    def _set_params(self):
        """
        Overriden to account for the fact the fit with volume**(2/3) instead
        of volume.
        """
        deriv0 = np.poly1d(self.eos_params)
        deriv1 = np.polyder(deriv0, 1)
        deriv2 = np.polyder(deriv1, 1)
        deriv3 = np.polyder(deriv2, 1)

        for x in np.roots(deriv1):
            if x > 0 and deriv2(x) > 0:
                v0 = x ** (-3. / 2.)
                break
        else:
            raise EOSError("No minimum could be found")

        derivV2 = 4. / 9. * x ** 5. * deriv2(x)
        derivV3 = (-20. / 9. * x ** (13. / 2.) * deriv2(x) - 8. / 27. *
                   x ** (15. / 2.) * deriv3(x))
        b0 = derivV2 / x ** (3. / 2.)
        b1 = -1 - x ** (-3. / 2.) * derivV3 / derivV2

        # e0, b0, b1, v0
        self._params = [deriv0(v0 ** (-2. / 3.)), b0, b1, v0]


class NumericalEOS(PolynomialEOS):

    def fit(self, min_ndata_factor=3, max_poly_order_factor=5, min_poly_order=2):
        """
        Fit the input data to the 'numerical eos', the equation of state employed
        in the quasiharmonic Debye model described in the paper:
        10.1103/PhysRevB.90.174107.

        credits: Cormac Toher

        Args:
            min_ndata_factor (int): parameter that controls the minimum number
                of data points that will be used for fitting.
                minimum number of data points =
                    total data points-2*min_ndata_factor
            max_poly_order_factor (int): parameter that limits the max order
                of the polynomial used for fitting.
                max_poly_order = number of data points used for fitting -
                                 max_poly_order_factor
            min_poly_order (int): minimum order of the polynomial to be
                considered for fitting.
        """
        warnings.simplefilter('ignore', np.RankWarning)

        def get_rms(x, y):
            return np.sqrt(np.sum((np.array(x) - np.array(y)) ** 2) / len(x))

        # list of (energy, volume) tuples
        e_v = [(i, j) for i, j in zip(self.energies, self.volumes)]
        ndata = len(e_v)
        # minimum number of data points used for fitting
        ndata_min = max(ndata - 2 * min_ndata_factor, min_poly_order + 1)
        rms_min = np.inf
        # number of data points available for fit in each iteration
        ndata_fit = ndata
        # store the fit polynomial coefficients and the rms in a dict,
        # where the key=(polynomial order, number of data points used for
        # fitting)
        all_coeffs = {}

        # sort by energy
        e_v = sorted(e_v, key=lambda x: x[0])
        # minimum energy tuple
        e_min = e_v[0]
        # sort by volume
        e_v = sorted(e_v, key=lambda x: x[1])
        # index of minimum energy tuple in the volume sorted list
        emin_idx = e_v.index(e_min)
        # the volume lower than the volume corresponding to minimum energy
        v_before = e_v[emin_idx - 1][1]
        # the volume higher than the volume corresponding to minimum energy
        v_after = e_v[emin_idx + 1][1]
        e_v_work = deepcopy(e_v)

        # loop over the data points.
        while (ndata_fit >= ndata_min) and (e_min in e_v_work):
            max_poly_order = ndata_fit - max_poly_order_factor
            e = [ei[0] for ei in e_v_work]
            v = [ei[1] for ei in e_v_work]
            # loop over polynomial order
            for i in range(min_poly_order, max_poly_order + 1):
                coeffs = np.polyfit(v, e, i)
                pder = np.polyder(coeffs)
                a = np.poly1d(pder)(v_before)
                b = np.poly1d(pder)(v_after)
                if a * b < 0:
                    rms = get_rms(e, np.poly1d(coeffs)(v))
                    rms_min = min(rms_min, rms * i / ndata_fit)
                    all_coeffs[(i, ndata_fit)] = [coeffs.tolist(), rms]
                    # store the fit coefficients small to large,
                    # i.e a0, a1, .. an
                    all_coeffs[(i, ndata_fit)][0].reverse()
            # remove 1 data point from each end.
            e_v_work.pop()
            e_v_work.pop(0)
            ndata_fit = len(e_v_work)

        logger.info("total number of polynomials: {}".format(len(all_coeffs)))

        norm = 0.
        fit_poly_order = ndata
        # weight average polynomial coefficients.
        weighted_avg_coeffs = np.zeros((fit_poly_order,))

        # combine all the filtered polynomial candidates to get the final fit.
        for k, v in all_coeffs.items():
            # weighted rms = rms * polynomial order / rms_min / ndata_fit
            weighted_rms = v[1] * k[0] / rms_min / k[1]
            weight = np.exp(-(weighted_rms ** 2))
            norm += weight
            coeffs = np.array(v[0])
            # pad the coefficient array with zeros
            coeffs = np.lib.pad(coeffs,
                                (0, max(fit_poly_order - len(coeffs), 0)),
                                'constant')
            weighted_avg_coeffs += weight * coeffs

        # normalization
        weighted_avg_coeffs /= norm
        weighted_avg_coeffs = weighted_avg_coeffs.tolist()
        # large to small(an, an-1, ..., a1, a0) as expected by np.poly1d
        weighted_avg_coeffs.reverse()

        self.eos_params = weighted_avg_coeffs
        self._set_params()


class EOS:
    """
    Convenient wrapper. Retained in its original state to ensure backward
    compatibility.

    Fit equation of state for bulk systems.

    The following equations are supported::

        murnaghan: PRB 28, 5480 (1983)

        birch: Intermetallic compounds: Principles and Practice, Vol I:
            Principles. pages 195-210

        birch_murnaghan: PRB 70, 224107

        pourier_tarantola: PRB 70, 224107

        vinet: PRB 70, 224107

        deltafactor

        numerical_eos: 10.1103/PhysRevB.90.174107.

    Usage::

       eos = EOS(eos_name='murnaghan')
       eos_fit = eos.fit(volumes, energies)
       eos_fit.plot()
    """

    MODELS = {
        "murnaghan": Murnaghan,
        "birch": Birch,
        "birch_murnaghan": BirchMurnaghan,
        "pourier_tarantola": PourierTarantola,
        "vinet": Vinet,
        "deltafactor": DeltaFactor,
        "numerical_eos": NumericalEOS
    }

    def __init__(self, eos_name='murnaghan'):
        if eos_name not in self.MODELS:
            raise EOSError("The equation of state '{}' is not supported. "
                           "Please choose one from the following list: {}".
                           format(eos_name, list(self.MODELS.keys())))
        self._eos_name = eos_name
        self.model = self.MODELS[eos_name]

    def fit(self, volumes, energies):
        """
        Fit energies as function of volumes.

        Args:
            volumes (list/np.array)
            energies (list/np.array)

        Returns:
            EOSBase: EOSBase object
        """
        eos_fit = self.model(np.array(volumes), np.array(energies))
        eos_fit.fit()
        return eos_fit


class EOSError(Exception):
    pass