Coverage for src / thunderlab / powerspectrum.py: 92%
161 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-02-14 12:12 +0000
1"""Powerspectra and spectrograms for a given frequency resolution
3## Computation of nfft
5- `next_power_of_two()`: round an integer up to the next power of two.
6- `nfft()`: compute nfft based on a given frequency resolution.
8## Decibel
10- `decibel()`: transform power to decibel.
11- `power()`: transform decibel to power.
13## Power spectra
15- `psd()`: power spectrum for a given frequency resolution.
16- `spectrogram()`: spectrogram of a given frequency resolution and overlap fraction.
18## Power spectrum analysis
20- `peak_freqs()`: peak frequencies computed from power spectra of data snippets.
22## Visualization
24- `plot_decibel_psd()`: plot power spectrum in decibel.
26## Configuration parameter
28- `add_spectrum_config()`: add parameters for psd() and spectrogram() to configuration.
29- `sepctrum_args()`: retrieve parameters for psd() and spectrogram() from configuration.
30"""
32import numpy as np
34from scipy.signal import get_window
35from matplotlib.mlab import psd as mpsd
36try:
37 from scipy.signal import welch as swelch
38 psdscipy = True
39except ImportError:
40 psdscipy = False
41from matplotlib.mlab import specgram as mspecgram
42try:
43 from scipy.signal import spectrogram as sspectrogram
44 specgramscipy = True
45except ImportError:
46 specgramscipy = False
48from .eventdetection import detect_peaks
51def next_power_of_two(n):
52 """The next integer power of two.
54 Parameters
55 ----------
56 n: int
57 A positive number.
59 Returns
60 -------
61 m: int
62 The next integer power of two equal or larger than `n`.
63 """
64 return int(2 ** np.floor(np.log(n) / np.log(2.0) + 1.0 - 1e-8))
67def nfft(rate, freq_resolution, min_nfft=16, max_nfft=None):
68 """Required number of samples for an FFT of a given frequency resolution.
70 Note that the returned number of FFT samples results
71 in frequency intervals that are smaller or equal to `freq_resolution`.
73 Parameters
74 ----------
75 rate: float
76 Sampling rate of the data in Hertz.
77 freq_resolution: float
78 Minimum frequency resolution in Hertz.
79 min_nfft: int
80 Smallest value of nfft to be used.
81 max_nfft: int or None
82 If not None, largest value of nfft to be used.
84 Returns
85 -------
86 nfft: int
87 Number of FFT points.
88 """
89 nfft = next_power_of_two(rate / freq_resolution)
90 if not max_nfft is None:
91 if nfft > max_nfft:
92 nfft = next_power_of_two(max_nfft//2 + 1)
93 if nfft < min_nfft:
94 nfft = min_nfft
95 return nfft
98def decibel(power, ref_power=1.0, min_power=1e-20):
99 """Transform power to decibel relative to ref_power.
101 \\[ decibel = 10 \\cdot \\log_{10}(power/ref\\_power) \\]
102 Power values smaller than `min_power` are set to `-np.inf`.
104 Parameters
105 ----------
106 power: float or array
107 Power values, for example from a power spectrum or spectrogram.
108 ref_power: float or None or 'peak'
109 Reference power for computing decibel.
110 If set to `None` or 'peak', the maximum power is used.
111 min_power: float
112 Power values smaller than `min_power` are set to `-np.inf`.
114 Returns
115 -------
116 decibel_psd: array
117 Power values in decibel relative to `ref_power`.
118 """
119 if np.isscalar(power):
120 tmp_power = np.array([power])
121 decibel_psd = np.array([power])
122 else:
123 tmp_power = power
124 decibel_psd = power.copy()
125 if ref_power is None or ref_power == 'peak':
126 ref_power = np.max(decibel_psd)
127 decibel_psd[tmp_power <= min_power] = float('-inf')
128 decibel_psd[tmp_power > min_power] = 10.0 * np.log10(decibel_psd[tmp_power > min_power]/ref_power)
129 if np.isscalar(power):
130 return decibel_psd[0]
131 else:
132 return decibel_psd
135def power(decibel, ref_power=1.0):
136 """Transform decibel back to power relative to `ref_power`.
138 \\[ power = ref\\_power \\cdot 10^{decibel/10} \\]
140 Parameters
141 ----------
142 decibel: array
143 Decibel values of the power spectrum or spectrogram.
144 ref_power: float
145 Reference power for computing power.
147 Returns
148 -------
149 power: array
150 Power values of the power spectrum or spectrogram.
151 """
152 return ref_power * 10.0 ** (0.1 * decibel)
155def psd(data, ratetime, freq_resolution=1, overlap_frac=0.5,
156 n_fft=None, n_overlap=None, min_nfft=16, max_nfft=None,
157 detrend='constant', window='hann'):
158 """Power spectrum density of a given frequency resolution.
160 If `freq_resolution` is given, the number of samples used for each
161 FFT segment is computed from the requested frequency resolution
162 and the sampling rate. Check the returned frequency array for the
163 actually used frequency resolution. The frequency intervals are
164 smaller or equal to `freq_resolution`. The actually used number
165 of samples used for each FFT segment can be retrieved by dividing
166 the sampling rate by the actual frequency resolution:
167 ```
168 freq, power = psd(data, samplingrate, 0.1)
169 df = np.mean(np.diff(freq)) # the actual frequency resolution
170 nfft = int(samplingrate/df)
171 ```
173 Uses `scipy signal.welch()` if available, otherwise
174 `matplotlib.mlab.psd()`.
176 Parameters
177 ----------
178 data: 1-D array
179 Data from which power spectra are computed.
180 ratetime: float or array
181 If float, sampling rate of the data in Hertz.
182 If array, assume `ratetime` to be the time array
183 corresponding to the data.
184 Compute sampling rate as `1/(ratetime[1]-ratetime[0])`.
185 freq_resolution: float or None
186 Desired frequency resolution of the power spectrum in Hertz.
187 See `nfft()` for details.
188 Alternatively, the number of samples used for computing FFTs
189 can be specified by the `n_fft` argument.
190 overlap_frac: float or None
191 Fraction of overlap between subsequent FFT segments
192 (0: no overlap, 1: complete overlap).
193 Alternatively, the overlap can be specified by the `n_overlap` argument.
194 n_fft: int or None
195 If `freq_resolution` is None, then this is the number of
196 samples used for each FFT segment.
197 n_overlap: int or None
198 If `overlap_frac` is None, then this is the number of
199 samples subsequent FFT segments overlap.
200 min_nfft: int
201 Smallest value of nfft to be used.
202 max_nfft: int or None
203 If not None, largest value of nfft to be used.
204 detrend: string
205 If 'constant' or 'mean' subtract mean of data.
206 If 'linear' subtract line fitted to the data.
207 If 'none' do not detrend the data.
208 window: string
209 Function used for windowing data segements.
210 One of hann, blackman, hamming, bartlett, boxcar, triang, parzen,
211 bohman, blackmanharris, nuttall, fattop, barthann
212 (see scipy.signal window functions).
214 Returns
215 -------
216 freq: 1-D array
217 Frequencies corresponding to `power` array.
218 power: 1-D array
219 Power spectral density in [data]^2/Hz.
221 Raises
222 ------
223 ValueError:
224 Both, `freq_resolution`and `n_fft` are not specified, or
225 both, `overlap_frac`and `n_overlap` are not specified.
226 """
227 rate = ratetime if np.isscalar(ratetime) else 1.0/(ratetime[1]-ratetime[0])
228 if freq_resolution is not None:
229 n_fft = nfft(rate, freq_resolution, min_nfft, max_nfft)
230 if n_fft is None:
231 raise ValueError('freq_resolution or n_fft needs to be specified')
232 if overlap_frac is not None:
233 n_overlap = int(n_fft*overlap_frac)
234 if n_overlap is None:
235 raise ValueError('overlap_frac or n_overlap needs to be specified')
236 if n_fft >= len(data):
237 n_overlap = len(data) - 1
238 if psdscipy:
239 if detrend == 'none':
240 detrend = False
241 elif detrend == 'mean':
242 detrend = 'constant'
243 freqs, power = swelch(data, fs=rate, nperseg=n_fft, nfft=None,
244 noverlap=n_overlap, detrend=detrend,
245 window=window, scaling='density')
246 else:
247 if detrend == 'constant':
248 detrend = 'mean'
249 power, freqs = mpsd(data, Fs=rate, NFFT=n_fft,
250 noverlap=n_overlap, detrend=detrend,
251 window=get_window(window, n_fft),
252 scale_by_freq=True)
253 # squeeze is necessary when n_fft is too large with respect to the data:
254 return freqs, np.squeeze(power)
257def spectrogram(data, ratetime, freq_resolution=0.2, overlap_frac=0.5,
258 n_fft=None, n_overlap=None, min_nfft=16, max_nfft=None,
259 detrend='constant', window='hann'):
260 """Spectrogram of a given frequency resolution.
262 Check the returned frequency array for the actually used frequency
263 resolution.
264 The actual frequency resolution is smaller or equal to `freq_resolution`.
265 The used number of data points per FFT segment (NFFT) is the
266 sampling rate divided by the actual frequency resolution:
268 ```
269 freqs, times, spec = spectrum(data, samplingrate, 0.1) # request 0.1Hz resolution
270 df = np.mean(np.diff(freq)) # the actual frequency resolution
271 df = freqs[1] # same
272 nfft = int(samplingrate/df)
273 ```
275 You can plot the spectrogram directly with `pcolormesh()` with
276 `shading='nearest'` (default):
278 ```
279 fig, ax = plt.subplots()
280 ax.pcolormesh(times, freqs, decibel(spec),
281 vmin=-100, shading='nearest')
282 ```
284 Uses `scipy signal.spectrogram()` if available, otherwise
285 `matplotlib.mlab.specgram()`.
287 Parameters
288 ----------
289 data: 1D or 2D array of floats
290 Data for the spectrograms. First dimension is time,
291 optional second dimension is channel.
292 ratetime: float or array
293 If float, sampling rate of the data in Hertz.
294 If array, assume `ratetime` to be the time array
295 corresponding to the data.
296 The sampling rate is then computed as `1/(ratetime[1]-ratetime[0])`.
297 freq_resolution: float
298 Desired frequency resolution of the spectrogram in Hertz.
299 See `nfft()` for details.
300 Alternatively, the number of samples used for computing FFTs
301 can be specified by the `n_fft` argument.
302 overlap_frac: float
303 Fraction of overlap between subsequent FFT segments
304 (0: no overlap, 1: complete overlap).
305 Alternatively, the overlap can be specified by the `n_overlap` argument.
306 n_fft: int or None
307 If `freq_resolution` is None, then this is the number of
308 samples used for each FFT segment.
309 n_overlap: int or None
310 If `overlap_frac` is None, then this is the number of
311 samples subsequent FFT segments overlap.
312 min_nfft: int
313 Smallest value of nfft to be used. See `nfft()` for details.
314 max_nfft: int or None
315 If not None, largest value of nfft to be used.
316 See `nfft()` for details.
317 detrend: string or False
318 If 'constant' subtract mean of each data segment.
319 If 'linear' subtract line fitted to each data segment.
320 If `False` do not detrend the data segments.
321 window: string
322 Function used for windowing data segements.
323 One of hann, blackman, hamming, bartlett, boxcar, triang, parzen,
324 bohman, blackmanharris, nuttall, fattop, barthann, tukey
325 (see scipy.signal window functions).
327 Returns
328 -------
329 freqs: array
330 Frequencies of the spectrogram.
331 First element is zero. Second element is frequency resolution.
332 time: array
333 Times of the nfft segments, centered on the segments.
334 spectrum: 2D or 3D array
335 Power spectral density for each frequency and time.
336 First dimension is frequency and second dimension is time.
337 Optional last dimension is channel.
339 Raises
340 ------
341 ValueError:
342 Both, `freq_resolution`and `n_fft` are not specified, or
343 both, `overlap_frac`and `n_overlap` are not specified.
344 """
345 rate = ratetime if np.isscalar(ratetime) else 1.0/(ratetime[1]-ratetime[0])
346 if freq_resolution is not None:
347 n_fft = nfft(rate, freq_resolution, min_nfft, max_nfft)
348 if n_fft is None:
349 raise ValueError('freq_resolution or n_fft needs to be specified')
350 if overlap_frac is not None:
351 n_overlap = int(n_fft*overlap_frac)
352 if n_overlap is None:
353 raise ValueError('overlap_frac or n_overlap needs to be specified')
354 if n_fft >= len(data):
355 n_overlap = len(data) - 1
356 if specgramscipy:
357 freqs, time, spec = sspectrogram(data, fs=rate, window=window,
358 nperseg=n_fft, noverlap=n_overlap,
359 detrend=detrend, scaling='density',
360 mode='psd', axis=0)
361 if data.ndim > 1:
362 # scipy spectrogram() returns f x n x t:
363 spec = np.transpose(spec, (0, 2, 1))
364 else:
365 if data.ndim > 1:
366 spec = None
367 for k in range(data.shape[1]):
368 try:
369 ssx, freqs, time = mspecgram(data[:,k], NFFT=n_fft, Fs=rate,
370 noverlap=n_overlap,
371 detrend=detrend,
372 scale_by_freq=True,
373 scale='linear',
374 mode='psd',
375 window=get_window(window, n_fft))
376 except TypeError:
377 ssx, freqs, time = mspecgram(data[:,k], NFFT=n_fft, Fs=rate,
378 noverlap=n_overlap,
379 detrend=detrend,
380 scale_by_freq=True,
381 window=get_window(window, n_fft))
382 if spec is None:
383 spec = np.zeros((len(freqs), len(time), data.shape[1]))
384 spec[:, :, k] = ssx
385 else:
386 try:
387 spec, freqs, time = mspecgram(data, NFFT=n_fft, Fs=rate,
388 noverlap=n_overlap,
389 detrend=detrend,
390 scale_by_freq=True,
391 scale='linear',
392 mode='psd',
393 window=get_window(window, n_fft))
394 except TypeError:
395 spec, freqs, time = mspecgram(data, NFFT=n_fft, Fs=rate,
396 noverlap=n_overlap,
397 detrend=detrend,
398 scale_by_freq=True,
399 window=get_window(window, n_fft))
400 return freqs, time, spec
403def plot_decibel_psd(ax, freqs, power, ref_power=1.0, min_power=1e-20,
404 log_freq=False, min_freq=0.0, max_freq=2000.0,
405 ymarg=0.0, sstyle=dict(color='tab:blue', lw=1)):
406 """Plot the powerspectum in decibel relative to `ref_power`.
408 Parameters
409 ----------
410 ax:
411 Axis for plot.
412 freqs: 1-D array
413 Frequency array of the power spectrum.
414 power: 1-D array
415 Power values of the power spectrum.
416 ref_power: float
417 Reference power for computing decibel. If set to `None` the maximum power is used.
418 min_power: float
419 Power values smaller than `min_power` are set to `np.nan`.
420 log_freq: boolean
421 Logarithmic (True) or linear (False) frequency axis.
422 min_freq: float
423 Limits of frequency axis are set to `(min_freq, max_freq)`
424 if `max_freq` is greater than zero
425 max_freq: float
426 Limits of frequency axis are set to `(min_freq, max_freq)`
427 and limits of power axis are computed from powers below max_freq
428 if `max_freq` is greater than zero
429 ymarg: float
430 Add this to the maximum decibel power for setting the ylim.
431 sstyle: dict
432 Plot parameter that are passed on to the `plot()` function.
433 """
434 decibel_psd = decibel(power, ref_power=ref_power, min_power=min_power)
435 ax.plot(freqs, decibel_psd, **sstyle)
436 ax.set_xlabel('Frequency [Hz]')
437 if max_freq > 0.0:
438 if log_freq and min_freq < 1e-8:
439 min_freq = 1.0
440 ax.set_xlim(min_freq, max_freq)
441 else:
442 max_freq = freqs[-1]
443 if log_freq:
444 ax.set_xscale('log')
445 dpmf = decibel_psd[freqs < max_freq]
446 pmin = np.min(dpmf[np.isfinite(dpmf)])
447 pmin = np.floor(pmin / 10.0) * 10.0
448 pmax = np.max(dpmf[np.isfinite(dpmf)])
449 pmax = np.ceil((pmax + ymarg) / 10.0) * 10.0
450 ax.set_ylim(pmin, pmax)
451 ax.set_ylabel('Power [dB]')
454def peak_freqs(onsets, offsets, data, rate, freq_resolution=0.2,
455 thresh=None, **kwargs):
456 """Peak frequencies computed from power spectra of data snippets.
458 Parameters
459 ----------
460 onsets: array of ints
461 Indices indicating the onsets of the snippets in `data`.
462 offsets: array of ints
463 Indices indicating the offsets of the snippets in `data`.
464 data: 1-D array
465 Data array that contains the data snippets defined by
466 `onsets` and `offsets`.
467 rate: float
468 Sampling rate of data in Hertz.
469 freq_resolution: float
470 Desired frequency resolution of the computed power spectra in Hertz.
471 thresh: None or float
472 If not None than this is the threshold required for the minimum height
473 of the peak in the decibel power spectrum. If the peak is too small
474 than the peak frequency of that snippet is set to NaN.
475 kwargs: dict
476 Further arguments passed on to psd().
478 Returns
479 -------
480 freqs: array of floats
481 For each data snippet the frequency of the maximum power.
482 """
483 freqs = []
484 for i0, i1 in zip(onsets, offsets):
485 if 'max_nfft' in kwargs:
486 del kwargs['max_nfft']
487 f, power = psd(data[i0:i1], rate, freq_resolution,
488 max_nfft=i1 - i0, **kwargs)
489 if thresh is None:
490 fpeak = f[np.argmax(power)]
491 else:
492 p, _ = detect_peaks(decibel(power, None), thresh)
493 if len(p) > 0:
494 ipeak = np.argmax(power[p])
495 fpeak = f[p[ipeak]]
496 else:
497 fpeak = float('NaN')
498 freqs.append(fpeak)
499 return np.array(freqs)
502def add_spectrum_config(cfg, freq_resolution=0.2, overlap_frac=0.5,
503 detrend='constant', window='hann'):
504 """Add all parameters needed for the psd() and spectrogram() functions as a new section to a configuration.
506 Parameters
507 ----------
508 cfg: ConfigFile
509 The configuration.
511 See psd() and spectrogram() for details on the remaining arguments.
512 """
513 cfg.add_section('Power spectra and spectrograms:')
514 cfg.add('frequencyResolution', freq_resolution, 'Hz', 'Frequency resolution of power spectra and spectrograms.')
515 cfg.add('overlapFraction', 100*overlap_frac, '%', 'Overlap of subsequent data segments in power spectra and spectrograms.')
516 cfg.add('detrendMethod', detrend, '', 'Detrend method to be applied on data segments for computing power spectra and spectrograms ("constant", "linear", or "none".')
517 cfg.add('windowFunction', window, '', 'Window function applied on data segements for computing power spectra and spectrograms (one of "hann", "blackman", "hamming", "bartlett", "boxcar", "triang", "parzen", "bohman", "blackmanharris", "nuttall", "fattop", "barthann").')
520def spectrum_args(cfg):
521 """Translates a configuration to the respective parameter names of the psd() and spectrogram() functions.
523 The return value can then be passed as key-word arguments to
524 these functions.
526 Parameters
527 ----------
528 cfg: ConfigFile
529 The configuration.
531 Returns
532 -------
533 a: dict
534 Dictionary with names of arguments of the psd() and spectrogram()
535 functions and their values as supplied by `cfg`.
536 """
537 a = cfg.map({'freq_resolution': 'frequencyResolution',
538 'overlap_frac': 'overlapFraction',
539 'detrend': 'detrendMethod',
540 'window': 'windowFunction'})
541 a['overlap_frac'] *= 0.01
542 return a
545def main():
546 import matplotlib.pyplot as plt
547 print('Compute powerspectra of two sine waves (300 and 450 Hz)')
549 # generate data:
550 fundamentals = [300, 450] # Hz
551 rate = 100_000.0 # Hz
552 time = np.arange(0.0, 8.0, 1.0/rate)
553 data = np.sin(2*np.pi*fundamentals[0]*time) + 0.5*np.sin(2*np.pi*fundamentals[1]*time)
555 # compute power spectrum:
556 freqs, power = psd(data, rate, freq_resolution=0.5,
557 detrend='none', window='hann')
558 df = np.mean(np.diff(freqs))
559 nfft = int(rate/df)
561 # plot power spectrum:
562 fig, (ax1, ax2) = plt.subplots(1, 2, layout='constrained')
563 plot_decibel_psd(ax1, freqs, power,
564 sstyle=dict(lw=2,
565 label=f'$\\Delta f={df:.1f}$Hz, nnft={nfft}'))
566 ax1.legend(loc='upper left')
568 # compute spectrogram:
569 freqs, times, power = spectrogram(data, rate, freq_resolution=2)
570 df = np.mean(np.diff(freqs))
571 dt = 1/df
572 # plot spectrogram:
573 ax2.pcolormesh(times, freqs, decibel(power, min_power=1e-40),
574 vmin=-100, shading='nearest')
575 ax2.text(0.05, 0.98, f'$\\Delta f={df:.2f}$Hz, $\\Delta t={dt:.2f}$s',
576 transform=ax2.transAxes, va='top',
577 bbox=dict(boxstyle='round', facecolor='white'))
578 ax2.set_xlim(time[0], time[-1] + time[1])
579 ax2.set_ylim(0, 600)
580 ax2.set_xlabel('Time [s]')
581 ax2.set_ylabel('Frequency [Hz]')
583 plt.show()
586if __name__ == '__main__':
587 main()