Source code for rhodent.spectrum

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()