Coverage for rhodent/density_matrices/time.py: 89%
378 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-08-01 16:57 +0000
« prev ^ index » next coverage.py v7.9.1, created at 2025-08-01 16:57 +0000
1from __future__ import annotations
3from typing import Generator, Collection
5from pathlib import Path
6import numpy as np
7from numpy.typing import NDArray
9from gpaw.lcaotddft.ksdecomposition import KohnShamDecomposition
10from gpaw.tddft.units import as_to_au, au_to_as, eV_to_au, au_to_eV
11from gpaw.mpi import world
13from .buffer import DensityMatrixBuffer
14from .distributed.pulse import PulseConvolver
15from .readers.gpaw import KohnShamRhoWfsReader
16from .readers.numpy import FrequencyDensityMatrixReader, TimeDensityMatrixReader
17from .density_matrix import DensityMatrix
18from .base import BaseDensityMatrices, WorkMetadata
19from ..perturbation import create_perturbation, Perturbation, DeltaKick, PerturbationLike, PulsePerturbation
20from ..utils import Logger, two_communicator_sizes, get_gaussian_pulse_values
21from ..utils.logging import format_times
22from ..utils.memory import MemoryEstimate
23from ..typing import Array1D, Communicator
26class ConvolutionDensityMatrixMetadata(WorkMetadata):
27 """ Metadata to the density matrices.
29 Properties
30 ----------
31 density_matrices
32 Parent of this object.
33 globalt
34 Time index.
35 localt
36 Time index on this rank.
37 globalp
38 Pulse index.
39 localp
40 Pulse index on this rank.
41 """
42 density_matrices: ConvolutionDensityMatrices
43 globalt: int
44 globalp: int
45 localt: int
46 localp: int
48 def __new__(cls,
49 density_matrices: ConvolutionDensityMatrices,
50 globalt: int,
51 globalp: int,
52 localt: int,
53 localp: int):
54 self = WorkMetadata.__new__(cls, density_matrices=density_matrices)
55 self.globalt = globalt
56 self.globalp = globalp
57 self.localt = localt
58 self.localp = localp
59 return self
61 @property
62 def global_indices(self):
63 return (self.globalp, self.globalt)
65 @property
66 def time(self) -> float:
67 """ Simulation time in units of as. """
68 return self.density_matrices.times[self.globalt]
70 @property
71 def pulse(self) -> Perturbation:
72 """ Pulse. """
73 return self.density_matrices.pulses[self.globalp]
75 @property
76 def desc(self) -> str:
77 timestr = f'{self.time:.1f}as'
78 if len(self.density_matrices.pulses) == 1:
79 return timestr
81 d = self.pulse.todict()
82 if d['name'] == 'GaussianPulse':
83 pulsestr = f'{d["frequency"]:.1f}eV'
84 else:
85 pulsestr = f'#{self.globalp}'
86 return f'{timestr} @ Pulse {pulsestr}'
89class ConvolutionDensityMatrices(BaseDensityMatrices[ConvolutionDensityMatrixMetadata]):
91 """
92 Collection of density matrices in the Kohn-Sham basis for different times,
93 after convolution with various pulses.
95 Plain density matrices and/or derivatives thereof may be represented.
97 Parameters
98 ----------
99 ksd
100 KohnShamDecomposition object or file name.
101 pulses
102 Convolute the density matrices with these pulses.
103 times
104 Compute density matrices for these times (or as close to them as possible). In units of as.
105 derivative_order_s
106 Compute density matrix derivatives of the following orders.
107 0 for plain density matrix and positive integers for derivatives.
108 real
109 Calculate the real part of density matrices.
110 imag
111 Calculate the imaginary part of density matrices.
112 calc_size
113 Size of the calculation communicator.
114 """
116 _derivative_order_s: list[int]
118 def __init__(self,
119 ksd: KohnShamDecomposition | str,
120 pulses: Collection[PerturbationLike],
121 times: list[float] | NDArray[np.float64],
122 derivative_order_s: list[int] = [0],
123 real: bool = True,
124 imag: bool = True,
125 calc_size: int = 1,
126 log: Logger | None = None):
127 self._time_t = np.array(times)
128 self._pulses = [create_perturbation(pulse) for pulse in pulses]
129 self._derivative_order_s = derivative_order_s
131 super().__init__(ksd=ksd,
132 real=real, imag=imag,
133 calc_size=calc_size, log=log)
135 tag_s = ['', '-Iomega', '-omega2']
136 return_keys_by_order = {0: 'pulserho', 1: 'pulsedrho', 2: 'pulseddrho'}
137 self.keys_by_order = {0: 'rho_ia', 1: 'drho_ia', 2: 'ddrho_ia'}
138 if 1 not in derivative_order_s:
139 tag_s.remove('-Iomega')
140 return_keys_by_order.pop(1)
141 if 2 not in derivative_order_s:
142 tag_s.remove('-omega2')
143 return_keys_by_order.pop(2)
144 self.tag_s = tag_s
145 self.return_keys_by_order = return_keys_by_order
147 @property
148 def times(self) -> Array1D[np.float64]:
149 """ Simulation time in units of as. """
150 return self._time_t # type: ignore
152 @property
153 def nt(self) -> int:
154 """ Number of times. """
155 return len(self.times)
157 @property
158 def pulses(self) -> list[Perturbation]:
159 """ Pulses with which density matrices are convoluted. """
160 return self._pulses
162 @property
163 def derivative_order_s(self) -> list[int]:
164 """
165 List of orders of the density matrix derivatives to be computed.
166 `0` for plain density matrix and positive integers for derivatives.
167 """
168 return self._derivative_order_s
170 def work_loop(self,
171 rank: int) -> Generator[ConvolutionDensityMatrixMetadata | None, None, None]:
172 nt = len(self.times)
173 ntperrank = (nt + self.loop_comm.size - 1) // self.loop_comm.size
175 # Do convolution pulse-by-pulse, time-by-time
176 for p, pulse in enumerate(self.pulses):
177 # Determine which times to compute on this loop_comm rank for good load balancing
178 shift = (p * nt + rank) % self.loop_comm.size
179 for localt in range(ntperrank):
180 globalt = shift + localt * self.loop_comm.size
181 if globalt < nt:
182 yield ConvolutionDensityMatrixMetadata(density_matrices=self, globalt=globalt,
183 localt=localt, globalp=p, localp=p)
184 else:
185 yield None
187 def write_to_disk(self,
188 pulserho_fmt: str):
189 """ Calculate the density matrices and save to disk.
191 Parameters
192 ----------
193 pulserho_fmt
194 Formatting string for the density matrices saved to disk.
196 The formatting string should be a plain string containing variable
197 placeholders within curly brackets `{}`. It should not be confused with
198 a formatted string literal (f-string).
200 Example:
202 * pulserho_fmt = `pulserho/t{time:09.1f}{tag}.npy`.
204 Accepts variables
206 * `{time}` - Time in units of as.
207 * `{tag}` - Derivative tag, `''`, `'-Iomega'`, or `'-omega2'`.
208 * `{pulsefreq}` - Pulse frequency in units of eV.
209 * `{pulsefwhm}` - Pulse FWHM in units of fs.
210 """
211 nlocaltot = len(self.local_work_plan)
213 tags_keys = [(tag, key) for s, (tag, key) in enumerate(
214 [('', 'rho_ia'), ('-Iomega', 'drho_ia'), ('-omega2', 'ddrho_ia')]) if s in self.derivative_order_s]
216 # Iterate over density matrices on all ranks
217 for ndone, (work, dm) in enumerate(self, 1):
218 avg = self.log.elapsed('read')/ndone
219 estrem = avg * (nlocaltot - ndone)
220 self.log(f'Obtained density matrix {ndone:4.0f}/{nlocaltot:4.0f} on this rank. '
221 f'Avg: {avg:10.3f}s, estimated remaining: {estrem:10.3f}s', who='Writer', flush=True)
222 for tag, key in tags_keys:
223 fname_kw = dict(time=work.time, tag=tag, **get_gaussian_pulse_values(work.pulse))
224 fpath = Path(pulserho_fmt.format(**fname_kw))
225 rho_ia = getattr(dm, key)
226 if self.calc_comm.rank == 0:
227 assert isinstance(rho_ia, np.ndarray)
228 fpath.parent.mkdir(parents=True, exist_ok=True)
229 if fpath.suffix == '.npz':
230 # Write numpy archive
231 np.savez(str(fpath), rho_ia=rho_ia)
232 else:
233 # Save numpy binary file
234 np.save(str(fpath), rho_ia)
235 self.log(f'Written {fpath}', who='Writer', flush=True)
236 world.barrier()
239class TimeDensityMatricesFromWaveFunctions(ConvolutionDensityMatrices):
241 """
242 Collection of density matrices in the Kohn-Sham basis for different times,
243 obtained by reading the time-dependent wave functions file.
245 Parameters
246 ----------
247 ksd
248 KohnShamDecomposition object or file name.
249 wfs_fname
250 File name of the time-dependent wave functions file.
251 times
252 Compute density matrices for these times (or as close to them as possible). In units of as.
253 real
254 Calculate the real part of density matrices.
255 imag
256 Calculate the imaginary part of density matrices.
257 calc_size
258 Size of the calculation communicator.
259 stridet
260 Skip this many steps when reading the time-dependent wave functions file.
261 """
263 def __init__(self,
264 ksd: KohnShamDecomposition | str,
265 wfs_fname: str,
266 times: list[float] | Array1D[np.float64],
267 real: bool = True,
268 imag: bool = True,
269 calc_size: int = 1,
270 log: Logger | None = None,
271 stridet: int = 1):
272 rho_nn_direct = KohnShamRhoWfsReader(wfs_fname=wfs_fname, ksd=ksd,
273 yield_re=real, yield_im=imag,
274 filter_times=np.array(times) * as_to_au, # type: ignore
275 stridet=stridet, log=log)
276 self.rho_nn_direct = rho_nn_direct
278 super().__init__(ksd=rho_nn_direct.ksd,
279 times=rho_nn_direct.time_t * au_to_as,
280 real=real,
281 imag=imag,
282 pulses=[None],
283 calc_size=1,
284 log=rho_nn_direct.log)
286 imin, imax, amin, amax = self.ksd.ialims()
288 # Read density matrices corresponding to ksd ialims
289 self._n1slice = slice(imin, imax + 1)
290 self._n2slice = slice(amin, amax + 1)
292 def __str__(self) -> str:
293 lines = ['Response from time-dependent wave functions']
295 lines.append('')
296 lines += [' ' + line for line in str(self.rho_nn_direct).split('\n')]
297 shape = (self._n1slice.stop - self._n1slice.start,
298 self._n2slice.stop - self._n2slice.start,)
299 if 'Re' in self.reim and 'Im' in self.reim:
300 reim_desc = 'real and imaginary parts'
301 elif 'Re' in self.reim:
302 reim_desc = 'real parts'
303 else:
304 reim_desc = 'imaginary parts'
306 lines.append('')
307 lines.append(f'Density matrix shape {shape}')
308 lines.append(f'reading {reim_desc}')
310 return '\n'.join(lines)
312 def get_memory_estimate(self) -> MemoryEstimate:
313 shape = (self._n1slice.stop - self._n1slice.start,
314 self._n2slice.stop - self._n2slice.start,)
315 narrays = 2 if 'Re' in self.reim and 'Im' in self.reim else 1
316 for t_r in self.rho_nn_direct.work_loop_by_ranks():
317 nreading = len(t_r)
318 break
319 memory_estimate = MemoryEstimate()
320 memory_estimate.add_child('Time-dependent wave functions reader',
321 self.rho_nn_direct.get_memory_estimate())
322 memory_estimate.add_key('Density matrix', (shape) + (narrays, ),
323 on_num_ranks=nreading)
325 return memory_estimate
327 def __iter__(self) -> Generator[tuple[ConvolutionDensityMatrixMetadata, DensityMatrix], None, None]:
328 assert self.calc_comm.size == 1 # TODO
330 for work, dm_buffer in zip(self.work_loop(self.loop_comm.rank),
331 self.rho_nn_direct.iread(0, 0, self._n1slice, self._n2slice)):
332 assert work is not None
333 if 'Re' in self.reim:
334 Rerho_ia = dm_buffer._get_real(0)
335 if 'Im' in self.reim:
336 Imrho_ia = dm_buffer._get_imag(0)
337 if 'Re' in self.reim and 'Im' in self.reim:
338 # Complex result
339 # Compared to numpy, we use another convention, hence the minus sign
340 rho_ia = Rerho_ia - 1j * Imrho_ia
341 elif 'Re' in self.reim:
342 # Real result
343 rho_ia = Rerho_ia
344 else:
345 rho_ia = -1j * Imrho_ia
346 matrices: dict[int, NDArray[np.complex128] | None] = {0: rho_ia}
347 dm = DensityMatrix(ksd=self.ksd, matrices=matrices, comm=self.calc_comm)
349 yield work, dm
351 def parallel_prepare(self):
352 self.rho_nn_direct.parallel_prepare()
355class ConvolutionDensityMatricesFromDisk(ConvolutionDensityMatrices):
357 """
358 Collection of density matrices in the Kohn-Sham basis after convolution with
359 one or several pulses, for different times. Read from disk.
361 Parameters
362 ----------
363 ksd
364 KohnShamDecomposition object or file name.
365 pulserho_fmt
366 Formatting string for the density matrices saved to disk.
368 The formatting string should be a plain string containing variable
369 placeholders within curly brackets `{}`. It should not be confused with
370 a formatted string literal (f-string).
372 Example:
374 * pulserho_fmt = `pulserho/t{time:09.1f}{tag}.npy`.
376 Accepts variables
378 * `{time}` - Time in units of as.
379 * `{tag}` - Derivative tag, `''`, `'-Iomega'`, or `'-omega2'`.
380 * `{pulsefreq}` - Pulse frequency in units of eV.
381 * `{pulsefwhm}` - Pulse FWHM in units of fs.
382 pulses
383 Convolute the density matrices with these pulses.
384 times
385 Compute density matrices for these times (or as close to them as possible). In units of as.
386 derivative_order_s
387 Compute density matrix derivatives of the following orders.
388 `0` for plain density matrix and positive integers for derivatives.
389 real
390 Calculate real part of density matrices.
391 imag
392 Calculate imaginary part of density matrices.
393 calc_size
394 Size of the calculation communicator.
395 """
397 def __init__(self,
398 ksd: KohnShamDecomposition | str,
399 pulserho_fmt: str,
400 pulses: Collection[PerturbationLike],
401 times: list[float] | Array1D[np.float64],
402 derivative_order_s: list[int] = [0],
403 log: Logger | None = None,
404 real: bool = True,
405 imag: bool = True,
406 calc_size: int = 1):
407 reader = TimeDensityMatrixReader(pulserho_fmt=pulserho_fmt,
408 ksd=ksd,
409 pulses=pulses,
410 filter_times=times,
411 derivative_order_s=derivative_order_s,
412 log=log)
413 self.reader = reader
414 super().__init__(ksd=reader.ksd, pulses=reader.pulses, times=reader.times,
415 real=real, imag=imag,
416 derivative_order_s=derivative_order_s, calc_size=calc_size, log=reader.log)
417 self.ksd.distribute(self.calc_comm)
418 self.pulserho_fmt = pulserho_fmt
420 def __str__(self) -> str:
421 return str(self.reader)
423 def __iter__(self) -> Generator[tuple[ConvolutionDensityMatrixMetadata, DensityMatrix], None, None]:
424 for work in self.local_work_plan:
425 self.log.start('read')
426 matrices: dict[int, NDArray[np.complex128] | None] = dict()
427 for derivative in self.derivative_order_s:
428 if self.calc_comm.rank > 0:
429 # Don't read on non calc comm root ranks
430 matrices[derivative] = None
431 continue
433 matrices[derivative] = self.reader.read(work.time, work.pulse, derivative)
435 dm = DensityMatrix(ksd=self.ksd, matrices=matrices, comm=self.calc_comm)
437 yield work, dm
440class ConvolutionDensityMatricesFromFrequency(ConvolutionDensityMatrices):
442 """
443 Collection of density matrices in the Kohn-Sham basis after convolution with
444 one or several pulses, for different times. Obtained from the the density
445 matrices in frequency space, which are read from disk.
447 Parameters
448 ----------
449 ksd
450 KohnShamDecomposition object or file name.
451 frho_fmt
452 Formatting string for the density matrices
453 in frequency space saved to disk.
455 The formatting string should be a plain string containing variable
456 placeholders within curly brackets `{}`. It should not be confused with
457 a formatted string literal (f-string).
459 Example:
461 * frho_fmt = `frho/w{freq:05.2f}-{reim}.npy'`.
463 Accepts variables:
465 * `{freq}` - Frequency in units of eV.
466 * `{reim}` - `'Re'` or `'Im'` for Fourier transform of real/imaginary
467 part of density matrix.
468 perturbation
469 Perturbation that was present during time propagation.
470 pulses
471 Convolute the density matrices with these pulses.
472 times
473 Compute density matrices for these times (or as close to them as possible). In units of as.
474 derivative_order_s
475 Compute density matrix derivatives of the following orders.
476 `0` for plain density matrix and positive integers for derivatives.
477 real
478 Calculate real part of density matrices.
479 imag
480 Calculate imaginary part of density matrices.
481 calc_size
482 Size of the calculation communicator.
483 """
485 def __init__(self,
486 ksd: KohnShamDecomposition | str,
487 frho_fmt: str,
488 perturbation: PerturbationLike,
489 pulses: Collection[PerturbationLike],
490 times: list[float] | NDArray[np.float64],
491 derivative_order_s: list[int] = [0],
492 log: Logger | None = None,
493 real: bool = True,
494 imag: bool = True,
495 calc_size: int = 1):
496 super().__init__(ksd=ksd, pulses=pulses, times=times,
497 real=real, imag=imag,
498 derivative_order_s=derivative_order_s, calc_size=calc_size, log=log)
499 # Frequency grid for convolution
500 if not all(pulse.todict()['name'] == 'GaussianPulse' for pulse in self.pulses):
501 raise NotImplementedError('Only Gaussian pulses implemented for ResponseFromFourierTransform')
503 freq_spacing = 0.05 # 50meV
504 freq_min = min((pulse.pulse.omega0 - 4 * pulse.pulse.sigma) * au_to_eV # type: ignore
505 for pulse in self.pulses)
506 freq_min = max(freq_min, freq_spacing)
507 freq_max = min((pulse.pulse.omega0 + 4 * pulse.pulse.sigma) * au_to_eV # type: ignore
508 for pulse in self.pulses)
510 transformer = ExactFourierTransformer.from_file(
511 frho_fmt=frho_fmt, ksd=ksd, perturbation=perturbation,
512 real=real, imag=imag, log=self.log, comm=self.calc_comm,
513 freq_spacing=freq_spacing, freq_min=freq_min, freq_max=freq_max)
514 self.transformer = transformer
516 def __str__(self) -> str:
517 lines = ['Response from Fourier transform of density matrices on disk']
518 lines.append('')
519 lines.append(f'Calculating response for {self.nt} times and {len(self.pulses)} pulses')
520 lines.append(f' times: {format_times(self.times)}')
522 return '\n'.join(lines)
524 def __iter__(self) -> Generator[tuple[ConvolutionDensityMatrixMetadata, DensityMatrix], None, None]:
525 self.transformer.parallel_prepare()
527 for work in self.local_work_plan:
528 matrices: dict[int, NDArray[np.complex128] | None] = dict()
530 for derivative in self.derivative_order_s:
531 buffer = self.transformer.convolve(work.time, work.pulse, derivative)
532 if self.calc_comm.rank > 0:
533 matrices[derivative] = None
534 continue
535 assert buffer is not None
536 # Buffer shape is i, a, pulses, times
537 if 'Re' in self.reim:
538 Rerho_ia = buffer._get_real(derivative)
539 if 'Im' in self.reim:
540 Imrho_ia = buffer._get_imag(derivative)
541 if 'Re' in self.reim and 'Im' in self.reim:
542 # Complex result
543 rho_ia = Rerho_ia + 1j * Imrho_ia
544 elif 'Re' in self.reim:
545 # Real result
546 rho_ia = Rerho_ia
547 else:
548 rho_ia = 1j * Imrho_ia
549 matrices[derivative] = rho_ia
551 dm = DensityMatrix(ksd=self.ksd, matrices=matrices, comm=self.calc_comm)
553 yield work, dm
556class ConvolutionDensityMatricesFromWaveFunctions(ConvolutionDensityMatrices):
558 """
559 Collection of density matrices in the Kohn-Sham basis after convolution with
560 one or several pulses, for different times. Obtained from the time-dependent wave functions file,
561 which is read from disk.
563 Parameters
564 ----------
565 ksd
566 KohnShamDecomposition object or file name.
567 wfs_fname
568 File name of the time-dependent wave functions file.
569 perturbation
570 Perturbation that was present during time propagation.
571 pulses
572 Convolute the density matrices with these pulses.
573 times
574 Compute density matrices for these times (or as close to them as possible). In units of as.
575 derivative_order_s
576 Compute density matrix derivatives of the following orders.
577 `0` for plain density matrix and positive integers for derivatives.
578 real
579 Calculate real part of density matrices.
580 imag
581 Calculate imaginary part of density matrices.
582 calc_size
583 Size of the calculation communicator.
584 stridet
585 Skip this many steps when reading the time-dependent wave functions file.
586 """
588 def __init__(self,
589 ksd: KohnShamDecomposition | str,
590 wfs_fname: str,
591 perturbation: PerturbationLike,
592 pulses: Collection[PerturbationLike],
593 times: list[float] | Array1D[np.float64],
594 derivative_order_s: list[int] = [0],
595 real: bool = True,
596 imag: bool = True,
597 calc_size: int = 1,
598 log: Logger | None = None,
599 stridet: int = 1):
600 _, calc_size = two_communicator_sizes(-1, calc_size)
601 # The calc_comm rank 0's are world ranks 0, with a spacing of calc_size
602 result_on_ranks = list(range(0, world.size, calc_size))
604 rho_nn_conv = PulseConvolver.from_parameters(
605 wfs_fname, ksd,
606 pulses=pulses,
607 perturbation=perturbation,
608 yield_re=real, yield_im=imag,
609 filter_times=np.array(times) * as_to_au,
610 derivative_order_s=list(sorted(derivative_order_s)),
611 stridet=stridet,
612 result_on_ranks=result_on_ranks,
613 log=log)
614 self.rho_nn_conv = rho_nn_conv
616 super().__init__(ksd=rho_nn_conv.ksd, pulses=pulses, times=rho_nn_conv.time_t * au_to_as,
617 real=real, imag=imag,
618 derivative_order_s=derivative_order_s, calc_size=calc_size,
619 log=rho_nn_conv.log)
621 def __str__(self) -> str:
622 lines = []
623 lines.append('Response from time-dependent wave functions ')
624 lines.append('')
625 lines += ['' + line for line in str(self.rho_nn_conv.rho_nn_reader.rho_wfs_reader).split('\n')]
626 lines.append('')
627 lines += ['' + line for line in str(self.rho_nn_conv.rho_nn_reader).split('\n')]
628 lines.append('')
629 lines += ['' + line for line in str(self.rho_nn_conv).split('\n')]
631 return '\n'.join(lines)
633 def get_memory_estimate(self) -> MemoryEstimate:
634 return self.rho_nn_conv.get_memory_estimate()
636 @property
637 def myt(self) -> list[int]:
638 """ List of indices corresponding to the time indices on held on this rank. """
639 return self.rho_nn_conv.my_work()
641 def work_loop(self,
642 rank: int) -> Generator[ConvolutionDensityMatrixMetadata | None, None, None]:
643 nt = len(self.times)
644 ntperrank = (nt + self.loop_comm.size - 1) // self.loop_comm.size
646 for p, pulse in enumerate(self.pulses):
647 for localt in range(ntperrank):
648 globalt = rank + localt * self.loop_comm.size
649 if globalt < nt:
650 yield ConvolutionDensityMatrixMetadata(density_matrices=self, globalt=globalt,
651 localt=localt, globalp=p, localp=p)
652 else:
653 yield None
655 def __iter__(self) -> Generator[tuple[ConvolutionDensityMatrixMetadata, DensityMatrix], None, None]:
656 parameters = self.rho_nn_conv.rho_nn_reader._parameters
657 flt = (slice(parameters.n1size), slice(parameters.n2size))
659 dist_buffer = self.rho_nn_conv.dist_buffer # Perform the redistribution
660 self.ksd.distribute(self.calc_comm)
662 if self.calc_comm.rank != 0:
663 assert len(self.myt) == 0, self.myt
665 for work in self.local_work_plan:
666 if self.calc_comm.rank == 0:
667 assert self.myt[work.localt] == work.globalt
668 localflt = flt + (work.localp, work.localt)
670 matrices: dict[int, NDArray[np.complex128] | None] = dict()
671 for derivative in self.derivative_order_s:
672 if self.calc_comm.rank > 0:
673 matrices[derivative] = None
674 continue
675 # Buffer shape is i, a, pulses, times
676 if 'Re' in self.reim:
677 Rerho_ia = dist_buffer._get_real(derivative)[localflt]
678 if 'Im' in self.reim:
679 Imrho_ia = dist_buffer._get_imag(derivative)[localflt]
680 if 'Re' in self.reim and 'Im' in self.reim:
681 # Complex result
682 # Compared to numpy, we use another convention, hence the minus sign
683 rho_ia = Rerho_ia - 1j * Imrho_ia
684 elif 'Re' in self.reim:
685 # Real result
686 rho_ia = Rerho_ia
687 else:
688 rho_ia = -1j * Imrho_ia
689 matrices[derivative] = rho_ia
690 dm = DensityMatrix(ksd=self.ksd, matrices=matrices, comm=self.calc_comm)
692 yield work, dm
694 def parallel_prepare(self):
695 self.rho_nn_conv.dist_buffer # Perform the redistribution
698class ExactFourierTransformer:
700 def __init__(self,
701 reader: FrequencyDensityMatrixReader,
702 perturbation: PerturbationLike,
703 real: bool = True,
704 imag: bool = True,
705 comm: Communicator | None = None,
706 log: Logger | None = None):
707 if comm is None:
708 comm = world
709 self._comm = comm
710 self.reader = reader
711 if log is None:
712 log = Logger()
713 self.log = log
715 self.real = real
716 self.imag = imag
718 self.frequency_buffer = DensityMatrixBuffer(self.nnshape, (self.reader.nw, ), np.complex128)
719 self._ready = False
721 omega_w = self.reader.frequencies * eV_to_au
723 self.perturbation = create_perturbation(perturbation)
724 assert isinstance(self.perturbation, DeltaKick)
726 # Need a uniform spacing for Simpson's integration rule
727 dw = omega_w[1] - omega_w[0]
728 nw = len(omega_w)
729 if not np.allclose(omega_w[1:] - dw, omega_w[:-1]):
730 raise ValueError('Variable frequency spacing.')
731 # integration weights from Simpson's integration rule
732 self._weight_w = dw / 3 * np.array([1] + [4, 2] * int((nw - 2) / 2)
733 + [4] * (nw % 2) + [1])
735 self._weight_w *= 2 / (2 * np.pi)
736 self._weight_w /= self.perturbation.strength
738 @property
739 def comm(self) -> Communicator:
740 """ MPI commonicator.
742 During convolution, the density matrices are distributed with
743 different chunks on different ranks of this communicator.
744 """
745 return self._comm
747 @property
748 def full_nnshape(self) -> tuple[int, int]:
749 imin, imax, amin, amax = [int(i) for i in self.reader.ksd.ialims()]
750 return (imax - imin + 1, amax - amin + 1)
752 @property
753 def nnshape(self) -> tuple[int, int]:
754 stridei = (self.full_nnshape[0] + self.comm.size - 1) // self.comm.size
755 return (stridei, self.full_nnshape[1])
757 @property
758 def myslice(self) -> slice:
759 return self.slice_of_rank(self.comm.rank)
761 def slice_of_rank(self, rank: int) -> slice:
762 stridei = (self.full_nnshape[0] + self.comm.size - 1) // self.comm.size
763 return slice(rank * stridei, (rank + 1) * stridei)
765 def parallel_prepare(self):
766 # Read all the Fourier transform of density matrices
767 self.frequency_buffer.zero_buffers(real=self.real, imag=self.imag, derivative_order_s=[0])
768 for w, freq in enumerate(self.reader.frequencies):
769 if self.real:
770 rho_ia = self.reader.read(frequency=freq, real=True)[self.myslice]
771 self.frequency_buffer[w].real[:len(rho_ia)] = rho_ia
772 if self.imag:
773 rho_ia = self.reader.read(frequency=freq, real=False)[self.myslice]
774 self.frequency_buffer[w].imag[:len(rho_ia)] = rho_ia
775 self._ready = True
777 def convolve(self,
778 time: float,
779 pulse: Perturbation,
780 derivative: int) -> DensityMatrixBuffer | None:
781 if not self._ready:
782 self.parallel_prepare()
784 assert isinstance(pulse, PulsePerturbation)
785 omega_w = self.reader.frequencies * eV_to_au
786 pulse_w = pulse.pulse.fourier(omega_w)
788 buffer = DensityMatrixBuffer(self.frequency_buffer.nnshape, (), dtype=np.float64)
790 self.log.start('convolve')
791 exp_w = np.exp(-1j * omega_w * time * as_to_au)
792 exp_w *= self._weight_w * pulse_w * (-1j * omega_w) ** derivative
794 if self.real:
795 # Real part
796 rho_iaw = self.frequency_buffer.real[:]
797 conv_rho_ia = rho_iaw.real @ exp_w.real
798 conv_rho_ia -= rho_iaw.imag @ exp_w.imag
800 buffer.store(True, derivative, conv_rho_ia)
802 if self.imag:
803 # Imaginary part
804 rho_iaw = self.frequency_buffer.imag[:]
805 conv_rho_ia = rho_iaw.real @ exp_w.real
806 conv_rho_ia -= rho_iaw.imag @ exp_w.imag
808 buffer.store(False, derivative, conv_rho_ia)
810 if self.comm.rank == 0:
811 full_buffer = DensityMatrixBuffer(self.frequency_buffer.nnshape, (), dtype=np.float64)
812 full_buffer.zero_buffers(real=self.real, imag=self.imag, derivative_order_s=[derivative])
814 # Send to root
815 for rank in range(self.comm.size):
816 if rank == self.comm.rank:
817 # Send own work to root
818 buffer.send_arrays(self.comm, 0)
820 if self.comm.rank == 0:
821 # Receive on root, fill buffer
822 buffer.recv_arrays(self.comm, rank)
823 for full_rho_ia, part_rho_ia in zip(full_buffer._iter_buffers(),
824 buffer._iter_buffers()):
825 full_slice_ia = full_rho_ia[self.slice_of_rank(rank)]
826 full_slice_ia[:] = part_rho_ia[:len(full_slice_ia)]
828 if self.comm.rank == 0:
829 self.log(f'Convolved density matrices in {self.log.elapsed("convolve"):.1f} s',
830 who='Response', flush=True)
832 if self.comm.rank == 0:
833 return full_buffer
835 return None
837 @classmethod
838 def from_file(cls,
839 frho_fmt: str,
840 ksd: KohnShamDecomposition | str,
841 perturbation: PerturbationLike,
842 freq_min: float,
843 freq_spacing: float,
844 freq_max: float,
845 real: bool = True,
846 imag: bool = True,
847 comm: Communicator | None = None,
848 log: Logger | None = None) -> ExactFourierTransformer:
849 frequencies = np.arange(freq_min, freq_max + freq_spacing * 0.1, freq_spacing)
851 reader = FrequencyDensityMatrixReader(frho_fmt=frho_fmt,
852 ksd=ksd,
853 filter_frequencies=frequencies,
854 real=real, imag=imag,
855 comm=None, # Look for files on world
856 log=log)
857 if reader.frequencies[0] > freq_min:
858 raise ValueError(f'Could not satisfy lower bound for frequency grid with files matching {frho_fmt}. '
859 f'Found {reader.frequencies[0]:.2f}eV, reqested {freq_min:.2f}eV.')
861 if reader.frequencies[-1] > freq_max:
862 raise ValueError(f'Could not satisfy upper bound for frequency grid with files matching {frho_fmt}. '
863 f'Found {reader.frequencies[-1]:.2f}eV, reqested {freq_max:.2f}eV.')
865 dw = reader.frequencies[1] - reader.frequencies[0]
866 if dw > freq_spacing + 1e-8:
867 raise ValueError(f'Could not construct sufficiently dense frequency grid with files matching {frho_fmt}. '
868 f'Found spacing of {dw:.4f}eV, reqested {freq_spacing:.4f}eV.')
870 return cls(reader=reader, perturbation=perturbation, real=real, imag=imag, comm=comm, log=log)