Source code for gpaw.tetrahedron

"""Tetrahedron method for Brillouin-zone integrations.

See::

    Improved tetrahedron method for Brillouin-zone integrations.

    Peter E. Blöchl, O. Jepsen, and O. K. Andersen
    Phys. Rev. B 49, 16223 – Published 15 June 1994

    DOI:https://doi.org/10.1103/PhysRevB.49.16223
"""

from math import nan
from typing import List, Tuple, cast
import numpy as np
from scipy.spatial import Delaunay

from gpaw.occupations import (ZeroWidth, findroot, collect_eigelvalues,
                              distribute_occupation_numbers,
                              OccupationNumberCalculator, ParallelLayout)
from gpaw.mpi import broadcast_float
from gpaw.typing import Array1D, Array2D, Array3D, ArrayLike1D, ArrayLike2D


def bja1(e1: Array1D, e2: Array1D, e3: Array1D, e4: Array1D
         ) -> Tuple[float, Array1D]:
    """Eq. (A2) and (C2) from Blöchl, Jepsen and Andersen."""
    x = 1.0 / ((e2 - e1) * (e3 - e1) * (e4 - e1))
    return (-(e1**3).dot(x),
            3 * e1**2 * x)


def bja2(e1: Array1D, e2: Array1D, e3: Array1D, e4: Array1D
         ) -> Tuple[float, Array1D]:
    """Eq. (A3) and (C3) from Blöchl, Jepsen and Andersen."""
    x = 1.0 / ((e3 - e1) * (e4 - e1))
    y = (e3 - e1 + e4 - e2) / ((e3 - e2) * (e4 - e2))
    return (x.dot((e2 - e1)**2
                  - 3 * (e2 - e1) * e2
                  + 3 * e2**2
                  + y * e2**3),
            x * (3 * (e2 - e1)
                 - 6 * e2
                 - 3 * y * e2**2))


def bja3(e1: Array1D, e2: Array1D, e3: Array1D, e4: Array1D
         ) -> Tuple[float, Array1D]:
    """Eq. (A4) and (C4) from Blöchl, Jepsen and Andersen."""
    x = 1.0 / ((e4 - e1) * (e4 - e2) * (e4 - e3))
    return (len(e1) - x.dot(e4**3),
            3 * x * e4**2)


def bja1b(e1: Array1D, e2: Array1D, e3: Array1D, e4: Array1D) -> Array2D:
    """Eq. (B2)-(B5) from Blöchl, Jepsen and Andersen."""
    C = -0.25 * e1**3 / ((e2 - e1) * (e3 - e1) * (e4 - e1))
    w2 = -C * e1 / (e2 - e1)
    w3 = -C * e1 / (e3 - e1)
    w4 = -C * e1 / (e4 - e1)
    w1 = 4 * C - w2 - w3 - w4
    return np.array([w1, w2, w3, w4])


def bja2b(e1: Array1D, e2: Array1D, e3: Array1D, e4: Array1D) -> Array2D:
    """Eq. (B7)-(B10) from Blöchl, Jepsen and Andersen."""
    C1 = 0.25 * e1**2 / ((e4 - e1) * (e3 - e1))
    C2 = 0.25 * e1 * e2 * e3 / ((e4 - e1) * (e3 - e2) * (e3 - e1))
    C3 = 0.25 * e2**2 * e4 / ((e4 - e2) * (e3 - e2) * (e4 - e1))
    w1 = C1 + (C1 + C2) * e3 / (e3 - e1) + (C1 + C2 + C3) * e4 / (e4 - e1)
    w2 = C1 + C2 + C3 + (C2 + C3) * e3 / (e3 - e2) + C3 * e4 / (e4 - e2)
    w3 = (C1 + C2) * e1 / (e1 - e3) - (C2 + C3) * e2 / (e3 - e2)
    w4 = (C1 + C2 + C3) * e1 / (e1 - e4) + C3 * e2 / (e2 - e4)
    return np.array([w1, w2, w3, w4])


def bja3b(e1: Array1D, e2: Array1D, e3: Array1D, e4: Array1D) -> Array2D:
    """Eq. (B14)-(B17) from Blöchl, Jepsen and Andersen."""
    C = 0.25 * e4**3 / ((e4 - e1) * (e4 - e2) * (e4 - e3))
    w1 = 0.25 - C * e4 / (e4 - e1)
    w2 = 0.25 - C * e4 / (e4 - e2)
    w3 = 0.25 - C * e4 / (e4 - e3)
    w4 = 1.0 - 4 * C - w1 - w2 - w3
    return np.array([w1, w2, w3, w4])


def triangulate_submesh(rcell_cv: Array2D) -> Array3D:
    """Find the 6 tetrahedra."""
    ABC_sc = np.array([[A, B, C]
                       for A in [0, 1] for B in [0, 1] for C in [0, 1]])
    dt = Delaunay(ABC_sc.dot(rcell_cv))
    s_tq = dt.simplices
    ABC_tqc = ABC_sc[s_tq]

    # Remove zero-volume slivers:
    ABC_tqc = ABC_tqc[np.linalg.det(ABC_tqc[:, 1:] - ABC_tqc[:, :1]) != 0]

    assert ABC_tqc.shape == (6, 4, 3)
    return ABC_tqc


def triangulate_everything(size_c: Array1D,
                           ABC_tqc: Array3D,
                           i_k: Array1D) -> Array3D:
    """Triangulate the whole BZ.

    Returns i_ktq ndarray mapping:

    * k: BZ k-point index (0, 1, ...,  nbzk - 1)
    * t: tetrahedron index (0, 1, ..., 5)
    * q: tetrahedron corner index (0, 1, 2, 3)

    to i: IBZ k-point index (0, 1, ...,  nibzk - 1).
    """
    nbzk = cast(int, size_c.prod())
    ABC_ck = np.unravel_index(np.arange(nbzk), size_c)
    ABC_tqck = ABC_tqc[..., np.newaxis] + ABC_ck
    ABC_cktq = np.transpose(ABC_tqck, (2, 3, 0, 1))
    k_ktq = np.ravel_multi_index(
        ABC_cktq.reshape((3, nbzk * 6 * 4)),
        size_c,
        mode='wrap').reshape((nbzk, 6, 4))  # type: ignore
    i_ktq = i_k[k_ktq]
    return i_ktq


[docs]class TetrahedronMethod(OccupationNumberCalculator): name = 'tetrahedron-method' extrapolate_factor = 0.0 def __init__(self, rcell: ArrayLike2D, size: ArrayLike1D, improved=False, bz2ibzmap: ArrayLike1D = None, parallel_layout: ParallelLayout = None): """Tetrahedron method for calculating occupation numbers. The reciprocal cell, *rcell*, can be given in arbitrary units (only the shape matters) and *size* is the size of the Monkhorst-Pack grid. If k-points have been symmetry-reduced the *bz2ibzmap* parameter mapping BZ k-point indizes to IBZ k-point indices must be given. """ OccupationNumberCalculator.__init__(self, parallel_layout) self.rcell_cv = np.asarray(rcell) self.size_c = np.asarray(size) self.improved = improved nbzk = self.size_c.prod() if bz2ibzmap is None: bz2ibzmap = np.arange(nbzk) self.i_k = np.asarray(bz2ibzmap) assert self.size_c.shape == (3,) assert self.rcell_cv.shape == (3, 3) assert self.i_k.shape == (nbzk,) ABC_tqc = triangulate_submesh( self.rcell_cv / self.size_c[:, np.newaxis]) self.i_ktq = triangulate_everything(self.size_c, ABC_tqc, self.i_k) self.nibzkpts = self.i_k.max() + 1 def __repr__(self): return ( 'TetrahedronMethod(' f'rcell={self.rcell_cv.tolist()}, ' f'size={self.size_c.tolist()}, ' f'bz2ibzmap={self.i_k.tolist()}, ' 'parallel_layout=<' f'{self.bd.comm.size}x{self.kpt_comm.size}x{self.domain_comm.size}' '>)') def copy(self, parallel_layout: ParallelLayout = None, bz2ibzmap: List[int] = None ) -> OccupationNumberCalculator: return TetrahedronMethod( self.rcell_cv, self.size_c, self.improved, self.i_k if bz2ibzmap is None else bz2ibzmap, parallel_layout or self.parallel_layout) def _calculate(self, nelectrons, eig_qn, weight_q, f_qn, fermi_level_guess=nan) -> Tuple[float, float]: if np.isnan(fermi_level_guess): zero = ZeroWidth(self.parallel_layout) fermi_level_guess, _ = zero._calculate( nelectrons, eig_qn, weight_q, f_qn) if np.isinf(fermi_level_guess): return fermi_level_guess, 0.0 x = fermi_level_guess eig_in, weight_i, nkpts_r = collect_eigelvalues(eig_qn, weight_q, self.bd, self.kpt_comm) if eig_in.size != 0: if len(eig_in) == self.nibzkpts: nspins = 1 else: nspins = 2 assert len(eig_in) == 2 * self.nibzkpts def func(x, eig_in=eig_in): """Return excess electrons and derivative.""" if nspins == 1: n, dn = count(x, eig_in, self.i_ktq) else: n1, dn1 = count(x, eig_in[::2], self.i_ktq) n2, dn2 = count(x, eig_in[1::2], self.i_ktq) n = n1 + n2 dn = dn1 + dn2 return n - nelectrons, dn fermi_level, niter = findroot(func, x) def w(de_in): return weights(de_in, self.i_ktq, self.improved) if nspins == 1: f_in = w(eig_in - fermi_level) else: f_in = np.zeros_like(eig_in) f_in[::2] = w(eig_in[::2] - fermi_level) f_in[1::2] = w(eig_in[1::2] - fermi_level) f_in *= 1 / (weight_i[:, np.newaxis] * len(self.i_k)) else: f_in = None fermi_level = nan distribute_occupation_numbers(f_in, f_qn, nkpts_r, self.bd, self.kpt_comm) if self.kpt_comm.rank == 0: fermi_level = broadcast_float(fermi_level, self.bd.comm) fermi_level = broadcast_float(fermi_level, self.kpt_comm) return fermi_level, 0.0
def count(fermi_level: float, eig_in: Array2D, i_ktq: Array3D) -> Tuple[float, float]: """Count electrons. Return number of electrons and derivative with respect to fermi level. """ eig_in = eig_in - fermi_level nocc_i = (eig_in < 0.0).sum(axis=1) n1 = nocc_i.min() n2 = nocc_i.max() ne = n1 dnedef = 0.0 if n1 == n2: return ne, dnedef ntetra = 6 * i_ktq.shape[0] eig_Tq = eig_in[i_ktq, n1:n2].transpose((0, 1, 3, 2)).reshape( (ntetra * (n2 - n1), 4)) eig_Tq.sort(axis=1) eig_Tq = eig_Tq[eig_Tq[:, 0] < 0.0] mask1_T = eig_Tq[:, 1] > 0.0 mask2_T = ~mask1_T & (eig_Tq[:, 2] > 0.0) mask3_T = ~mask1_T & ~mask2_T & (eig_Tq[:, 3] > 0.0) for mask_T, bjaa in [(mask1_T, bja1), (mask2_T, bja2), (mask3_T, bja3)]: n, dn_T = bjaa(*eig_Tq[mask_T].T) ne += n / ntetra dnedef += dn_T.sum() / ntetra # type: ignore mask4_T = ~mask1_T & ~mask2_T & ~mask3_T ne += mask4_T.sum() / ntetra return ne, dnedef def weights(eig_in: Array2D, i_ktq: Array3D, improved=False) -> Array2D: """Calculate occupation numbers.""" nocc_i = (eig_in < 0.0).sum(axis=1) n1 = nocc_i.min() n2 = nocc_i.max() f_in = np.zeros_like(eig_in) for i in i_ktq[:, 0, 0]: f_in[i, :n1] += 6.0 if n1 == n2: return f_in / 6 ntetra = 6 * i_ktq.shape[0] eig_Tq = eig_in[i_ktq, n1:n2].transpose((0, 1, 3, 2)).reshape( (ntetra * (n2 - n1), 4)) q_Tq = eig_Tq.argsort(axis=1) eig_Tq = np.take_along_axis(eig_Tq, q_Tq, 1) f_Tq = np.zeros_like(eig_Tq) mask0_T = eig_Tq[:, 0] > 0.0 mask1_T = ~mask0_T & (eig_Tq[:, 1] > 0.0) mask2_T = ~mask0_T & ~mask1_T & (eig_Tq[:, 2] > 0.0) mask3_T = ~mask0_T & ~mask1_T & ~mask2_T & (eig_Tq[:, 3] > 0.0) for mask_T, bjab in [(mask1_T, bja1b), (mask2_T, bja2b), (mask3_T, bja3b)]: w_qT = bjab(*eig_Tq[mask_T].T) f_Tq[mask_T] += w_qT.T if improved: for mask_T, bja in [(mask1_T, bja1), (mask2_T, bja2), (mask3_T, bja3)]: e_Tq = eig_Tq[mask_T] _, d_T = bja(*e_Tq.T) f_Tq[mask_T] += (d_T * (e_Tq.sum(1) - 4 * e_Tq.T)).T / 40 mask4_T = ~mask0_T & ~mask1_T & ~mask2_T & ~mask3_T f_Tq[mask4_T] += 0.25 ktn_T = np.array(np.unravel_index(np.arange(len(eig_Tq)), (len(i_ktq), 6, n2 - n1))).T for f_q, q_q, (k, t, n) in zip(f_Tq, q_Tq, ktn_T): for q, f in zip(q_q, f_q): f_in[i_ktq[k, t, q], n1 + n] += f f_in *= 1 / 6 return f_in