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

1r""" 

2Kramers-Kronig relation: 

3 

4.. math:: 

5 

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' 

9 

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. 

13""" 

14 

15import numpy as np 

16from pandas import DataFrame 

17 

18 

19def _hilbert(x): 

20 """ Discrete Hilbert transform via FFT (one-sided spectrum approach). """ 

21 N = len(x) 

22 Xf = np.fft.fft(x) 

23 h = np.zeros(N) 

24 if N % 2 == 0: 

25 h[0] = h[N // 2] = 1 

26 h[1:N // 2] = 2 

27 else: 

28 h[0] = 1 

29 h[1:(N + 1) // 2] = 2 

30 return np.fft.ifft(Xf * h) 

31 

32 

33def _kramers_kronig_vectorized(omega: np.ndarray, chi_imag: np.ndarray) -> np.ndarray: 

34 """ Compute Re[chi] via a fully vectorized (n x n) matrix approach. 

35 

36 Parameters 

37 ---------- 

38 omega 

39 Non-negative angular frequency grid in rad/s, uniform or non-uniform. 

40 chi_imag 

41 Imaginary part of the susceptibility at each point in :attr:`omega`. 

42 

43 Returns 

44 ------- 

45 np.ndarray 

46 Real part of the susceptibility at each point in :attr:`omega`. 

47 

48 Notes 

49 ----- 

50 Complexity: O(n²) time, O(n²) memory. 

51 Accuracy: exact trapezoid rule over the supplied angular frequency grid. 

52 """ 

53 omega = np.asarray(omega, dtype=float) 

54 chi_imag = np.asarray(chi_imag, dtype=float) 

55 

56 w = omega[:, np.newaxis] 

57 om = omega[np.newaxis, :] 

58 

59 denom = om ** 2 - w ** 2 

60 np.fill_diagonal(denom, 1.0) 

61 

62 integrand = om * chi_imag[np.newaxis, :] / denom 

63 np.fill_diagonal(integrand, 0.0) 

64 

65 return (2.0 / np.pi) * np.trapezoid(integrand, omega, axis=1) 

66 

67 

68def _kramers_kronig_fft(omega: np.ndarray, chi_imag: np.ndarray) -> np.ndarray: 

69 """ Compute Re[chi] via an FFT-based Hilbert transform. 

70 

71 Parameters 

72 ---------- 

73 omega 

74 Non-negative, uniformly spaced angular frequency grid in rad/s. 

75 chi_imag 

76 Imaginary part of the susceptibility at each point in :attr:`omega`. 

77 

78 Returns 

79 ------- 

80 np.ndarray 

81 Real part of the susceptibility at each point in :attr:`omega`. 

82 

83 Notes 

84 ----- 

85 Complexity: O(n log n) time, O(n) memory. 

86 Accuracy: approximation; differs from direct quadrature by roughly 5-10% of 

87 peak amplitude on typical grids. 

88 Requirement: :attr:`omega` must be uniformly spaced and :attr:`chi_imag` must be 

89 negligibly small at both endpoints. 

90 """ 

91 chi_full = np.concatenate([-chi_imag[::-1], chi_imag]) # odd extension 

92 return -_hilbert(chi_full).imag[len(omega):] 

93 

94 

95def apply_kramers_kronig(df: DataFrame, method: str = 'vectorized') -> DataFrame: 

96 r""" Apply the Kramers-Kronig relation to add the real part of the dielectric function. 

97 

98 Takes the output of :func:`~calorine.tools.get_dielectric_function`, computes 

99 :math:`\epsilon_1(\omega)` from :math:`\epsilon_2(\omega)` via the Kramers-Kronig 

100 relation, and returns the DataFrame with new ``epsilon_real*`` columns appended. 

101 

102 Each ``epsilon_imag{suffix}`` column in :attr:`df` produces a corresponding 

103 ``epsilon_real{suffix}`` column (e.g. ``epsilon_imag_xx`` → ``epsilon_real_xx``). 

104 

105 Parameters 

106 ---------- 

107 df 

108 DataFrame as returned by :func:`~calorine.tools.get_dielectric_function` or 

109 :func:`~calorine.tools.get_ir_spectrum`; 

110 must contain ``angular_frequency`` and at least one column whose name starts 

111 with ``epsilon_imag``. 

112 method 

113 Integration method. ``'vectorized'`` (default) uses an exact trapezoid rule 

114 over an :math:`n \times n` matrix (:math:`\mathcal{O}(n^2)` time and memory); 

115 ``'fft'`` uses an :math:`\mathcal{O}(n \log n)` Hilbert transform and is faster 

116 for large arrays but approximates the principal-value integral. 

117 The two methods agree to within a few percent on typical grids. 

118 

119 Returns 

120 ------- 

121 DataFrame 

122 Input DataFrame with additional ``epsilon_real*`` columns. 

123 """ 

124 if method not in ('fft', 'vectorized'): 

125 raise ValueError(f"method must be 'fft' or 'vectorized', got {method!r}") 

126 

127 kk = _kramers_kronig_fft if method == 'fft' else _kramers_kronig_vectorized 

128 omega = df['angular_frequency'].to_numpy() 

129 imag_cols = [c for c in df.columns if c.startswith('epsilon_imag')] 

130 

131 df = df.copy() 

132 for col in imag_cols: 

133 suffix = col[len('epsilon_imag'):] 

134 df['epsilon_real' + suffix] = kk(omega, df[col].to_numpy()) 

135 return df