Coverage for tbkit/dos.py: 100%

1""" 

2Density of states from a set of eigenenergies, real-space 

3(:class:`tbkit.system.System`) or reciprocal-space 

4(:class:`tbkit.kspace.KSpace`, sampled over a k-mesh). 

5""" 

6from __future__ import annotations 

7 

8import numpy as np 

9from numpy.typing import ArrayLike, NDArray 

10 

11import tbkit.error_handling as error_handling 

12 

13 

14def density_of_states( 

15 energies: ArrayLike, 

16 e_grid: ArrayLike | None = None, 

17 broadening: float = 0.05, 

18 kernel: str = 'gaussian', 

19) -> tuple[NDArray[np.float64], NDArray[np.float64]]: 

20 r''' 

21 Get the density of states, broadened by a Gaussian or Lorentzian 

22 kernel of width *broadening*: 

23 

24 .. math:: 

25 

26 \rho(E) = \sum_n g(E-E_n)\, ,\quad 

27 g(x) = \frac{1}{\sqrt{2\pi}\sigma}e^{-x^2/2\sigma^2}\ \text{(gaussian)} 

28 \ \text{or}\ 

29 g(x) = \frac{1}{\pi}\frac{\sigma}{x^2+\sigma^2}\ \text{(lorentzian)} 

30 

31 Each level contributes a kernel of unit area, so 

32 :math:`\int\rho(E)dE` equals the number of levels in *energies*, 

33 for an *e_grid* wide enough to contain the tails. 

34 

35 :param energies: Array of (real) eigenenergies. Any shape (e.g. the 

36 *en* attribute of **System**, or of **KSpace** after *get_bands* 

37 over a k-mesh -- flattened automatically). 

38 :param e_grid: Real ndarray. Default value None. Energies at which to 

39 evaluate the density of states. If None, a grid of 401 points 

40 spanning ``[min(energies)-3*broadening, max(energies)+3*broadening]`` 

41 is used. 

42 :param broadening: Positive real number. Default value 0.05. Kernel width 

43 :math:`\sigma`. 

44 :param kernel: String. Default value 'gaussian'. 'gaussian' or 'lorentzian'. 

45 

46 :returns: 

47 * **e_grid** -- Real ndarray. The energy grid used. 

48 * **dos** -- Real ndarray, same shape as *e_grid*. Density of states. 

49 ''' 

50 error_handling.ndarray_empty(np.asarray(energies), 'energies') 

51 error_handling.positive_real(broadening, 'broadening') 

52 error_handling.dos_kernel(kernel) 

53 energies = np.asarray(energies).real.astype('f8').ravel() 

54 if e_grid is None: 

55 pad = 3 * broadening 

56 e_grid = np.linspace(energies.min() - pad, energies.max() + pad, 401) 

57 else: 

58 error_handling.ndarray_empty(np.asarray(e_grid), 'e_grid') 

59 e_grid = np.asarray(e_grid, dtype='f8') 

60 diff = e_grid[:, None] - energies[None, :] 

61 if kernel == 'gaussian': 

62 weight = np.exp(-diff**2 / (2*broadening**2)) / (broadening*np.sqrt(2*np.pi)) 

63 else: 

64 weight = (broadening/np.pi) / (diff**2 + broadening**2) 

65 return e_grid, weight.sum(axis=1)