# This module contains components adapted from GPAW:
# gpaw.tddft.spectrum
# https://gitlab.com/gpaw/gpaw/-/blob/aca9ed6f520d6a855013247119b50c630f5eebc9/gpaw/tddft/spectrum.py
"""This module provides functionality for transforming time and
frequency-series including Lorentzian and Gaussian broadening.
"""
import numpy as np
from scipy.special import dawsn
from ..units import au_to_eV, eV_to_au
class Broadening:
"""Class representing broadening types of spectra."""
def __repr__(self) -> str:
return f'{self.__class__.__name__}'
[docs]
class NoArtificialBroadening(Broadening):
"""Class representing no artificial broadening."""
[docs]
class ArtificialBroadening(Broadening):
"""Class representing artificial broadenings.
Parameters
----------
width
Broadening width; meaning interpreted by the subclasses;
must be strictly positive.
Units are eV by default, atomic units can be selected via :attr:`units`.
units
`eV` to specify :attr:`width` in eV or `au` to specify :attr:`width` in
atomic units.
This parameter determines whether conversion should be performed during
initialization and has no effect on instance methods and variables.
"""
def __init__(self, width: float, units: str = 'eV') -> None:
if width <= 0:
raise ValueError('width must be strictly positive')
if units == 'eV':
width = width * eV_to_au
elif units != 'au':
raise ValueError(f"units has to be 'eV' or 'au', not '{units}'")
self._width = width
def __repr__(self) -> str:
return f"{self.__class__.__name__}({self.width}, units='eV')"
@property
def width(self) -> float:
"""Broadening width in units of eV."""
return self._width * au_to_eV
[docs]
class GaussianBroadening(ArtificialBroadening):
r"""Class for representing artificial Gaussian broadening of spectra.
Parameters
----------
width
Broadening width (:math:`\sigma`);
must be strictly positive.
Units are eV by default, atomic units can be selected via :attr:`units`.
units
`eV` to specify :attr:`width` in eV or `au` to specify :attr:`width` in
atomic units.
This parameter determines whether conversion should be performed during
initialization and has no effect on instance methods and variables.
Examples
-------
>>> from strongcoca.response.utilities import GaussianBroadening
>>> GaussianBroadening(0.1)
GaussianBroadening(0.1, units='eV')
>>> GaussianBroadening(0.005, units='au')
GaussianBroadening(0.136..., units='eV')
"""
def __init__(self, width: float, units: str = 'eV') -> None:
super().__init__(width=width, units=units)
[docs]
class LorentzianBroadening(ArtificialBroadening):
r"""Class for representing artificial Lorentzian broadening of spectra.
Parameters
----------
width
Broadening width (:math:`\eta`);
must be strictly positive.
Units are eV by default, atomic units can be selected via :attr:`units`.
units
`eV` to specify :attr:`width` in eV or `au` to specify :attr:`width` in
atomic units.
This parameter determines whether conversion should be performed during
initialization and has no effect on instance methods and variables.
Examples
-------
>>> from strongcoca.response.utilities import LorentzianBroadening
>>> LorentzianBroadening(0.1)
LorentzianBroadening(0.1, units='eV')
>>> LorentzianBroadening(0.005, units='au')
LorentzianBroadening(0.136..., units='eV')
"""
def __init__(self, width: float, units: str = 'eV') -> None:
super().__init__(width=width, units=units)
def fourier_transform(time_t: np.ndarray,
data_tX: np.ndarray,
freq_w: np.ndarray,
freq_axis: str = 'real',
broadening: Broadening = NoArtificialBroadening()) -> np.ndarray:
r"""Calculates the Fourier transform of a time-series data according to
.. math::
y_X(\omega) = \int_{t_\text{min}}^{t_\text{max}} y_X(t) e^{i \omega t} \mathrm{d}t
The integral is evaluated as such when broadening is not specified.
Lorentzian broadening corresponds to
.. math::
e^{i \omega t} \to e^{i (\omega + i \eta) t} = e^{i \omega t} e^{- \eta t}
Gaussian broadening corresponds to
.. math::
e^{i \omega t} \to e^{i \omega t} e^{- \frac{1}{2}\sigma^2 t^2}
This function does not carry out unit conversions, but all parameters are
assumed to be in mutually compatible units.
Parameters
----------
time_t
Equally spaced time grid.
data_tX
Data to be transformed.
The first dimension must be of same size as `time_t`. Other
dimensions can be of any size.
freq_w
Frequency values for the transform.
freq_axis : `'real'` or `'imag'`
Interpretation of frequency values as real or imaginary frequencies.
broadening
Broadening to be used.
Notes
-----
When working with imaginary axes
`fourier_transform(..., freq_w=1j * freq_w, freq_axis='real')` and
`fourier_transform(..., freq_w=freq_w, freq_axis='imag')`
yield similar results, yet the latter is computationally
more efficient when `freq_w` is real.
Returns
-------
:class:`np.ndarray`
Transform of input data to frequency.
The first dimension is of the same size as `freq_w`. Other dimensions
are the same as in input data `data_tX`.
"""
if time_t.ndim != 1:
raise ValueError('time_t must be one-dimensional array')
if data_tX.shape[0] != time_t.shape[0]:
raise ValueError('data_tX must have compatible shape with time_t')
if freq_axis == 'real':
prefactor = 1.0j
if isinstance(broadening, NoArtificialBroadening):
weight_t = np.ones_like(time_t)
elif isinstance(broadening, LorentzianBroadening):
weight_t = np.exp(-broadening._width * time_t)
elif isinstance(broadening, GaussianBroadening):
weight_t = np.exp(-0.5 * broadening._width**2 * time_t**2)
else:
raise ValueError(f'Unknown broadening: {broadening}')
elif freq_axis == 'imag':
prefactor = -1.0
weight_t = np.ones_like(time_t)
if not isinstance(broadening, NoArtificialBroadening):
raise ValueError('Artificial broadening is not applied to imaginary frequencies')
else:
raise ValueError(f'Unknown frequency axis: {freq_axis}')
# 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.')
# integration weights from Simpson's integration rule
weight_t *= dt / 3 * np.array([1] + [4, 2] * int((len(time_t) - 2) / 2)
+ [4] * (len(time_t) % 2) + [1])
# transform
exp_tw = np.exp(np.outer(prefactor * time_t, freq_w))
data_wX = np.einsum('t...,tw,t->w...', data_tX, exp_tw, weight_t, optimize=True)
return data_wX # type: ignore
def broaden(freq_i: np.ndarray,
data_iX: np.ndarray,
freq_w: np.ndarray,
freq_axis: str = 'real',
broadening: Broadening = NoArtificialBroadening()) -> np.ndarray:
r"""Broaden discrete data and returns it on a continuous frequency grid.
Specifically this function performs the summation
.. math::
y_{X}(\omega) = \sum_I \frac{y_{I}^{(X)}}{\omega_I^2 - \omega^2}
This sum is evaluated as such when no broadening is specified.
Lorentzian broadening corresponds to
.. math::
\frac{1}{\omega_I^2 - \omega^2} \to \frac{1}{\omega_I^2 - (\omega + i \eta)^2}
Gaussian broadening corresponds to
.. math::
\frac{1}{\omega_I^2 - \omega^2} \to
-\frac{\pi}{2 \omega_I}
\left\{
\frac{\sqrt{2}}{\pi\sigma}
\left[
D\left(\frac{\omega - \omega_I}{\sqrt{2}\sigma}\right)
- D\left(\frac{\omega + \omega_I}{\sqrt{2}\sigma}\right)
\right] \\
-\frac{i}{\sqrt{2\pi}\sigma}
\left[
G\left(\frac{\omega - \omega_I}{\sqrt{2}\sigma}\right)
- G\left(\frac{\omega + \omega_I}{\sqrt{2}\sigma}\right)
\right]
\right\}
where :math:`D(\omega)` and :math:`G(\omega)` are Dawson and Gaussian functions:
.. math::
D(\omega) &= e^{-\omega^2} \int_0^\omega e^{s^2} \mathrm{d}s \\
G(\omega) &= e^{-\omega^2}
Parameters
----------
freq_i
Discrete frequency values of the data (:math:`\omega_I`)
data_iX
Data to be broadened (:math:`y_{I}^{(X)}`)
freq_w
Continuous frequency values for broadening (:math:`\omega`)
freq_axis : `'real'` or `'imag'`
Interpretation of frequency values as real or imaginary frequencies.
broadening
Broadening to be used.
Notes
-----
When working with imaginary axes
`broaden(..., freq_w=1j * freq_w, freq_axis='real')` and
`broaden(..., freq_w=freq_w, freq_axis='imag')`
yield similar results, yet the latter is computationally
more efficient when `freq_w` is real.
Returns
-------
:class:`np.ndarray`
Broadened data along continuous frequency axis.
The first dimension is of the same size as `freq_w`. Other dimensions
are the same as in input data `data_iX`.
"""
if freq_i.ndim != 1:
raise ValueError('freq_i must be one-dimensional array')
if data_iX.shape[0] != freq_i.shape[0]:
raise ValueError('data_iX must have compatible shape with freq_i')
if freq_axis == 'real':
if isinstance(broadening, NoArtificialBroadening):
weight_iw = 1. / (freq_i[:, np.newaxis]**2 - freq_w[np.newaxis, :]**2)
elif isinstance(broadening, LorentzianBroadening):
width = broadening._width
weight_iw = 1. / (freq_i[:, np.newaxis]**2 - (freq_w[np.newaxis, :] + 1j * width)**2)
elif isinstance(broadening, GaussianBroadening):
width = broadening._width
xm_iw = (freq_w[np.newaxis, :] - freq_i[:, np.newaxis]) / (np.sqrt(2) * width)
xp_iw = (freq_w[np.newaxis, :] + freq_i[:, np.newaxis]) / (np.sqrt(2) * width)
weight_iw = -np.pi / (2 * freq_i[:, np.newaxis]) * (
np.sqrt(2) / (np.pi * width) * (dawsn(xm_iw) - dawsn(xp_iw))
- 1j / (np.sqrt(2 * np.pi) * width) * (np.exp(-xm_iw**2) - np.exp(-xp_iw**2))
)
else:
raise ValueError(f'Unknown broadening: {broadening}')
elif freq_axis == 'imag':
weight_iw = 1. / (freq_i[:, np.newaxis]**2 + freq_w[np.newaxis, :]**2)
if not isinstance(broadening, NoArtificialBroadening):
raise ValueError('Artificial broadening is not applied to imaginary frequencies')
else:
raise ValueError(f'Unknown frequency axis: {freq_axis}')
# transform
data_wX = np.einsum('i...,iw->w...', data_iX, weight_iw, optimize=True)
return data_wX # type: ignore
