"""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