# 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
```