Coverage for rhodent/calculators/energy.py: 92%
143 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-08-01 16:57 +0000
1from __future__ import annotations
3from collections.abc import Sequence
4import numpy as np
5from typing import Any, Collection, Generator
6from ase.units import Hartree, Bohr
7from gpaw.tddft.units import as_to_au, au_to_eV, au_to_eA, eV_to_au
9from .base import BaseObservableCalculator
10from ..density_matrices.base import WorkMetadata
11from ..density_matrices.time import ConvolutionDensityMatrixMetadata
12from ..density_matrices.distributed.pulse import PulsePerturbation
13from ..perturbation import PerturbationLike
14from ..response import BaseResponse
15from ..voronoi import VoronoiWeights
16from ..typing import Array1D
17from ..utils import ResultKeys, Result, broaden_n2e
20class EnergyCalculator(BaseObservableCalculator):
22 r""" Calculate energy contributions in the time domain.
24 The total energy can be written
26 .. math::
28 E_\text{tot}(t) = E^{(0)}_\text{tot} + \sum_{ia}^\text{eh} E_{ia}(t) + E_\text{pulse}(t).
30 The contributions to the total energy are
32 .. math::
34 E_{ia} = \frac{1}{2} \left[
35 p_{ia}\dot{q}_{ia} - q_{ia} \dot{p}_{ia} - v_{ia} q_{ia} \right],
37 the contributions to the Hartree-xc energy are
39 .. math::
41 E_{ia}^\text{Hxc} = -\frac{1}{2} \left[
42 \omega_{ia} q_{ia}^2 - q_{ia} \dot{p}_{ia} - v_{ia} q_{ia} \right],
44 and the rate of energy change is
46 .. math::
48 \dot{E}_{ia} = \frac{1}{2} \left[
49 p_{ia}\ddot{q}_{ia} - q_{ia} \ddot{p}_{ia}
50 - v_{ia} \dot{q}_{ia} - \dot{v}_{ia} q_{ia} \right].
52 The matrix element :math:`v_{ia}` can be computed from the dipole matrix element
54 .. math::
56 \boldsymbol{\mu}_{ia} = \int \psi^{(0)}_i(\boldsymbol{r}) \boldsymbol{r}
57 \psi^{(0)}_a(\boldsymbol{r}) \mathrm{d}\boldsymbol{r}
59 projected on the direction of the perturbation :math:`\hat{\boldsymbol{e}}`,
60 the occupation number difference :math:`f_{ia}` and
61 the pulse amplitude :math:`v_\text{pulse}(t)`
63 .. math::
65 v_{ia} = \sqrt{2 f_{ia}}
66 \boldsymbol{\mu}_{ia} \cdot\hat{\boldsymbol{e}}
67 v_\text{pulse}.
70 Parameters
71 ----------
72 response
73 Response object.
74 voronoi
75 Voronoi weights object.
76 filter_pair
77 Filter electron-hole pairs (occupied-unoccupied transitions) in summation of energies.
78 Provide a tuples (low, high) to compute the sum of energies
79 :math:`E_{ia}` and :math:`E_{ia}^\text{Hxc}` for pairs with transition energy
80 :math:`\varepsilon_a-\varepsilon_{i}` in the interval low-high (in units of eV).
81 energies_occ
82 Energy grid in units of eV for occupied levels (holes).
83 energies_unocc
84 Energy grid in units of eV for unoccupied levels (electrons).
85 sigma
86 Gaussian broadening width for energy grid in units of eV.
87 times
88 Compute energies in the time domain, for these times (or as close to them as possible).
89 In units of as.
90 pulses
91 Compute energies in the time domain, in response to these pulses.
92 If none, then no pulse convolution is performed.
93 """
94 def __init__(self,
95 response: BaseResponse,
96 voronoi: VoronoiWeights | None = None,
97 *,
98 filter_pair: tuple[float, float] = (0, np.inf),
99 energies_occ: list[float] | Array1D[np.float64] | None = None,
100 energies_unocc: list[float] | Array1D[np.float64] | None = None,
101 sigma: float | None = None,
102 times: list[float] | Array1D[np.float64] | None = None,
103 pulses: Collection[PerturbationLike] | None = None,
104 ):
105 super().__init__(response=response,
106 voronoi=voronoi,
107 energies_occ=energies_occ,
108 energies_unocc=energies_unocc,
109 sigma=sigma,
110 times=times,
111 pulses=pulses)
112 if len(filter_pair) != 2:
113 raise ValueError('filter_pair must be tuple of two floats')
114 self._filter_pair_low = float(filter_pair[0]) * eV_to_au
115 self._filter_pair_high = float(filter_pair[1]) * eV_to_au
117 def get_result_keys(self,
118 yield_total_E_ia: bool = False,
119 yield_proj_E_ia: bool = False,
120 yield_total_E_ou: bool = False,
121 yield_total_dists: bool = False,
122 direction: int | Sequence[int] = 2,
123 ) -> ResultKeys:
124 r""" Get the keys that each result will contain, and dimensions thereof.
126 Parameters
127 ----------
128 yield_total_E_ia
129 The results should include the contributions in the electron-hole basis to
130 the total energy :math:`E_{ia}` and Hartree-xc energy :math:`E_{ia}^\text{Hxc}`
131 yield_proj_E_ia
132 The results should include the contributions in the electron-hole basis to
133 the total energy projected on the occupied and unoccupied Voronoi weights
134 :math:`E_{ia} w_i` and :math:`E_{ia} w_a`.
135 yield_total_E_ou
136 The results should include the contributions the total energy broadened on the
137 occupied and unoccupied energy grids
138 :math:`\sum_{ia} E_{ia}\delta(\varepsilon_\text{occ}-\varepsilon_{i})
139 \delta(\varepsilon_\text{unocc}-\varepsilon_{a})` and
140 yield_total_dists
141 The results should include the contributions the total energy and Hartree-xc energy
142 broadened by electronic transition energy onto the unoccupied energies grid
143 :math:`\sum_{ia} E_{ia} \delta(\varepsilon-\omega_{ia})` and
144 :math:`\sum_{ia} E_{ia}^\text{Hxc} \delta(\varepsilon-\omega_{ia})`
145 direction
146 Direction :math:`\hat{\boldsymbol{e}}` of the polarization of the
147 pulse. Integer 0, 1 or 2 to specify x, y or z, or the direction vector specified as
148 a list of three values. Default: polarization along z.
149 """
150 nI = self.nproj
151 imin, imax, amin, amax = self.ksd.ialims()
152 ni = int(imax - imin + 1)
153 na = int(amax - amin + 1)
154 no = len(self.energies_occ)
155 nu = len(self.energies_unocc)
157 assert direction in [0, 1, 2] or (isinstance(direction, Sequence) and len(direction) == 3)
159 resultkeys = ResultKeys('dm', 'total', 'total_Hxc', 'field',
160 'total_resonant', 'total_resonant_Hxc', 'Epulse')
162 if yield_total_E_ia:
163 resultkeys.add_key('E_ia', (ni, na))
164 resultkeys.add_key('Ec_ia', (ni, na))
165 if yield_total_dists:
166 resultkeys.add_key('E_transition_u', nu)
167 resultkeys.add_key('Ec_transition_u', nu)
168 if yield_total_E_ou:
169 resultkeys.add_key('E_ou', (no, nu))
170 resultkeys.add_key('Ec_ou', (no, nu))
172 resultkeys.add_key('total_proj_II', (nI, nI))
173 resultkeys.add_key('total_Hxc_proj_II', (nI, nI))
174 if yield_proj_E_ia:
175 resultkeys.add_key('E_occ_proj_Iia', (nI, ni, na))
176 resultkeys.add_key('E_unocc_proj_Iia', (nI, ni, na))
178 return resultkeys
180 def icalculate(self,
181 yield_total_E_ia: bool = False,
182 yield_proj_E_ia: bool = False,
183 yield_total_E_ou: bool = False,
184 yield_total_dists: bool = False,
185 direction: int | Sequence[int] = 2,
186 ) -> Generator[tuple[WorkMetadata, Result], None, None]:
187 """ Iteratively calculate energies.
189 Parameters
190 ----------
191 yield_total_E_ia
192 The results should include the contributions in the electron-hole basis to
193 the total energy :math:`E_{ia}` and Hartree-xc energy :math:`E_{ia}^\\text{Hxc}`
194 yield_proj_E_ia
195 The results should include the contributions in the electron-hole basis to
196 the total energy projected on the occupied and unoccupied Voronoi weights
197 :math:`E_{ia} w_i` and :math:`E_{ia} w_a`.
198 yield_total_E_ou
199 The results should include the contributions the total energy broadened on the
200 occupied and unoccupied energy grids
201 :math:`\\sum_{ia} E_{ia}\\delta(\\varepsilon_\\text{occ}-\\varepsilon_{i})
202 \\delta(\\varepsilon_\\text{unocc}-\\varepsilon_{a})` and
203 yield_total_dists
204 The results should include the contributions the total energy and Hartree-xc energy
205 broadened by electronic transition energy onto the unoccupied energies grid
206 :math:`\\sum_{ia} E_{ia} \\delta(\\varepsilon-\\omega_{ia})` and
207 :math:`\\sum_{ia} E_{ia}^\\text{Hxc} \\delta(\\varepsilon-\\omega_{ia})`
208 direction
209 Direction :math:`\\hat{\\boldsymbol{e}}` of the polarization of the
210 pulse. Integer 0, 1 or 2 to specify x, y or z, or the direction vector specified as
211 a list of three values. Default: polarization along z.
213 Yields
214 ------
215 Tuple (work, result) on the root rank of the calculation communicator. \
216 Does not yield on non-root ranks of the calculation communicator.
218 work
219 An object representing the metadata (time or pulse) for the work done.
220 result
221 Object containg the calculation results for this time and pulse.
222 """
223 include_energy_dists = (yield_total_dists or yield_total_E_ou)
224 if include_energy_dists:
225 assert self.sigma is not None
227 assert direction in [0, 1, 2] or (isinstance(direction, Sequence) and len(direction) == 3)
228 direction_v = np.zeros(3)
229 if direction in [0, 1, 2]:
230 direction_v[direction] = 1
231 else:
232 direction_v[:] = direction
233 direction_v /= np.linalg.norm(direction_v)
235 dm_p = direction_v @ self.ksd.dm_vp
236 v0_p = dm_p * np.sqrt(2 * self.ksd.f_p)
237 dm_ia = self.ksd.M_ia_from_M_p(dm_p)
238 v0_ia = self.ksd.M_ia_from_M_p(v0_p)
239 w_ia = self.ksd.M_ia_from_M_p(self.ksd.w_p)
241 # List all keys that this method computes
242 # This is necessary to safely send and receive data across ranks
243 resultkeys = self.get_result_keys(yield_total_dists=yield_total_dists,
244 yield_total_E_ia=yield_total_E_ia,
245 yield_proj_E_ia=yield_proj_E_ia,
246 yield_total_E_ou=yield_total_E_ou,
247 )
249 self._read_weights_diagonal()
251 # Iterate over the pulses and times
252 for work, dm in self.density_matrices:
253 if dm.rank > 0:
254 continue
255 self.log.start('calculate')
257 assert isinstance(work, ConvolutionDensityMatrixMetadata)
258 assert isinstance(work.pulse, PulsePerturbation)
259 pulsestr = work.pulse.pulse.strength(work.time * as_to_au)
261 dipmom = - 2 * np.sum(dm_ia * dm.rho_ia.real)
262 resonant_filter_ia = (w_ia > self._filter_pair_low) & (w_ia < self._filter_pair_high)
264 Epulse = dipmom * pulsestr * au_to_eV
266 # Calculate v_ia
267 v_ia = v0_ia * pulsestr
269 E_ia = -v_ia * dm.Q_ia
270 E_ia -= dm.Q_ia * dm.dP_ia
272 Ec_ia = E_ia.copy()
274 E_ia += dm.P_ia * dm.dQ_ia
275 E_ia *= 0.5 * au_to_eV
277 Ec_ia -= w_ia * dm.Q_ia ** 2
278 Ec_ia *= 0.5 * au_to_eV
280 result = Result()
281 if yield_total_E_ia:
282 result['E_ia'] = E_ia
283 result['Ec_ia'] = Ec_ia
285 result['dm'] = dipmom * au_to_eA
286 result['Epulse'] = Epulse
287 result['field'] = pulsestr * Hartree / Bohr # V/Å
288 result['total'] = np.sum(E_ia)
289 result['total_Hxc'] = np.sum(Ec_ia)
290 result['total_resonant'] = E_ia.ravel() @ resonant_filter_ia.ravel()
291 result['total_resonant_Hxc'] = Ec_ia.ravel() @ resonant_filter_ia.ravel()
293 # (Optional) Broaden transitions by transition energy
294 if yield_total_dists:
295 assert self.sigma is not None
296 result['E_transition_u'] = broaden_n2e(E_ia.ravel(), w_ia.ravel() * au_to_eV,
297 self.energies_unocc, self.sigma)
298 result['Ec_transition_u'] = broaden_n2e(Ec_ia.ravel(), w_ia.ravel() * au_to_eV,
299 self.energies_unocc, self.sigma)
301 # (Optional) Compute energy contribution matrix
302 if yield_total_E_ou:
303 result['E_ou'] = self.broaden_ia2ou(E_ia)
304 result['Ec_ou'] = self.broaden_ia2ou(Ec_ia)
306 # Initialize the remaining empty arrays
307 result.create_all_empty(resultkeys)
309 # Iterate over projections
310 for iI, weight_n in enumerate(self._iterate_weights_diagonal):
311 assert weight_n is not None
312 weight_i = weight_n[self.flti]
313 weight_a = weight_n[self.flta]
315 for iI2, weight2_n in enumerate(self._iterate_weights_diagonal):
316 assert weight2_n is not None
317 weight2_a = weight2_n[self.flta]
319 result['total_proj_II'][iI, iI2] = np.einsum(
320 'ia,i,a->', E_ia, weight_i, weight2_a, optimize=True)
321 result['total_Hxc_proj_II'][iI, iI2] = np.einsum(
322 'ia,i,a->', Ec_ia, weight_i, weight2_a, optimize=True)
324 if yield_proj_E_ia:
325 E_occ_proj_ia = E_ia * weight_i[:, None]
326 E_unocc_proj_ia = E_ia * weight_a[None, :]
328 result['E_occ_proj_Iia'][iI] = E_occ_proj_ia
329 result['E_unocc_proj_Iia'][iI] = E_unocc_proj_ia
331 self.log_parallel(f'Calculated and broadened energies in {self.log.elapsed("calculate"):.2f}s '
332 f'for {work.desc}', flush=True)
334 yield work, result
336 if self.calc_comm.rank == 0:
337 self.log_parallel('Finished calculating energies', flush=True)
339 @property
340 def _need_derivatives_real_imag(self) -> tuple[list[int], bool, bool]:
341 return ([0, 1], True, True)
343 @property
344 def _voronoi_shapes(self) -> dict[str, tuple[int, ...]]:
345 nI = self.voronoi.nproj
346 if nI == 0:
347 return {}
348 Nn = self.voronoi.nn
349 # Diagonal weights only
350 return {'diagonal': (nI, Nn)}
352 def calculate_and_write(self,
353 out_fname: str,
354 write_extra: dict[str, Any] = dict(),
355 include_tcm: bool = False,
356 save_dist: bool = False,
357 only_one_pulse: bool = True):
358 """ Calculate energy contributions.
360 Energies are saved in a numpy archive if the file extension is `.npz`
361 or in an ULM file if the file extension is `.ulm`.
363 Parameters
364 ----------
365 out_fname
366 File name of the resulting data file.
367 write_extra
368 Dictionary of extra key-value pairs to write to the data file.
369 include_tcm
370 Whether the transition contribution map (TCM) to the energies should be computed and saved.
371 save_dist
372 Whether the transition energy distributions should be computed and saved.
373 only_one_pulse
374 If False, group arrays by pulse.
375 """
376 from ..writers.energy import EnergyWriter
377 from ..writers.writer import TimeResultsCollector, PulseConvolutionResultsCollector
379 out_fname = str(out_fname)
380 if out_fname[-4:] == '.npz':
381 calc_kwargs = dict(yield_total_E_ou=include_tcm, yield_total_dists=save_dist)
382 cls = TimeResultsCollector if only_one_pulse else PulseConvolutionResultsCollector
383 writer = EnergyWriter(cls(self, calc_kwargs), only_one_pulse=only_one_pulse)
384 writer.calculate_and_save_npz(out_fname, write_extra=write_extra)
385 elif out_fname[-4:] == '.ulm':
386 assert only_one_pulse
387 calc_kwargs = dict(yield_total_E_ou=include_tcm, yield_total_dists=save_dist)
388 writer = EnergyWriter(TimeResultsCollector(self, calc_kwargs, exclude=['E_ou']), only_one_pulse=True)
389 writer.calculate_and_save_ulm(out_fname, write_extra=write_extra)
390 else:
391 raise ValueError(f'output-file must have ending .npz or .ulm, is {out_fname}')