Coverage for rhodent/calculators/dipole.py: 57%
158 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 typing import Any, Generator
5from numpy.typing import NDArray
7from gpaw.tddft.units import au_to_eA, au_to_eV
9from .base import BaseObservableCalculator
10from ..density_matrices.base import WorkMetadata
11from ..utils import ResultKeys, Result
14class DipoleCalculator(BaseObservableCalculator):
16 r""" Calculate the induced dipole moment in the time or frequency domain.
18 The induced dipole moment (i.e. the dipole moment minus the permanent
19 component) is to first order given by
21 .. math::
23 \delta\boldsymbol{\mu} = -2 \sum_{ia}^\text{eh}
24 \boldsymbol{\mu}_{ia} \mathrm{Re}\:\delta\rho_{ia},
26 where :math:`\boldsymbol{\mu}_{ia}` is the dipole matrix element of ground state Kohn-Sham
27 pair :math:`ia`
29 .. math::
31 \boldsymbol{\mu}_{ia} = \int \psi^{(0)}_i(\boldsymbol{r}) \boldsymbol{r}
32 \psi^{(0)}_a(\boldsymbol{r}) \mathrm{d}\boldsymbol{r},
34 and :math:`\delta\rho_{ia}` the induced Kohn-Sham density matrix.
36 In the frequency domain, this calculator calculates the polarizability, i.e. the Fourier
37 transform of the dipole moment divided by the perturbation.
39 .. math::
41 \boldsymbol{\alpha}(\omega) = -2 \sum_{ia}^\text{eh}
42 \boldsymbol{\mu}_{ia} \frac{\mathcal{F}\left[\mathrm{Re}\:\delta\rho_{ia}\right](\omega)}{v(\omega)}.
44 The absorption spectrum in units of dipole strength function is the imaginary part
45 of the polarizability times a prefactor
47 .. math::
49 \boldsymbol{S}(\omega) = \frac{2\omega}{\pi} \mathrm{Im}\:\boldsymbol{\alpha}(\omega).
51 This class can also compute projections of the above on Voronoi weights :math:`w_{ia}`.
53 Parameters
54 ----------
55 response
56 Response object.
57 voronoi
58 Voronoi weights object.
59 energies_occ
60 Energy grid in units of eV for occupied levels (holes).
61 energies_unocc
62 Energy grid in units of eV for unoccupied levels (electrons)
63 sigma
64 Gaussian broadening width for energy grid in units of eV.
65 times
66 Compute induced dipole in the time domain, for these times (or as close to them as possible).
67 In units of as.
69 May not be used together with :attr:`frequencies` or :attr:`frequency_broadening`.
70 pulses
71 Compute induced dipole in the time domain, in response to these pulses.
72 If none, then no pulse convolution is performed.
74 May not be used together with :attr:`frequencies` or :attr:`frequency_broadening`.
75 frequencies
76 Compute polarizability in the frequency domain, for these frequencies. In units of eV.
78 May not be used together with :attr:`times` or :attr:`pulses`.
79 frequency_broadening
80 Compute polarizability in the frequency domain, with Gaussian broadening of this width.
81 In units of eV.
83 May not be used together with :attr:`times` or :attr:`pulses`.
84 """
86 def get_result_keys(self,
87 yield_total_ia: bool = False,
88 yield_proj_ia: bool = False,
89 yield_total_ou: bool = False,
90 yield_proj_ou: bool = False,
91 decompose_v: bool = True,
92 v: int | None = None,
93 ) -> ResultKeys:
94 r""" Get the keys that each result will contain, and dimensions thereof.
96 Parameters
97 ----------
98 yield_total_ia
99 The results should include the total dipole contributions in the electron-hole basis
100 :math:`-2 \boldsymbol{\mu}_{ia} \mathrm{Re}\:\delta\rho_{ia}`.
101 yield_proj_ia
102 The results should include projections of the dipole contributions in the electron-hole basis
103 :math:`-2 \boldsymbol{\mu}_{ia} \mathrm{Re}\:\delta\rho_{ia} w_{ia}`.
104 yield_total_ou
105 The results should include the total dipole contributions on the energy grid.
106 yield_proj_ou
107 The results should include projections of the dipole contributions on the energy grid.
108 decompose_v
109 The results should include the dipole moment and/or its contributions decomposed
110 by Cartesian direction.
111 v
112 If not None, then the results should include the v:th Cartesian component
113 of the dipole moment and its contributions.
114 """
115 assert decompose_v or v is not None
117 nI = self.nproj
118 imin, imax, amin, amax = self.ksd.ialims()
119 ni = imax - imin + 1
120 na = amax - amin + 1
121 no = len(self.energies_occ)
122 nu = len(self.energies_unocc)
124 # Time domain: save dipole moment, which is real
125 # Frequency domain: save polarizability, which is complex
126 dtype = float if self._is_time_density_matrices else complex
128 resultkeys = ResultKeys()
129 if v is not None:
130 resultkeys.add_key('dm', dtype=dtype)
131 if yield_total_ia:
132 resultkeys.add_key('dm_ia', (ni, na), dtype=dtype)
133 if yield_total_ou:
134 resultkeys.add_key('dm_ou', (no, nu), dtype=dtype)
136 resultkeys.add_key('dm_proj_II', (nI, nI), dtype=dtype)
137 if yield_proj_ia:
138 resultkeys.add_key('dm_occ_proj_Iia', (nI, ni, na), dtype=dtype)
139 resultkeys.add_key('dm_unocc_proj_Iia', (nI, ni, na), dtype=dtype)
140 if yield_proj_ou:
141 resultkeys.add_key('dm_occ_proj_Iou', (nI, no, nu), dtype=dtype)
142 resultkeys.add_key('dm_unocc_proj_Iou', (nI, no, nu), dtype=dtype)
143 if decompose_v:
144 resultkeys.add_key('dm_v', 3, dtype=dtype)
145 if yield_total_ia:
146 resultkeys.add_key('dm_iav', (ni, na, 3), dtype=dtype)
147 if yield_total_ou:
148 resultkeys.add_key('dm_ouv', (no, nu, 3), dtype=dtype)
150 resultkeys.add_key('dm_proj_IIv', (nI, nI, 3), dtype=dtype)
151 if yield_proj_ia:
152 resultkeys.add_key('dm_occ_proj_Iiav', (nI, ni, na, 3), dtype=dtype)
153 resultkeys.add_key('dm_unocc_proj_Iiav', (nI, ni, na, 3), dtype=dtype)
154 if yield_proj_ou:
155 resultkeys.add_key('dm_occ_proj_Iouv', (nI, no, nu, 3), dtype=dtype)
156 resultkeys.add_key('dm_unocc_proj_Iouv', (nI, no, nu, 3), dtype=dtype)
158 return resultkeys
160 def icalculate(self,
161 yield_total_ia: bool = False,
162 yield_proj_ia: bool = False,
163 yield_total_ou: bool = False,
164 yield_proj_ou: bool = False,
165 decompose_v: bool = True,
166 v: int | None = None,
167 ) -> Generator[tuple[WorkMetadata, Result], None, None]:
168 """ Iteratively calculate dipole contributions.
170 Parameters
171 ----------
172 yield_total_ia
173 The results should include the total dipole contributions in the electron-hole basis
174 :math:`-2 \\boldsymbol{\\mu}_{ia} \\mathrm{Re}\\:\\delta\\rho_{ia}`.
175 yield_proj_ia
176 The results should include projections of the dipole contributions in the electron-hole basis
177 :math:`-2 \\boldsymbol{\\mu}_{ia} \\mathrm{Re}\\:\\delta\\rho_{ia} w_{ia}`.
178 yield_total_ou
179 The results should include the total dipole contributions on the energy grid.
180 yield_proj_ou
181 The results should include projections of the dipole contributions on the energy grid.
182 decompose_v
183 The results should include the dipole moment and/or its contributions decomposed
184 by Cartesian direction.
185 v
186 If not `None`, then the results should include the v:th Cartesian component \
187 of the dipole moment and its contributions.
189 Yields
190 ------
191 Tuple (work, result) on the root rank of the calculation communicator. \
192 Does not yield on non-root ranks of the calculation communicator.
194 work
195 An object representing the metadata (time, frequency or pulse) for the work done.
196 result
197 Object containg the calculation results for this time, frequency or pulse.
198 """
199 include_energy_dists = (yield_total_ou or yield_proj_ou)
200 if include_energy_dists:
201 assert self.sigma is not None
202 need_entire_matrix = (yield_total_ou or yield_proj_ou
203 or yield_total_ia or yield_proj_ia
204 or self.nproj > 0)
206 # Time domain: dipole moment in units of eÅ
207 # Frequency domain: polarizability in units of (eÅ)**2/eV
208 unit = au_to_eA if self._is_time_density_matrices else au_to_eA**2 / au_to_eV
209 dm0_iav = -2 * np.moveaxis(np.array([self.ksd.M_ia_from_M_p(dm_p) for dm_p in self.ksd.dm_vp]), 0, -1) * unit
211 # List all keys that this method computes
212 # This is necessary to safely send and receive data across ranks
213 resultkeys = self.get_result_keys(yield_total_ia=yield_total_ia,
214 yield_proj_ia=yield_proj_ia,
215 yield_total_ou=yield_total_ou,
216 yield_proj_ou=yield_proj_ou,
217 decompose_v=decompose_v,
218 v=v,
219 )
221 self._read_weights_diagonal()
223 # Iterate over the pulses and times
224 for work, dm in self.density_matrices:
225 if dm.rank > 0:
226 continue
228 self.log.start('calculate')
230 if need_entire_matrix:
231 if self._is_time_density_matrices:
232 # Real dipole
233 dm_iav = dm0_iav * dm.rho_ia[..., None].real
234 else:
235 # Complex polarizability
236 dm_iav = dm0_iav * dm.rho_ia[..., None]
237 dm_v = np.sum(dm_iav, axis=(0, 1))
239 if yield_total_ou:
240 dm_ouv = self.broaden_ia2ou(dm_iav)
241 else:
242 if self._is_time_density_matrices:
243 # Real dipole
244 dm_v = np.einsum('iav,ia->v', dm0_iav, dm.rho_ia.real, optimize=True)
245 else:
246 # Complex polarizability
247 dm_v = np.einsum('iav,ia->v', dm0_iav, dm.rho_ia, optimize=True)
249 result = Result()
250 if decompose_v:
251 result['dm_v'] = dm_v
252 if yield_total_ia:
253 result['dm_iav'] = dm_iav
254 if yield_total_ou:
255 result['dm_ouv'] = dm_ouv
256 if v is not None:
257 result['dm'] = dm_v[v]
258 if yield_total_ia:
259 result['dm_ia'] = dm_iav[..., v]
260 if yield_total_ou:
261 result['dm_ou'] = dm_ouv[..., v]
263 # Initialize the remaining empty arrays
264 result.create_all_empty(resultkeys)
266 # Iterate over projections
267 for iI, weight_n in enumerate(self._iterate_weights_diagonal):
268 assert weight_n is not None
269 weight_i = weight_n[self.flti]
270 weight_a = weight_n[self.flta]
272 for iI2, weight2_n in enumerate(self._iterate_weights_diagonal):
273 assert weight2_n is not None
274 weight2_a = weight2_n[self.flta]
276 dm_proj_v = np.einsum('iav,i,a->v', dm_iav, weight_i, weight2_a, optimize=True)
277 result['dm_proj_IIv'][iI, iI2] = dm_proj_v
279 if yield_proj_ia or yield_proj_ou:
280 dm_occ_proj_iav = dm_iav * weight_i[:, None, None]
281 dm_unocc_proj_iav = dm_iav * weight_a[None, :, None]
283 if yield_proj_ou:
284 dm_occ_proj_ouv = self.broaden_ia2ou(dm_occ_proj_iav)
285 dm_unocc_proj_ouv = self.broaden_ia2ou(dm_unocc_proj_iav)
287 if decompose_v:
288 if yield_proj_ia:
289 result['dm_occ_proj_Iiav'][iI] = dm_occ_proj_iav
290 result['dm_unocc_proj_Iiav'][iI] = dm_unocc_proj_iav
291 if yield_proj_ou:
292 result['dm_occ_proj_Iouv'][iI] = dm_occ_proj_ouv
293 result['dm_unocc_proj_Iouv'][iI] = dm_unocc_proj_ouv
294 if v is not None:
295 if yield_proj_ia:
296 result['dm_occ_proj_Iia'][iI] = dm_occ_proj_iav[..., v]
297 result['dm_unocc_proj_Iia'][iI] = dm_unocc_proj_iav[..., v]
298 if yield_proj_ou:
299 dm_occ_proj_ouv = self.broaden_ia2ou(dm_occ_proj_iav)
300 dm_unocc_proj_ouv = self.broaden_ia2ou(dm_unocc_proj_iav)
302 if decompose_v:
303 if yield_proj_ia:
304 result['dm_occ_proj_Iiav'][iI] = dm_occ_proj_iav
305 result['dm_unocc_proj_Iiav'][iI] = dm_unocc_proj_iav
306 if yield_proj_ou:
307 result['dm_occ_proj_Iouv'][iI] = dm_occ_proj_ouv
308 result['dm_unocc_proj_Iouv'][iI] = dm_unocc_proj_ouv
309 if v is not None:
310 if yield_proj_ia:
311 result['dm_occ_proj_Iia'][iI] = dm_occ_proj_iav[..., v]
312 result['dm_unocc_proj_Iia'][iI] = dm_unocc_proj_iav[..., v]
313 if yield_proj_ou:
314 result['dm_occ_proj_Iou'][iI] = dm_occ_proj_ouv[..., v]
315 result['dm_unocc_proj_Iou'][iI] = dm_unocc_proj_ouv[..., v]
317 self.log_parallel(f'Calculated and broadened dipoles contributions in {self.log.elapsed("calculate"):.2f}s '
318 f'for {work.desc}', flush=True)
320 yield work, result
322 if self.calc_comm.rank == 0:
323 self.log_parallel('Finished calculating dipoles contributions', flush=True)
325 @property
326 def _need_derivatives_real_imag(self) -> tuple[list[int], bool, bool]:
327 # Time domain: We only need the real part of the density matrix.
328 # Frequency domain: We need the (complex) Fourier transform of
329 # the real part of the density matrix.
330 return ([0], True, False)
332 @property
333 def _voronoi_shapes(self) -> dict[str, tuple[int, ...]]:
334 nI = self.voronoi.nproj
335 if nI == 0:
336 return {}
337 Nn = self.voronoi.nn
338 # Diagonal weights only
339 return {'diagonal': (nI, Nn)}
341 @property
342 def oscillator_strength_prefactor(self) -> NDArray[np.float64]:
343 """ Conversion factor from polarizability to dipole strength function.
345 """
346 from gpaw.tddft.units import eA_to_au, eV_to_au
347 # Convert polarizability (eÅ**2/eV) to atomic units
348 # Multiply by 2 omega / pi in atomic units
349 # Convert to units of dipole strength function (1/eV)
350 prefactor_w = 2 * (self.frequencies * eV_to_au) / np.pi * eA_to_au ** 2
351 return prefactor_w
353 def calculate_and_write(self,
354 out_fname: str,
355 write_extra: dict[str, Any] = dict(),
356 include_tcm: bool = False,
357 only_one_pulse: bool = True):
358 """ Calculate induced dipole moments and transition contribution maps.
360 Dipole moments and contributions are saved in a numpy archive if
361 the file extension is `.npz` 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 dipole moments should be computed and saved.
371 only_one_pulse
372 If False, group arrays by pulse. Only valid in time domain.
373 """
374 from ..writers.tcm import DipoleWriter
375 from ..writers.writer import FrequencyResultsCollector, TimeResultsCollector, PulseConvolutionResultsCollector
377 cls = ((TimeResultsCollector if only_one_pulse else PulseConvolutionResultsCollector)
378 if self._is_time_density_matrices else FrequencyResultsCollector)
379 calc_kwargs = dict(yield_total_ou=include_tcm, yield_proj_ou=include_tcm)
381 out_fname = str(out_fname)
382 if out_fname[-4:] == '.npz':
383 writer = DipoleWriter(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 writer = DipoleWriter(cls(self, calc_kwargs, exclude=['dm_ouv']), only_one_pulse=only_one_pulse)
387 writer.calculate_and_save_ulm(out_fname, write_extra=write_extra)
388 else:
389 raise ValueError(f'output-file must have ending .npz or .ulm, is {out_fname}')