diff --git a/neurodsp/spectral/checks.py b/neurodsp/spectral/checks.py index 4dc0ac1f..3673df52 100644 --- a/neurodsp/spectral/checks.py +++ b/neurodsp/spectral/checks.py @@ -1,10 +1,12 @@ """Checker functions for neurodsp.spectral.""" +from scipy.fft import next_fast_len + ################################################################################################### ################################################################################################### -def check_spg_settings(fs, window, nperseg, noverlap): - """Check settings used for calculating spectrogram. +def check_windowing_settings(fs, window, nperseg, noverlap, fast_len=False): + """Check settings used for defining windows. Parameters ---------- @@ -17,6 +19,9 @@ def check_spg_settings(fs, window, nperseg, noverlap): Length of each segment, in number of samples. noverlap : int or None Number of points to overlap between segments. + fast_len : bool, optional, default: False + If True, updates nperseg to the next fastest length to reduce computation time. + See scipy.fft.next_fast_len for details. Returns ------- @@ -35,8 +40,10 @@ def check_spg_settings(fs, window, nperseg, noverlap): # If the window is an array, defaults to window length else: nperseg = len(window) - else: - nperseg = int(nperseg) + assert not fast_len, "Cannot use fast_len if setting length to a defined window." + + nperseg = int(nperseg) + nperseg = next_fast_len(nperseg) if noverlap is not None: noverlap = int(noverlap) diff --git a/neurodsp/spectral/power.py b/neurodsp/spectral/power.py index 07284ff4..cb7095d2 100644 --- a/neurodsp/spectral/power.py +++ b/neurodsp/spectral/power.py @@ -7,8 +7,7 @@ """ import numpy as np -from scipy.signal import spectrogram, medfilt -from scipy.fft import next_fast_len +from scipy.signal import welch, spectrogram, medfilt from neurodsp.utils.core import get_avg_func from neurodsp.utils.data import create_freqs @@ -16,8 +15,8 @@ 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, window_pad -from neurodsp.spectral.checks import check_spg_settings, check_mt_settings +from neurodsp.spectral.utils import trim_spectrum +from neurodsp.spectral.checks import check_windowing_settings, check_mt_settings ################################################################################################### ################################################################################################### @@ -71,8 +70,8 @@ def compute_spectrum(sig, fs, method='welch', **kwargs): SPECTRUM_INPUTS = { - 'welch' : ['avg_type', 'window', 'nperseg', 'noverlap', 'nfft', \ - 'fast_len', 'f_range', 'outlier_percent'], + 'welch' : ['avg_type', 'window', 'nperseg', 'noverlap', \ + 'nfft', 'fast_len', 'f_range'], 'wavelet' : ['freqs', 'avg_type', 'n_cycles', 'scaling', 'norm'], 'medfilt' : ['filt_len', 'f_range'], } @@ -137,9 +136,8 @@ def compute_spectrum_wavelet(sig, fs, freqs, avg_type='mean', **kwargs): return freqs, spectrum -def compute_spectrum_welch(sig, fs, avg_type='mean', window='hann', - nperseg=None, noverlap=None, nfft=None, - fast_len=False, f_range=None, outlier_percent=None): +def compute_spectrum_welch(sig, fs, avg_type='mean', window='hann', nperseg=None, + noverlap=None, nfft=None, fast_len=False, f_range=None): """Compute the power spectral density using Welch's method. Parameters @@ -167,12 +165,10 @@ def compute_spectrum_welch(sig, fs, avg_type='mean', window='hann', Number of samples per window. Requires nfft > nperseg. Windows are zero-padded by the difference, nfft - nperseg. fast_len : bool, optional, default: False - Moves nperseg to the fastest length to reduce computation. + If True, updates nperseg to the next fastest length to reduce computation time. See scipy.fft.next_fast_len for details. 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 ------- @@ -184,6 +180,7 @@ def compute_spectrum_welch(sig, fs, avg_type='mean', window='hann', Notes ----- - Welch's method ([1]_) computes a power spectra by averaging over windowed FFTs. + - This function uses scipy.signal.welch to compute the Welch's estimate. References ---------- @@ -202,30 +199,12 @@ def compute_spectrum_welch(sig, fs, avg_type='mean', window='hann', >>> 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) + nperseg, noverlap = check_windowing_settings(fs, window, nperseg, noverlap, fast_len) - # Pad signal if requested - if nfft is not None and nfft < nperseg: - raise ValueError('nfft must be greater than nperseg.') - elif nfft is not None: - npad = nfft - nperseg - noverlap = nperseg // 8 if noverlap is None else noverlap - sig, nperseg, noverlap = window_pad(sig, nperseg, noverlap, npad, fast_len) - elif fast_len: - nperseg = next_fast_len(nperseg) + freqs, spectrum = welch(sig, fs, window, nperseg, noverlap, nfft, + detrend=False, return_onesided=True, scaling='density', + average=avg_type) - # Compute spectrogram - 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) diff --git a/neurodsp/spectral/utils.py b/neurodsp/spectral/utils.py index b32614bc..4002e4a8 100644 --- a/neurodsp/spectral/utils.py +++ b/neurodsp/spectral/utils.py @@ -128,119 +128,3 @@ def trim_spectrogram(freqs, times, spg, f_range=None, t_range=None): times_ext = times return freqs_ext, times_ext, spg_ext - - -def window_pad(sig, nperseg, noverlap, npad, fast_len, - nwindows=None, nsamples=None, pad_left=None, pad_right=None): - """Pads windows (for Welch's PSD) with zeros. - - Parameters - ---------- - sig : 1d or 2d array - Time series. - nperseg : int - Length of each segment, in number of samples, at the beginning and end of each window. - noverlap : int - Number of points to overlap between segments, applied prior to zero padding. - npad : int - Number of samples to zero pad windows per side. - fast_len : bool, optional - Moves nperseg to the fastest length to reduce computation. - Adjusts zero-padding to account for the new nperseg. - See scipy.fft.next_fast_len for details. - nwindows, nsamples, pad_left, pad_right : int, optional, default: None - Prevents redundant computation when sig is 2d. - - Returns - ------- - sig_windowed : 1d or 2d array - Windowed signal, with zeros padded at the around each window. - """ - - if sig.ndim == 2: - # Determine the number of samples and padding once, - # to prevent redundant computation in the loop - nwindows = int(np.ceil(len(sig[0])/nperseg)) - if nsamples is None or pad_left is None or pad_right is None: - nsamples, pad_left, pad_right = _find_pad_size( - nperseg, npad, fast_len - ) - - # Recursively call window_pad on each signal - for sind, csig in enumerate(sig): - - _sig_win, _nperseg, _noverlap = window_pad( - # Required arguments - csig, nperseg, noverlap, npad, fast_len, - # Optional arguments to prevent redundant computation - nwindows, nsamples, pad_left, pad_right, - ) - - if sind == 0: - # Initialize windowed array - sig_windowed = np.zeros((len(sig), len(_sig_win))) - - sig_windowed[sind] = _sig_win - - # Update nperseg and noverlap - nperseg, noverlap = _nperseg, _noverlap - - else: - - # Compute the number of windows, samples, and padding. - # Do not recompute if called from the 2d case - if nwindows is None: - nwindows = int(np.ceil(len(sig) / nperseg)) - - if nsamples is None or pad_left is None or pad_right is None: - # Skipped if called from the 2d case - nsamples, pad_left, pad_right = _find_pad_size( - nperseg, npad, fast_len - ) - - # Window signal - sig_windowed = np.zeros((nwindows, nsamples)) - - for wind in range(nwindows): - - # Signal indices - start = max(0, (wind * nperseg) - noverlap) - end = min(len(sig), start + nperseg) - - if end - start != nperseg: - # Stop if a full window can't be created at end of signal - break - - # Pad - sig_windowed[wind] = np.pad(sig[start:end], (pad_left, pad_right)) - - # Removed incomplete windows and flatten - sig_windowed = sig_windowed[:wind].flatten() - - # Update nperseg - nperseg += (pad_left + pad_right) - - # Overlap is zero since overlapping segments was applied prior to padding each window - noverlap = 0 - - return sig_windowed, nperseg, noverlap - - -def _find_pad_size(nperseg, npad, fast_len): - """Determine pad size and number of samples required.""" - - nsamples = nperseg + npad - - pad_left = npad // 2 - pad_right = npad - pad_left - - if fast_len: - # Increase nsamples to the next fastest length and update for zero-padding size - nsamples = next_fast_len(nsamples) - - # New padding - npad = nsamples - nperseg - pad_left = npad // 2 - pad_right = npad - pad_left - - return nsamples, pad_left, pad_right diff --git a/neurodsp/spectral/variance.py b/neurodsp/spectral/variance.py index f9e9071a..68026614 100644 --- a/neurodsp/spectral/variance.py +++ b/neurodsp/spectral/variance.py @@ -7,7 +7,7 @@ from neurodsp.utils.decorators import multidim from neurodsp.utils.checks import check_param_options from neurodsp.spectral.utils import trim_spectrum -from neurodsp.spectral.checks import check_spg_settings +from neurodsp.spectral.checks import check_windowing_settings ################################################################################################### ################################################################################################### @@ -57,7 +57,7 @@ def compute_scv(sig, fs, window='hann', nperseg=None, noverlap=0, outlier_pct=No """ # Compute spectrogram of data - nperseg, noverlap = check_spg_settings(fs, window, nperseg, noverlap) + nperseg, noverlap = check_windowing_settings(fs, window, nperseg, noverlap) freqs, _, spg = spectrogram(sig, fs, window, nperseg, noverlap) if outlier_pct is not None: @@ -134,7 +134,7 @@ def compute_scv_rs(sig, fs, window='hann', nperseg=None, noverlap=0, """ check_param_options(method, 'method', ['bootstrap', 'rolling']) - nperseg, noverlap = check_spg_settings(fs, window, nperseg, noverlap) + nperseg, noverlap = check_windowing_settings(fs, window, nperseg, noverlap) # Compute spectrogram of data freqs, ts, spg = spectrogram(sig, fs, window, nperseg, noverlap) @@ -229,7 +229,7 @@ def compute_spectral_hist(sig, fs, window='hann', nperseg=None, noverlap=None, """ # Compute spectrogram of data - nperseg, noverlap = check_spg_settings(fs, window, nperseg, noverlap) + nperseg, noverlap = check_windowing_settings(fs, window, nperseg, noverlap) freqs, _, spg = spectrogram(sig, fs, window, nperseg, noverlap, return_onesided=True) # Get log10 power & limit to frequency range of interest before binning diff --git a/neurodsp/tests/spectral/test_checks.py b/neurodsp/tests/spectral/test_checks.py index 951a1f44..a8b78c98 100644 --- a/neurodsp/tests/spectral/test_checks.py +++ b/neurodsp/tests/spectral/test_checks.py @@ -9,14 +9,18 @@ ################################################################################################### ################################################################################################### -def test_spg_settings(): +def test_check_windowing_settings(): - nperseg, noverlap = check_spg_settings(FS, 'hann', None, None) + nperseg, noverlap = check_windowing_settings(FS, 'hann', None, None) assert nperseg == FS assert noverlap == None window = np.array([1, 2, 3, 4]) noverlap_in = 20 - nperseg, noverlap_out = check_spg_settings(FS, window, None, noverlap_in) + nperseg, noverlap_out = check_windowing_settings(FS, window, None, noverlap_in) assert nperseg == len(window) assert noverlap_out == noverlap_in + + # Check next fast len + nperseg, noverlap = check_windowing_settings(FS, 'hann', 101, None, fast_len=True) + nperseg == 105 diff --git a/neurodsp/tests/spectral/test_utils.py b/neurodsp/tests/spectral/test_utils.py index 430c90e0..16da0420 100644 --- a/neurodsp/tests/spectral/test_utils.py +++ b/neurodsp/tests/spectral/test_utils.py @@ -42,26 +42,3 @@ def test_trim_spectrogram(): f_ext, t_ext, p_ext = trim_spectrogram(freqs, times, pows, f_range=[6, 8], t_range=None) assert_equal(f_ext, np.array([6, 7, 8])) assert_equal(t_ext, times) - - -@pytest.mark.parametrize("fast_len", [True, False]) -def test_window_pad(fast_len): - - nperseg = 100 - noverlap = 10 - npad = 1000 - - sig = np.random.rand(1000) - - sig_windowed, _nperseg, _noverlap = window_pad(sig, nperseg, noverlap, npad, fast_len) - - # Overlap was handled correctly b/w the first two windows - assert np.all(sig_windowed[npad:npad+nperseg][-noverlap:] == - sig_windowed[(3*npad)+nperseg:(3*npad)+nperseg+noverlap]) - - # Updated nperseg has no remainder - nwin = (len(sig_windowed) / nperseg) - assert nwin == int(nwin) - - # Ensure updated nperseg is correct - assert _nperseg == nperseg + npad