Coverage for rhodent/response/base.py: 94%
49 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 abc import abstractmethod, ABC
4from typing import Generator, Collection
6import numpy as np
8from gpaw.lcaotddft.ksdecomposition import KohnShamDecomposition
10from ..perturbation import create_perturbation, Perturbation, PerturbationLike
11from ..density_matrices.density_matrix import DensityMatrix
12from ..density_matrices.frequency import (FrequencyDensityMatrixMetadata,
13 FrequencyDensityMatrices)
14from ..density_matrices.time import (ConvolutionDensityMatrixMetadata,
15 ConvolutionDensityMatrices)
16from ..typing import Array1D
17from ..utils import add_fake_kpts, Logger
20class BaseResponse(ABC):
22 """ Object describing response; obtained from :term:`TDDFT` calculation.
24 Parameters
25 ----------
26 ksd
27 KohnShamDecomposition object or file name.
28 perturbation
29 The perturbation that was present during the TDDFT calculation.
30 None to mark it as an ne perturbation.
31 calc_size
32 Size of the calculation communicator.
33 """
34 def __init__(self,
35 ksd: KohnShamDecomposition | str,
36 perturbation: PerturbationLike = None,
37 calc_size: int = 1):
38 if isinstance(ksd, KohnShamDecomposition):
39 self._ksd = ksd
40 else:
41 self._ksd = KohnShamDecomposition(filename=ksd)
42 add_fake_kpts(self._ksd)
44 self._perturbation = create_perturbation(perturbation)
45 self.calc_size = calc_size
47 def __str__(self) -> str:
48 lines = [f'{self.__class__.__name__}']
49 lines += [f' ksd: {self.ksd.filename if self.ksd.filename is not None else "From calc"}']
50 lines += [' perturbation:']
51 lines += [' ' + line for line in str(self.perturbation).split('\n')]
52 return '\n'.join(lines)
54 @property
55 def ksd(self) -> KohnShamDecomposition:
56 """ Kohn-Sham decomposition object. """
57 return self._ksd
59 @property
60 def perturbation(self) -> Perturbation:
61 """ The perturbation that caused this response. """
62 return self._perturbation
64 @abstractmethod
65 def _get_time_density_matrices(self,
66 times: list[float] | Array1D[np.float64],
67 pulses: Collection[PerturbationLike],
68 derivative_order_s: list[int] = [0],
69 real: bool = True,
70 imag: bool = True,
71 log: Logger | None = None,
72 ) -> ConvolutionDensityMatrices:
73 raise NotImplementedError
75 @abstractmethod
76 def _get_frequency_density_matrices(self,
77 frequencies: list[float] | Array1D[np.float64],
78 frequency_broadening: float = 0,
79 real: bool = True,
80 imag: bool = True,
81 log: Logger | None = None,
82 ) -> FrequencyDensityMatrices:
83 raise NotImplementedError
85 def iterate_density_matrices_in_time(self,
86 times: list[float] | Array1D[np.float64],
87 pulses: Collection[PerturbationLike],
88 derivative_order_s: list[int] = [0],
89 real: bool = True,
90 imag: bool = True,
91 log: Logger | None = None,
92 ) -> Generator[tuple[ConvolutionDensityMatrixMetadata,
93 DensityMatrix], None, None]:
94 """ Obtain density matrices at the given times in response to the given pulses.
96 If the given pulse(s) differ from the perturbation that caused this response,
97 then the pulse convolution trick is applied to obtain the response to the given
98 pulse(s).
100 Parameters
101 ----------
102 times
103 Calculate density matrices for these times (or as close to them as possible). In units of as.
104 pulses
105 Calculate density matrices in response to these pulses.
106 derivative_order_s
107 Calculate density matrix derivatives of the following orders.
108 ``0`` for plain density matrix and positive integers for derivatives.
109 real
110 Calculate the real part of density matrices.
111 imag
112 Calculate the imaginary part of density matrices.
113 log
114 Logger object.
116 Yields
117 ------
118 Tuple (work, dm) on the root rank of the calculation communicator:
120 work
121 An object representing the metadata (time and pulse) for the work done.
122 dm
123 Density matrix for this time and pulse.
124 """
125 density_matrices = self._get_time_density_matrices(
126 times, pulses, derivative_order_s, real, imag, log)
127 yield from density_matrices
129 def iterate_density_matrices_in_frequency(self,
130 frequencies: list[float] | Array1D[np.float64],
131 frequency_broadening: float = 0,
132 real: bool = True,
133 imag: bool = True,
134 log: Logger | None = None,
135 ) -> Generator[tuple[FrequencyDensityMatrixMetadata,
136 DensityMatrix], None, None]:
137 """ Obtain density matrices at the given frequencies.
139 Parameters
140 ----------
141 frequencies
142 Compute density matrices for these frequencies (or as close to them as possible). In units of eV.
143 frequency_broadening
144 Gaussian broadening width in atomic units. Default (0) is no broadening.
145 real
146 Calculate the Fourier transform of the real part of the density matrix.
147 imag
148 Calculate the Fourier transform of the imaginary part of the density matrix.
150 Yields
151 ------
152 Tuple (work, dm) on the root rank of the calculation communicator:
154 work
155 An object representing the metadata (frequency) for the work done.
156 dm
157 Density matrix for this frequency.
158 """
159 density_matrices = self._get_frequency_density_matrices(
160 frequencies, frequency_broadening, real, imag, log)
161 yield from density_matrices
163 def write_in_time(self,
164 pulserho_fmt: str,
165 times: list[float] | Array1D[np.float64],
166 pulses: Collection[PerturbationLike],
167 derivative_order_s: list[int] = [0],
168 real: bool = True,
169 imag: bool = True,
170 log: Logger | None = None):
171 density_matrices = self._get_time_density_matrices(
172 times, pulses, derivative_order_s, real, imag, log)
173 density_matrices.write_to_disk(pulserho_fmt)
175 def write_in_frequency(self,
176 frho_fmt: str,
177 frequencies: list[float] | Array1D[np.float64],
178 frequency_broadening: float = 0,
179 real: bool = True,
180 imag: bool = True,
181 log: Logger | None = None):
182 density_matrices = self._get_frequency_density_matrices(
183 frequencies, frequency_broadening, real, imag, log)
184 density_matrices.write_to_disk(frho_fmt)