Source code for gpaw.dos

from __future__ import annotations
from pathlib import Path
from typing import Union, List, Optional, Sequence, TYPE_CHECKING

import numpy as np
from ase.dft.dos import linear_tetrahedron_integration as lti

from gpaw.setup import Setup
from gpaw.spinorbit import soc_eigenstates, BZWaveFunctions
from gpaw.typing import Array1D, Array2D, Array3D, ArrayLike1D

if TYPE_CHECKING:
    from gpaw.calculator import GPAW
    from gpaw.new.ase_interface import ASECalculator


class IBZWaveFunctions:
    """Container for eigenvalues and PAW projections (only IBZ)."""
    def __init__(self, calc: ASECalculator | GPAW):
        self.calc = calc
        self.fermi_level = self.calc.get_fermi_level()
        self.size = calc.wfs.kd.N_c
        self.bz2ibz_map = calc.wfs.kd.bz2ibz_k

    def weights(self) -> Array1D:
        """Weigths of IBZ k-points (adds to 1.0)."""
        return self.calc.wfs.kd.weight_k

    def eigenvalues(self) -> Array3D:
        """All eigenvalues."""
        kd = self.calc.wfs.kd
        eigs = np.array([[self.calc.get_eigenvalues(kpt=k, spin=s)
                          for k in range(kd.nibzkpts)]
                         for s in range(kd.nspins)])
        return eigs

    def pdos_weights(self,
                     a: int,
                     indices: List[int]
                     ) -> Array3D:
        """Projections for PDOS.

        Returns (nibzkpts, nbands, nspins)-shaped ndarray
        of the square of absolute value of the projections.  The *indices*
        list contains projector-indices.
        """
        kd = self.calc.wfs.kd
        dos_kns = np.zeros((kd.nibzkpts,
                            self.calc.wfs.bd.nbands,
                            kd.nspins))
        bands = self.calc.wfs.bd.get_slice()

        for wf in self.calc.wfs.kpt_u:
            P_ani = wf.projections
            if a in P_ani:
                P_ni = P_ani[a][:, indices]
                dos_kns[wf.k, bands, wf.s] = (abs(P_ni)**2).sum(1)

        self.calc.world.sum(dos_kns)
        return dos_kns


def get_projector_numbers(setup: Setup, ell: int) -> List[int]:
    """Find indices of bound-state PAW projector functions.

    >>> from gpaw.setup import create_setup
    >>> setup = create_setup('Li')
    >>> get_projector_numbers(setup, 0)
    [0]
    >>> get_projector_numbers(setup, 1)
    [1, 2, 3]
    """
    indices = []
    i1 = 0
    for n, l in zip(setup.n_j, setup.l_j):
        i2 = i1 + 2 * l + 1
        if l == ell and n >= 0:
            indices += list(range(i1, i2))
        i1 = i2
    return indices


def gaussian_dos(eig_kn,
                 weight_kn,
                 weight_k,
                 energies,
                 width: float) -> Array1D:
    """Simple broadening with a Gaussian."""
    dos = np.zeros_like(energies)
    if weight_kn is None:
        for e_n, w in zip(eig_kn, weight_k):
            for e in e_n:
                dos += w * np.exp(-((energies - e) / width)**2)
    else:
        for e_n, w, w_n in zip(eig_kn, weight_k, weight_kn):
            for e, w2 in zip(e_n, w_n):
                dos += w * w2 * np.exp(-((energies - e) / width)**2)
    return dos / (np.pi**0.5 * width)


def linear_tetrahedron_dos(eig_kn,
                           weight_kn,
                           energies,
                           cell,
                           size,
                           bz2ibz_map=None) -> Array1D:
    """Linear-tetrahedron method."""
    if len(eig_kn) != np.prod(size):
        eig_kn = eig_kn[bz2ibz_map]
        if weight_kn is not None:
            weight_kn = weight_kn[bz2ibz_map]

    shape = tuple(size) + (-1,)
    eig_kn = eig_kn.reshape(shape)
    if weight_kn is not None:
        weight_kn = weight_kn.reshape(shape)

    dos = lti(cell, eig_kn, energies, weight_kn)
    return dos


[docs]class DOSCalculator: def __init__(self, wfs, setups=None, cell=None, shift_fermi_level=True): self.wfs = wfs self.setups = setups self.cell = cell self.eig_skn = wfs.eigenvalues() self.fermi_level = wfs.fermi_level if shift_fermi_level: self.eig_skn -= wfs.fermi_level self.collinear = (self.eig_skn.ndim == 3) if self.collinear: self.degeneracy = 2 / len(self.eig_skn) else: self.eig_skn = np.array([self.eig_skn, self.eig_skn]) self.degeneracy = 0.5 self.nspins = len(self.eig_skn) self.weight_k = wfs.weights() def get_energies(self, emin: Optional[float] = None, emax: Optional[float] = None, npoints: int = 100): emin = emin if emin is not None else self.eig_skn.min() emax = emax if emax is not None else self.eig_skn.max() return np.linspace(emin, emax, npoints)
[docs] @classmethod def from_calculator(cls, filename: ASECalculator | GPAW | Path | str, soc=False, theta=0.0, phi=0.0, shift_fermi_level=True): """Create DOSCalculator from a GPAW calculation. filename: str Name of restart-file or GPAW calculator object. """ if not isinstance(filename, (str, Path)): calc = filename else: calc = GPAW(filename, txt=None) wfs: Union[BZWaveFunctions, IBZWaveFunctions] if soc: wfs = soc_eigenstates(calc, theta=theta, phi=phi) else: wfs = IBZWaveFunctions(calc) return DOSCalculator(wfs, calc.setups, calc.atoms.cell, shift_fermi_level)
def calculate(self, energies: ArrayLike1D, eig_kn: Array2D, weight_kn: Array2D = None, width: float = 0.1): energies = np.asarray(energies) if width > 0.0: return gaussian_dos(eig_kn, weight_kn, self.weight_k, energies, width) else: return linear_tetrahedron_dos( eig_kn, weight_kn, energies, self.cell, self.wfs.size, self.wfs.bz2ibz_map)
[docs] def raw_dos(self, energies: Sequence[float], spin: Union[int, None] = None, width: float = 0.1) -> Array1D: """Calculate density of states. width: float Width of Gaussians in eV. Use width=0.0 to use the linear-tetrahedron-interpolation method. """ if spin is None: dos = sum(self.calculate(energies, eig_kn, width=width) for eig_kn in self.eig_skn) dos *= self.degeneracy else: dos = self.calculate(energies, self.eig_skn[spin], width=width) return dos
[docs] def raw_pdos(self, energies: Sequence[float], a: int, l: int, m: Optional[int] = None, spin: int = None, width: float = 0.1) -> Array1D: """Calculate projected density of states. a: Atom index. l: Angular momentum quantum number. m: Magnetic quantum number. Default is None meaning sum over all m. For p-orbitals, m=0,1,2 translates to y, z and x. For d-orbitals, m=0,1,2,3,4 translates to xy, yz, 3z2-r2, zx and x2-y2. spin: Must be 0, 1 or None meaning spin-up, down or total respectively. width: float Width of Gaussians in eV. Use width=0.0 to use the linear-tetrahedron-interpolation method. """ indices = get_projector_numbers(self.setups[a], l) if m is not None: indices = indices[m::(2 * l) + 1] weight_kns = self.wfs.pdos_weights(a, indices) if spin is None: dos = sum(self.calculate(energies, eig_kn, weight_nk.T, width=width) for eig_kn, weight_nk in zip(self.eig_skn, weight_kns.T)) dos *= self.degeneracy else: dos = self.calculate(energies, self.eig_skn[spin], weight_kns[:, :, spin], width=width) return dos