Coverage for dynasor/tools/acfs.py: 98%
80 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-10 08:27 +0000
1"""
2A number of utility functions, for example for dealing with
3autocorrelation functions, Fourier transforms, and smoothing.
4"""
6from typing import Optional
7import numpy as np
8from scipy.signal import correlate
9from numpy.typing import NDArray
10import pandas as pd
13def psd_from_acf(
14 acf: NDArray[float],
15 dt: Optional[float] = 1,
16 even: Optional[bool] = True,
17) -> tuple[NDArray[float], NDArray[float]]:
18 """Computes the power spectral density (PSD) from an auto-correlation function (ACF).
20 Let x(t) be a time signal and define its auto-correlation function
22 C(τ) = ⟨ x(t) x(t+τ) ⟩
24 where ⟨·⟩ denotes a time or ensemble average. According to the
25 Wiener–Khinchin theorem, the power spectral density (PSD) of x(t) is
26 the Fourier transform of its auto-correlation function:
28 S(ω) = ∫ C(τ) e^{-i ω τ} dτ
30 This function computes the PSD by performing a discrete Fourier
31 transform of the provided ACF.
33 Parameters
34 ----------
35 acf
36 The auto-correlation function C(τ) evaluated at discrete time lags.
37 dt
38 Time spacing between consecutive samples in the ACF.
39 even
40 Whether the ACF is assumed to be even in time. If True, the ACF
41 is mirrored to construct a two-sided correlation function before
42 computing the Fourier transform.
44 Returns
45 -------
46 freqs, psd
47 Angular frequencies ω and corresponding PSD values.
48 """
49 if dt <= 0:
50 raise ValueError('dt must be positive.')
51 assert acf.ndim == 1
52 if np.iscomplexobj(acf):
53 raise ValueError('acf must be real.')
55 signal = np.asarray(acf)
57 if even:
58 signal = np.hstack((signal, signal[:0:-1]))
60 fft = dt * np.fft.rfft(signal)
61 psd = fft.real if even else fft
63 freqs = np.fft.rfftfreq(len(signal), dt)
64 freqs = 2 * np.pi * freqs
66 return freqs, psd
69def psd_from_acf_2d(
70 acf: NDArray[float],
71 dt: Optional[float] = 1,
72 even: Optional[bool] = True,
73) -> tuple[NDArray[float], NDArray[float]]:
74 """Computes the power spectral density (PSD) from an auto-correlation function (ACF).
76 This is the 2D analogue of `psd_from_acf`, where the ACF is provided for
77 multiple signals, e.g. F(q, t). The Fourier transform is performed along
78 the last axis.
80 Parameters
81 ----------
82 acf
83 The ACF as a 2D array with shape (N, Nt), where the last axis corresponds
84 to time lags.
85 dt
86 The time step between samples.
87 even
88 Whether the ACF is even in time and should be mirrored before computing
89 the Fourier transform.
91 Returns
92 -------
93 freqs, psd
94 Angular frequencies ω and corresponding PSD values with shape (N, Nω).
95 """
96 if dt <= 0:
97 raise ValueError('dt must be positive.')
98 assert acf.ndim == 2
99 if np.iscomplexobj(acf):
100 raise ValueError('acf must be real.')
102 signal = np.asarray(acf)
104 if even:
105 signal = np.hstack((signal, signal[:, :0:-1]))
107 fft = dt * np.fft.rfft(signal, axis=1)
108 psd = fft.real if even else fft
110 freqs = np.fft.rfftfreq(signal.shape[1], dt)
111 freqs = 2 * np.pi * freqs
113 return freqs, psd
116def psd_from_time_signal(
117 x: NDArray[float],
118 dt: Optional[float] = 1,
119 complex: Optional[bool] = False,
120) -> tuple[NDArray[float], NDArray[float]]:
121 """Computes the power spectral density (PSD) directly from a time signal.
123 Let x(t) be a time signal. The power spectral density (PSD) can be
124 computed directly from the Fourier transform of the signal:
126 ``S(ω) = |F[x(t)]|²``
128 where F denotes the Fourier transform.
130 This function computes the PSD by performing a discrete Fourier
131 transform of the time signal and taking the squared magnitude of
132 the transform.
134 This is equivalent to computing the PSD from the Fourier transform of the
135 auto-correlation function.
137 Parameters
138 ----------
139 x
140 Time signal as a 1D array.
141 dt
142 Time spacing between consecutive samples.
143 complex
144 Whether the time signal is complex. If False (default), the signal is
145 assumed to be real and the PSD is computed using a real FFT (`rfft`),
146 returning only non-negative frequencies. If True, the full FFT (`fft`)
147 is used and both positive and negative frequencies are returned.
149 Returns
150 -------
151 freqs, psd
152 Angular frequencies ω and corresponding PSD values.
153 """
154 if dt <= 0:
155 raise ValueError('dt must be positive.')
156 assert x.ndim == 1
158 signal = np.asarray(x)
159 N = len(signal)
161 if complex:
162 fft = np.fft.fft(signal)
163 freqs = np.fft.fftfreq(N, dt)
164 else:
165 fft = np.fft.rfft(signal)
166 freqs = np.fft.rfftfreq(N, dt)
168 psd = dt * np.abs(fft) ** 2 / N
169 freqs = 2 * np.pi * freqs
171 return freqs, psd
174def compute_acf(
175 Z: NDArray[float],
176 delta_t: Optional[float] = 1.0,
177 method: Optional[str] = 'scipy',
178) -> tuple[NDArray[float], NDArray[float]]:
179 r"""
180 Computes the autocorrelation function (ACF) for a one-dimensional signal :math:`Z` in time as
182 .. math::
184 \text{ACF}(\tau) = \frac{\left < Z(t) Z^*(t+\tau) \right >}{\left < Z(t) Z^*(t) \right >}
186 Here, only the real part of the ACF is returned since if :math:`Z` is complex
187 the imaginary part should average out to zero for any stationary signal.
189 Parameters
190 ----------
191 Z
192 Complex time signal.
193 delta_t
194 Spacing in time between two consecutive values in :math:`Z`.
195 method
196 Implementation to use; possible values: `numpy` and `scipy` (default and usually faster).
198 Returns
199 -------
200 time_lags, acf
201 Time lags τ and the corresponding normalized ACF values.
202 """
204 if delta_t <= 0:
205 raise ValueError('delta_t must be positive.')
207 # keep only real part and normalize
208 acf = _compute_correlation_function(Z, Z, method)
209 acf = np.real(acf)
210 if acf[0] == 0:
211 raise ValueError('Cannot normalize ACF; zero-lag autocorrelation (acf[0]) is zero')
213 acf /= acf[0]
215 time_lags = delta_t * np.arange(0, len(acf), 1)
216 return time_lags, acf
219def _compute_correlation_function(Z1, Z2, method: Optional[str] = 'scipy'):
220 N = len(Z1)
221 assert len(Z1) == len(Z2)
222 if method == 'scipy':
223 cf = correlate(Z1, Z2, mode='full')[N - 1:] / np.arange(N, 0, -1)
224 elif method == 'numpy':
225 cf = np.correlate(Z1, Z2, mode='full')[N - 1:] / np.arange(N, 0, -1)
226 else:
227 raise ValueError('method must be either numpy or scipy')
228 return cf
231# smoothing functions / FFT filters
232# -------------------------------------
233def gaussian_decay(t: NDArray[float], t_sigma: float) -> NDArray[float]:
234 r"""
235 Evaluates a gaussian distribution in time :math:`f(t)`, which can be applied to an ACF in time
236 to artificially damp it, i.e., forcing it to go to zero for long times.
238 .. math::
240 f(t) = \exp{\left [-\frac{1}{2} \left (\frac{t}{t_\mathrm{sigma}}\right )^2 \right ] }
242 Parameters
243 ----------
244 t
245 Time array.
246 t_sigma
247 Width (standard deviation of the gaussian) of the decay.
248 """
250 if t_sigma <= 0: 250 ↛ 252line 250 didn't jump to line 252 because the condition on line 250 was always true
251 raise ValueError('t_sigma must be positive.')
252 return np.exp(- 1 / 2 * (t / t_sigma) ** 2)
255def fermi_dirac(t: NDArray[float], t_0: float, t_width: float) -> NDArray[float]:
256 r"""
257 Evaluates a Fermi-Dirac-like function in time :math:`f(t)`, which can be applied to an
258 auto-correlation function (ACF) in time to artificially dampen it, i.e., forcing it to
259 go to zero for long times without affecting the short-time correlations too much.
261 .. math::
263 f(t) = \frac{1}{\exp{[(t-t_0)/t_\mathrm{width}}] + 1}
265 Parameters
266 ----------
267 t
268 Time array.
269 t_0
270 Starting time for decay.
271 t_width
272 Width of the decay.
274 """
275 if t_width <= 0:
276 raise ValueError('t_width must be positive.')
277 return 1.0 / (np.exp((t - t_0) / t_width) + 1)
280def smoothing_function(
281 data: NDArray[float],
282 window_size: int,
283 window_type: Optional[str] = 'hamming',
284) -> NDArray[float]:
285 """
286 Smoothing function for 1D arrays.
287 This functions employs the pandas rolling window average function.
289 Parameters
290 ----------
291 data
292 1D data array.
293 window_size
294 The size of smoothing/smearing window.
295 window_type
296 What type of window-shape to use, e.g. ``'blackman'``, ``'hamming'``, ``'boxcar'``
297 (see pandas and SciPy documentation for more details).
299 """
300 if window_size <= 0:
301 raise ValueError('window_size must be a positive integer.')
302 series = pd.Series(data)
303 new_data = series.rolling(window_size, win_type=window_type, center=True, min_periods=1).mean()
304 return np.array(new_data)