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