Coverage for calorine/tools/spectra.py: 100%

183 statements  

« prev     ^ index     » next       coverage.py v7.16.0, created at 2026-10-08 12:52 +0000

1r""" 

2Tools for computing optical and vibrational spectra from MD trajectories. 

3 

4Covers three routes depending on the model and GPUMD output: 

5 

6* **Dielectric tensor** (:func:`get_dielectric_function`): full 3×3 tensor 

7 from ``dpdt.out`` (qNEP, ``column='dP'`` or ``column='P'``) or ``dipole.out`` 

8 (TNEP, ``column='mu'``). 

9* **IR spectrum** (:func:`get_ir_spectrum`): from ``dpdt.out`` 

10 (qNEP, ``column='dP'`` or ``column='P'``) or ``dipole.out`` 

11 (TNEP, ``column='mu'``). 

12* **Raman spectrum** (:func:`get_raman_spectrum`): from ``polarizability.out`` 

13 (TNEP polarizability/susceptibility model). 

14 

15Nomenclature 

16------------ 

17- **Molecules / nanoparticles (localized systems)**: dipole moment :math:`\mu` 

18 (e·Å) and polarizability :math:`\alpha` (bohr³ or training-data units). 

19- **Solids / liquids (extended systems)**: polarization :math:`P` (e·Å, total 

20 cell dipole) and susceptibility :math:`\chi` (same unit as training data per 

21 cell). 

22 

23Both cases use the same GPUMD output files and formulas; the distinction is 

24whether :attr:`volume` is supplied for normalization. 

25 

26Quantum corrections 

27------------------- 

28The spectra from classical MD can be corrected for quantum statistics via 

29:func:`apply_quantum_correction`. The correction factors differ by scattering 

30order (M. Cardona, *Topics in Applied Physics*, Vol. 50 (Springer, 1982); 

31Rosander *et al.*, PRB **111**, 064107 (2025)). 

32""" 

33 

34import warnings 

35 

36import numpy as np 

37from pandas import DataFrame 

38from scipy.signal import correlate 

39from ase.units import _k as kB_SI, _e as e_SI, _eps0 as eps0_SI, _c as c_SI, _hbar as hbar_SI 

40 

41_rad_s_to_thz = 1e-12 

42_rad_s_to_invcm = 1 / (2 * np.pi * c_SI * 100) 

43 

44 

45def _compute_correlation_function(Z1, Z2): 

46 """Normalized cross-correlation; normalization divides by the number of contributing pairs.""" 

47 N = len(Z1) 

48 # scipy.signal.correlate switches to an FFT for long signals, np.correlate is always O(N^2). 

49 return correlate(Z1, Z2, mode='full')[N - 1:] / np.arange(N, 0, -1) 

50 

51 

52def _gaussian_decay(t, t_sigma): 

53 r"""Gaussian envelope :math:`\exp(-t^2 / 2\sigma^2)` used as a window applied to the ACF.""" 

54 return np.exp(-0.5 * (t / t_sigma) ** 2) 

55 

56 

57def _psd_from_acf(acf, dt): 

58 """Power spectral density via even extension of the ACF followed by rfft. 

59 

60 The even extension makes this the two-sided transform, so a Green-Kubo integral, which 

61 takes the one-sided one, needs a factor of one half. 

62 """ 

63 signal = np.hstack((acf, acf[:0:-1])) 

64 fft = dt * np.fft.rfft(signal) 

65 freqs = 2 * np.pi * np.fft.rfftfreq(len(signal), dt) 

66 return freqs, fft.real 

67 

68 

69def _parse_polarization_pair(polarization_in, polarization_out): 

70 """Validate and parse a pair of polarization unit vectors, or return (None, None).""" 

71 if (polarization_in is None) != (polarization_out is None): 

72 raise ValueError('polarization_in and polarization_out must be given together') 

73 if polarization_in is None: 

74 return None, None 

75 n_in = np.asarray(polarization_in, dtype=float) 

76 n_out = np.asarray(polarization_out, dtype=float) 

77 if n_in.shape != (3,) or n_out.shape != (3,): 

78 raise ValueError('polarization_in and polarization_out must have shape (3,)') 

79 return n_in, n_out 

80 

81 

82# ── Signal column mapping ────────────────────────────────────────────────────── 

83 

84_SIGNAL_COLS = { 

85 'dP': (['dPx', 'dPy', 'dPz'], True), 

86 'P': (['Px', 'Py', 'Pz'], False), 

87 'mu': (['mu_x', 'mu_y', 'mu_z'], False), 

88} 

89 

90 

91def _extract_dt(df, dt): 

92 """Return dt (ps), auto-extracting from the DataFrame's time column when not given.""" 

93 if dt is None: 

94 if 'time' not in df.columns: 

95 raise ValueError( 

96 'dt must be provided when the DataFrame has no time column ' 

97 f'(got columns: {list(df.columns)!r})' 

98 ) 

99 diffs = df['time'].diff().dropna() 

100 if diffs.empty: 

101 raise ValueError( 

102 'dt must be provided when the DataFrame has fewer than two rows ' 

103 f'(got {len(df)} row(s))' 

104 ) 

105 dt = float(diffs.iloc[0]) 

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

107 raise ValueError( 

108 'dt must be provided when the time column is not uniformly spaced ' 

109 f'(spacing ranges from {diffs.min()} to {diffs.max()} ps)' 

110 ) 

111 if dt <= 0: 

112 raise ValueError(f'dt must be positive, got {dt} ps') 

113 return dt 

114 

115 

116def _resolve_signal(signal, column, dt): 

117 """Select signal columns from a DataFrame, infer derivative flag, auto-extract dt.""" 

118 if column not in _SIGNAL_COLS: 

119 raise ValueError( 

120 f'column must be one of {list(_SIGNAL_COLS)!r}, got {column!r}' 

121 ) 

122 col_names, derivative = _SIGNAL_COLS[column] 

123 missing = [c for c in col_names if c not in signal.columns] 

124 if missing: 

125 raise ValueError( 

126 f'signal is missing columns {missing!r} for column={column!r}' 

127 ) 

128 dt = _extract_dt(signal, dt) 

129 return signal[col_names].to_numpy(dtype=float), dt, derivative 

130 

131 

132# ── IR spectrum ──────────────────────────────────────────────────────────────── 

133 

134def _acf_to_psd(acf, dt, t_sigma, window_size): 

135 """Truncate ACF, apply optional Gaussian window, and compute PSD. dt, t_sigma, and 

136 window_size in ps.""" 

137 if window_size is None: 

138 window_size = 5 * t_sigma if t_sigma is not None else None 

139 if window_size is not None: 

140 if t_sigma is not None and window_size < 3 * t_sigma: 

141 warnings.warn( 

142 f'window_size ({window_size:.3g} ps) is less than 3 * t_sigma ' 

143 f'({3 * t_sigma:.3g} ps); the Gaussian window is truncated before it has ' 

144 'meaningfully decayed, which can introduce spectral artifacts. Consider ' 

145 'increasing window_size or decreasing t_sigma.' 

146 ) 

147 max_available = len(acf) * dt 

148 if window_size > max_available: 

149 raise ValueError( 

150 f'window_size ({window_size:.3g} ps) exceeds the {max_available:.3g} ps ' 

151 'of autocorrelation data available from the input trajectory. Use a ' 

152 'smaller window_size (or a larger t_sigma, since window_size defaults to ' 

153 '5 * t_sigma), or provide a longer trajectory.' 

154 ) 

155 acf = acf if window_size is None else acf[:int(window_size / dt)] 

156 t_acf = dt * np.arange(len(acf)) # ps 

157 if t_sigma is not None: 

158 acf = acf * _gaussian_decay(t_acf, t_sigma=t_sigma) # both ps 

159 return _psd_from_acf(acf, dt * 1e-12) # _psd_from_acf uses seconds 

160 

161 

162def _ir_spectrum_1d(dt, signal, volume_SI, temperature, t_sigma, window_size, derivative=False): 

163 """Compute IR spectrum for a single dipole/polarization component in SI units.""" 

164 D = e_SI * 1e-10 / (1e-15 if derivative else 1.0) 

165 signal_SI = (signal - np.mean(signal)) * D 

166 acf = _compute_correlation_function(signal_SI, signal_SI) 

167 w, S = _acf_to_psd(acf, dt, t_sigma, window_size) 

168 # When input is the signal itself, multiply by omega**2 to get PSD of its derivative 

169 # (PSD[dx/dt] = omega**2 * PSD[x]). When input is already the derivative, use PSD directly. 

170 S_dpdt = S if derivative else w ** 2 * S 

171 

172 if volume_SI is not None: 

173 beta = 1.0 / (kB_SI * temperature) 

174 # One-sided transform. 

175 conductivity = (beta / (2 * volume_SI)) * S_dpdt 

176 with np.errstate(divide='ignore', invalid='ignore'): 

177 eps_imag = conductivity / (eps0_SI * w) 

178 return w, eps_imag, conductivity 

179 else: 

180 return w, S_dpdt 

181 

182 

183def get_ir_spectrum( 

184 signal: DataFrame, 

185 volume: float = None, 

186 temperature: float = None, 

187 column: str = 'dP', 

188 dt: float = None, 

189 polarization: np.ndarray = None, 

190 t_sigma: float = None, 

191 window_size: float = None, 

192) -> DataFrame: 

193 r"""Compute the IR spectrum from a time series of dipole moments or polarizations. 

194 

195 With :attr:`volume` and :attr:`temperature` given, the returned 

196 ``epsilon_imag`` is the isotropic average of the dielectric tensor, 

197 

198 .. math:: 

199 

200 \epsilon_2(\omega) = \frac{1}{3} \sum_\alpha \frac{\beta}{\epsilon_0 \omega V} 

201 \,\mathrm{Re}\!\left[\int_0^\infty 

202 \langle \dot{P}_\alpha(0)\,\dot{P}_\alpha(t) \rangle 

203 e^{-i\omega t}\,dt\right], 

204 

205 where :math:`\beta = 1/(k_\mathrm{B}T)` and :math:`V` is the cell volume. 

206 A given :attr:`polarization` :math:`\hat{n}` replaces the average by the single 

207 projected component :math:`\hat{n} \cdot \boldsymbol{\epsilon}_2 \cdot \hat{n}`. 

208 ``conductivity`` is :math:`\epsilon_0 \omega \epsilon_2(\omega)`. 

209 

210 Parameters 

211 ---------- 

212 signal 

213 ``DataFrame`` as returned by :func:`~calorine.gpumd.read_dpdt` or 

214 :func:`~calorine.gpumd.read_dipole`. The columns to use are selected 

215 via :attr:`column`. 

216 volume 

217 Simulation cell volume in ų. Required for extended systems; when 

218 provided the function returns the imaginary dielectric function and 

219 conductivity. Pass ``None`` for molecules to obtain an unnormalized 

220 line shape in arbitrary units. 

221 temperature 

222 Temperature in K. Required when :attr:`volume` is not ``None``. 

223 column 

224 Selects which columns of :attr:`signal` to use and implies whether 

225 the input is a time derivative: 

226 

227 ``'dP'`` (default) 

228 Uses ``dPx``, ``dPy``, ``dPz`` from :func:`~calorine.gpumd.read_dpdt` 

229 (qNEP); signal is :math:`\mathrm{d}P/\mathrm{d}t` in e·Å/fs. 

230 ``'P'`` 

231 Uses ``Px``, ``Py``, ``Pz`` from :func:`~calorine.gpumd.read_dpdt` 

232 (qNEP); signal is the polarization :math:`P(t)` in e·Å. 

233 ``'mu'`` 

234 Uses ``mu_x``, ``mu_y``, ``mu_z`` from 

235 :func:`~calorine.gpumd.read_dipole` (TNEP); signal is the dipole 

236 moment :math:`\mu(t)` in e·Å. 

237 

238 dt 

239 Time between consecutive frames in ps. Auto-extracted from the 

240 ``time`` column when present (i.e. for :func:`~calorine.gpumd.read_dpdt` 

241 output). Must be supplied explicitly for step-indexed inputs such as 

242 :func:`~calorine.gpumd.read_dipole` output. 

243 polarization 

244 Unit vector ``(3,)`` defining the electric-field polarization direction. 

245 When given, the spectrum is computed from the projected signal 

246 :math:`s(t) = \hat{n} \cdot \mathrm{signal}(t)` instead of the 

247 isotropic average. 

248 t_sigma 

249 Width of the Gaussian window applied to the ACF in ps. 

250 ``None`` uses no windowing. 

251 window_size 

252 Length of the ACF to retain, in ps. ``None`` (default) uses 

253 :math:`5 \times` :attr:`t_sigma` when :attr:`t_sigma` is given, otherwise the 

254 full ACF. A warning is issued if :attr:`window_size` is less than 

255 :math:`3 \times` :attr:`t_sigma`, since the Gaussian window is then truncated 

256 before it has meaningfully decayed. Raises :class:`ValueError` if 

257 :attr:`window_size` exceeds the autocorrelation data available from 

258 :attr:`signal`. 

259 

260 Returns 

261 ------- 

262 DataFrame 

263 Always contains ``angular_frequency`` (THz) and 

264 ``wavenumber_invcm`` (:math:`\mathrm{cm}^{-1}`). 

265 When :attr:`volume` and :attr:`temperature` are given: 

266 ``epsilon_imag`` (dimensionless) and ``conductivity`` (S/m). 

267 When :attr:`volume` is ``None``: ``ir_intensity`` (arbitrary units, 

268 proportional to :math:`\mathrm{PSD}[\dot{\mu}]`). 

269 """ 

270 if volume is not None and temperature is None: 

271 raise ValueError('temperature must be provided when volume is given') 

272 

273 arr, dt, derivative = _resolve_signal(signal, column, dt) 

274 volume_SI = volume * 1e-30 if volume is not None else None 

275 

276 if polarization is not None: 

277 n = np.asarray(polarization, dtype=float) 

278 if n.shape != (3,): 

279 raise ValueError('polarization must have shape (3,)') 

280 components = [arr @ n] 

281 else: 

282 components = [arr[:, i] for i in range(3)] 

283 

284 results = [_ir_spectrum_1d(dt, c, volume_SI, temperature, t_sigma, window_size, derivative) 

285 for c in components] 

286 

287 if volume_SI is not None: 

288 w = results[0][0] 

289 eps_imag = np.mean([r[1] for r in results], axis=0) 

290 conductivity = np.mean([r[2] for r in results], axis=0) 

291 mask = np.isfinite(eps_imag) 

292 cols = ['angular_frequency', 'wavenumber_invcm', 'epsilon_imag', 'conductivity'] 

293 data = np.column_stack((w[mask] * _rad_s_to_thz, w[mask] * _rad_s_to_invcm, 

294 eps_imag[mask], conductivity[mask])) 

295 return DataFrame(data, columns=cols) 

296 else: 

297 w = results[0][0] 

298 ir_intensity = np.mean([r[1] for r in results], axis=0) 

299 cols = ['angular_frequency', 'wavenumber_invcm', 'ir_intensity'] 

300 data = np.column_stack((w * _rad_s_to_thz, w * _rad_s_to_invcm, ir_intensity)) 

301 return DataFrame(data, columns=cols) 

302 

303 

304# ── Raman spectrum ───────────────────────────────────────────────────────────── 

305 

306_POL_COLS = ['xx', 'yy', 'zz', 'yz', 'xz', 'xy'] 

307 

308 

309def get_raman_spectrum( 

310 polarizability: DataFrame, 

311 dt: float = None, 

312 polarization_in: np.ndarray = None, 

313 polarization_out: np.ndarray = None, 

314 t_sigma: float = None, 

315 window_size: float = None, 

316) -> DataFrame: 

317 r"""Compute the Raman spectrum from a time series of polarizability tensors. 

318 

319 Parameters 

320 ---------- 

321 polarizability 

322 DataFrame with columns ``xx``, ``yy``, ``zz``, ``xy``, ``yz``, ``xz`` 

323 containing the polarizability :math:`\alpha` (molecules) or susceptibility 

324 :math:`\chi` (extended systems), as returned by 

325 :func:`~calorine.gpumd.read_polarizability`. Units are those of the 

326 TNEP training data (typically bohr³). 

327 dt 

328 Time between consecutive frames in ps. Auto-extracted from the 

329 ``time`` column when present (i.e. for 

330 :func:`~calorine.gpumd.read_polarizability` output with a header). 

331 Must be supplied explicitly for header-less/step-indexed input. 

332 polarization_in 

333 Unit vector ``(3,)`` for the polarization of the incoming light. If 

334 both :attr:`polarization_in` and :attr:`polarization_out` are given, 

335 the polarization-resolved intensity 

336 :math:`I(\omega) \propto \mathrm{FT}[\langle s(0)s(t)\rangle]` with 

337 :math:`s(t) = \hat{n}^\mathrm{out} \cdot \alpha(t) \cdot \hat{n}^\mathrm{in}` 

338 is added as the column ``raman_polarized``. 

339 polarization_out 

340 Unit vector ``(3,)`` for the polarization of the outgoing (scattered) 

341 light. 

342 t_sigma 

343 Width of the Gaussian window applied to the ACF in ps. ``None`` uses no 

344 windowing. 

345 window_size 

346 Length of the ACF to retain, in ps. ``None`` (default) uses 

347 :math:`5 \times` :attr:`t_sigma` when :attr:`t_sigma` is given, otherwise the 

348 full ACF. A warning is issued if :attr:`window_size` is less than 

349 :math:`3 \times` :attr:`t_sigma`, since the Gaussian window is then truncated 

350 before it has meaningfully decayed. Raises :class:`ValueError` if 

351 :attr:`window_size` exceeds the autocorrelation data available from 

352 :attr:`polarizability`. 

353 

354 Returns 

355 ------- 

356 DataFrame 

357 Always contains ``angular_frequency`` (THz), 

358 ``wavenumber_invcm`` (:math:`\mathrm{cm}^{-1}`), 

359 ``raman_isotropic`` (proportional to 

360 :math:`\mathrm{FT}[\langle\gamma(0)\gamma(t)\rangle]`), and 

361 ``raman_anisotropic`` (proportional to 

362 :math:`\mathrm{FT}[\langle\mathrm{Tr}[\beta(0)\beta(t)]\rangle]`). 

363 If both polarization vectors are given, also contains 

364 ``raman_polarized``. 

365 """ 

366 missing = [c for c in _POL_COLS if c not in polarizability.columns] 

367 if missing: 

368 raise ValueError(f'polarizability is missing columns: {missing}') 

369 dt = _extract_dt(polarizability, dt) 

370 

371 # Remove mean (Rayleigh component) before computing fluctuation spectrum. 

372 pol = polarizability[_POL_COLS] - polarizability[_POL_COLS].mean() 

373 

374 n_in, n_out = _parse_polarization_pair(polarization_in, polarization_out) 

375 

376 # Isotropic component gamma(t) = Tr[alpha(t)] / 3. 

377 gamma = (pol['xx'] + pol['yy'] + pol['zz']).to_numpy() / 3.0 

378 

379 # Traceless part beta = alpha - gamma*I. 

380 beta_xx = pol['xx'].to_numpy() - gamma 

381 beta_yy = pol['yy'].to_numpy() - gamma 

382 beta_zz = pol['zz'].to_numpy() - gamma 

383 # Off-diagonal elements are unchanged: beta_xy = alpha_xy, etc. 

384 

385 with np.errstate(invalid='ignore', divide='ignore'): 

386 # Isotropic Raman: FT[<gamma(0)*gamma(t)>]. 

387 acf_iso = _compute_correlation_function(gamma, gamma) 

388 w, L_iso = _acf_to_psd(acf_iso, dt, t_sigma, window_size) 

389 

390 # Anisotropic Raman: FT[<Tr[beta(0)*beta(t)]>]. 

391 # Tr[beta(0)*beta(t)] = sum_ij beta_ij(0)*beta_ij(t), the sum of per-element 

392 # ACFs with factor 2 for off-diagonal (since beta_ij = beta_ji). 

393 acf_aniso = ( 

394 _compute_correlation_function(beta_xx, beta_xx) 

395 + _compute_correlation_function(beta_yy, beta_yy) 

396 + _compute_correlation_function(beta_zz, beta_zz) 

397 + 2.0 * _compute_correlation_function(pol['yz'].to_numpy(), pol['yz'].to_numpy()) 

398 + 2.0 * _compute_correlation_function(pol['xz'].to_numpy(), pol['xz'].to_numpy()) 

399 + 2.0 * _compute_correlation_function(pol['xy'].to_numpy(), pol['xy'].to_numpy()) 

400 ) 

401 _, L_aniso = _acf_to_psd(acf_aniso, dt, t_sigma, window_size) 

402 

403 cols = ['angular_frequency', 'wavenumber_invcm', 'raman_isotropic', 'raman_anisotropic'] 

404 data = np.column_stack((w * _rad_s_to_thz, w * _rad_s_to_invcm, L_iso, L_aniso)) 

405 df = DataFrame(data, columns=cols) 

406 

407 # Polarization-resolved spectrum: s(t) = n_out . alpha(t) . n_in. 

408 # s = sum_ij n_out_i * alpha_ij * n_in_j; use symmetry alpha_ij = alpha_ji. 

409 if n_in is not None: 

410 s = (n_in[0]*n_out[0] * pol['xx'].to_numpy() 

411 + n_in[1]*n_out[1] * pol['yy'].to_numpy() 

412 + n_in[2]*n_out[2] * pol['zz'].to_numpy() 

413 + (n_in[1]*n_out[2] + n_in[2]*n_out[1]) * pol['yz'].to_numpy() 

414 + (n_in[0]*n_out[2] + n_in[2]*n_out[0]) * pol['xz'].to_numpy() 

415 + (n_in[0]*n_out[1] + n_in[1]*n_out[0]) * pol['xy'].to_numpy()) 

416 with np.errstate(invalid='ignore', divide='ignore'): 

417 acf_pol = _compute_correlation_function(s, s) 

418 _, L_pol = _acf_to_psd(acf_pol, dt, t_sigma, window_size) 

419 df['raman_polarized'] = L_pol 

420 

421 return df 

422 

423 

424# ── Quantum correction ───────────────────────────────────────────────────────── 

425 

426def apply_quantum_correction( 

427 df: DataFrame, 

428 temperature: float, 

429 column: str, 

430 order: str = 'first', 

431 force: bool = False, 

432) -> DataFrame: 

433 r"""Apply a quantum correction to a classically computed spectrum. 

434 

435 Classical MD underestimates spectral intensities because it samples the 

436 classical Boltzmann distribution rather than the Bose-Einstein distribution. 

437 The correction factor depends on the scattering order and mode type 

438 [Cardona1982]_ [Rosander2025]_. 

439 

440 Parameters 

441 ---------- 

442 df 

443 DataFrame with an ``angular_frequency`` column (THz) and the column 

444 to correct. 

445 temperature 

446 Temperature in K. 

447 column 

448 Name of the spectral column to correct (e.g. ``'raman_isotropic'``, 

449 ``'ir_intensity'``, ``'epsilon_imag'``). 

450 order 

451 Correction type: 

452 

453 ``'first'`` 

454 First-order scattering (IR absorption and first-order Raman): 

455 :math:`f = \beta\hbar\omega / (1 - e^{-\beta\hbar\omega})`. 

456 ``'overtone'`` 

457 Second-order overtone (same mode, :math:`\omega_1 = \omega/2`): 

458 :math:`f = [y/(1 - e^{-y})]^2 (2 - e^{-y})` with 

459 :math:`y = \beta\hbar\omega/2`. 

460 ``'combination'`` 

461 Second-order combination band upper bound 

462 (:math:`\omega_1 = \omega_2 = \omega/2`): 

463 :math:`f = [y/(1 - e^{-y})]^2` with :math:`y = \beta\hbar\omega/2`. 

464 force 

465 If ``True``, overwrite an existing ``column + '_qm'`` column. 

466 Default ``False`` raises :class:`ValueError` to prevent accidental 

467 double-correction. 

468 

469 Returns 

470 ------- 

471 DataFrame 

472 Copy of :attr:`df` with a new column ``column + '_qm'`` containing 

473 the quantum-corrected intensities. The input is not mutated. 

474 

475 References 

476 ---------- 

477 .. [Cardona1982] M. Cardona, *Resonance phenomena*, in *Topics in Applied 

478 Physics*, Vol. 50, edited by M. Cardona and G. Güntherodt 

479 (Springer, Berlin, 1982). 

480 .. [Rosander2025] Rosander *et al.*, Phys. Rev. B **111**, 064107 (2025), 

481 Supp. Eqs. S5, S7, S10. 

482 """ 

483 if order not in ('first', 'overtone', 'combination'): 

484 raise ValueError("order must be 'first', 'overtone', or 'combination'") 

485 

486 if not force and column + '_qm' in df.columns: 

487 raise ValueError( 

488 f"column '{column}_qm' already exists; pass force=True to overwrite" 

489 ) 

490 

491 w = df['angular_frequency'].to_numpy() * 1e12 # THz → rad/s 

492 beta = 1.0 / (kB_SI * temperature) 

493 x = beta * hbar_SI * w # beta*hbar*omega (dimensionless) 

494 

495 with np.errstate(divide='ignore', invalid='ignore', over='ignore'): 

496 if order == 'first': 

497 # f = beta*hbar*omega / (1 - exp(-beta*hbar*omega)) 

498 factor = np.where(x > 0, x / (1.0 - np.exp(-x)), 1.0) 

499 elif order == 'overtone': 

500 # f = [y/(1-exp(-y))]**2 * (2-exp(-y)), y = beta*hbar*omega/2 

501 y = x / 2.0 

502 base = np.where(y > 0, y / (1.0 - np.exp(-y)), 1.0) 

503 factor = base ** 2 * np.where(y > 0, 2.0 - np.exp(-y), 1.0) 

504 else: # combination 

505 # f = [y/(1-exp(-y))]**2, y = beta*hbar*omega/2 

506 y = x / 2.0 

507 base = np.where(y > 0, y / (1.0 - np.exp(-y)), 1.0) 

508 factor = base ** 2 

509 

510 result = df.copy() 

511 result[column + '_qm'] = df[column].to_numpy() * factor 

512 return result 

513 

514 

515# ── Dielectric tensor ───────────────────────────────────────────────────────── 

516 

517_TENSOR_COMPONENTS = [(0, 0, 'xx'), (1, 1, 'yy'), (2, 2, 'zz'), 

518 (1, 2, 'yz'), (0, 2, 'xz'), (0, 1, 'xy')] 

519 

520 

521def get_dielectric_function( 

522 signal: DataFrame, 

523 volume: float, 

524 temperature: float, 

525 column: str = 'dP', 

526 dt: float = None, 

527 t_sigma: float = None, 

528 window_size: float = None, 

529 return_real_part: bool = True, 

530 kk_method: str = 'vectorized', 

531 kk_max_memory_gb: float = 4.0, 

532) -> DataFrame: 

533 r"""Compute the dielectric function from a three-component polarization time series. 

534 

535 Computes all six unique Voigt components of :math:`\epsilon_2^{\alpha\beta}(\omega)` 

536 (xx, yy, zz, yz, xz, xy) from the three-component time series written by GPUMD. 

537 The real part :math:`\epsilon_1^{\alpha\beta}(\omega)` is obtained via the 

538 Kramers-Kronig relation when :attr:`return_real_part` is ``True`` (the default). 

539 

540 The imaginary part is related to the symmetrized cross-correlation via 

541 

542 .. math:: 

543 

544 \epsilon_2^{\alpha\beta}(\omega) = \frac{\beta}{\epsilon_0 \omega V} 

545 \,\mathrm{Re}\!\left[\int_0^\infty 

546 \langle \dot{P}_\alpha(0)\,\dot{P}_\beta(t) \rangle 

547 e^{-i\omega t}\,dt\right], 

548 

549 where :math:`\beta = 1/(k_\mathrm{B}T)` and :math:`V` is the cell volume. 

550 

551 Parameters 

552 ---------- 

553 signal 

554 ``DataFrame`` as returned by :func:`~calorine.gpumd.read_dpdt` or 

555 :func:`~calorine.gpumd.read_dipole`. The columns to use are selected 

556 via :attr:`column`. 

557 volume 

558 Simulation cell volume in ų. 

559 temperature 

560 Temperature in K. 

561 column 

562 Selects which columns of :attr:`signal` to use and implies whether 

563 the input is a time derivative: 

564 

565 ``'dP'`` (default) 

566 Uses ``dPx``, ``dPy``, ``dPz`` from :func:`~calorine.gpumd.read_dpdt` 

567 (qNEP); input is :math:`\dot{\mathbf{P}}(t)` in e·Å/fs. 

568 ``'P'`` 

569 Uses ``Px``, ``Py``, ``Pz`` from :func:`~calorine.gpumd.read_dpdt` 

570 (qNEP); input is the polarization :math:`\mathbf{P}(t)` in e·Å. 

571 ``'mu'`` 

572 Uses ``mu_x``, ``mu_y``, ``mu_z`` from 

573 :func:`~calorine.gpumd.read_dipole` (TNEP); input is the dipole 

574 moment :math:`\boldsymbol{\mu}(t)` in e·Å. 

575 

576 dt 

577 Time between consecutive frames in ps. Auto-extracted from the 

578 ``time`` column when present (i.e. for :func:`~calorine.gpumd.read_dpdt` 

579 output). Must be supplied explicitly for step-indexed inputs such as 

580 :func:`~calorine.gpumd.read_dipole` output. 

581 t_sigma 

582 Width of the Gaussian window applied to the ACF in ps. 

583 ``None`` uses no windowing. 

584 window_size 

585 Length of the ACF to retain, in ps. ``None`` (default) uses 

586 :math:`5 \times` :attr:`t_sigma` when :attr:`t_sigma` is given, otherwise the 

587 full ACF (i.e. the full trajectory length). This directly 

588 sets the size of the frequency grid (:math:`n \approx` :attr:`window_size` 

589 :math:`/\,dt`, or the number of frames in :attr:`signal` when unset) passed to 

590 the Kramers-Kronig step below; see :attr:`kk_method` for how the cost of that 

591 step depends on :math:`n`. A warning is issued if :attr:`window_size` is less 

592 than :math:`3 \times` :attr:`t_sigma`, since the Gaussian window is then 

593 truncated before it has meaningfully decayed. Raises :class:`ValueError` if 

594 :attr:`window_size` exceeds the autocorrelation data available from 

595 :attr:`signal`. 

596 return_real_part 

597 When ``True`` (default), the real part :math:`\epsilon_1^{\alpha\beta}(\omega)` 

598 is computed via the Kramers-Kronig relation and appended as 

599 ``epsilon_real_{xx,yy,zz,yz,xz,xy}`` columns. 

600 kk_method 

601 Integration method passed to :func:`~calorine.tools.apply_kramers_kronig` 

602 when :attr:`return_real_part` is ``True``. ``'vectorized'`` (default) uses 

603 an exact trapezoid rule over an :math:`n \times n` matrix 

604 (:math:`\mathcal{O}(n^2)` time and memory, where :math:`n` is the size of the 

605 frequency grid set by :attr:`window_size`); ``'fft'`` uses a faster 

606 (:math:`\mathcal{O}(n \log n)` time, :math:`\mathcal{O}(n)` memory) Hilbert 

607 transform approximation. Ignored when :attr:`return_real_part` is ``False``. 

608 kk_max_memory_gb 

609 Memory limit (in GB) passed to :func:`~calorine.tools.apply_kramers_kronig` as 

610 ``max_memory_gb`` when :attr:`kk_method` is ``'vectorized'``. If the estimated 

611 peak memory for the :math:`n \times n` matrix would exceed this, :class:`ValueError` 

612 is raised before attempting the allocation, rather than letting the process hang 

613 or get OOM-killed; pass ``kk_method='fft'``, a smaller :attr:`window_size`, or a 

614 larger :attr:`kk_max_memory_gb` instead. Set to ``None`` to disable the check. 

615 Ignored when :attr:`return_real_part` is ``False`` or :attr:`kk_method` is 

616 ``'fft'``. 

617 

618 Returns 

619 ------- 

620 DataFrame 

621 Contains ``angular_frequency`` (THz), ``wavenumber_invcm`` 

622 (:math:`\mathrm{cm}^{-1}`), 

623 ``epsilon_imag_{xx,yy,zz,yz,xz,xy}`` (dimensionless, imaginary dielectric 

624 tensor components), and ``conductivity_{xx,yy,zz,yz,xz,xy}`` (S/m). 

625 When :attr:`return_real_part` is ``True``, ``epsilon_real_{xx,yy,zz,yz,xz,xy}`` 

626 columns are appended. 

627 """ 

628 arr, dt, derivative = _resolve_signal(signal, column, dt) 

629 

630 D = e_SI * 1e-10 / (1e-15 if derivative else 1.0) 

631 volume_SI = volume * 1e-30 

632 beta = 1.0 / (kB_SI * temperature) 

633 

634 signal_SI = (arr - arr.mean(axis=0)) * D 

635 

636 def _psd(z1, z2=None): 

637 if z2 is None: 

638 acf = _compute_correlation_function(z1, z1) 

639 else: 

640 acf = (_compute_correlation_function(z1, z2) 

641 + _compute_correlation_function(z2, z1)) / 2 

642 w, S = _acf_to_psd(acf, dt, t_sigma, window_size) 

643 # PSD[dx/dt] = omega**2 * PSD[x] when input is the signal itself. 

644 return w, S if derivative else w ** 2 * S 

645 

646 w = None 

647 conductivities = {} 

648 eps_imags = {} 

649 for i, j, lbl in _TENSOR_COMPONENTS: 

650 zi = signal_SI[:, i] 

651 zj = signal_SI[:, j] 

652 w, S_dpdt = _psd(zi) if i == j else _psd(zi, zj) 

653 # One-sided transform. 

654 conductivity = (beta / (2 * volume_SI)) * S_dpdt 

655 with np.errstate(divide='ignore', invalid='ignore'): 

656 eps_imag = conductivity / (eps0_SI * w) 

657 conductivities[lbl] = conductivity 

658 eps_imags[lbl] = eps_imag 

659 

660 mask = w > 0 

661 cols = (['angular_frequency', 'wavenumber_invcm'] 

662 + [f'epsilon_imag_{lbl}' for _, _, lbl in _TENSOR_COMPONENTS] 

663 + [f'conductivity_{lbl}' for _, _, lbl in _TENSOR_COMPONENTS]) 

664 arrays = ([w[mask] * _rad_s_to_thz, w[mask] * _rad_s_to_invcm] 

665 + [eps_imags[lbl][mask] for _, _, lbl in _TENSOR_COMPONENTS] 

666 + [conductivities[lbl][mask] for _, _, lbl in _TENSOR_COMPONENTS]) 

667 result = DataFrame(np.column_stack(arrays), columns=cols) 

668 if return_real_part: 

669 from .kramers_kronig import apply_kramers_kronig 

670 result = apply_kramers_kronig(result, method=kk_method, max_memory_gb=kk_max_memory_gb) 

671 return result