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 def _get_time_weight(self, time_t: np.ndarray) -> np.ndarray: 

26 raise ValueError(f'Unknown broadening: {self}') 

27 

28 def _get_frequency_weight(self, freq_i: np.ndarray, freq_w: np.ndarray) -> np.ndarray: 

29 raise ValueError(f'Unknown broadening: {self}') 

30 

31 

32class NoArtificialBroadening(Broadening): 

33 """Class representing no artificial broadening.""" 

34 

35 def _get_time_weight(self, time_t: np.ndarray) -> np.ndarray: 

36 return np.ones_like(time_t) 

37 

38 def _get_frequency_weight(self, freq_i: np.ndarray, freq_w: np.ndarray) -> np.ndarray: 

39 return 1. / (freq_i[:, np.newaxis]**2 - freq_w[np.newaxis, :]**2) # type: ignore 

40 

41 

42class ArtificialBroadening(Broadening): 

43 """Class representing artificial broadenings. 

44 

45 Parameters 

46 ---------- 

47 width 

48 Broadening width; meaning interpreted by the subclasses; 

49 must be strictly positive. 

50 

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

52 units 

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

54 atomic units. 

55 

56 This parameter determines whether conversion should be performed during 

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

58 """ 

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

60 if width <= 0: 

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

62 if units == 'eV': 

63 width = width * eV_to_au 

64 elif units != 'au': 

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

66 self._width = width 

67 

68 def __repr__(self) -> str: 

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

70 

71 def _repr_html_(self) -> str: 

72 """HTML representation for Jupyter notebooks.""" 

73 rows = [ 

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

75 ] 

76 return ( 

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

78 '<table>' 

79 '<thead><tr>' 

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

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

82 '</tr></thead>' 

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

84 '</table>' 

85 ) 

86 

87 @property 

88 def width(self) -> float: 

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

90 return self._width * au_to_eV 

91 

92 

93class GaussianBroadening(ArtificialBroadening): 

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

95 

96 Parameters 

97 ---------- 

98 width 

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

100 must be strictly positive. 

101 

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

103 units 

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

105 atomic units. 

106 

107 This parameter determines whether conversion should be performed during 

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

109 

110 Examples 

111 ------- 

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

113 >>> GaussianBroadening(0.1) 

114 GaussianBroadening(0.1, units='eV') 

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

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

117 """ 

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

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

120 

121 def _get_time_weight(self, time_t: np.ndarray) -> np.ndarray: 

122 return np.exp(-0.5 * self._width**2 * time_t**2) 

123 

124 def _get_frequency_weight(self, freq_i: np.ndarray, freq_w: np.ndarray) -> np.ndarray: 

125 width = self._width 

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

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

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

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

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

131 ) 

132 return weight_iw # type: ignore 

133 

134 

135class LorentzianBroadening(ArtificialBroadening): 

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

137 

138 Parameters 

139 ---------- 

140 width 

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

142 must be strictly positive. 

143 

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

145 units 

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

147 atomic units. 

148 

149 This parameter determines whether conversion should be performed during 

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

151 

152 Examples 

153 ------- 

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

155 >>> LorentzianBroadening(0.1) 

156 LorentzianBroadening(0.1, units='eV') 

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

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

159 """ 

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

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

162 

163 def _get_time_weight(self, time_t: np.ndarray) -> np.ndarray: 

164 return np.exp(-self._width * time_t) 

165 

166 def _get_frequency_weight(self, freq_i: np.ndarray, freq_w: np.ndarray) -> np.ndarray: 

167 width = self._width 

168 return 1. / (freq_i[:, np.newaxis]**2 - (freq_w[np.newaxis, :] + 1j * width)**2) 

169 

170 

171def fourier_transform(time_t: np.ndarray, 

172 data_tX: np.ndarray, 

173 freq_w: np.ndarray, 

174 freq_axis: str = 'real', 

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

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

177 

178 .. math:: 

179 

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

181 

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

183 

184 Lorentzian broadening corresponds to 

185 

186 .. math:: 

187 

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

189 

190 Gaussian broadening corresponds to 

191 

192 .. math:: 

193 

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

195 

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

197 assumed to be in mutually compatible units. 

198 

199 Parameters 

200 ---------- 

201 time_t 

202 Equally spaced time grid. 

203 data_tX 

204 Data to be transformed. 

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

206 dimensions can be of any size. 

207 freq_w 

208 Frequency values for the transform. 

209 freq_axis : `'real'` or `'imag'` 

210 Interpretation of frequency values as real or imaginary frequencies. 

211 broadening 

212 Broadening to be used. 

213 

214 Notes 

215 ----- 

216 When working with imaginary axes 

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

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

219 yield similar results, yet the latter is computationally 

220 more efficient when `freq_w` is real. 

221 

222 Returns 

223 ------- 

224 :class:`np.ndarray` 

225 Transform of input data to frequency. 

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

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

228 """ 

229 if time_t.ndim != 1: 

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

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

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

233 

234 if freq_axis == 'real': 

235 prefactor = 1.0j 

236 weight_t = broadening._get_time_weight(time_t) 

237 elif freq_axis == 'imag': 

238 prefactor = -1.0 

239 weight_t = np.ones_like(time_t) 

240 if not isinstance(broadening, NoArtificialBroadening): 

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

242 else: 

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

244 

245 # check time step 

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

247 dt = dt_t[0] 

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

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

250 

251 # integration weights from Simpson's integration rule 

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

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

254 

255 # transform 

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

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

258 return data_wX # type: ignore 

259 

260 

261def broaden(freq_i: np.ndarray, 

262 data_iX: np.ndarray, 

263 freq_w: np.ndarray, 

264 freq_axis: str = 'real', 

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

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

267 Specifically this function performs the summation 

268 

269 .. math:: 

270 

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

272 

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

274 

275 Lorentzian broadening corresponds to 

276 

277 .. math:: 

278 

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

280 

281 Gaussian broadening corresponds to 

282 

283 .. math:: 

284 

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

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

287 \left\{ 

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

289 \left[ 

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

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

292 \right] \\ 

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

294 \left[ 

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

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

297 \right] 

298 \right\} 

299 

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

301 

302 .. math:: 

303 

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

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

306 

307 

308 Parameters 

309 ---------- 

310 freq_i 

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

312 data_iX 

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

314 freq_w 

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

316 freq_axis : `'real'` or `'imag'` 

317 Interpretation of frequency values as real or imaginary frequencies. 

318 broadening 

319 Broadening to be used. 

320 

321 Notes 

322 ----- 

323 When working with imaginary axes 

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

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

326 yield similar results, yet the latter is computationally 

327 more efficient when `freq_w` is real. 

328 

329 Returns 

330 ------- 

331 :class:`np.ndarray` 

332 Broadened data along continuous frequency axis. 

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

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

335 """ 

336 if freq_i.ndim != 1: 

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

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

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

340 

341 if freq_axis == 'real': 

342 weight_iw = broadening._get_frequency_weight(freq_i, freq_w) 

343 elif freq_axis == 'imag': 

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

345 if not isinstance(broadening, NoArtificialBroadening): 

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

347 else: 

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

349 

350 # transform 

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

352 return data_wX # type: ignore