Coverage for rhodent/calculators/hotcarriers.py: 96%
127 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-08-01 16:57 +0000
1from __future__ import annotations
3import numpy as np
4from numpy.typing import NDArray
5from typing import Any, Generator
7from .base import BaseObservableCalculator
8from ..density_matrices.base import WorkMetadata
9from ..utils import ResultKeys, Result
12class HotCarriersCalculator(BaseObservableCalculator):
14 r""" Calculate hot-carrier distributions in the time domain.
16 For weak perturbations, the response of the density matrix is
17 to first order non-zero only in the occupied-unoccupied space,
18 i.e. the block off-diagonals
20 .. math::
22 \delta\rho = \begin{bmatrix}
23 0 & [\delta\rho_{ai}^*] \\
24 [\delta\rho_{ia}] & 0
25 \end{bmatrix}.
27 The unoccupied-occupied, or electron-hole, part of the density matrix is thus
28 linear in perturbation and can by transformed using Fourier transforms.
30 From the first-order response, the second order response, i.e. the hole-hole
31 (:math:`\delta\rho_{ii'}`) and electron-electron (:math:`\delta\rho_{aa'}`) parts
32 can be obtained.
34 The hole-hole part is
36 .. math::
38 \delta\rho_{ii'} = - \frac{1}{2} \sum_n^{f_n > f_i, f_n > f_{i'}}
39 P_{ni} P_{ni'} + Q_{ni} Q_{ni'}
41 and the electron-hole part
43 .. math::
45 \delta\rho_{aa'} = \frac{1}{2} \sum_n^{f_n < f_a, f_n < f_a'}
46 P_{ia} P_{ia'} + Q_{ia} Q_{ia'}
48 where
50 .. math::
52 \begin{align}
53 P_{ia} &= \frac{2 \mathrm{Im}\:\delta\rho_{ia}}{\sqrt{2 f_{ia}}} \\
54 Q_{ia} &= \frac{2 \mathrm{Re}\:\delta\rho_{ia}}{\sqrt{2 f_{ia}}} ,
55 \end{align}
57 where :math:`f_{ia}` is the occupation number difference of pair :math:`ia`.
59 Hot-carrier distributions are calculated by convolution of :math:`\delta\rho_{ii'}` (holes)
60 and :math:`\delta\rho_{aa'}` (electrons) by Gaussian broadening functions on the energy
61 grid.
63 .. math::
65 \begin{align}
66 P^\text{holes}(\varepsilon) &=
67 \sum_i \delta\rho_{ii'} G(\varepsilon - \varepsilon_i) \\
68 P^\text{electrons}(\varepsilon) &=
69 \sum_a \delta\rho_{aa'} G(\varepsilon - \varepsilon_a)
70 \end{align}
73 where :math:`\varepsilon_n` are Kohn-Sham energies and
75 .. math::
77 G(\varepsilon - \varepsilon_n)
78 = \frac{1}{\sqrt{2 \pi \sigma^2}}
79 \exp\left(-\frac{
80 \left(\varepsilon_i - \varepsilon\right)^2
81 }{
82 2 \sigma^2
83 }\right).
85 Projected hot-carrier distributions are defined
87 .. math::
89 \begin{align}
90 P^\text{holes}(\varepsilon) &=
91 \sum_{ii'} \delta\rho_{ii'} w_{ii'} G(\varepsilon - \varepsilon_i) \\
92 P^\text{electrons}(\varepsilon) &=
93 \sum_{aa'} \delta\rho_{aa'} w_{aa'} G(\varepsilon - \varepsilon_a).
94 \end{align}
96 The Voronoi weights are
98 .. math::
100 W_{nn}
101 = \left<\psi_n|\hat{w}|\psi_{n}\right>
102 = \int w(\boldsymbol{r}) \psi_n^*(\boldsymbol{r}) \psi_{n}(\boldsymbol{r}) d\boldsymbol{r}
104 where the operator :math:`\hat{w} = w(\boldsymbol{r})` is 1 in the Voronoi
105 region of the atomic projections and 0 outside.
107 Parameters
108 ----------
109 response
110 Response object.
111 voronoi
112 Voronoi weights object.
113 energies_occ
114 Energy grid in units of eV for occupied levels (holes).
115 energies_unocc
116 Energy grid in units of eV for unoccupied levels (electrons)
117 sigma
118 Gaussian broadening width for energy grid in units of eV.
119 times
120 Compute hot carriers in the time domain, for these times (or as close to them as possible).
121 In units of as.
122 pulses
123 Compute hot carriers in the time domain, in response to these pulses.
124 If `None` no pulse convolution is performed.
125 """
127 def get_result_keys(self,
128 yield_total_hcdists: bool = False,
129 yield_proj_hcdists: bool = False,
130 yield_total_P: bool = False,
131 yield_proj_P: bool = False,
132 yield_total_P_ou: bool = False,
133 ):
134 """ Get the keys that each result will contain, and dimensions thereof.
136 Parameters
137 ----------
138 yield_total_hcdists
139 The results should include the total hot-carrier distributions on the energy grid.
140 yield_proj_hcdists
141 The results should include the projections of the hot-carrier distributions on the energy grid.
142 yield_total_P
143 The results should include the total hot-carrier distributions in the electron-hole basis.
144 yield_proj_P
145 The results should include the projections of the hot-carrier distributions in the electron-hole basis.
146 yield_total_P_ou
147 The results should include the transition matrix broadened on the energy grid.
148 """
149 nI = self.nproj
150 imin, imax, amin, amax = self.ksd.ialims()
151 ni = int(imax - imin + 1)
152 na = int(amax - amin + 1)
153 no = len(self.energies_occ)
154 nu = len(self.energies_unocc)
156 resultkeys = ResultKeys('sumocc', 'sumunocc')
158 if yield_total_P:
159 resultkeys.add_key('P_i', ni)
160 resultkeys.add_key('P_a', na)
161 if yield_total_hcdists:
162 resultkeys.add_key('hcdist_o', no)
163 resultkeys.add_key('hcdist_u', nu)
165 if yield_total_P_ou:
166 resultkeys.add_key('P_ou', (no, nu))
168 resultkeys.add_key('sumocc_proj_I', nI)
169 resultkeys.add_key('sumunocc_proj_I', nI)
170 if yield_proj_P:
171 resultkeys.add_key('P_proj_Ii', (nI, ni))
172 resultkeys.add_key('P_proj_Ia', (nI, na))
173 if yield_proj_hcdists:
174 resultkeys.add_key('hcdist_proj_Io', (nI, no))
175 resultkeys.add_key('hcdist_proj_Iu', (nI, nu))
177 return resultkeys
179 def icalculate(self,
180 yield_total_hcdists: bool = False,
181 yield_proj_hcdists: bool = False,
182 yield_total_P: bool = False,
183 yield_proj_P: bool = False,
184 yield_total_P_ou: bool = False,
185 ) -> Generator[tuple[WorkMetadata, Result], None, None]:
186 """ Iteratively calculate second order density matrices and hot-carrier distributions.
188 Parameters
189 ----------
190 yield_total_hcdists
191 The results should include the total hot-carrier distributions on the energy grid.
192 yield_proj_hcdists
193 The results should include the projections of the hot-carrier distributions on the energy grid.
194 yield_total_P
195 The results should include the total hot-carrier distributions in the electron-hole basis.
196 yield_proj_P
197 The results should include the projections of the hot-carrier distributions in the electron-hole basis.
198 yield_total_P_ou
199 The results should include the transition matrix broadened on the energy grid.
201 Yields
202 ------
203 Tuple (work, result) on the root rank of the calculation communicator. \
204 Does not yield on non-root ranks of the calculation communicator.
206 work
207 An object representing the metadata (time and pulse) for the work done.
208 result
209 Object containg the calculation results for this time and pulse.
210 """
211 include_energy_dists = (yield_proj_hcdists or yield_total_hcdists)
212 if include_energy_dists:
213 assert self.sigma is not None
215 # List all keys that this method computes
216 # This is necessary to safely send and receive data across ranks
217 resultkeys = self.get_result_keys(yield_total_hcdists=yield_total_hcdists,
218 yield_proj_hcdists=yield_proj_hcdists,
219 yield_total_P=yield_total_P,
220 yield_proj_P=yield_proj_P,
221 yield_total_P_ou=yield_total_P_ou,
222 )
224 self._read_weights_eh()
226 # Iterate over the pulses and times
227 for work, dm in self.density_matrices:
228 Q_ia = dm.Q_ia
229 P_ia = dm.P_ia
231 if dm.rank > 0:
232 continue
234 # Holes
235 M_ii = 0.5 * (Q_ia @ Q_ia.T + P_ia @ P_ia.T)
237 # Electrons
238 M_aa = 0.5 * (Q_ia.T @ Q_ia + P_ia.T @ P_ia)
240 result = Result()
242 # (Optional) Compute broadened transition matrix
243 if yield_total_P_ou:
244 transition_ia = 0.5 * (Q_ia**2 + P_ia**2)
245 result['P_ou'] = self.broaden_ia2ou(transition_ia)
247 # Compute quantities in all space
248 P_i = calculate_hcdist(None, M_ii)
249 P_a = calculate_hcdist(None, M_aa)
250 result['sumocc'] = np.sum(P_i)
251 result['sumunocc'] = np.sum(P_a)
252 if yield_total_hcdists:
253 result['hcdist_o'] = self.broaden_occ(P_i)
254 result['hcdist_u'] = self.broaden_unocc(P_a)
255 if yield_total_P:
256 result['P_i'] = P_i
257 result['P_a'] = P_a
259 result.create_all_empty(resultkeys)
261 # Iterate over projections
262 for iI, (weight_ii, weight_aa) in enumerate(self._iterate_weights_eh):
263 P_proj_i = calculate_hcdist(weight_ii, M_ii)
264 P_proj_a = calculate_hcdist(weight_aa, M_aa)
265 result['sumocc_proj_I'][iI] = np.sum(P_proj_i)
266 result['sumunocc_proj_I'][iI] = np.sum(P_proj_a)
267 if yield_proj_hcdists:
268 result['hcdist_proj_Io'][iI] = self.broaden_occ(P_proj_i)
269 result['hcdist_proj_Iu'][iI] = self.broaden_unocc(P_proj_a)
270 if yield_proj_P:
271 result['P_proj_Ii'][iI] = P_proj_i
272 result['P_proj_Ia'][iI] = P_proj_a
274 self.log_parallel(f'Calculated and broadened HC distributions in {self.log.elapsed("calculate"):.2f}s '
275 f'for {work.desc}', flush=True)
277 yield work, result
279 if self.calc_comm.rank == 0 and self.density_matrices.localn > 0:
280 self.log_parallel('Finished calculating hot-carrier matrices', flush=True)
282 @property
283 def _need_derivatives_real_imag(self) -> tuple[list[int], bool, bool]:
284 return ([0], True, True)
286 @property
287 def _voronoi_shapes(self) -> dict[str, tuple[int, ...]]:
288 nI = self.voronoi.nproj
289 if nI == 0:
290 return {}
291 # Hole-hole part and electron-electron part
292 Ni, Na = len(self.eig_i), len(self.eig_a)
293 return {'ii': (nI, Ni, Ni), 'aa': (nI, Na, Na)}
295 def calculate_totals_by_pulse_and_write(self,
296 out_fname: str,
297 write_extra: dict[str, Any] = dict()):
298 """ Calculate the number of generated hot carriers, projected on groups of atoms, for
299 a list of pulses.
301 HC numbers are saved in a numpy archive if the file extension is `.npz`
302 or in a comma separated file if the file extension is `.dat`.
304 Parameters
305 ----------
306 out_fname
307 File name of the resulting data file.
308 """
309 from ..writers.hcdist import HotCarriersWriter, write_hot_carrier_totals_by_pulse
310 from ..writers.writer import PulseConvolutionAverageResultsCollector
312 calc_kwargs = dict(yield_total_hcdists=False, yield_proj_hcdists=False)
313 collector = PulseConvolutionAverageResultsCollector(self, calc_kwargs)
314 writer = HotCarriersWriter(collector, only_one_pulse=False)
316 out_fname = str(out_fname)
317 if out_fname[-4:] == '.npz':
318 writer.calculate_and_save_npz(out_fname, write_extra=write_extra)
319 elif out_fname[-4:] == '.dat':
320 write_hot_carrier_totals_by_pulse(out_fname, writer)
321 else:
322 raise ValueError(f'output-file must have ending .npz or .dat, is {out_fname}')
324 def calculate_totals_by_time_and_write(self,
325 out_fname: str,
326 write_extra: dict[str, Any] = dict()):
327 """ Calculate the number of generated hot carriers, projected on groups of atoms, for
328 a list of times.
330 HC numbers are saved in a numpy archive if the file extension is `.npz`
331 or in a comma separated file if the file extension is `.dat`.
333 Parameters
334 ----------
335 out_fname
336 File name of the resulting data file.
337 """
338 from ..writers.hcdist import HotCarriersWriter, write_hot_carrier_totals_by_time
339 from ..writers.writer import TimeResultsCollector
341 calc_kwargs = dict(yield_total_hcdists=False, yield_proj_hcdists=False)
342 collector = TimeResultsCollector(self, calc_kwargs)
343 writer = HotCarriersWriter(collector, only_one_pulse=True)
345 if out_fname[-4:] == '.npz':
346 writer.calculate_and_save_npz(out_fname, write_extra=write_extra)
347 elif out_fname[-4:] == '.dat':
348 write_hot_carrier_totals_by_time(out_fname, writer)
349 else:
350 raise ValueError(f'output-file must have ending .npz or .dat, is {out_fname}')
352 def calculate_distributions_and_write(self,
353 out_fname: str,
354 write_extra: dict[str, Any] = dict(),
355 average_times: bool = True,
356 only_one_pulse: bool = True):
357 """ Calculate broadened hot-carrier energy distributions, optionally averaged over time.
359 HC distributions are saved in a numpy archive if the file extension is `.npz`
360 or in a comma separated file if the file extension is `.dat`.
362 Parameters
363 ----------
364 out_fname
365 File name of the resulting data file.
366 write_extra
367 Dictionary of extra key-value pairs to write to the data file.
368 average_times
369 If `True`, an average over the given times will be taken. If false, then
370 hot-carrier distributions are computed separately over the times, and
371 each output is written separately for each time.
372 only_one_pulse
373 There is only one pulse, don't group by pulse.
374 """
375 from ..writers.hcdist import HotCarriersWriter, write_hot_carrier_distributions
376 from ..writers.writer import TimeResultsCollector, TimeAverageResultsCollector
378 if len(self.energies_occ) == 0 and len(self.energies_unocc) == 0:
379 raise ValueError('Either occupied or unoccupied energies grid must be given')
381 calc_kwargs = dict(yield_total_hcdists=True, yield_proj_hcdists=True)
382 cls = TimeAverageResultsCollector if average_times else TimeResultsCollector
383 writer = HotCarriersWriter(cls(self, calc_kwargs), only_one_pulse=only_one_pulse)
385 if out_fname[-4:] == '.npz':
386 writer.calculate_and_save_npz(out_fname, write_extra=write_extra)
387 elif out_fname[-4:] == '.dat':
388 write_hot_carrier_distributions(out_fname, writer)
389 else:
390 raise ValueError(f'output-file must have ending .npz or .dat, is {out_fname}')
393def calculate_hcdist(weight_xx: NDArray[np.float64] | None,
394 M_xx: NDArray[np.float64],
395 ) -> NDArray[np.float64]:
396 r""" Calculate row-wise summed hot carrier distribution.
398 .. math::
400 \rho_n = \sum_{n'} \rho_{nn'} w_{nn'}
402 Parameters
403 ----------
404 weight_xx
405 Voronoi weights :math:`w_{ii}` or :math:`w_{aa}`.
406 Specify None to let the weights be the identity matrix
407 M_xx
408 Matrix :math:`M_{ii}` or :math:`M_{aa}`
410 Returns
411 -------
412 Hot carrier distribution by eigenvalue :math:`\rho_n`
413 """
414 if weight_xx is None:
415 P_x = np.diag(M_xx)
416 else:
417 P_x = np.sum(weight_xx*M_xx, axis=0)
419 return P_x