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

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 def _repr_html_(self) -> str: 

20 """HTML representation for Jupyter notebooks.""" 

21 return ( 

22 f'<h4>{self.__class__.__name__}</h4>' 

23 ) 

24 

25 

26class NoArtificialBroadening(Broadening): 

27 """Class representing no artificial broadening.""" 

28 

29 

30class ArtificialBroadening(Broadening): 

31 """Class representing artificial broadenings. 

32 

33 Parameters 

34 ---------- 

35 width 

36 Broadening width; meaning interpreted by the subclasses; 

37 must be strictly positive. 

38 

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

40 units 

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

42 atomic units. 

43 

44 This parameter determines whether conversion should be performed during 

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

46 """ 

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

48 if width <= 0: 

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

50 if units == 'eV': 

51 width = width * eV_to_au 

52 elif units != 'au': 

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

54 self._width = width 

55 

56 def __repr__(self) -> str: 

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

58 

59 def _repr_html_(self) -> str: 

60 """HTML representation for Jupyter notebooks.""" 

61 rows = [ 

62 f'<tr><td style="text-align: left;">width (eV)</td><td>{self.width:.4f}</td></tr>', 

63 ] 

64 return ( 

65 f'<h4>{self.__class__.__name__}</h4>' 

66 '<table>' 

67 '<thead><tr>' 

68 '<th style="text-align: left;">field</th>' 

69 '<th style="text-align: left;">value</th>' 

70 '</tr></thead>' 

71 '<tbody>' + ''.join(rows) + '</tbody>' 

72 '</table>' 

73 ) 

74 

75 @property 

76 def width(self) -> float: 

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

78 return self._width * au_to_eV 

79 

80 

81class GaussianBroadening(ArtificialBroadening): 

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

83 

84 Parameters 

85 ---------- 

86 width 

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

88 must be strictly positive. 

89 

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

91 units 

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

93 atomic units. 

94 

95 This parameter determines whether conversion should be performed during 

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

97 

98 Examples 

99 ------- 

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

101 >>> GaussianBroadening(0.1) 

102 GaussianBroadening(0.1, units='eV') 

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

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

105 """ 

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

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

108 

109 

110class LorentzianBroadening(ArtificialBroadening): 

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

112 

113 Parameters 

114 ---------- 

115 width 

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

117 must be strictly positive. 

118 

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

120 units 

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

122 atomic units. 

123 

124 This parameter determines whether conversion should be performed during 

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

126 

127 Examples 

128 ------- 

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

130 >>> LorentzianBroadening(0.1) 

131 LorentzianBroadening(0.1, units='eV') 

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

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

134 """ 

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

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

137 

138 

139def fourier_transform(time_t: np.ndarray, 

140 data_tX: np.ndarray, 

141 freq_w: np.ndarray, 

142 freq_axis: str = 'real', 

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

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

145 

146 .. math:: 

147 

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

149 

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

151 

152 Lorentzian broadening corresponds to 

153 

154 .. math:: 

155 

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

157 

158 Gaussian broadening corresponds to 

159 

160 .. math:: 

161 

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

163 

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

165 assumed to be in mutually compatible units. 

166 

167 Parameters 

168 ---------- 

169 time_t 

170 Equally spaced time grid. 

171 data_tX 

172 Data to be transformed. 

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

174 dimensions can be of any size. 

175 freq_w 

176 Frequency values for the transform. 

177 freq_axis : `'real'` or `'imag'` 

178 Interpretation of frequency values as real or imaginary frequencies. 

179 broadening 

180 Broadening to be used. 

181 

182 Notes 

183 ----- 

184 When working with imaginary axes 

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

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

187 yield similar results, yet the latter is computationally 

188 more efficient when `freq_w` is real. 

189 

190 Returns 

191 ------- 

192 :class:`np.ndarray` 

193 Transform of input data to frequency. 

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

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

196 """ 

197 if time_t.ndim != 1: 

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

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

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

201 

202 if freq_axis == 'real': 

203 prefactor = 1.0j 

204 if isinstance(broadening, NoArtificialBroadening): 

205 weight_t = np.ones_like(time_t) 

206 elif isinstance(broadening, LorentzianBroadening): 

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

208 elif isinstance(broadening, GaussianBroadening): 

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

210 else: 

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

212 elif freq_axis == 'imag': 

213 prefactor = -1.0 

214 weight_t = np.ones_like(time_t) 

215 if not isinstance(broadening, NoArtificialBroadening): 

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

217 else: 

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

219 

220 # check time step 

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

222 dt = dt_t[0] 

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

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

225 

226 # integration weights from Simpson's integration rule 

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

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

229 

230 # transform 

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

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

233 return data_wX # type: ignore 

234 

235 

236def broaden(freq_i: np.ndarray, 

237 data_iX: np.ndarray, 

238 freq_w: np.ndarray, 

239 freq_axis: str = 'real', 

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

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

242 Specifically this function performs the summation 

243 

244 .. math:: 

245 

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

247 

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

249 

250 Lorentzian broadening corresponds to 

251 

252 .. math:: 

253 

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

255 

256 Gaussian broadening corresponds to 

257 

258 .. math:: 

259 

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

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

262 \left\{ 

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

264 \left[ 

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

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

267 \right] \\ 

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

269 \left[ 

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

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

272 \right] 

273 \right\} 

274 

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

276 

277 .. math:: 

278 

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

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

281 

282 

283 Parameters 

284 ---------- 

285 freq_i 

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

287 data_iX 

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

289 freq_w 

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

291 freq_axis : `'real'` or `'imag'` 

292 Interpretation of frequency values as real or imaginary frequencies. 

293 broadening 

294 Broadening to be used. 

295 

296 Notes 

297 ----- 

298 When working with imaginary axes 

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

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

301 yield similar results, yet the latter is computationally 

302 more efficient when `freq_w` is real. 

303 

304 Returns 

305 ------- 

306 :class:`np.ndarray` 

307 Broadened data along continuous frequency axis. 

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

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

310 """ 

311 if freq_i.ndim != 1: 

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

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

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

315 

316 if freq_axis == 'real': 

317 if isinstance(broadening, NoArtificialBroadening): 

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

319 elif isinstance(broadening, LorentzianBroadening): 

320 width = broadening._width 

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

322 elif isinstance(broadening, GaussianBroadening): 

323 width = broadening._width 

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

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

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

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

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

329 ) 

330 else: 

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

332 elif freq_axis == 'imag': 

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

334 if not isinstance(broadening, NoArtificialBroadening): 

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

336 else: 

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

338 

339 # transform 

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

341 return data_wX # type: ignore