Coverage for rhodent/dos.py: 88%
120 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-08-01 16:57 +0000
1from __future__ import annotations
3from typing import Any, Generator
4from pathlib import Path
6import numpy as np
7from numpy.typing import NDArray
9from ase.io.ulm import Reader
10from gpaw.mpi import world
12from .typing import GPAWCalculator
13from .utils import gauss_ij, Logger
14from .voronoi import VoronoiWeights, atom_projections_to_numpy
17class DOSCalculator:
19 r""" Density of states (:term:`DOS`) and partial DOS (:term:`PDOS`) calculator.
21 Calculates DOS
23 .. math::
25 \mathrm{DOS}(\varepsilon) = \Sigma_n G(\varepsilon - \varepsilon_n)
27 and PDOS
29 .. math::
31 \mathrm{PDOS}(\varepsilon) = \Sigma_n w_{nn} G(\varepsilon - \varepsilon_n)
33 where :math:`\varepsilon_n` are Kohn-Sham energies and :math:`G(\varepsilon - \varepsilon_n)`
34 a Gaussian broadening function
36 .. math::
38 G(\varepsilon - \varepsilon_n)
39 = \frac{1}{\sqrt{2 \pi \sigma^2}}
40 \exp\left(-\frac{
41 \left(\varepsilon_i - \varepsilon\right)^2
42 }{
43 2 \sigma^2
44 }\right).
46 The Voronoi weights are defined
48 .. math::
50 W_{nn}
51 = \left<\psi_n|\hat{w}|\psi_{n}\right>
52 = \int w(\boldsymbol{r}) \psi_n^*(\boldsymbol{r}) \psi_{n}(\boldsymbol{r}) d\boldsymbol{r}
54 where the operator :math:`\hat{w} = w(\boldsymbol{r})` is 1 in the Voronoi
55 region of the atomic projections and 0 outside.
57 Parameters
58 ----------
59 eigenvalues
60 List of eigenvalues :math:`\varepsilon_n` (relative to Fermi level).
61 fermilevel
62 Fermi level.
63 voronoi
64 Voronoi weights object defining the atomic projections for :term:`PDOS`.
66 Leave out if only :term:`DOS` is desired.
67 energies
68 Array of energies (in units of eV) for which the broadened :term:`PDOS` is computed.
69 sigma
70 Gaussian broadening width :math:`\sigma` in units of eV.
71 zerofermi
72 Eigenvalues relative to Fermi level if `True`, else relative to vacuum
73 """
75 def __init__(self,
76 eigenvalues: list[float] | NDArray[np.float64],
77 fermilevel: float,
78 voronoi: VoronoiWeights | None,
79 energies: list[float] | NDArray[np.float64],
80 sigma: float,
81 zerofermi: bool = False):
82 self._voronoi = voronoi
83 self._eig_n = np.array(eigenvalues)
84 self._fermilevel = fermilevel
85 self._energies = np.array(energies)
86 self._sigma = sigma
87 self._zerofermi = zerofermi
88 if voronoi is None:
89 self._log = Logger()
90 else:
91 self._log = voronoi.log
93 @property
94 def voronoi(self) -> VoronoiWeights | None:
95 """ Voronoi weights object. """
96 return self._voronoi
98 @property
99 def log(self) -> Logger:
100 """ Logger. """
101 return self._log
103 @property
104 def zero(self) -> float:
105 if self._zerofermi:
106 return self._fermilevel
107 else:
108 return 0
110 @property
111 def eig_n(self) -> NDArray[np.float64]:
112 return self._eig_n - self.zero
114 @property
115 def energies(self) -> NDArray[np.float64]:
116 """ Energy grid in units of eV. """
117 return self._energies
119 @property
120 def sigma(self) -> float:
121 """ Gaussian broadening width in units of eV. """
122 return self._sigma
124 def calculate_dos(self) -> NDArray[np.float64] | None:
125 """ Calculate :term:`DOS`.
127 Returns
128 ------
129 DOS on the root rank, None on other ranks.
130 """
131 if world.rank != 0:
132 return None
134 # Construct gaussians
135 gauss_en = gauss_ij(self.energies, self.eig_n, self.sigma)
136 dos_e = np.sum(gauss_en, axis=1)
137 self.log('Computed DOS', who='Calculator', comm=world, flush=True)
138 return dos_e
140 def icalculate_pdos(self) -> Generator[dict[str, NDArray[np.float64] | None], None, None]:
141 """ Calculate :term:`PDOS` for each of the atomic projections in the :attr:`voronoi` object..
143 Yields
144 ------
145 Once per set of Voronoi weights a dictionary with keys
146 * `weight_n` - Array of dimensions `(Nn)` of projections. `None` on non-root ranks.
147 * `pdos_e` - Broadened PDOS. `None` on non-root ranks.
148 """
149 if self.voronoi is None:
150 raise ValueError('Voronoi must be given to the calculator')
152 if world.rank == 0:
153 # Construct gaussians
154 gauss_en = gauss_ij(self.energies, self.eig_n, self.sigma)
155 self.log('Computed gaussians', who='Calculator', comm=world, flush=True)
157 for i, weight_nn in enumerate(self.voronoi):
158 if world.rank == 0:
159 assert weight_nn is not None
160 weight_n = weight_nn.diagonal()
161 pdos_e = gauss_en @ weight_n
162 self.log(f'Computed PDOS for projection {self.voronoi.atom_projections[i]}',
163 who='Calculator', comm=world, flush=True)
164 yield dict(weight_n=weight_n, pdos_e=pdos_e)
165 else:
166 yield dict(weight_n=None, pdos_e=None)
168 @classmethod
169 def from_gpw(cls,
170 gpw_file: Path | str,
171 voronoi: VoronoiWeights | None,
172 energies: list[float] | NDArray[np.float64],
173 sigma: float,
174 zerofermi: bool = False):
175 r"""
176 Initialize from `.gpw` file.
178 Parameters
179 ----------
180 gpw_file
181 File name of GPAW ground state file.
182 voronoi
183 Voronoi weights object defining the atomic projections for :term:`PDOS`.
185 Leave out if only :term:`DOS` is desired.
186 energies
187 Array of energies (in units of eV) for which the broadened :term:`PDOS` is computed.
188 sigma
189 Gaussian broadening width :math:`\sigma` in units of eV.
190 zerofermi
191 Eigenvalues relative to Fermi level if `True`, else relative to vacuum.
192 """
194 reader = Reader(gpw_file)
195 eig_skn = reader.wave_functions.eigenvalues
196 # Assume only one spin channel and one k point
197 assert eig_skn.shape[:2] == (1, 1), 'Many spins or kpoints'
198 eig_n = eig_skn[0, 0]
199 fermilevel = reader.wave_functions.fermi_levels[0]
201 return cls(eig_n, fermilevel, voronoi, energies, sigma, zerofermi)
203 @classmethod
204 def from_calc(cls,
205 calc: GPAWCalculator,
206 voronoi: VoronoiWeights | None,
207 energies: list[float] | NDArray[np.float64],
208 sigma: float,
209 zerofermi: bool = False):
210 r"""
211 Initialize from GPAW calculator.
213 Parameters
214 ----------
215 calc
216 GPAW calculator.
217 voronoi
218 Voronoi weights object defining the atomic projections for :term:`PDOS`.
220 Leave out if only :term:`DOS` is desired.
221 energies
222 Array of energies (in units of eV) for which the broadened :term:`PDOS` is computed.
223 sigma
224 Gaussian broadening width :math:`\sigma` in units of eV.
225 zerofermi
226 Eigenvalues relative to Fermi level if `True`, else relative to vacuum.
227 """
228 eig_n = calc.get_eigenvalues()
229 fermilevel = calc.get_fermi_level()
231 return cls(eig_n, fermilevel, voronoi, energies, sigma, zerofermi)
233 def calculate_dos_and_write(self,
234 out_fname: Path | str,
235 write_extra: dict[str, Any] = dict()):
236 """ Calculate the :term:`DOS` and write to file.
238 The DOS is saved in a numpy archive if the file extension is `.npz`
239 or in a comma separated file if the file extension is `.dat`.
241 Parameters
242 ----------
243 out_fname
244 File name of the resulting data file.
245 write_extra
246 Dictionary of additional data written to numpy archive (ignored for `.dat`) files.
247 """
248 from .writers.dos import write_density_of_states
250 # Calculate
251 if world.rank == 0:
252 dos_e = self.calculate_dos()
254 # Write to file on root
255 if world.rank > 0:
256 world.barrier()
257 return
258 assert dos_e is not None
260 out_fname = str(out_fname)
261 if out_fname[-4:] == '.npz':
262 write: dict[str, Any] = dict(
263 energy_e=self.energies,
264 dos_e=dos_e,
265 sigma=self.sigma,
266 zerofermi=self._zerofermi,
267 )
268 write.update(sigma=self.sigma, zerofermi=self._zerofermi)
269 write.update(write_extra)
271 np.savez_compressed(out_fname, **write)
272 elif out_fname[-4:] == '.dat':
273 write_density_of_states(out_fname=out_fname,
274 energies=self.energies,
275 dos=dos_e,
276 sigma=self.sigma,
277 zerofermi=self._zerofermi)
278 else:
279 raise ValueError(f'output-file must have ending .npz or .dat, is {out_fname}')
281 self.log(f'Written {out_fname}', flush=True, who='Calculator', rank=0)
282 world.barrier()
284 def calculate_pdos_and_write(self,
285 out_fname: Path | str,
286 write_extra: dict[str, Any] = dict(),
287 write_extra_from_voronoi: bool = False):
288 """ Calculate the :term:`PDOS` and write to file.
290 The PDOS is saved in a numpy archive if the file extension is `.npz`
291 or in a comma separated file if the file extension is `.dat`.
293 Parameters
294 ----------
295 out_fname
296 File name of the resulting data file.
297 write_extra
298 Dictionary of additional data written to numpy archive (ignored for `.dat`) files.
299 write_extra_from_voronoi
300 If true, and voronoi is a ULM reader, extra key-value pairs are read from
301 voronoi and written to the `.npz` file (ignored for `.dat`) files.
302 """
303 from .writers.dos import write_partial_density_of_states
305 if self.voronoi is None:
306 raise ValueError('Voronoi must be given to the calculator')
308 # Calculate
309 if world.rank == 0:
310 pdos_ei = np.zeros((len(self.energies), len(self.voronoi)))
311 for i, ret in enumerate(self.icalculate_pdos()):
312 if world.rank != 0:
313 continue
314 pdos_ei[:, i] = ret['pdos_e']
316 # Write to file on root
317 if world.rank > 0:
318 world.barrier()
319 return
321 out_fname = str(out_fname)
322 if out_fname[-4:] == '.npz':
323 write: dict[str, Any] = dict(
324 energy_e=self.energies,
325 pdos_ei=pdos_ei,
326 atom_projections=atom_projections_to_numpy(self.voronoi.atom_projections),
327 sigma=self.sigma,
328 zerofermi=self._zerofermi,
329 )
330 write.update(sigma=self.sigma, zerofermi=self._zerofermi)
331 if write_extra_from_voronoi:
332 write.update(self.voronoi.saved_fields)
333 write.update(write_extra)
335 np.savez_compressed(out_fname, **write)
336 elif out_fname[-4:] == '.dat':
337 write_partial_density_of_states(out_fname=out_fname,
338 energies=self.energies,
339 pdos=pdos_ei,
340 atom_projections=self.voronoi.atom_projections,
341 sigma=self.sigma,
342 zerofermi=self._zerofermi)
343 else:
344 raise ValueError(f'output-file must have ending .npz or .dat, is {out_fname}')
346 self.log(f'Written {out_fname}', flush=True, who='Calculator', rank=0)
347 world.barrier()