Coverage for rhodent/spectrum.py: 90%
89 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-08-01 16:57 +0000
1from __future__ import annotations
3from typing import Any
4from pathlib import Path
5import numpy as np
6from numpy.typing import NDArray
8from ase.parallel import parprint
9from gpaw.mpi import world
10from gpaw.tddft.units import as_to_au, au_to_as, eV_to_au, au_to_eV, eA_to_au, au_to_eA
11from gpaw.tddft.spectrum import read_dipole_moment_file
13from .perturbation import create_perturbation, Perturbation, PerturbationLike, NoPerturbation
16class SpectrumCalculator:
18 r""" Spectrum calculator.
20 Calculates the dynamic polarizability and dipole strength function (the spectrum).
22 The polarizability :math:`\alpha(\omega)` is related to the perturbation
23 :math:`\mathcal{E}(\omega)` and the Fouier transform of the dipole moment
24 :math:`\boldsymbol\mu(t)` as
26 .. math::
28 \boldsymbol\alpha(\omega) \mathcal{E}(\omega) = \mathcal{F}[\boldsymbol\mu(t)](\omega).
30 The dipole strength function is
32 .. math::
34 \boldsymbol{S}(\omega) = \frac{2}{\pi} \omega\:\mathrm{Im}[\boldsymbol\alpha(\omega)].
36 Both quantities can be broadened by supplying a non-zero value :math:`\sigma`. Then the
37 convolution
39 .. math::
41 \frac{1}{\sqrt{2\pi/\sigma^2}}
42 \int_{-\infty}^{\infty} \boldsymbol\alpha(\omega')
43 \exp\left(- \frac{(\omega - \omega')^2}{2\sigma^2}\right) \mathrm{d}\omega'
45 is computed. When the perturbation is a delta-kick, this can efficiently be computed
46 by multiplying the dipole moment by a Gaussian envelope before Fourier transformation
48 .. math::
50 \boldsymbol\mu(t) \exp(-\sigma^2 t^2 / 2).
52 For other perturbations, a :term:`FFT` and :term:`IFFT` is first
53 performed to obtain the delta-kick response.
56 Parameters
57 ----------
58 times
59 Array (length T) of times in units of as.
60 dipole_moments
61 Array (shape (T, 3)) of dipole moments in units of eÅ.
62 perturbation
63 The perturbation that was applied in the :term:`TDDFT` calculation.
64 """
66 def __init__(self,
67 times: NDArray[np.float64],
68 dipole_moments: NDArray[np.float64],
69 perturbation: PerturbationLike):
70 time_t = np.array(times) * as_to_au
71 dm_tv = np.array(dipole_moments) * eA_to_au
72 dm_tv -= dm_tv[0] # Remove static dipole
74 # Remove duplicates due to stopped and restarted calculations, and delta kick
75 flt_t = np.ones_like(time_t, dtype=bool)
76 maxtime = time_t[0]
77 for t in range(1, len(time_t)):
78 if time_t[t] > maxtime:
79 maxtime = time_t[t]
80 else:
81 flt_t[t] = False
82 time_t = time_t[flt_t]
83 dm_tv = dm_tv[flt_t]
84 ndup = len(flt_t) - flt_t.sum()
85 if ndup > 0:
86 parprint(f'Removed {ndup} duplicates')
88 # check time step
89 dt_t = time_t[1:] - time_t[:-1]
90 dt = dt_t[0]
91 if not np.allclose(dt_t, dt, rtol=1e-6, atol=0):
92 raise ValueError('Time grid must be equally spaced.')
94 self._time_t = time_t
95 self._dm_tv = dm_tv
96 self._perturbation = create_perturbation(perturbation)
98 if isinstance(self.perturbation, NoPerturbation):
99 raise ValueError('A perturbation must be given')
101 @property
102 def times(self) -> NDArray[np.float64]:
103 """ Array of times corresponding to dipole moments, in units of as. """
104 return self._time_t * au_to_as
106 @property
107 def dipole_moments(self) -> NDArray[np.float64]:
108 """ Array of dipole moments, in units of eÅ. """
109 return self._dm_tv * au_to_eA
111 @property
112 def perturbation(self) -> Perturbation:
113 """ The perturbation that was applied in the :term:`TDDFT` calculation. """
114 return self._perturbation
116 def _get_dynamic_polarizability(self,
117 frequencies: list[float] | NDArray[np.float64],
118 frequency_broadening: float = 0):
119 """ Calculate the dynamic polarizability in atomic units.
121 Parameters
122 ----------
123 frequencies
124 Array of frequencies for which to calculate the polarizability; in atomic units.
125 frequency_broadening
126 Gaussian broadening width in atomic units. Default (0) is no broadening.
128 Returns
129 -------
130 Array of dynamic polarizabilities in atomic units. \
131 The array has shape (N, 3), where N is the length of :attr:`frequencies`.
132 """
133 dt = self._time_t[1] - self._time_t[0]
134 nt = len(self._time_t)
136 # Get a response equivalent to a unity-strength delta-kick
137 dm_tv = self.perturbation.normalize_time_response(self._dm_tv, self._time_t, axis=0)
139 if frequency_broadening == 0:
140 # No broadening
141 weight_t = np.ones_like(self._time_t)
142 else:
143 # Gaussian broadening
144 weight_t = np.exp(-0.5 * frequency_broadening ** 2 * self._time_t**2)
146 # integration weights from Simpson's integration rule
147 weight_t *= dt / 3 * np.array([1] + [4, 2] * int((nt - 2) / 2)
148 + [4] * (nt % 2) + [1])
150 # transform
151 exp_tw = np.exp(np.outer(1.0j * self._time_t, frequencies))
152 dm_wv = np.einsum('t...,tw,t->w...', dm_tv, exp_tw, weight_t, optimize=True)
154 return dm_wv
156 def get_dynamic_polarizability(self,
157 frequencies: list[float] | NDArray[np.float64],
158 frequency_broadening: float = 0):
159 """ Calculate the dynamic polarizability.
161 Parameters
162 ----------
163 frequencies
164 Array of frequencies for which to calculate the polarizability; in units of eV.
165 frequency_broadening
166 Gaussian broadening width in atomic units. Default (0) is no broadening.
168 Returns
169 -------
170 Array of dynamic polarizabilities in (eÅ)**2/eV. \
171 The array has shape (N, 3), where N is the length of :attr:`frequencies`.
172 """
174 dm_wv = self._get_dynamic_polarizability(np.array(frequencies) * eV_to_au,
175 frequency_broadening * eV_to_au)
176 dm_wv = dm_wv * au_to_eA**2 / au_to_eV
177 return dm_wv
179 def _get_dipole_strength_function(self,
180 frequencies: list[float] | NDArray[np.float64],
181 frequency_broadening: float = 0):
182 """ Calculate the dipole strength function (spectrum) in atomic units.
184 Parameters
185 ----------
186 frequencies
187 Array of frequencies for which to calculate the spectrum; in atomic units.
188 frequency_broadening
189 Gaussian broadening width in atomic units. Default (0) is no broadening.
191 Returns
192 -------
193 Array of dipole strength function in atomic units. \
194 The array has shape (N, 3), where N is length of frequencies.
195 """
196 freq_w = np.array(frequencies)
197 pol_wv = self._get_dynamic_polarizability(frequencies, frequency_broadening)
198 osc_wv = 2 / np.pi * freq_w[:, np.newaxis] * pol_wv.imag
199 return osc_wv # type: ignore
201 def get_dipole_strength_function(self,
202 frequencies: list[float] | NDArray[np.float64],
203 frequency_broadening: float = 0):
204 """ Calculate the dipole strength function (spectrum) in units of 1/eV.
206 Parameters
207 ----------
208 frequencies
209 Array of frequencies for which to calculate the spectrum; in units of eV.
210 frequency_broadening
211 Gaussian broadening width in units of eV. Default (0) is no broadening.
213 Returns
214 -------
215 Array of dipole strength function in units of 1/eV. \
216 The array has shape (N, 3), where N is length of frequencies.
217 """
218 osc_wv = self._get_dipole_strength_function(np.array(frequencies) * eV_to_au,
219 frequency_broadening * eV_to_au)
220 osc_wv = osc_wv / au_to_eV
221 return osc_wv
223 @classmethod
224 def from_file(cls,
225 dipolefile: str | Path,
226 perturbation: PerturbationLike) -> SpectrumCalculator:
227 _, time_t, _, dm_tv = read_dipole_moment_file(str(dipolefile))
228 return cls(time_t * au_to_as, dm_tv * au_to_eA, perturbation)
230 def calculate_spectrum_and_write(self,
231 out_fname: Path | str,
232 frequencies: list[float] | NDArray[np.float64],
233 frequency_broadening: float = 0,
234 write_extra: dict[str, Any] = dict()):
235 """ Calculate the dipole strength function (spectrum) and write to file.
237 The spectrum is saved in a numpy archive if the file extension is `.npz`
238 or in a comma separated file if the file extension is `.dat`.
240 Parameters
241 ----------
242 out_fname
243 File name of the resulting data file.
244 frequencies
245 Array of frequencies for which to calculate the spectrum; in units of eV.
246 frequency_broadening
247 Gaussian broadening width in units of eV. Default (0) is no broadening.
248 write_extra
249 Dictionary of additional data written to numpy archive (ignored for `.dat`) files.
250 """
251 from .writers.spectrum import write_spectrum
253 out_fname = str(out_fname)
254 if world.rank > 0:
255 world.barrier()
256 return
258 osc_wv = self.get_dipole_strength_function(frequencies, frequency_broadening)
259 if out_fname[-4:] == '.npz':
260 d = dict(freq_w=frequencies,
261 osc_wv=osc_wv,
262 frequency_broadening=frequency_broadening)
263 d.update({f'perturbation_{key}': value
264 for key, value in self.perturbation.todict().items()})
265 d.update(write_extra)
266 np.savez_compressed(str(out_fname), **d)
267 elif out_fname[-4:] == '.dat':
268 write_spectrum(out_fname,
269 frequencies=frequencies,
270 spectrum=osc_wv,
271 frequency_broadening=frequency_broadening,
272 total_time=self.times[-1] - self.times[0],
273 timestep=self.times[1] - self.times[0],
274 perturbation=self.perturbation)
275 else:
276 raise ValueError(f'output-file must have ending .npz or .dat, is {out_fname}')
278 print(f'Written {out_fname}', flush=True)
279 world.barrier()