Merge pull request #354 from neurodsp-tools/fft

[ENH] - Extend spectral power with more FFT options

Author: TomDonoghue

doc/api.rst

@@ -111,6 +111,7 @@ Spectral Power
     :toctree: generated/
 
     compute_spectrum
+    compute_spectrum_fft
     compute_spectrum_welch
     compute_spectrum_wavelet
     compute_spectrum_medfilt

neurodsp/spectral/__init__.py

@@ -1,7 +1,8 @@
 """Spectral module, for calculating power spectra, spectral variance, etc."""
 
-from .power import (compute_spectrum, compute_spectrum_welch, compute_spectrum_wavelet,
-                    compute_spectrum_medfilt, compute_spectrum_multitaper)
+from .power import (compute_spectrum, compute_spectrum_fft, compute_spectrum_welch,
+                    compute_spectrum_wavelet, compute_spectrum_medfilt,
+                    compute_spectrum_multitaper)
 from .measures import compute_absolute_power, compute_relative_power, compute_band_ratio
 from .variance import compute_scv, compute_scv_rs, compute_spectral_hist
 from .utils import trim_spectrum, trim_spectrogram

neurodsp/spectral/power.py

@@ -7,15 +7,17 @@
 """
 
 import numpy as np
+from scipy.signal import spectrogram, medfilt
 from scipy.signal import welch, spectrogram, medfilt
+from scipy.signal.windows import get_window
 
 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.utils import trim_spectrum, pad_signal
 from neurodsp.spectral.checks import check_windowing_settings, check_mt_settings
 
 ###################################################################################################
@@ -70,9 +72,10 @@ def compute_spectrum(sig, fs, method='welch', **kwargs):
 
 
 SPECTRUM_INPUTS = {
+    'wavelet' : ['freqs', 'avg_type', 'n_cycles', 'scaling', 'norm'],
+    'fft' : ['window', 'f_range'],
     'welch' : ['avg_type', 'window', 'nperseg', 'noverlap', \
                'nfft', 'fast_len', 'f_range'],
-    'wavelet' : ['freqs', 'avg_type', 'n_cycles', 'scaling', 'norm'],
     'medfilt' : ['filt_len', 'f_range'],
 }
 
@@ -136,6 +139,50 @@ def compute_spectrum_wavelet(sig, fs, freqs, avg_type='mean', **kwargs):
     return freqs, spectrum
 
 
+@multidim(select=[0])
+def compute_spectrum_fft(sig, fs, window=None, nfft=None, f_range=None):
+    """Compute the power spectrum based on a single FFT.
+
+    Parameters
+    ----------
+    sig : array
+        Time series.
+    fs : float
+        Sampling rate, in Hz.
+    window : str or tuple or float, optional
+        Window function to apply to signal.
+        Typically, this is a string of the name of the window to use (e.g. 'hann' or 'hamming').
+        See `scipy.signal.windows.get_window` for details.
+    nfft : int, optional
+        Number of samples per for the FFT estimation.
+        If provided and nfft > len(sig), then the signal is zero-padded to this length.
+    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 : array
+        Power spectral density.
+    """
+
+    if window is not None:
+        sig = sig * get_window(window, len(sig))
+
+    if nfft is not None:
+        sig = pad_signal(sig, nfft)
+
+    # Compute the FFT and convert to power & compute corresponding frequency vector
+    spectrum = np.abs(np.fft.rfft(sig)) ** 2.
+    freqs = np.fft.rfftfreq(len(sig), 1. / fs)
+
+    if f_range:
+        freqs, spectrum = trim_spectrum(freqs, spectrum, f_range)
+
+    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):
     """Compute the power spectral density using Welch's method.
@@ -243,16 +290,15 @@ def compute_spectrum_medfilt(sig, fs, filt_len=1., f_range=None):
     >>> 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.))]
+    # Compute spectrum estimate as a single FFT
+    freqs, spectrum = compute_spectrum_fft(sig, fs)
 
     # 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)
+    spectrum = medfilt(spectrum / (fs * len(sig)), filt_len_samp)
 
     if f_range:
         freqs, spectrum = trim_spectrum(freqs, spectrum, f_range)
@@ -319,7 +365,7 @@ def compute_spectrum_multitaper(sig, fs, bandwidth=None, n_tapers=None,
                              "Could not compute spectrum with low_bias=True.")
 
     # Compute Fourier transform on signal weighted by each slepian sequence
-    freqs = np.fft.rfftfreq(sig_len, 1. /fs)
+    freqs = np.fft.rfftfreq(sig_len, 1. / fs)
     spectra = np.abs(np.fft.rfft(slepian_sequences[:, np.newaxis] * sig)) ** 2
 
     # combine estimates to compute final spectrum

neurodsp/spectral/utils.py

@@ -128,3 +128,38 @@ def trim_spectrogram(freqs, times, spg, f_range=None, t_range=None):
         times_ext = times
 
     return freqs_ext, times_ext, spg_ext
+
+
+def pad_signal(sig, length, fast_len=False):
+    """Pad a signal to a desired length.
+
+    Parameters
+    ----------
+    sig : 1d array
+        Signal to pad.
+    length : int
+        Output length to pad the signal to.
+    fast_len : bool, optional, default: False
+        If True, updates length to the next fastest length to reduce computation time.
+        See scipy.fft.next_fast_len for details.
+
+    Returns
+    -------
+    sig : 1d array
+        Padded signal.
+
+    Notes
+    -----
+    This approach pads the signal evenly on the left and right side with 0s.
+    If the padding length ends up being odd, this approach will split to padding to
+    have one less at the front / left side pad, and one more on the right / end side pad.
+    """
+
+    if length > len(sig):
+        if fast_len:
+            length = next_fast_len(length)
+        npad_total = length - len(sig)
+        npad_left, npad_right = int(np.floor(npad_total / 2)), int(np.ceil(npad_total / 2))
+        sig = np.pad(sig, (npad_left, npad_right), mode='constant', constant_values=0)
+
+    return sig

neurodsp/tests/spectral/test_power.py

@@ -55,6 +55,20 @@ def test_compute_spectrum_2d(tsig2d):
     assert freqs.shape[-1] == spectrum.shape[-1]
     assert spectrum.ndim == 2
 
+def test_compute_spectrum_fft(tsig, tsig_sine):
+
+    freqs1, spectrum1 = compute_spectrum_fft(tsig, FS)
+    assert freqs1.shape == spectrum1.shape
+
+    # Test applying a window function
+    freqs2, spectrum2 = compute_spectrum_fft(tsig, FS, window='hann')
+    assert freqs2.shape == spectrum2.shape
+
+    # Test padding signal
+    freqs3, spectrum3 = compute_spectrum_fft(tsig, FS, nfft=1.5*len(tsig))
+    assert freqs3.shape == spectrum3.shape
+    assert freqs2.shape != freqs3.shape
+
 def test_compute_spectrum_welch(tsig, tsig_sine):
 
     freqs, spectrum = compute_spectrum_welch(tsig, FS, avg_type='mean')
@@ -66,8 +80,7 @@ def test_compute_spectrum_welch(tsig, tsig_sine):
     # Use a rectangular window with a width of one period/cycle and no overlap
     #   The spectrum should just be a dirac spike at the first frequency
     window = np.ones(FS)
-    _, psd_welch = compute_spectrum(tsig_sine, FS, method='welch',
-                                    nperseg=FS, noverlap=0, window=window)
+    _, psd_welch = compute_spectrum_welch(tsig_sine, FS, nperseg=FS, noverlap=0, window=window)
 
     # Spike at frequency 1
     assert np.abs(psd_welch[FREQ_SINE] - 0.5) < EPS
@@ -81,7 +94,7 @@ def test_compute_spectrum_welch(tsig, tsig_sine):
     assert np.allclose(psd_welch[0:FREQ_SINE], expected_answer, atol=EPS)
 
     # Test zero padding
-    freqs, spectrum = compute_spectrum(
+    freqs, spectrum = compute_spectrum_welch(
         np.tile(tsig, (2, 1)), FS, nperseg=100, noverlap=0, nfft=1000, f_range=(1, 200)
     )
     assert np.all(spectrum[0] == spectrum[1])
@@ -99,16 +112,13 @@ def test_compute_spectrum_medfilt(tsig, tsig_sine):
     freqs, spectrum = compute_spectrum_medfilt(tsig, FS)
     assert freqs.shape == spectrum.shape
 
-    # Compute raw estimate of psd using fourier transform
-    #   Only look at the spectrum up to the Nyquist frequency
+    # Compute raw estimate of psd using FFT
     sig_len = len(tsig_sine)
-    nyq_freq = sig_len//2
-    sig_ft = np.fft.fft(tsig_sine)[:nyq_freq]
-    psd = np.abs(sig_ft)**2/(FS * sig_len)
+    psd = np.abs(np.fft.rfft(tsig_sine))**2 / (FS * sig_len)
 
     # The medfilt here should be taking the median of a window with one sample
     #   Therefore, it should match the estimate of psd from above
-    _, psd_medfilt = compute_spectrum(tsig_sine, FS, method='medfilt', filt_len=0.1)
+    _, psd_medfilt = compute_spectrum_medfilt(tsig_sine, FS, filt_len=0.1)
     assert np.allclose(psd, psd_medfilt, atol=EPS)
 
 def test_compute_spectrum_multitaper(tsig_sine, tsig2d):
@@ -126,4 +136,3 @@ def test_compute_spectrum_multitaper(tsig_sine, tsig2d):
     idx_freq_sine = np.argmin(np.abs(freqs - FREQ_SINE))
     idx_peak = np.argmax(spectrum)
     assert idx_freq_sine == idx_peak
-

neurodsp/tests/spectral/test_utils.py

@@ -42,3 +42,25 @@ 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)
+
+def test_pad_signal():
+
+    # Test case: odd length, even number added per side
+    length = 5
+    out1 = pad_signal(np.array([1, 2, 3]), length)
+    assert len(out1) == length
+
+    # Test case: even length, even number added per side
+    length = 6
+    out2 = pad_signal(np.array([1, 2]), length)
+    assert len(out2) == length
+
+    # Test case: odd length, uneven number added per side
+    length = 5
+    out3 = pad_signal(np.array([1, 2]), length)
+    assert len(out3) == length
+
+    # Test case: even length, uneven number added per side
+    length = 6
+    out4 = pad_signal(np.array([1, 2, 3]), length)
+    assert len(out4) == length