Coverage for calorine/tools/kramers_kronig.py: 100%
41 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-23 12:50 +0000
1r"""
2Kramers-Kronig relation:
4.. math::
6 \chi_\mathrm{real}(\omega) = \frac{2}{\pi} \mathcal{P}
7 \int_0^\infty \frac{\omega' \chi_\mathrm{imag}(\omega')}
8 {\omega'^2 - \omega^2} \, d\omega'
10Two private implementations are provided for validation purposes.
11The public entry point is :func:`apply_kramers_kronig`, which defaults to the O(n²) vectorized
12variant for accuracy; pass ``method='fft'`` for the O(n log n) approximation. Since the
13vectorized variant keeps multiple n x n float64 matrices alive at once, it guards against
14accidental huge allocations via ``max_memory_gb``.
15"""
17import numpy as np
18from pandas import DataFrame
21def _hilbert(x):
22 """ Discrete Hilbert transform via FFT (one-sided spectrum approach). """
23 N = len(x)
24 Xf = np.fft.fft(x)
25 h = np.zeros(N)
26 if N % 2 == 0:
27 h[0] = h[N // 2] = 1
28 h[1:N // 2] = 2
29 else:
30 h[0] = 1
31 h[1:(N + 1) // 2] = 2
32 return np.fft.ifft(Xf * h)
35def _kramers_kronig_vectorized(omega: np.ndarray, chi_imag: np.ndarray) -> np.ndarray:
36 """ Compute Re[chi] via a fully vectorized (n x n) matrix approach.
38 Parameters
39 ----------
40 omega
41 Non-negative angular frequency grid in rad/s, uniform or non-uniform.
42 chi_imag
43 Imaginary part of the susceptibility at each point in :attr:`omega`.
45 Returns
46 -------
47 np.ndarray
48 Real part of the susceptibility at each point in :attr:`omega`.
50 Notes
51 -----
52 Complexity: O(n²) time, O(n²) memory.
53 Accuracy: exact trapezoid rule over the supplied angular frequency grid.
54 """
55 omega = np.asarray(omega, dtype=float)
56 chi_imag = np.asarray(chi_imag, dtype=float)
58 w = omega[:, np.newaxis]
59 om = omega[np.newaxis, :]
61 denom = om ** 2 - w ** 2
62 np.fill_diagonal(denom, 1.0)
64 integrand = om * chi_imag[np.newaxis, :] / denom
65 np.fill_diagonal(integrand, 0.0)
67 return (2.0 / np.pi) * np.trapezoid(integrand, omega, axis=1)
70def _kramers_kronig_fft(omega: np.ndarray, chi_imag: np.ndarray) -> np.ndarray:
71 """ Compute Re[chi] via an FFT-based Hilbert transform.
73 Parameters
74 ----------
75 omega
76 Non-negative, uniformly spaced angular frequency grid in rad/s.
77 chi_imag
78 Imaginary part of the susceptibility at each point in :attr:`omega`.
80 Returns
81 -------
82 np.ndarray
83 Real part of the susceptibility at each point in :attr:`omega`.
85 Notes
86 -----
87 Complexity: O(n log n) time, O(n) memory.
88 Accuracy: approximation; differs from direct quadrature by roughly 5-10% of
89 peak amplitude on typical grids.
90 Requirement: :attr:`omega` must be uniformly spaced and :attr:`chi_imag` must be
91 negligibly small at both endpoints.
92 """
93 chi_full = np.concatenate([-chi_imag[::-1], chi_imag]) # odd extension
94 return -_hilbert(chi_full).imag[len(omega):]
97def apply_kramers_kronig(
98 df: DataFrame,
99 method: str = 'vectorized',
100 max_memory_gb: float = 4.0,
101) -> DataFrame:
102 r""" Apply the Kramers-Kronig relation to add the real part of the dielectric function.
104 Takes the output of :func:`~calorine.tools.get_dielectric_function`, computes
105 :math:`\epsilon_1(\omega)` from :math:`\epsilon_2(\omega)` via the Kramers-Kronig
106 relation, and returns the DataFrame with new ``epsilon_real*`` columns appended.
108 Each ``epsilon_imag{suffix}`` column in :attr:`df` produces a corresponding
109 ``epsilon_real{suffix}`` column (e.g. ``epsilon_imag_xx`` → ``epsilon_real_xx``).
111 Parameters
112 ----------
113 df
114 DataFrame as returned by :func:`~calorine.tools.get_dielectric_function` or
115 :func:`~calorine.tools.get_ir_spectrum`;
116 must contain ``angular_frequency`` and at least one column whose name starts
117 with ``epsilon_imag``.
118 method
119 Integration method. ``'vectorized'`` (default) uses an exact trapezoid rule
120 over an :math:`n \times n` matrix (:math:`\mathcal{O}(n^2)` time and memory,
121 where :math:`n` is the number of rows of :attr:`df`); ``'fft'`` uses an
122 :math:`\mathcal{O}(n \log n)` Hilbert transform (:math:`\mathcal{O}(n)` memory)
123 and is faster for large arrays but approximates the principal-value integral.
124 The two methods agree to within a few percent on typical grids.
125 max_memory_gb
126 When :attr:`method` is ``'vectorized'``, the peak memory footprint is
127 approximately :math:`3 \times 8n^2` bytes (three :math:`n \times n` ``float64``
128 buffers alive at once). If this estimate exceeds :attr:`max_memory_gb`,
129 :class:`ValueError` is raised before attempting the allocation, rather than
130 letting the process hang or get OOM-killed. Pass ``method='fft'``, a smaller
131 frequency grid (e.g. a smaller ``window_size`` in
132 :func:`~calorine.tools.get_dielectric_function`), or a larger
133 :attr:`max_memory_gb` to proceed anyway. Set to ``None`` to disable the check.
134 Ignored when :attr:`method` is ``'fft'``.
136 Returns
137 -------
138 DataFrame
139 Input DataFrame with additional ``epsilon_real*`` columns.
140 """
141 if method not in ('fft', 'vectorized'):
142 raise ValueError(f"method must be 'fft' or 'vectorized', got {method!r}")
144 omega = df['angular_frequency'].to_numpy()
145 imag_cols = [c for c in df.columns if c.startswith('epsilon_imag')]
147 if method == 'vectorized':
148 n = len(omega)
149 estimated_gb = 3 * n ** 2 * 8 / 1e9
150 if max_memory_gb is not None and estimated_gb > max_memory_gb:
151 raise ValueError(
152 f"method='vectorized' would build a {n}x{n} matrix, requiring an "
153 f'estimated {estimated_gb:.1f} GB of memory (limit: {max_memory_gb} GB). '
154 "Use method='fft' (O(n log n) time and memory), reduce the frequency "
155 'grid size (e.g. a smaller window_size in get_dielectric_function), or '
156 'pass a larger max_memory_gb to proceed anyway.'
157 )
159 kk = _kramers_kronig_fft if method == 'fft' else _kramers_kronig_vectorized
160 df = df.copy()
161 for col in imag_cols:
162 suffix = col[len('epsilon_imag'):]
163 df['epsilon_real' + suffix] = kk(omega, df[col].to_numpy())
164 return df