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
« 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.
4Covers three routes depending on the model and GPUMD output:
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).
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).
23Both cases use the same GPUMD output files and formulas; the distinction is
24whether :attr:`volume` is supplied for normalization.
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"""
34import warnings
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
41_rad_s_to_thz = 1e-12
42_rad_s_to_invcm = 1 / (2 * np.pi * c_SI * 100)
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)
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)
57def _psd_from_acf(acf, dt):
58 """Power spectral density via even extension of the ACF followed by rfft.
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
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
82# ── Signal column mapping ──────────────────────────────────────────────────────
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}
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
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
132# ── IR spectrum ────────────────────────────────────────────────────────────────
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
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
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
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.
195 With :attr:`volume` and :attr:`temperature` given, the returned
196 ``epsilon_imag`` is the isotropic average of the dielectric tensor,
198 .. math::
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],
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)`.
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:
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·Å.
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`.
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')
273 arr, dt, derivative = _resolve_signal(signal, column, dt)
274 volume_SI = volume * 1e-30 if volume is not None else None
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)]
284 results = [_ir_spectrum_1d(dt, c, volume_SI, temperature, t_sigma, window_size, derivative)
285 for c in components]
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)
304# ── Raman spectrum ─────────────────────────────────────────────────────────────
306_POL_COLS = ['xx', 'yy', 'zz', 'yz', 'xz', 'xy']
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.
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`.
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)
371 # Remove mean (Rayleigh component) before computing fluctuation spectrum.
372 pol = polarizability[_POL_COLS] - polarizability[_POL_COLS].mean()
374 n_in, n_out = _parse_polarization_pair(polarization_in, polarization_out)
376 # Isotropic component gamma(t) = Tr[alpha(t)] / 3.
377 gamma = (pol['xx'] + pol['yy'] + pol['zz']).to_numpy() / 3.0
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.
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)
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)
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)
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
421 return df
424# ── Quantum correction ─────────────────────────────────────────────────────────
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.
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]_.
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:
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.
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.
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'")
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 )
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)
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
510 result = df.copy()
511 result[column + '_qm'] = df[column].to_numpy() * factor
512 return result
515# ── Dielectric tensor ─────────────────────────────────────────────────────────
517_TENSOR_COMPONENTS = [(0, 0, 'xx'), (1, 1, 'yy'), (2, 2, 'zz'),
518 (1, 2, 'yz'), (0, 2, 'xz'), (0, 1, 'xy')]
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.
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).
540 The imaginary part is related to the symmetrized cross-correlation via
542 .. math::
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],
549 where :math:`\beta = 1/(k_\mathrm{B}T)` and :math:`V` is the cell volume.
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:
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·Å.
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'``.
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)
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)
634 signal_SI = (arr - arr.mean(axis=0)) * D
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
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
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