Coverage for tbkit/occupation.py: 100%

1""" 

2Occupation of single-particle states: the Fermi-Dirac distribution and the 

3Fermi level that holds a given number of electrons, at zero or finite 

4temperature. Used by :class:`tbkit.system.System` (a finite spectrum) and 

5:class:`tbkit.kspace.KSpace` (a Brillouin-zone mesh, with weights). 

6 

7Energies and temperatures are in the same units (:math:`k_B = 1`). Each 

8state holds one electron: for spin-degenerate levels of a spinless model, 

9count two electrons per level. 

10""" 

11from __future__ import annotations 

12 

13import numpy as np 

14from numpy.typing import ArrayLike, NDArray 

15from scipy.optimize import brentq 

16from scipy.special import expit 

17 

18import tbkit.error_handling as error_handling 

19 

20 

21def fermi_dirac(energies: ArrayLike, mu: float, temperature: float = 0.) -> NDArray[np.float64]: 

22 r''' 

23 Get the Fermi-Dirac occupations 

24 

25 .. math:: 

26 

27 f(E) = \frac{1}{e^{(E-\mu)/T} + 1}\, , 

28 

29 the step :math:`\Theta(\mu - E)` at :math:`T = 0` (with :math:`f = 1/2` 

30 exactly at :math:`E = \mu`). 

31 

32 Parameters 

33 ---------- 

34 energies : Array of real energies (for complex ones, the real part is used) 

35 mu : Real number 

36 Chemical potential. 

37 temperature : Positive real number or zero, default 0 

38 

39 Returns 

40 ------- 

41 f : Real ndarray, same shape as *energies* 

42 ''' 

43 error_handling.real_number(mu, 'mu') 

44 error_handling.positive_real_zero(temperature, 'temperature') 

45 energies = np.asarray(energies).real 

46 if temperature == 0: 

47 return np.heaviside(mu - energies, 0.5) 

48 return expit(-(energies - mu) / temperature) 

49 

50 

51def fermi_level( 

52 energies: ArrayLike, n_electrons: float, temperature: float = 0., 

53 weights: ArrayLike | None = None, 

54) -> float: 

55 r''' 

56 Get the Fermi level :math:`\mu` at which :math:`\sum_n w_n f(E_n) = N_e`. 

57 

58 At :math:`T = 0` with a gap above the last occupied level, :math:`\mu` is 

59 placed midway in that gap (one unit above the top level if every level 

60 is filled); if the last electrons only partly fill a (degenerate) 

61 level, :math:`\mu` is that level's energy, where *fermi_dirac* gives 

62 each state :math:`f = 1/2` -- the electron count is then exact only for 

63 a half-filled shell. For partly filled degenerate shells, use a small 

64 temperature. 

65 

66 Parameters 

67 ---------- 

68 energies 

69 Array of real energies, any shape (e.g. *System.en*, or 

70 the bands of *KSpace.mesh_bands*). 

71 n_electrons : Positive real number or zero 

72 Number of electrons 

73 :math:`N_e` (per unit cell, with the weights of a k-mesh). 

74 temperature : Positive real number or zero, default 0 

75 weights : Array of positive reals, same shape as *energies* 

76 Default value None (all 1). Weight of each level (e.g. 1/nk on a k-mesh). 

77 

78 Returns 

79 ------- 

80 mu : Real number 

81 ''' 

82 energies = np.asarray(energies).real.ravel() 

83 error_handling.ndarray_empty(energies, 'energies') 

84 weights = np.ones_like(energies) if weights is None else np.asarray(weights, dtype='f8').ravel() 

85 error_handling.weights(weights, len(energies)) 

86 error_handling.positive_real_zero(temperature, 'temperature') 

87 error_handling.electrons(n_electrons, weights.sum()) 

88 order = np.argsort(energies, kind='stable') 

89 en, w = energies[order], weights[order] 

90 if temperature == 0: 

91 cum = np.cumsum(w) 

92 tol = 1e-9 * max(1., cum[-1]) 

93 if n_electrons <= tol: 

94 return float(en[0] - 1.) 

95 m = int(np.searchsorted(cum, n_electrons - tol)) 

96 if abs(cum[m] - n_electrons) > tol or m == len(en) - 1: 

97 # a partly filled level, or every level filled 

98 return float(en[m] if abs(cum[m] - n_electrons) > tol else en[m] + 1.) 

99 above = en[m + 1:][en[m + 1:] > en[m] + 1e-12] 

100 if len(above) == 0 or np.any(np.isclose(en[m + 1:], en[m], atol=1e-12)): 

101 # the next level is degenerate with the last occupied one 

102 return float(en[m]) 

103 return float((en[m] + above[0]) / 2) 

104 pad = 50. * temperature 

105 return float(brentq(lambda mu: np.sum(w * fermi_dirac(en, mu, temperature)) - n_electrons, 

106 en[0] - pad, en[-1] + pad, xtol=1e-14))