Coverage for rhodent/response/numpy.py: 69%
51 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-08-01 16:57 +0000
1from __future__ import annotations
3from typing import Collection
5import numpy as np
7from gpaw.lcaotddft.ksdecomposition import KohnShamDecomposition
9from .base import BaseResponse
10from ..perturbation import DeltaKick, PerturbationLike, NoPerturbation
11from ..density_matrices.frequency import (FrequencyDensityMatrices,
12 FrequencyDensityMatricesFromDisk)
13from ..density_matrices.time import (ConvolutionDensityMatricesFromDisk,
14 ConvolutionDensityMatricesFromFrequency,
15 ConvolutionDensityMatrices)
16from ..utils import Logger
17from ..typing import Array1D
20class ResponseFromDensityMatrices(BaseResponse):
22 """ Response from density matrices saved on disk.
24 Parameters
25 ----------
26 pulserho_fmt
27 Formatting string for the density matrices saved to disk.
29 The formatting string should be a plain string containing variable
30 placeholders within curly brackets `{}`. It should not be confused with
31 a formatted string literal (f-string).
33 Example:
35 * pulserho_fmt = `pulserho/t{time:09.1f}{tag}.npy`.
37 Accepts variables
39 * `{time}` - Time in units of as.
40 * `{tag}` - Derivative tag, `''`, `'-Iomega'`, or `'-omega2'`.
41 * `{pulsefreq}` - Pulse frequency in units of eV.
42 * `{pulsefwhm}` - Pulse FWHM in units of fs.
43 ksd
44 KohnShamDecomposition object or file name.
45 perturbation
46 Perturbation that was present during time propagation.
47 calc_size
48 Size of the calculation communicator.
49 """
51 def __init__(self,
52 pulserho_fmt: str,
53 ksd: KohnShamDecomposition | str,
54 perturbation: PerturbationLike = None,
55 calc_size: int = 1):
56 super().__init__(ksd=ksd,
57 perturbation=perturbation,
58 calc_size=calc_size)
59 self.pulserho_fmt = pulserho_fmt
61 def __str__(self) -> str:
62 lines = [f'{self.__class__.__name__}']
63 lines += [f' ksd: {self.ksd.filename if self.ksd.filename is not None else "From calc"}']
64 lines += [f' pulserho_fmt: {self.pulserho_fmt}']
65 lines += [' perturbation:']
66 lines += [' ' + line for line in str(self.perturbation).split('\n')]
67 return '\n'.join(lines)
69 def _get_time_density_matrices(self,
70 times: list[float] | Array1D[np.float64],
71 pulses: Collection[PerturbationLike],
72 derivative_order_s: list[int] = [0],
73 real: bool = True,
74 imag: bool = True,
75 log: Logger | None = None,
76 ) -> ConvolutionDensityMatrices:
77 if (len(pulses) == 1 and list(pulses)[0] == self.perturbation) or isinstance(self.perturbation, NoPerturbation):
78 # Yield the density matrices without performing convolution
79 density_matrices = ConvolutionDensityMatricesFromDisk(
80 ksd=self.ksd,
81 pulserho_fmt=self.pulserho_fmt,
82 times=times,
83 pulses=pulses,
84 derivative_order_s=derivative_order_s,
85 real=real,
86 imag=imag,
87 log=log,
88 calc_size=self.calc_size)
89 else:
90 raise NotImplementedError('Pulse convolution of density matrices on disk is not implemented')
92 return density_matrices
94 def _get_frequency_density_matrices(self,
95 frequencies: list[float] | Array1D[np.float64],
96 frequency_broadening: float = 0,
97 real: bool = True,
98 imag: bool = True,
99 log: Logger | None = None,
100 ) -> FrequencyDensityMatrices:
101 raise NotImplementedError('Fourier transformation of density matrices on disk is not implemented')
104class ResponseFromFourierTransform(BaseResponse):
106 """ Response from Fourier transform of density matrices save on disk.
108 Parameters
109 ----------
110 frho_fmt
111 Formatting string for the density matrices
112 in frequency space saved to disk.
114 The formatting string should be a plain string containing variable
115 placeholders within curly brackets `{}`. It should not be confused with
116 a formatted string literal (f-string).
118 Example:
120 * frho_fmt = `frho/w{freq:05.2f}-{reim}.npy`.
122 Accepts variables:
124 * `{freq}` - Frequency in units of eV.
125 * `{reim}` - `'Re'` or `'Im'` for Fourier transform of real/imaginary
126 part of density matrix.
127 ksd
128 KohnShamDecomposition object or file name.
129 perturbation
130 Perturbation that was present during time propagation.
131 calc_size
132 Size of the calculation communicator.
133 """
135 def __init__(self,
136 frho_fmt: str,
137 ksd: KohnShamDecomposition | str,
138 perturbation: PerturbationLike = None,
139 calc_size: int = 1):
140 super().__init__(ksd=ksd,
141 perturbation=perturbation,
142 calc_size=calc_size)
144 if not isinstance(self.perturbation, (DeltaKick, NoPerturbation)):
145 raise NotImplementedError('Only delta kick implemented')
147 self.frho_fmt = frho_fmt
149 def __str__(self) -> str:
150 lines = [f'{self.__class__.__name__}']
151 lines += [f' ksd: {self.ksd.filename if self.ksd.filename is not None else "From calc"}']
152 lines += [f' frho_fmt: {self.frho_fmt}']
153 lines += [' perturbation:']
154 lines += [' ' + line for line in str(self.perturbation).split('\n')]
155 return '\n'.join(lines)
157 def _get_time_density_matrices(self,
158 times: list[float] | Array1D[np.float64],
159 pulses: Collection[PerturbationLike],
160 derivative_order_s: list[int] = [0],
161 real: bool = True,
162 imag: bool = True,
163 log: Logger | None = None,
164 ) -> ConvolutionDensityMatrices:
165 if not isinstance(self.perturbation, DeltaKick):
166 raise NotImplementedError('Only delta kick implemented')
168 density_matrices = ConvolutionDensityMatricesFromFrequency(
169 ksd=self.ksd,
170 frho_fmt=self.frho_fmt,
171 perturbation=self.perturbation,
172 pulses=pulses,
173 times=times,
174 derivative_order_s=derivative_order_s,
175 real=real,
176 imag=imag,
177 log=log,
178 calc_size=self.calc_size)
180 return density_matrices
182 def _get_frequency_density_matrices(self,
183 frequencies: list[float] | Array1D[np.float64],
184 frequency_broadening: float = 0,
185 real: bool = True,
186 imag: bool = True,
187 log: Logger | None = None,
188 ) -> FrequencyDensityMatrices:
189 if frequency_broadening > 0:
190 raise NotImplementedError(f'Frequency broadening not implemented for {self.__class__.__name__}')
192 density_matrices = FrequencyDensityMatricesFromDisk(
193 ksd=self.ksd,
194 frho_fmt=self.frho_fmt,
195 perturbation=self.perturbation,
196 frequencies=frequencies,
197 real=real,
198 imag=imag,
199 log=log,
200 calc_size=self.calc_size)
202 return density_matrices