Coverage for tbkit/dos.py: 100%
91 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-15 06:13 +0000
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
8from itertools import permutations
10import numpy as np
11from numpy.typing import ArrayLike, NDArray
12import matplotlib.pyplot as plt
13from matplotlib.figure import Figure
15import tbkit.error_handling as error_handling
18def density_of_states(
19 energies: ArrayLike,
20 e_grid: ArrayLike | None = None,
21 broadening: float = 0.05,
22 kernel: str = 'gaussian',
23) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
24 r'''
25 Get the density of states, broadened by a Gaussian or Lorentzian
26 kernel of width *broadening*:
28 .. math::
30 \rho(E) = \sum_n g(E-E_n)\, ,\quad
31 g(x) = \frac{1}{\sqrt{2\pi}\sigma}e^{-x^2/2\sigma^2}\ \text{(gaussian)}
32 \ \text{or}\
33 g(x) = \frac{1}{\pi}\frac{\sigma}{x^2+\sigma^2}\ \text{(lorentzian)}
35 Each level contributes a kernel of unit area, so
36 :math:`\int\rho(E)dE` equals the number of levels in *energies*,
37 for an *e_grid* wide enough to contain the tails.
39 :param energies: Array of (real) eigenenergies. Any shape (e.g. the
40 *en* attribute of **System**, or of **KSpace** after *get_bands*
41 over a k-mesh -- flattened automatically).
42 :param e_grid: Real ndarray. Default value None. Energies at which to
43 evaluate the density of states. If None, a grid of 401 points
44 spanning ``[min(energies)-3*broadening, max(energies)+3*broadening]``
45 is used.
46 :param broadening: Positive real number. Default value 0.05. Kernel width
47 :math:`\sigma`.
48 :param kernel: String. Default value 'gaussian'. 'gaussian' or 'lorentzian'.
50 :returns:
51 * **e_grid** -- Real ndarray. The energy grid used.
52 * **dos** -- Real ndarray, same shape as *e_grid*. Density of states.
53 '''
54 error_handling.ndarray_empty(np.asarray(energies), 'energies')
55 error_handling.positive_real(broadening, 'broadening')
56 error_handling.dos_kernel(kernel)
57 energies = np.asarray(energies).real.astype('f8').ravel()
58 if e_grid is None:
59 pad = 3 * broadening
60 e_grid = np.linspace(energies.min() - pad, energies.max() + pad, 401)
61 else:
62 error_handling.ndarray_empty(np.asarray(e_grid), 'e_grid')
63 e_grid = np.asarray(e_grid, dtype='f8')
64 diff = e_grid[:, None] - energies[None, :]
65 if kernel == 'gaussian':
66 weight = np.exp(-diff**2 / (2*broadening**2)) / (broadening*np.sqrt(2*np.pi))
67 else:
68 weight = (broadening/np.pi) / (diff**2 + broadening**2)
69 return e_grid, weight.sum(axis=1)
72def _mesh_simplices(dim: int) -> NDArray[np.int64]:
73 '''
74 Private function. Split each cell of a periodic *dim*-dimensional mesh
75 (*dim* = 1, 2, 3) into *dim*! simplices of equal volume (segments,
76 triangles, tetrahedra), all sharing the cell's main diagonal (the
77 decomposition of Blochl et al.). The vertices are given as integer
78 offsets from the cell's lowest corner: shape (dim!, dim+1, dim), the
79 vertex m of simplex p being the sum of the first m unit vectors of
80 the permutation p of the axes.
81 '''
82 eye = np.eye(dim, dtype=int)
83 return np.array([np.cumsum([np.zeros(dim, int)] + [eye[a] for a in perm], axis=0)
84 for perm in permutations(range(dim))])
87def _simplex_values(values: NDArray, offsets: NDArray[np.int64]) -> NDArray:
88 '''
89 Private function. Values of *values* (shape (n1, ..., n_dim, ...),
90 periodic over its first dim axes) at the vertices of every simplex of
91 *_mesh_simplices*: shape (dim+1, dim!, n1, ..., n_dim, ...).
92 '''
93 dim = offsets.shape[-1]
94 axes = tuple(range(dim))
95 return np.stack([np.stack([np.roll(values, tuple(-o), axis=axes) for o in simplex])
96 for simplex in offsets], axis=1)
99def _simplex_fraction(v: NDArray[np.float64], e: float) -> NDArray[np.float64]:
100 r'''
101 Private function. Fraction of the volume of each simplex where the
102 linear interpolation of its vertex values is below *e*, for simplices
103 that *e* cuts (``v[0] <= e < v[-1]``): *v* has shape (dim+1, n), sorted
104 along the first axis. In 3D, the integrated density of states of the
105 linear tetrahedron method (Blochl et al. 1994, appendix).
106 '''
107 dim = len(v) - 1
108 if dim == 1:
109 return (e - v[0]) / (v[1] - v[0])
110 out = np.empty(v.shape[1])
111 if dim == 2:
112 low = e < v[1]
113 a, b, c = v[:, low]
114 out[low] = (e - a) ** 2 / ((b - a) * (c - a))
115 a, b, c = v[:, ~low]
116 out[~low] = 1. - (c - e) ** 2 / ((c - a) * (c - b))
117 return out
118 low, high = e < v[1], e >= v[2]
119 mid = ~low & ~high
120 a, b, c, d = v[:, low]
121 out[low] = (e - a) ** 3 / ((b - a) * (c - a) * (d - a))
122 a, b, c, d = v[:, high]
123 out[high] = 1. - (d - e) ** 3 / ((d - a) * (d - b) * (d - c))
124 a, b, c, d = v[:, mid]
125 x = e - b
126 out[mid] = ((b - a) ** 2 + 3 * (b - a) * x + 3 * x ** 2
127 - (c - a + d - b) / ((c - b) * (d - b)) * x ** 3) / ((c - a) * (d - a))
128 return out
131def tetrahedron_dos(
132 energies: ArrayLike,
133 e_grid: ArrayLike | None = None,
134) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
135 r'''
136 Get the density of states by the linear tetrahedron method (Jepsen and
137 Andersen 1971, Lehmann and Taut 1972, Blochl et al. 1994): the bands
138 are interpolated linearly inside the simplices (segments in 1D,
139 triangles in 2D, tetrahedra in 3D) of a uniform mesh of the Brillouin
140 zone, and the density of states of the interpolated bands is
141 integrated exactly. There is no broadening: band edges and van Hove
142 singularities (the logarithmic peak of the 2D square lattice, the
143 :math:`\sqrt{E}` edges of 3D bands) come out sharp, and converge with
144 the mesh instead of with a kernel width.
146 The value at ``e_grid[i]`` is the exact average of that density of
147 states over the bin between the midpoints to the neighbouring grid
148 points, so a flat band (a delta function) still carries its full
149 weight, in one or two bins. The normalization is that of
150 *density_of_states* on the same energies: :math:`\int\rho(E)dE` equals
151 the number of levels, ``energies.size``, for an *e_grid* that spans
152 them all. This is the linear method, without Blochl's curvature
153 correction. Complex energies (non-Hermitian bands) are replaced by
154 their real part.
156 :param energies: Real array, shape (n1, n2, ..., nbands): the bands on
157 a periodic mesh of 1, 2 or 3 dimensions, ``energies[i1, i2, ..., n]``
158 at the k-point :math:`\sum_d (i_d/n_d)\mathbf{b}_d` (as given by
159 ``KSpace.mesh_bands(nk).reshape(*nk, norb)``).
160 :param e_grid: Real ndarray, increasing. Default value None. Energies at
161 which to evaluate the density of states. If None, 401 points
162 spanning the bands, with a 2% margin on each side.
164 :returns:
165 * **e_grid** -- Real ndarray. The energy grid used.
166 * **dos** -- Real ndarray, same shape as *e_grid*. Density of states.
168 Example usage::
170 en = ks.mesh_bands((40, 40)).reshape(40, 40, ks.norb)
171 e_grid, rho = dos.tetrahedron_dos(en)
172 '''
173 error_handling.mesh_energies(energies)
174 energies = np.asarray(energies).real.astype('f8')
175 dim = energies.ndim - 1
176 if e_grid is None:
177 span = energies.max() - energies.min()
178 pad = 0.02 * (span if span > 0 else 1.)
179 e_grid = np.linspace(energies.min() - pad, energies.max() + pad, 401)
180 else:
181 error_handling.increasing_grid(e_grid, 'e_grid')
182 e_grid = np.asarray(e_grid, dtype='f8')
183 offsets = _mesh_simplices(dim)
184 v = np.sort(_simplex_values(energies, offsets).reshape(dim + 1, -1), axis=0)
185 weight = energies.shape[-1] * energies[..., 0].size / v.shape[1] # levels per simplex
186 # integrated density of states at the bin edges, halfway between grid points
187 mid = (e_grid[1:] + e_grid[:-1]) / 2
188 edges = np.concatenate([[2 * e_grid[0] - mid[0]], mid, [2 * e_grid[-1] - mid[-1]]])
189 tops = np.sort(v[-1])
190 count = np.empty(len(edges))
191 for n, e in enumerate(edges):
192 cut = (v[0] <= e) & (e < v[-1])
193 count[n] = np.searchsorted(tops, e, side='right') + _simplex_fraction(v[:, cut], e).sum()
194 return e_grid, weight * np.diff(count) / np.diff(edges)
197def _plot_density_of_states(
198 energies: ArrayLike,
199 e_grid: ArrayLike | None,
200 broadening: float,
201 kernel: str,
202 fs: float,
203 lw: float,
204 figsize: tuple[float, float] | None,
205) -> Figure:
206 '''
207 Private function. Plot *density_of_states* (shared by *Plot.dos* and
208 *KSpace.plot_dos*, which validate the plotting parameters).
209 '''
210 e_grid, rho = density_of_states(energies, e_grid=e_grid, broadening=broadening, kernel=kernel)
211 return _plot_dos_curve(e_grid, rho, fs, lw, figsize)
214def _plot_dos_curve(
215 e_grid: NDArray[np.float64],
216 rho: NDArray[np.float64],
217 fs: float,
218 lw: float,
219 figsize: tuple[float, float] | None,
220) -> Figure:
221 '''
222 Private function. Plot a density of states *rho* on *e_grid*.
223 '''
224 fig, ax = plt.subplots(figsize=figsize)
225 ax.plot(e_grid, rho, 'b', lw=lw)
226 ax.fill_between(e_grid, rho, color='b', alpha=0.2)
227 ax.set_xlim(e_grid[0], e_grid[-1])
228 ax.set_ylim(0., None)
229 ax.set_title('Density of states', fontsize=fs)
230 ax.set_xlabel('$E$', fontsize=fs)
231 ax.set_ylabel(r'$\rho(E)$', fontsize=fs)
232 for label in ax.xaxis.get_majorticklabels():
233 label.set_fontsize(fs)
234 for label in ax.yaxis.get_majorticklabels():
235 label.set_fontsize(fs)
236 fig.set_layout_engine('tight')
237 plt.draw()
238 return fig