Source code for neurodsp.spectral.power

"""Compute spectral power.

Notes
-----
Mostly relies on scipy.signal.spectrogram and numpy.fft
https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.spectrogram.html
"""

import numpy as np
from scipy.signal import spectrogram, medfilt

from neurodsp.utils.core import get_avg_func
from neurodsp.utils.data import create_freqs
from neurodsp.utils.decorators import multidim
from neurodsp.utils.checks import check_param_options
from neurodsp.utils.outliers import discard_outliers
from neurodsp.timefrequency.wavelets import compute_wavelet_transform
from neurodsp.spectral.utils import trim_spectrum
from neurodsp.spectral.checks import check_spg_settings

###################################################################################################
###################################################################################################

[docs]def compute_spectrum(sig, fs, method='welch', avg_type='mean', **kwargs): """Compute the power spectral density of a time series. Parameters ---------- sig : 1d or 2d array Time series. fs : float Sampling rate, in Hz. method : {'welch', 'wavelet', 'medfilt'}, optional Method to use to estimate the power spectrum. avg_type : {'mean', 'median'}, optional If relevant, the method to average across windows to create the spectrum. **kwargs Keyword arguments to pass through to the function that calculates the spectrum. Returns ------- freqs : 1d array Frequencies at which the measure was calculated. spectrum : 1d or 2d array Power spectral density. Examples -------- Compute the power spectrum of a simulated time series: >>> from neurodsp.sim import sim_combined >>> sig = sim_combined(n_seconds=10, fs=500, ... components={'sim_powerlaw': {}, 'sim_oscillation' : {'freq': 10}}) >>> freqs, spectrum = compute_spectrum(sig, fs=500) """ check_param_options(method, 'method', ['welch', 'wavelet', 'medfilt']) if method == 'welch': return compute_spectrum_welch(sig, fs, avg_type=avg_type, **kwargs) elif method == 'wavelet': return compute_spectrum_wavelet(sig, fs, avg_type=avg_type, **kwargs) elif method == 'medfilt': return compute_spectrum_medfilt(sig, fs, **kwargs)
[docs]@multidim(select=[0]) def compute_spectrum_wavelet(sig, fs, freqs, avg_type='mean', **kwargs): """Compute the power spectral density using wavelets. Parameters ---------- sig : 1d or 2d array Time series. fs : float Sampling rate, in Hz. freqs : 1d array or list of float If array, frequency values to estimate with morlet wavelets. If list, define the frequency range, as [freq_start, freq_stop, freq_step]. The `freq_step` is optional, and defaults to 1. Range is inclusive of `freq_stop` value. avg_type : {'mean', 'median'}, optional Method to average across the windows. **kwargs Optional inputs for using wavelets. Returns ------- freqs : 1d array Frequencies at which the measure was calculated. spectrum : 1d or 2d array Power spectral density. Examples -------- Compute the power spectrum of a simulated time series using wavelets: >>> from neurodsp.sim import sim_combined >>> sig = sim_combined(n_seconds=10, fs=500, ... components={'sim_powerlaw': {}, 'sim_oscillation' : {'freq': 10}}) >>> freqs, spectrum = compute_spectrum_wavelet(sig, fs=500, freqs=[1, 30]) """ if isinstance(freqs, (tuple, list)): freqs = create_freqs(*freqs) # Compute the wavelet transform mwt = compute_wavelet_transform(sig, fs, freqs, **kwargs) # Convert the wavelet coefficient outputs to units of power mwt_power = abs(mwt)**2 # Create the power spectrum by averaging across the time dimension spectrum = get_avg_func(avg_type)(mwt_power, axis=1) return freqs, spectrum
[docs]def compute_spectrum_welch(sig, fs, avg_type='mean', window='hann', nperseg=None, noverlap=None, f_range=None, outlier_percent=None): """Compute the power spectral density using Welch's method. Parameters ---------- sig : 1d or 2d array Time series. fs : float Sampling rate, in Hz. avg_type : {'mean', 'median'}, optional Method to average across the windows: * 'mean' is the same as Welch's method, taking the mean across FFT windows. * 'median' uses median across FFT windows instead of the mean, to minimize outlier effects. window : str or tuple or array_like, optional, default: 'hann' Desired window to use. See scipy.signal.get_window for a list of available windows. If array_like, the array will be used as the window and its length must be nperseg. nperseg : int, optional Length of each segment, in number of samples. If None, and window is str or tuple, is set to 1 second of data. If None, and window is array_like, is set to the length of the window. noverlap : int, optional Number of points to overlap between segments. If None, noverlap = nperseg // 8. f_range : list of [float, float], optional Frequency range to sub-select from the power spectrum. outlier_percent : float, optional The percentage of outlier values to be removed. Must be between 0 and 100. Returns ------- freqs : 1d array Frequencies at which the measure was calculated. spectrum : 1d or 2d array Power spectral density. Notes ----- - Welch's method ([1]_) computes a power spectra by averaging over windowed FFTs. References ---------- .. [1] Welch, P. (1967). The use of fast Fourier transform for the estimation of power spectra: A method based on time averaging over short, modified periodograms. IEEE Transactions on Audio and Electroacoustics, 15(2), 70–73. DOI: https://doi.org/10.1109/TAU.1967.1161901 Examples -------- Compute the power spectrum of a simulated time series using Welch's method: >>> from neurodsp.sim import sim_combined >>> sig = sim_combined(n_seconds=10, fs=500, ... components={'sim_powerlaw': {}, 'sim_oscillation': {'freq': 10}}) >>> freqs, spec = compute_spectrum_welch(sig, fs=500) """ # Calculate the short time Fourier transform with signal.spectrogram nperseg, noverlap = check_spg_settings(fs, window, nperseg, noverlap) freqs, _, spg = spectrogram(sig, fs, window, nperseg, noverlap) # Throw out outliers if indicated if outlier_percent is not None: spg = discard_outliers(spg, outlier_percent) # Average across windows spectrum = get_avg_func(avg_type)(spg, axis=-1) # Trim spectrum, if requested if f_range: freqs, spectrum = trim_spectrum(freqs, spectrum, f_range) return freqs, spectrum
[docs]@multidim(select=[0]) def compute_spectrum_medfilt(sig, fs, filt_len=1., f_range=None): """Compute the power spectral density as a smoothed FFT. Parameters ---------- sig : 1d or 2d array Time series. fs : float Sampling rate, in Hz. filt_len : float, optional, default: 1 Length of the median filter, in Hz. f_range : list of [float, float], optional Frequency range to sub-select from the power spectrum. Returns ------- freqs : 1d array Frequencies at which the measure was calculated. spectrum : 1d or 2d array Power spectral density. Examples -------- Compute the power spectrum of a simulated time series as a smoothed FFT: >>> from neurodsp.sim import sim_combined >>> sig = sim_combined(n_seconds=10, fs=500, ... components={'sim_powerlaw': {}, 'sim_oscillation' : {'freq': 10}}) >>> freqs, spec = compute_spectrum_medfilt(sig, fs=500) """ # Take the positive half of the spectrum, since it's symmetrical ft = np.fft.fft(sig)[:int(np.ceil(len(sig) / 2.))] freqs = np.fft.fftfreq(len(sig), 1. / fs)[:int(np.ceil(len(sig) / 2.))] # Convert median filter length from Hz to samples, and make sure it is odd filt_len_samp = int(filt_len / (freqs[1] - freqs[0])) if filt_len_samp % 2 == 0: filt_len_samp += 1 spectrum = medfilt(np.abs(ft)**2. / (fs * len(sig)), filt_len_samp) if f_range: freqs, spectrum = trim_spectrum(freqs, spectrum, f_range) return freqs, spectrum