from __future__ import annotations
from typing import Any
from pathlib import Path
import numpy as np
from numpy.typing import NDArray
from ase.parallel import parprint
from gpaw.mpi import world
from gpaw.tddft.units import as_to_au, au_to_as, eV_to_au, au_to_eV, eA_to_au, au_to_eA
from gpaw.tddft.spectrum import read_dipole_moment_file
from .perturbation import create_perturbation, Perturbation, PerturbationLike, NoPerturbation
[docs]
class SpectrumCalculator:
r""" Spectrum calculator.
Calculates the dynamic polarizability and dipole strength function (the spectrum).
The polarizability :math:`\alpha(\omega)` is related to the perturbation
:math:`\mathcal{E}(\omega)` and the Fouier transform of the dipole moment
:math:`\boldsymbol\mu(t)` as
.. math::
\boldsymbol\alpha(\omega) \mathcal{E}(\omega) = \mathcal{F}[\boldsymbol\mu(t)](\omega).
The dipole strength function is
.. math::
\boldsymbol{S}(\omega) = \frac{2}{\pi} \omega\:\mathrm{Im}[\boldsymbol\alpha(\omega)].
Both quantities can be broadened by supplying a non-zero value :math:`\sigma`. Then the
convolution
.. math::
\frac{1}{\sqrt{2\pi/\sigma^2}}
\int_{-\infty}^{\infty} \boldsymbol\alpha(\omega')
\exp\left(- \frac{(\omega - \omega')^2}{2\sigma^2}\right) \mathrm{d}\omega'
is computed. When the perturbation is a delta-kick, this can efficiently be computed
by multiplying the dipole moment by a Gaussian envelope before Fourier transformation
.. math::
\boldsymbol\mu(t) \exp(-\sigma^2 t^2 / 2).
For other perturbations, a :term:`FFT` and :term:`IFFT` is first
performed to obtain the delta-kick response.
Parameters
----------
times
Array (length T) of times in units of as.
dipole_moments
Array (shape (T, 3)) of dipole moments in units of eÅ.
perturbation
The perturbation that was applied in the :term:`TDDFT` calculation.
"""
def __init__(self,
times: NDArray[np.float64],
dipole_moments: NDArray[np.float64],
perturbation: PerturbationLike):
time_t = np.array(times) * as_to_au
dm_tv = np.array(dipole_moments) * eA_to_au
dm_tv -= dm_tv[0] # Remove static dipole
# Remove duplicates due to stopped and restarted calculations, and delta kick
flt_t = np.ones_like(time_t, dtype=bool)
maxtime = time_t[0]
for t in range(1, len(time_t)):
if time_t[t] > maxtime:
maxtime = time_t[t]
else:
flt_t[t] = False
time_t = time_t[flt_t]
dm_tv = dm_tv[flt_t]
ndup = len(flt_t) - flt_t.sum()
if ndup > 0:
parprint(f'Removed {ndup} duplicates')
# check time step
dt_t = time_t[1:] - time_t[:-1]
dt = dt_t[0]
if not np.allclose(dt_t, dt, rtol=1e-6, atol=0):
raise ValueError('Time grid must be equally spaced.')
self._time_t = time_t
self._dm_tv = dm_tv
self._perturbation = create_perturbation(perturbation)
if isinstance(self.perturbation, NoPerturbation):
raise ValueError('A perturbation must be given')
@property
def times(self) -> NDArray[np.float64]:
""" Array of times corresponding to dipole moments, in units of as. """
return self._time_t * au_to_as
@property
def dipole_moments(self) -> NDArray[np.float64]:
""" Array of dipole moments, in units of eÅ. """
return self._dm_tv * au_to_eA
@property
def perturbation(self) -> Perturbation:
""" The perturbation that was applied in the :term:`TDDFT` calculation. """
return self._perturbation
def _get_dynamic_polarizability(self,
frequencies: list[float] | NDArray[np.float64],
frequency_broadening: float = 0):
""" Calculate the dynamic polarizability in atomic units.
Parameters
----------
frequencies
Array of frequencies for which to calculate the polarizability; in atomic units.
frequency_broadening
Gaussian broadening width in atomic units. Default (0) is no broadening.
Returns
-------
Array of dynamic polarizabilities in atomic units. \
The array has shape (N, 3), where N is the length of :attr:`frequencies`.
"""
dt = self._time_t[1] - self._time_t[0]
nt = len(self._time_t)
# Get a response equivalent to a unity-strength delta-kick
dm_tv = self.perturbation.normalize_time_response(self._dm_tv, self._time_t, axis=0)
if frequency_broadening == 0:
# No broadening
weight_t = np.ones_like(self._time_t)
else:
# Gaussian broadening
weight_t = np.exp(-0.5 * frequency_broadening ** 2 * self._time_t**2)
# integration weights from Simpson's integration rule
weight_t *= dt / 3 * np.array([1] + [4, 2] * int((nt - 2) / 2)
+ [4] * (nt % 2) + [1])
# transform
exp_tw = np.exp(np.outer(1.0j * self._time_t, frequencies))
dm_wv = np.einsum('t...,tw,t->w...', dm_tv, exp_tw, weight_t, optimize=True)
return dm_wv
[docs]
def get_dynamic_polarizability(self,
frequencies: list[float] | NDArray[np.float64],
frequency_broadening: float = 0):
""" Calculate the dynamic polarizability.
Parameters
----------
frequencies
Array of frequencies for which to calculate the polarizability; in units of eV.
frequency_broadening
Gaussian broadening width in atomic units. Default (0) is no broadening.
Returns
-------
Array of dynamic polarizabilities in (eÅ)**2/eV. \
The array has shape (N, 3), where N is the length of :attr:`frequencies`.
"""
dm_wv = self._get_dynamic_polarizability(np.array(frequencies) * eV_to_au,
frequency_broadening * eV_to_au)
dm_wv = dm_wv * au_to_eA**2 / au_to_eV
return dm_wv
def _get_dipole_strength_function(self,
frequencies: list[float] | NDArray[np.float64],
frequency_broadening: float = 0):
""" Calculate the dipole strength function (spectrum) in atomic units.
Parameters
----------
frequencies
Array of frequencies for which to calculate the spectrum; in atomic units.
frequency_broadening
Gaussian broadening width in atomic units. Default (0) is no broadening.
Returns
-------
Array of dipole strength function in atomic units. \
The array has shape (N, 3), where N is length of frequencies.
"""
freq_w = np.array(frequencies)
pol_wv = self._get_dynamic_polarizability(frequencies, frequency_broadening)
osc_wv = 2 / np.pi * freq_w[:, np.newaxis] * pol_wv.imag
return osc_wv # type: ignore
[docs]
def get_dipole_strength_function(self,
frequencies: list[float] | NDArray[np.float64],
frequency_broadening: float = 0):
""" Calculate the dipole strength function (spectrum) in units of 1/eV.
Parameters
----------
frequencies
Array of frequencies for which to calculate the spectrum; in units of eV.
frequency_broadening
Gaussian broadening width in units of eV. Default (0) is no broadening.
Returns
-------
Array of dipole strength function in units of 1/eV. \
The array has shape (N, 3), where N is length of frequencies.
"""
osc_wv = self._get_dipole_strength_function(np.array(frequencies) * eV_to_au,
frequency_broadening * eV_to_au)
osc_wv = osc_wv / au_to_eV
return osc_wv
@classmethod
def from_file(cls,
dipolefile: str | Path,
perturbation: PerturbationLike) -> SpectrumCalculator:
_, time_t, _, dm_tv = read_dipole_moment_file(str(dipolefile))
return cls(time_t * au_to_as, dm_tv * au_to_eA, perturbation)
[docs]
def calculate_spectrum_and_write(self,
out_fname: Path | str,
frequencies: list[float] | NDArray[np.float64],
frequency_broadening: float = 0,
write_extra: dict[str, Any] = dict()):
""" Calculate the dipole strength function (spectrum) and write to file.
The spectrum is saved in a numpy archive if the file extension is ``.npz``
or in a comma separated file if the file extension is ``.dat``.
Parameters
----------
out_fname
File name of the resulting data file.
frequencies
Array of frequencies for which to calculate the spectrum; in units of eV.
frequency_broadening
Gaussian broadening width in units of eV. Default (0) is no broadening.
write_extra
Dictionary of additional data written to numpy archive (ignored for ``.dat``) files.
"""
from .writers.spectrum import write_spectrum
out_fname = str(out_fname)
if world.rank > 0:
world.barrier()
return
osc_wv = self.get_dipole_strength_function(frequencies, frequency_broadening)
if out_fname[-4:] == '.npz':
d = dict(freq_w=frequencies,
osc_wv=osc_wv,
frequency_broadening=frequency_broadening)
d.update({f'perturbation_{key}': value
for key, value in self.perturbation.todict().items()})
d.update(write_extra)
np.savez_compressed(str(out_fname), **d)
elif out_fname[-4:] == '.dat':
write_spectrum(out_fname,
frequencies=frequencies,
spectrum=osc_wv,
frequency_broadening=frequency_broadening,
total_time=self.times[-1] - self.times[0],
timestep=self.times[1] - self.times[0],
perturbation=self.perturbation)
else:
raise ValueError(f'output-file must have ending .npz or .dat, is {out_fname}')
print(f'Written {out_fname}', flush=True)
world.barrier()