Coverage for rhodent/calculators/density.py: 84%
194 statements
« prev ^ index » next coverage.py v7.10.6, created at 2026-09-04 15:50 +0000
1from __future__ import annotations
3from pathlib import Path
4import numpy as np
5from numpy.typing import NDArray
6from typing import Any, Generator, Sequence, Collection
8from ase.units import Bohr
9from gpaw import GPAW
10try:
11 # gpaw 26.7 moved the module to this location
12 from gpaw.old.grid_descriptor import GridDescriptor
13except ImportError:
14 from gpaw.grid_descriptor import GridDescriptor
15from gpaw.lcaotddft.densitymatrix import get_density
17from .base import BaseObservableCalculator
18from ..typing import GPAWCalculator
19from ..density_matrices.base import WorkMetadata
20from ..density_matrices.frequency import FrequencyDensityMatrixMetadata
21from ..density_matrices.time import ConvolutionDensityMatrixMetadata
22from ..perturbation import PerturbationLike
23from ..response import BaseResponse
24from ..typing import ArrayIsOnRootRank, Array1D, DistributedArray
25from ..utils import ResultKeys, Result, get_gaussian_pulse_values, ParallelMatrix
28class DensityCalculator(BaseObservableCalculator):
30 r""" Calculate densities in the time or frequency domain.
32 The induced density (i.e. the density minus the ground state density) is to first
33 order given by
35 .. math::
37 \delta n(\boldsymbol{r}) = -2 \sum_{ia}^\text{eh}
38 n_{ia}(\boldsymbol{r}) \mathrm{Re}\:\delta\rho_{ia}
40 plus PAW corrections, where :math:`n_{ia}(\boldsymbol{r})` is the density of
41 ground state Kohn-Sham pair :math:`ia`
43 .. math::
45 n_{ia}(\boldsymbol{r}) = \psi^{(0)}_i(\boldsymbol{r}) \psi^{(0)}_a(\boldsymbol{r}).
47 In the time domain, electrons and holes densities can be computed.
49 .. math::
51 \begin{align}
52 n^\text{holes}(\boldsymbol{r}) &= \sum_{ii'}
53 n_{ii'}(\boldsymbol{r}) \delta\rho_{ii'} \\
54 n^\text{electrons}(\boldsymbol{r}) &= \sum_{aa'}
55 n_{aa'}(\boldsymbol{r}) \delta\rho_{aa'}.
56 \end{align}
58 Refer to the documentation of
59 :class:`HotCarriersCalculator <rhodent.calculators.HotCarriersCalculator>` for definitions
60 of :math:`\delta\rho_{ii'}` and :math:`\delta\rho_{aa'}`.
62 Parameters
63 ----------
64 gpw_file
65 File name of GPAW ground state file.
66 response
67 Response object.
68 filter_occ
69 Filters for occupied states (holes). Provide a list of tuples (low, high)
70 to compute the density of holes with energies within the interval low-high.
71 filter_unocc
72 Filters for unoccupied states (electrons). Provide a list of tuples (low, high)
73 to compute the density of excited electrons with energies within the interval low-high.
74 times
75 Compute densities in the time domain, for these times (or as close to them as possible).
76 In units of as.
78 May not be used together with :attr:`frequencies` or :attr:`frequency_broadening`.
79 pulses
80 Compute densities in the time domain, in response to these pulses.
81 If none, then no pulse convolution is performed.
83 May not be used together with :attr:`frequencies` or :attr:`frequency_broadening`.
84 frequencies
85 Compute densities in the frequency domain, for these frequencies. In units of eV.
87 May not be used together with :attr:`times` or :attr:`pulses`.
88 frequency_broadening
89 Compute densities in the frequency domain, with Gaussian broadening of this width.
90 In units of eV.
92 May not be used together with :attr:`times` or :attr:`pulses`.
93 """
95 def __init__(self,
96 gpw_file: str,
97 response: BaseResponse,
98 filter_occ: Sequence[tuple[float, float]] = [],
99 filter_unocc: Sequence[tuple[float, float]] = [],
100 *,
101 times: list[float] | Array1D[np.float64] | None = None,
102 pulses: Collection[PerturbationLike] | None = None,
103 frequencies: list[float] | Array1D[np.float64] | None = None,
104 frequency_broadening: float = 0,
105 ):
106 super().__init__(response=response,
107 times=times,
108 pulses=pulses,
109 frequencies=frequencies,
110 frequency_broadening=frequency_broadening)
111 self._occ_filters = [self._build_single_filter('o', low, high) for low, high in filter_occ]
112 self._unocc_filters = [self._build_single_filter('u', low, high) for low, high in filter_unocc]
114 self.log.start('load_gpaw')
115 self._calc = GPAW(gpw_file, txt=None, communicator=self.calc_comm,
116 parallel={'domain': self.calc_comm.size})
117 msg = f'Loaded/initialized GPAW in {self.log.elapsed("load_gpaw"):.1f}'
118 self.log.start('init_gpaw')
120 self.calc.initialize_positions() # Initialize in order to calculate density
121 msg += f'/{self.log.elapsed("init_gpaw"):.1f} s'
122 if self.calc_comm.rank == 0:
123 self.log_parallel(msg)
124 self.ksd.density = self.calc.density
126 @property
127 def gdshape(self) -> tuple[int, int, int]:
128 """ Shape of the real space grid.
129 """
130 shape = tuple(int(N) - 1 for N in self.N_c)
131 return shape # type: ignore
133 @property
134 def gd(self) -> GridDescriptor:
135 """ Real space grid. """
136 return self.ksd.density.finegd
138 @property
139 def N_c(self) -> NDArray[np.int_]:
140 """ Number of points in each Cartesian direction of the grid.
141 """
142 return self.gd.N_c
144 @property
145 def cell_cv(self) -> NDArray[np.float64]:
146 """ Cell vectors. """
147 return self.gd.cell_cv * Bohr
149 @property
150 def occ_filters(self) -> list[slice]:
151 """ List of energy filters for occupied states. """
152 return self._occ_filters
154 @property
155 def unocc_filters(self) -> list[slice]:
156 """ List of energy filters for unoccupied states. """
157 return self._unocc_filters
159 @property
160 def calc(self) -> GPAWCalculator:
161 """ GPAW calculator instance. """
162 return self._calc # type: ignore
164 def get_result_keys(self,
165 yield_total: bool = True,
166 yield_electrons: bool = False,
167 yield_holes: bool = False) -> ResultKeys:
168 noccf = len(self.occ_filters)
169 nunoccf = len(self.unocc_filters)
170 if (yield_electrons or yield_holes) and not self._is_time_density_matrices:
171 raise ValueError('Electron or hole densities can only be computed in the time domain.')
173 resultkeys = ResultKeys()
174 if yield_total:
175 resultkeys.add_key('rho_g', self.gdshape)
177 if yield_holes:
178 resultkeys.add_key('occ_rho_g', self.gdshape)
179 if noccf > 0:
180 resultkeys.add_key('occ_rho_rows_fg', (noccf, ) + self.gdshape)
181 resultkeys.add_key('occ_rho_diag_fg', (noccf, ) + self.gdshape)
183 if yield_electrons:
184 resultkeys.add_key('unocc_rho_g', self.gdshape)
185 if nunoccf > 0:
186 resultkeys.add_key('unocc_rho_rows_fg', (nunoccf, ) + self.gdshape)
187 resultkeys.add_key('unocc_rho_diag_fg', (nunoccf, ) + self.gdshape)
189 return resultkeys
191 @property
192 def _need_derivatives_real_imag(self) -> tuple[list[int], bool, bool]:
193 # Time domain: We only need the real part of the density matrix.
194 # Frequency domain: We need the (complex) Fourier transform of
195 # the real part of the density matrix.
196 return ([0], True, False)
198 def _find_limit(self,
199 lim: float) -> int:
200 """ Find the first eigenvalue larger than :attr:`lim`.
202 Parameters
203 ----------
204 lim
205 Threshold value in units of eV.
207 Returns
208 -------
209 The index of the first eigenvalue larger than :attr:`lim`.
210 Returns `len(eig_n)` if :attr:`lim` is larger than all eigenvalues.
211 """
212 if lim > self.eig_n[-1]:
213 return len(self.eig_n)
214 return int(np.argmax(self.eig_n > lim))
216 def _build_single_filter(self,
217 key: str,
218 low: float,
219 high: float) -> slice:
220 imin, imax, amin, amax = self.ksd.ialims()
222 if key == 'o':
223 nlow = min(self._find_limit(low), imax) - imin
224 nhigh = min(self._find_limit(high), imax) - imin
225 elif key == 'u':
226 nlow = min(self._find_limit(low), amax) - amin
227 nhigh = min(self._find_limit(high), amax) - amin
228 else:
229 raise RuntimeError(f'Unknown key {key}. Key must be "o" or "u"')
230 return slice(nlow, nhigh)
232 def get_density(self,
233 rho_nn: DistributedArray,
234 nn_indices: str,
235 fltn1: slice | NDArray[np.bool_] = slice(None),
236 fltn2: slice | NDArray[np.bool_] = slice(None),
237 u: int = 0) -> DistributedArray:
238 r""" Calculate a real space density from a density matrix in the Kohn-Sham basis.
240 Parameters
241 ----------
242 rho_nn
243 Density matrix :math:`\delta\rho_{ia}`, :math:`\delta\rho_{ii'}`, or
244 :math:`\delta\rho_{aa'}`.
245 nn_indices
246 Indices describing the density matrices :attr:`rho_nn`. One of
248 - `ia` for induced density :math:`\delta\rho_{ia'}`.
249 - `ii` for holes density :math:`\delta\rho_{ii'}`.
250 - `aa` for electrons density :math:`\delta\rho_{aa'}`.
251 flt_n1
252 Filter selecting rows of the density matrix.
253 flt_n2
254 Filter selecting columns of the density matrix.
255 u
256 k-point index.
257 Returns
258 -------
259 Distributed array with the density in real space on the root rank.
260 """
261 imin, imax, amin, amax = self.ksd.ialims()
262 if nn_indices not in ['ia', 'ii', 'aa']:
263 raise ValueError(f'Parameter nn_indices must be either "ia", "ii" or "aa". Is {nn_indices}.')
264 n1 = slice(imin, imax + 1) if nn_indices[0] == 'i' else slice(amin, amax + 1)
265 n2 = slice(imin, imax + 1) if nn_indices[1] == 'i' else slice(amin, amax + 1)
267 nn1, nn2 = n1.stop - n1.start, n2.stop - n2.start
268 nM = self.ksd.C0_unM[0].shape[-1]
270 if self.calc_comm.rank == 0:
271 C0_nM = self.ksd.C0_unM[u]
272 rho_n1n2 = ParallelMatrix((nn1, nn2), np.float64, comm=self.calc_comm,
273 array=rho_nn[fltn1][:, fltn2])
274 C0_n1M = ParallelMatrix((nn1, nM), np.float64, comm=self.calc_comm,
275 array=C0_nM[n1][fltn1])
276 C0_n2M = ParallelMatrix((nn2, nM), np.float64, comm=self.calc_comm,
277 array=C0_nM[n2][fltn2])
278 else:
279 rho_n1n2 = ParallelMatrix((nn1, nn2), np.float64, comm=self.calc_comm)
280 C0_n1M = ParallelMatrix((nn1, nM), np.float64, comm=self.calc_comm)
281 C0_n2M = ParallelMatrix((nn2, nM), np.float64, comm=self.calc_comm)
283 # Transform to LCAO basis C0_n1M.T @ rho_n1n2 @ C0_n2M
284 self.log.start('transform_dm')
286 rho_MM = (C0_n1M.T @ rho_n1n2 @ C0_n2M).broadcast()
287 # assert np.issubdtype(rho_nn.dtype, float)
288 rho_MM = 0.5 * (rho_MM + rho_MM.T)
290 msg = f'Transformed DM and constructed density in {self.log.elapsed("transform_dm"):.1f}s'
291 self.log.start('get_density')
292 rho_g = get_density(rho_MM, self.calc.wfs, self.calc.density, u=u)
293 msg += f'+{self.log.elapsed("get_density"):.1f}s'
294 if self.calc_comm.rank == 0:
295 self.log_parallel(msg, flush=True)
297 big_rho_g = self.gd.collect(rho_g)
299 if self.calc_comm.rank == 0:
300 return big_rho_g
301 else:
302 return ArrayIsOnRootRank()
304 def icalculate(self,
305 yield_total: bool = True,
306 yield_electrons: bool = False,
307 yield_holes: bool = False) -> Generator[tuple[WorkMetadata, Result], None, None]:
308 """ Iteratively calculate results.
310 Parameters
311 ----------
312 yield_total
313 The results should include the total induced density.
314 yield_holes
315 The results should include the holes densities, optionally decomposed by `filter_occ`.
316 yield_electrons
317 The results should include the electrons densities, optionally decomposed by `filter_unocc`.
319 Yields
320 ------
321 Tuple (work, result) on the root rank of the calculation communicator. \
322 Does not yield on non-root ranks of the calculation communicator.
324 work
325 An object representing the metadata (time, frequency or pulse) for the work done.
326 result
327 Object containg the calculation results for this time, frequency or pulse.
328 """
329 noccf = len(self.occ_filters)
330 nunoccf = len(self.unocc_filters)
332 if (yield_electrons or yield_holes) and not self._is_time_density_matrices:
333 raise ValueError('Electron or hole densities can only be computed in the time domain.')
335 # Iterate over the pulses and times, or frequencies
336 for work, dm in self.density_matrices:
337 if self._is_time_density_matrices:
338 # Real part contributes to density
339 rho_ia = dm.rho_ia.real
340 else:
341 # Imaginary part gives absorption contribution
342 rho_ia = -dm.rho_ia.imag
344 self.log.start('calculate')
346 # Non-root ranks on calc_comm will write empty arrays to result, but will not be yielded
347 result = Result()
349 if yield_total:
350 result['rho_g'] = self.get_density(rho_ia.real, 'ia') * Bohr ** -3
352 if yield_holes:
353 # Holes
354 M_ii = 0.5 * (dm.Q_ia @ dm.Q_ia.T + dm.P_ia @ dm.P_ia.T)
356 result['occ_rho_g'] = self.get_density(M_ii, 'ii') * Bohr ** -3
358 if noccf > 0:
359 result['occ_rho_rows_fg'] = np.array([self.get_density(M_ii, 'ii', fltn1=flt)
360 for flt in self.occ_filters]) * Bohr ** -3
361 result['occ_rho_diag_fg'] = np.array([self.get_density(M_ii, 'ii', fltn1=flt, fltn2=flt)
362 for flt in self.occ_filters]) * Bohr ** -3
364 if yield_electrons:
365 # Electrons
366 M_aa = 0.5 * (dm.Q_ia.T @ dm.Q_ia + dm.P_ia.T @ dm.P_ia)
368 result['unocc_rho_g'] = self.get_density(M_aa, 'aa') * Bohr ** -3
370 if nunoccf > 0:
371 result['unocc_rho_rows_fg'] = np.array([self.get_density(M_aa, 'aa', fltn1=flt)
372 for flt in self.unocc_filters]) * Bohr ** -3
373 result['unocc_rho_diag_fg'] = np.array([self.get_density(M_aa, 'aa', fltn1=flt, fltn2=flt)
374 for flt in self.unocc_filters]) * Bohr ** -3
375 if dm.rank > 0:
376 continue
378 self.log_parallel(f'Calculated density in {self.log.elapsed("calculate"):.2f}s '
379 f'for {work.desc}', flush=True)
381 yield work, result
383 if self.calc_comm.rank == 0:
384 self.log_parallel('Finished calculating density contributions', flush=True)
386 def calculate_and_write(self,
387 out_fname: str,
388 which: str | list[str] = 'induced',
389 write_extra: dict[str, Any] = dict()):
390 """ Calculate density contributions.
392 Densities are saved in a numpy archive, ULM file or cube file depending on
393 whether the file extension is `.npz`, `.ulm`, or `.cube`.
395 If the file extension is `.cube` then :attr:`out_fname` is taken to be a formatting string.
397 The formatting string should be a plain string containing variable
398 placeholders within curly brackets `{}`. It should not be confused with
399 a formatted string literal (f-string).
401 It acccepts the variables:
403 * `{time}` - Time in units of as (time domain only).
404 * `{freq}` - Frequency in units of eV (frequency domain only).
405 * `{which}` - The :attr:`which` argument.
406 * `{pulsefreq}` - Pulse frequency in units of eV (time domain only).
407 * `{pulsefwhm}` - Pulse FWHM in units of fs (time domain only).
409 Examples:
411 * out_fname = `{which}_density_t{time:09.1f}.cube`.
412 * out_fname = `{which}_density_w{freq:05.2f}.cube`.
414 Parameters
415 ----------
416 out_fname
417 File name of the resulting data file.
418 which
419 String, or list of strings specifying the types of density to compute:
421 * `induced` - Induced density.
422 * `holes` - Holes density (only allowed in the time domain).
423 * `electrons` - Electrons density (only allowed in the time domain).
424 write_extra
425 Dictionary of extra key-value pairs to write to the data file.
426 """
427 from ..writers.density import DensityWriter, write_density
428 from ..writers.writer import FrequencyResultsCollector, TimeResultsCollector
430 cls = TimeResultsCollector if self._is_time_density_matrices else FrequencyResultsCollector
432 if isinstance(which, str):
433 which = [which]
435 for which_key in which:
436 if which_key in ['holes', 'electrons'] and not self._is_time_density_matrices:
437 raise ValueError(f'Option which={which_key} not allowed in the frequency domain.')
438 if which_key not in ['induced', 'holes', 'electrons']:
439 raise ValueError(f'Option which={which} not recognized. '
440 'Must be one of: induced, holes, electrons')
442 calc_kwargs = dict(yield_total='induced' in which,
443 yield_holes='holes' in which,
444 yield_electrons='electrons' in which)
446 keys = {'induced': 'rho_g',
447 'holes': 'occ_rho_g',
448 'electrons': 'unocc_rho_g'}
450 out_fname = str(out_fname)
451 if out_fname.endswith('.npz'):
452 exclude = ['occ_rho_rows_fg', 'occ_rho_diag_fg', 'unocc_rho_rows_fg', 'unocc_rho_diag_fg']
453 writer = DensityWriter(cls(self, calc_kwargs=calc_kwargs, exclude=exclude))
454 writer.calculate_and_save_npz(out_fname=out_fname, write_extra=write_extra)
455 elif out_fname.endswith('.ulm'):
456 exclude = ['rho_g', 'occ_rho_rows_fg', 'occ_rho_diag_fg', 'unocc_rho_rows_fg', 'unocc_rho_diag_fg']
457 writer = DensityWriter(cls(self, calc_kwargs=calc_kwargs, exclude=exclude))
458 writer.calculate_and_save_ulm(out_fname=out_fname, write_extra=write_extra)
459 elif out_fname.endswith('.cube'):
460 atoms = self.calc.atoms
461 for work, res in self.icalculate(**calc_kwargs):
462 if self.calc_comm.rank > 0:
463 continue
464 for which_key in which:
465 key = keys[which_key]
466 fname_kw: dict[str, float | str] = dict(which=which_key)
467 data = res[key]
468 if self._is_time_density_matrices:
469 assert isinstance(work, ConvolutionDensityMatrixMetadata)
470 fname_kw.update(time=work.time, **get_gaussian_pulse_values(work.pulse))
471 else:
472 assert isinstance(work, FrequencyDensityMatrixMetadata)
473 fname_kw.update(freq=work.freq)
474 fpath = Path(out_fname.format(**fname_kw))
475 fpath.parent.mkdir(parents=True, exist_ok=True)
476 write_density(str(fpath), atoms, data)
477 self.log_parallel(f'Written {fpath}', flush=True)
479 else:
480 raise ValueError(f'output-file must have ending .npz or .ulm, is {out_fname}')