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 :param energies: Array of real energies (for complex ones, the real part is used). 

33 :param mu: Real number. Chemical potential. 

34 :param temperature: Positive real number or zero. Default value 0. 

35 

36 :returns: 

37 * **f** -- Real ndarray, same shape as *energies*. 

38 ''' 

39 error_handling.real_number(mu, 'mu') 

40 error_handling.positive_real_zero(temperature, 'temperature') 

41 energies = np.asarray(energies).real 

42 if temperature == 0: 

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

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

45 

46 

47def fermi_level( 

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

49 weights: ArrayLike | None = None, 

50) -> float: 

51 r''' 

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

53 

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

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

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

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

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

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

60 temperature. 

61 

62 :param energies: Array of real energies, any shape (e.g. *System.en*, or 

63 the bands of *KSpace.mesh_bands*). 

64 :param n_electrons: Positive real number or zero. Number of electrons 

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

66 :param temperature: Positive real number or zero. Default value 0. 

67 :param weights: Array of positive reals, same shape as *energies*. 

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

69 

70 :returns: 

71 * **mu** -- Real number. 

72 ''' 

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

74 error_handling.ndarray_empty(energies, 'energies') 

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

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

77 error_handling.positive_real_zero(temperature, 'temperature') 

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

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

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

81 if temperature == 0: 

82 cum = np.cumsum(w) 

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

84 if n_electrons <= tol: 

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

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

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

88 # a partly filled level, or every level filled 

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

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

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

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

93 return float(en[m]) 

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

95 pad = 50. * temperature 

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

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