Coverage for strongcoca/response/utilities.py: 97%

1# This module contains components adapted from GPAW: 

2# gpaw.tddft.spectrum 

3# https://gitlab.com/gpaw/gpaw/-/blob/aca9ed6f520d6a855013247119b50c630f5eebc9/gpaw/tddft/spectrum.py 

4 

5"""This module provides functionality for transforming time and 

6frequency-series including Lorentzian and Gaussian broadening. 

7""" 

8import numpy as np 

9from scipy.special import dawsn 

10 

11from ..units import au_to_eV, eV_to_au 

12 

13 

14class Broadening: 

15 """Class representing broadening types of spectra.""" 

16 def __repr__(self) -> str: 

17 return f'{self.__class__.__name__}' 

18 

19 

20class NoArtificialBroadening(Broadening): 

21 """Class representing no artificial broadening.""" 

22 

23 

24class ArtificialBroadening(Broadening): 

25 """Class representing artificial broadenings. 

26 

27 Parameters 

28 ---------- 

29 width 

30 Broadening width; meaning interpreted by the subclasses; 

31 must be strictly positive. 

32 

33 Units are eV by default, atomic units can be selected via :attr:`units`. 

34 units 

35 `eV` to specify :attr:`width` in eV or `au` to specify :attr:`width` in 

36 atomic units. 

37 

38 This parameter determines whether conversion should be performed during 

39 initialization and has no effect on instance methods and variables. 

40 """ 

41 def __init__(self, width: float, units: str = 'eV') -> None: 

42 if width <= 0: 

43 raise ValueError('width must be strictly positive') 

44 if units == 'eV': 

45 width = width * eV_to_au 

46 elif units != 'au': 

47 raise ValueError(f"units has to be 'eV' or 'au', not '{units}'") 

48 self._width = width 

49 

50 def __repr__(self) -> str: 

51 return f"{self.__class__.__name__}({self.width}, units='eV')" 

52 

53 @property 

54 def width(self) -> float: 

55 """Broadening width in units of eV.""" 

56 return self._width * au_to_eV 

57 

58 

59class GaussianBroadening(ArtificialBroadening): 

60 r"""Class for representing artificial Gaussian broadening of spectra. 

61 

62 Parameters 

63 ---------- 

64 width 

65 Broadening width (:math:`\sigma`); 

66 must be strictly positive. 

67 

68 Units are eV by default, atomic units can be selected via :attr:`units`. 

69 units 

70 `eV` to specify :attr:`width` in eV or `au` to specify :attr:`width` in 

71 atomic units. 

72 

73 This parameter determines whether conversion should be performed during 

74 initialization and has no effect on instance methods and variables. 

75 

76 Examples 

77 ------- 

78 >>> from strongcoca.response.utilities import GaussianBroadening 

79 >>> GaussianBroadening(0.1) 

80 GaussianBroadening(0.1, units='eV') 

81 >>> GaussianBroadening(0.005, units='au') 

82 GaussianBroadening(0.136..., units='eV') 

83 """ 

84 def __init__(self, width: float, units: str = 'eV') -> None: 

85 super().__init__(width=width, units=units) 

86 

87 

88class LorentzianBroadening(ArtificialBroadening): 

89 r"""Class for representing artificial Lorentzian broadening of spectra. 

90 

91 Parameters 

92 ---------- 

93 width 

94 Broadening width (:math:`\eta`); 

95 must be strictly positive. 

96 

97 Units are eV by default, atomic units can be selected via :attr:`units`. 

98 units 

99 `eV` to specify :attr:`width` in eV or `au` to specify :attr:`width` in 

100 atomic units. 

101 

102 This parameter determines whether conversion should be performed during 

103 initialization and has no effect on instance methods and variables. 

104 

105 Examples 

106 ------- 

107 >>> from strongcoca.response.utilities import LorentzianBroadening 

108 >>> LorentzianBroadening(0.1) 

109 LorentzianBroadening(0.1, units='eV') 

110 >>> LorentzianBroadening(0.005, units='au') 

111 LorentzianBroadening(0.136..., units='eV') 

112 """ 

113 def __init__(self, width: float, units: str = 'eV') -> None: 

114 super().__init__(width=width, units=units) 

115 

116 

117def fourier_transform(time_t: np.ndarray, 

118 data_tX: np.ndarray, 

119 freq_w: np.ndarray, 

120 freq_axis: str = 'real', 

121 broadening: Broadening = NoArtificialBroadening()) -> np.ndarray: 

122 r"""Calculates the Fourier transform of a time-series data according to 

123 

124 .. math:: 

125 

126 y_X(\omega) = \int_{t_\text{min}}^{t_\text{max}} y_X(t) e^{i \omega t} \mathrm{d}t 

127 

128 The integral is evaluated as such when broadening is not specified. 

129 

130 Lorentzian broadening corresponds to 

131 

132 .. math:: 

133 

134 e^{i \omega t} \to e^{i (\omega + i \eta) t} = e^{i \omega t} e^{- \eta t} 

135 

136 Gaussian broadening corresponds to 

137 

138 .. math:: 

139 

140 e^{i \omega t} \to e^{i \omega t} e^{- \frac{1}{2}\sigma^2 t^2} 

141 

142 This function does not carry out unit conversions, but all parameters are 

143 assumed to be in mutually compatible units. 

144 

145 Parameters 

146 ---------- 

147 time_t 

148 Equally spaced time grid. 

149 data_tX 

150 Data to be transformed. 

151 The first dimension must be of same size as `time_t`. Other 

152 dimensions can be of any size. 

153 freq_w 

154 Frequency values for the transform. 

155 freq_axis : `'real'` or `'imag'` 

156 Interpretation of frequency values as real or imaginary frequencies. 

157 broadening 

158 Broadening to be used. 

159 

160 Notes 

161 ----- 

162 When working with imaginary axes 

163 `fourier_transform(..., freq_w=1j * freq_w, freq_axis='real')` and 

164 `fourier_transform(..., freq_w=freq_w, freq_axis='imag')` 

165 yield similar results, yet the latter is computationally 

166 more efficient when `freq_w` is real. 

167 

168 Returns 

169 ------- 

170 :class:`np.ndarray` 

171 Transform of input data to frequency. 

172 The first dimension is of the same size as `freq_w`. Other dimensions 

173 are the same as in input data `data_tX`. 

174 """ 

175 if time_t.ndim != 1: 

176 raise ValueError('time_t must be one-dimensional array') 

177 if data_tX.shape[0] != time_t.shape[0]: 

178 raise ValueError('data_tX must have compatible shape with time_t') 

179 

180 if freq_axis == 'real': 

181 prefactor = 1.0j 

182 if isinstance(broadening, NoArtificialBroadening): 

183 weight_t = np.ones_like(time_t) 

184 elif isinstance(broadening, LorentzianBroadening): 

185 weight_t = np.exp(-broadening._width * time_t) 

186 elif isinstance(broadening, GaussianBroadening): 186 ↛ 189line 186 didn't jump to line 189 because the condition on line 186 was always true

187 weight_t = np.exp(-0.5 * broadening._width**2 * time_t**2) 

188 else: 

189 raise ValueError(f'Unknown broadening: {broadening}') 

190 elif freq_axis == 'imag': 

191 prefactor = -1.0 

192 weight_t = np.ones_like(time_t) 

193 if not isinstance(broadening, NoArtificialBroadening): 

194 raise ValueError('Artificial broadening is not applied to imaginary frequencies') 

195 else: 

196 raise ValueError(f'Unknown frequency axis: {freq_axis}') 

197 

198 # check time step 

199 dt_t = time_t[1:] - time_t[:-1] 

200 dt = dt_t[0] 

201 if not np.allclose(dt_t, dt, rtol=1e-6, atol=0): 

202 raise ValueError('Time grid must be equally spaced.') 

203 

204 # integration weights from Simpson's integration rule 

205 weight_t *= dt / 3 * np.array([1] + [4, 2] * int((len(time_t) - 2) / 2) 

206 + [4] * (len(time_t) % 2) + [1]) 

207 

208 # transform 

209 exp_tw = np.exp(np.outer(prefactor * time_t, freq_w)) 

210 data_wX = np.einsum('t...,tw,t->w...', data_tX, exp_tw, weight_t, optimize=True) 

211 return data_wX # type: ignore 

212 

213 

214def broaden(freq_i: np.ndarray, 

215 data_iX: np.ndarray, 

216 freq_w: np.ndarray, 

217 freq_axis: str = 'real', 

218 broadening: Broadening = NoArtificialBroadening()) -> np.ndarray: 

219 r"""Broaden discrete data and returns it on a continuous frequency grid. 

220 Specifically this function performs the summation 

221 

222 .. math:: 

223 

224 y_{X}(\omega) = \sum_I \frac{y_{I}^{(X)}}{\omega_I^2 - \omega^2} 

225 

226 This sum is evaluated as such when no broadening is specified. 

227 

228 Lorentzian broadening corresponds to 

229 

230 .. math:: 

231 

232 \frac{1}{\omega_I^2 - \omega^2} \to \frac{1}{\omega_I^2 - (\omega + i \eta)^2} 

233 

234 Gaussian broadening corresponds to 

235 

236 .. math:: 

237 

238 \frac{1}{\omega_I^2 - \omega^2} \to 

239 -\frac{\pi}{2 \omega_I} 

240 \left\{ 

241 \frac{\sqrt{2}}{\pi\sigma} 

242 \left[ 

243 D\left(\frac{\omega - \omega_I}{\sqrt{2}\sigma}\right) 

244 - D\left(\frac{\omega + \omega_I}{\sqrt{2}\sigma}\right) 

245 \right] \\ 

246 -\frac{i}{\sqrt{2\pi}\sigma} 

247 \left[ 

248 G\left(\frac{\omega - \omega_I}{\sqrt{2}\sigma}\right) 

249 - G\left(\frac{\omega + \omega_I}{\sqrt{2}\sigma}\right) 

250 \right] 

251 \right\} 

252 

253 where :math:`D(\omega)` and :math:`G(\omega)` are Dawson and Gaussian functions: 

254 

255 .. math:: 

256 

257 D(\omega) &= e^{-\omega^2} \int_0^\omega e^{s^2} \mathrm{d}s \\ 

258 G(\omega) &= e^{-\omega^2} 

259 

260 

261 Parameters 

262 ---------- 

263 freq_i 

264 Discrete frequency values of the data (:math:`\omega_I`) 

265 data_iX 

266 Data to be broadened (:math:`y_{I}^{(X)}`) 

267 freq_w 

268 Continuous frequency values for broadening (:math:`\omega`) 

269 freq_axis : `'real'` or `'imag'` 

270 Interpretation of frequency values as real or imaginary frequencies. 

271 broadening 

272 Broadening to be used. 

273 

274 Notes 

275 ----- 

276 When working with imaginary axes 

277 `broaden(..., freq_w=1j * freq_w, freq_axis='real')` and 

278 `broaden(..., freq_w=freq_w, freq_axis='imag')` 

279 yield similar results, yet the latter is computationally 

280 more efficient when `freq_w` is real. 

281 

282 Returns 

283 ------- 

284 :class:`np.ndarray` 

285 Broadened data along continuous frequency axis. 

286 The first dimension is of the same size as `freq_w`. Other dimensions 

287 are the same as in input data `data_iX`. 

288 """ 

289 if freq_i.ndim != 1: 

290 raise ValueError('freq_i must be one-dimensional array') 

291 if data_iX.shape[0] != freq_i.shape[0]: 

292 raise ValueError('data_iX must have compatible shape with freq_i') 

293 

294 if freq_axis == 'real': 

295 if isinstance(broadening, NoArtificialBroadening): 

296 weight_iw = 1. / (freq_i[:, np.newaxis]**2 - freq_w[np.newaxis, :]**2) 

297 elif isinstance(broadening, LorentzianBroadening): 

298 width = broadening._width 

299 weight_iw = 1. / (freq_i[:, np.newaxis]**2 - (freq_w[np.newaxis, :] + 1j * width)**2) 

300 elif isinstance(broadening, GaussianBroadening): 300 ↛ 309line 300 didn't jump to line 309 because the condition on line 300 was always true

301 width = broadening._width 

302 xm_iw = (freq_w[np.newaxis, :] - freq_i[:, np.newaxis]) / (np.sqrt(2) * width) 

303 xp_iw = (freq_w[np.newaxis, :] + freq_i[:, np.newaxis]) / (np.sqrt(2) * width) 

304 weight_iw = -np.pi / (2 * freq_i[:, np.newaxis]) * ( 

305 np.sqrt(2) / (np.pi * width) * (dawsn(xm_iw) - dawsn(xp_iw)) 

306 - 1j / (np.sqrt(2 * np.pi) * width) * (np.exp(-xm_iw**2) - np.exp(-xp_iw**2)) 

307 ) 

308 else: 

309 raise ValueError(f'Unknown broadening: {broadening}') 

310 elif freq_axis == 'imag': 

311 weight_iw = 1. / (freq_i[:, np.newaxis]**2 + freq_w[np.newaxis, :]**2) 

312 if not isinstance(broadening, NoArtificialBroadening): 

313 raise ValueError('Artificial broadening is not applied to imaginary frequencies') 

314 else: 

315 raise ValueError(f'Unknown frequency axis: {freq_axis}') 

316 

317 # transform 

318 data_wX = np.einsum('i...,iw->w...', data_iX, weight_iw, optimize=True) 

319 return data_wX # type: ignore