[MNT] - Update Welch implementation

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)

neurodsp/spectral/power.py

@@ -7,17 +7,16 @@
 """
 
 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
 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, 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)
 

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

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

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

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