<reponame>kavanase/py-sc-fermi
import numpy as np
from pymatgen.io.vasp import Vasprun # type: ignore
from pymatgen.electronic_structure.core import Spin # type: ignore

class DOS(object):
    """Class for handling density-of-states data and its integration"""

    def __init__(self, dos, edos, egap, nelect, spin_polarised=False, normalise=True):
        """Initialise a DOS instance.

        Args:
            dos (np.array): Density of states data.
            edos (np.array): Energies for the density of states (in eV).
            egap (float): Width of the band gap (in eV).
            nelect (int): Number of electrons.
            spin_polarise(`bool`, optional): is the calculated DOS spin polarisised? Default
                is `False`
            normalise (:obj:`bool`, optional): Normalise the DOS so that the integral
                from e_min --> 0.0 is equal to the `nelect`. Default is `True`.

        Returns:
            None

        Raises:
            ValueError: If `egap` > `max(edos)`.

        """
        if egap > edos[-1]:
            raise ValueError('ERROR: Your gap extends beyond your maximum energy')
        self._dos = dos
        self._edos = edos
        self._egap = egap
        self._nelect = nelect 
        if normalise:
            self.normalise_dos()

    @property
    def dos(self):
        """Get the dos array."""
        return self._dos

    @property
    def edos(self):
        """Get the edos array."""
        return self._edos

    @property
    def egap(self):
        """Get the bandgap."""
        return self._egap

    @property
    def nelect(self):
        """Get the number of electrons in the unit cell."""
        return self._nelect

    def sum_dos(self):
        """get integrated density of states"""
        vbm_index = np.where(self._edos <= 0)[0][-1]
        sum1 = np.trapz(self._dos[:vbm_index+1], self._edos[:vbm_index+1])
        return sum1
    
    def normalise_dos(self):
        """normalise the density of states w.r.t. number of electrons (self.nelect)""" 
        integrated_dos = self.sum_dos()  
        self._dos = self._dos / integrated_dos * self._nelect 

    def emin(self):
        """get lowest energy in the dos data"""
        return self._edos[0]

    def emax(self):
        """get highest energy in the dos data"""
        return self._edos[-1]

    def carrier_concentrations(self, e_fermi, kT):
        """get n0 and p0 using integrals (equations 28.9 in Ashcroft Mermin)
           for an abitrary value of kT and Fermi energy. Typically used internally 
           to calculate carrier concentrations at the self-cosistent Fermi energy
        Args:
            efermi (float): Fermi energy
            kT (float): k * temperature

        Returns:
            p0 (float): concentration of holes
            n0 (float): concentration of electrons
        """
        p0_index = np.where(self._edos <= 0)[0][-1]
        n0_index = np.where(self._edos > self.egap)[0][0]
        p0 = np.trapz( p_func(e_fermi, self._dos[:p0_index+1], self._edos[:p0_index+1], kT ),
                       self._edos[:p0_index+1])
        n0 = np.trapz( n_func(e_fermi, self._dos[n0_index:], self._edos[n0_index:], kT ),
                       self._edos[n0_index:])
        return p0, n0

# Possibly better as class methods that integrate over a certain energy range.
def p_func(e_fermi, dos, edos, kT):
    return dos / (1.0 + np.exp((e_fermi - edos)/kT))

def n_func(e_fermi, dos, edos, kT):
    return dos / (1.0 + np.exp((edos - e_fermi)/kT))

def DOS_from_vasprun(vasprun, nelect, bandgap = None):
    """
    generate DOS object (py_sc_fermi.dos) from a vasp 
    doscar given number of electrons in calculation
    
    Args:
        vasprun (string): path to vasprun to parse
        nelect (float): number of electrons in vasp calculation
        bandgap (float): user supplied band-gap (defualts to `None`),
        in which case, the bandgap will be determined from the DOS provided
    
    Returns:
        dos (DOS): DOS object 
    """
    vr = Vasprun(vasprun)
    densities = vr.complete_dos.densities
    cbm = vr.eigenvalue_band_properties[2]
    edos = vr.complete_dos.energies - vr.parameters['SIGMA'] - cbm
    if len(densities) == 2:
        tdos_data = np.stack([edos,densities[Spin.up], densities[Spin.down]], axis=1)
    else:
        tdos_data = np.stack([edos,densities[Spin.up]], axis = 1)
    edos = tdos_data[:,0]
    dos = np.sum(tdos_data[:,1:], axis=1)
    if bandgap == None:
        bandgap = vr.eigenvalue_band_properties[0]
    dos = DOS(dos, edos, bandgap, nelect=nelect)
    return dos
    


import numpy as np
from pymatgen.io.vasp import Vasprun # type: ignore
from pymatgen.electronic_structure.core import Spin # type: ignore

class DOS(object):
    """Class for handling density-of-states data and its integration"""

    def __init__(self, dos, edos, egap, nelect, spin_polarised=False, normalise=True):
        """Initialise a DOS instance.

        Args:
            dos (np.array): Density of states data.
            edos (np.array): Energies for the density of states (in eV).
            egap (float): Width of the band gap (in eV).
            nelect (int): Number of electrons.
            spin_polarise(`bool`, optional): is the calculated DOS spin polarisised? Default
                is `False`
            normalise (:obj:`bool`, optional): Normalise the DOS so that the integral
                from e_min --> 0.0 is equal to the `nelect`. Default is `True`.

        Returns:
            None

        Raises:
            ValueError: If `egap` > `max(edos)`.

        """
        if egap > edos[-1]:
            raise ValueError('ERROR: Your gap extends beyond your maximum energy')
        self._dos = dos
        self._edos = edos
        self._egap = egap
        self._nelect = nelect 
        if normalise:
            self.normalise_dos()

    @property
    def dos(self):
        """Get the dos array."""
        return self._dos

    @property
    def edos(self):
        """Get the edos array."""
        return self._edos

    @property
    def egap(self):
        """Get the bandgap."""
        return self._egap

    @property
    def nelect(self):
        """Get the number of electrons in the unit cell."""
        return self._nelect

    def sum_dos(self):
        """get integrated density of states"""
        vbm_index = np.where(self._edos <= 0)[0][-1]
        sum1 = np.trapz(self._dos[:vbm_index+1], self._edos[:vbm_index+1])
        return sum1
    
    def normalise_dos(self):
        """normalise the density of states w.r.t. number of electrons (self.nelect)""" 
        integrated_dos = self.sum_dos()  
        self._dos = self._dos / integrated_dos * self._nelect 

    def emin(self):
        """get lowest energy in the dos data"""
        return self._edos[0]

    def emax(self):
        """get highest energy in the dos data"""
        return self._edos[-1]

    def carrier_concentrations(self, e_fermi, kT):
        """get n0 and p0 using integrals (equations 28.9 in Ashcroft Mermin)
           for an abitrary value of kT and Fermi energy. Typically used internally 
           to calculate carrier concentrations at the self-cosistent Fermi energy
        Args:
            efermi (float): Fermi energy
            kT (float): k * temperature

        Returns:
            p0 (float): concentration of holes
            n0 (float): concentration of electrons
        """
        p0_index = np.where(self._edos <= 0)[0][-1]
        n0_index = np.where(self._edos > self.egap)[0][0]
        p0 = np.trapz( p_func(e_fermi, self._dos[:p0_index+1], self._edos[:p0_index+1], kT ),
                       self._edos[:p0_index+1])
        n0 = np.trapz( n_func(e_fermi, self._dos[n0_index:], self._edos[n0_index:], kT ),
                       self._edos[n0_index:])
        return p0, n0

# Possibly better as class methods that integrate over a certain energy range.
def p_func(e_fermi, dos, edos, kT):
    return dos / (1.0 + np.exp((e_fermi - edos)/kT))

def n_func(e_fermi, dos, edos, kT):
    return dos / (1.0 + np.exp((edos - e_fermi)/kT))

def DOS_from_vasprun(vasprun, nelect, bandgap = None):
    """
    generate DOS object (py_sc_fermi.dos) from a vasp 
    doscar given number of electrons in calculation
    
    Args:
        vasprun (string): path to vasprun to parse
        nelect (float): number of electrons in vasp calculation
        bandgap (float): user supplied band-gap (defualts to `None`),
        in which case, the bandgap will be determined from the DOS provided
    
    Returns:
        dos (DOS): DOS object 
    """
    vr = Vasprun(vasprun)
    densities = vr.complete_dos.densities
    cbm = vr.eigenvalue_band_properties[2]
    edos = vr.complete_dos.energies - vr.parameters['SIGMA'] - cbm
    if len(densities) == 2:
        tdos_data = np.stack([edos,densities[Spin.up], densities[Spin.down]], axis=1)
    else:
        tdos_data = np.stack([edos,densities[Spin.up]], axis = 1)
    edos = tdos_data[:,0]
    dos = np.sum(tdos_data[:,1:], axis=1)
    if bandgap == None:
        bandgap = vr.eigenvalue_band_properties[0]
    dos = DOS(dos, edos, bandgap, nelect=nelect)
    return dos
# type: ignore # type: ignore Class for handling density-of-states data and its integration Initialise a DOS instance.          Args:             dos (np.array): Density of states data.             edos (np.array): Energies for the density of states (in eV).             egap (float): Width of the band gap (in eV).             nelect (int): Number of electrons.             spin_polarise(`bool`, optional): is the calculated DOS spin polarisised? Default                 is `False`             normalise (:obj:`bool`, optional): Normalise the DOS so that the integral                 from e_min --> 0.0 is equal to the `nelect`. Default is `True`.          Returns:             None          Raises:             ValueError: If `egap` > `max(edos)`. Get the dos array. Get the edos array. Get the bandgap. Get the number of electrons in the unit cell. get integrated density of states normalise the density of states w.r.t. number of electrons (self.nelect) get lowest energy in the dos data get highest energy in the dos data get n0 and p0 using integrals (equations 28.9 in Ashcroft Mermin)            for an abitrary value of kT and Fermi energy. Typically used internally             to calculate carrier concentrations at the self-cosistent Fermi energy         Args:             efermi (float): Fermi energy             kT (float): k * temperature          Returns:             p0 (float): concentration of holes             n0 (float): concentration of electrons # Possibly better as class methods that integrate over a certain energy range. generate DOS object (py_sc_fermi.dos) from a vasp      doscar given number of electrons in calculation          Args:         vasprun (string): path to vasprun to parse         nelect (float): number of electrons in vasp calculation         bandgap (float): user supplied band-gap (defualts to `None`),         in which case, the bandgap will be determined from the DOS provided          Returns:         dos (DOS): DOS object