diff --git a/doc/api.rst b/doc/api.rst index e76e040b..4b7464b5 100644 --- a/doc/api.rst +++ b/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 diff --git a/neurodsp/spectral/__init__.py b/neurodsp/spectral/__init__.py index c63c4e41..74d81b90 100644 --- a/neurodsp/spectral/__init__.py +++ b/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 diff --git a/neurodsp/spectral/power.py b/neurodsp/spectral/power.py index cb7095d2..a2faa0c7 100644 --- a/neurodsp/spectral/power.py +++ b/neurodsp/spectral/power.py @@ -7,7 +7,9 @@ """ 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 @@ -15,7 +17,7 @@ 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 diff --git a/neurodsp/spectral/utils.py b/neurodsp/spectral/utils.py index 4002e4a8..5db1454e 100644 --- a/neurodsp/spectral/utils.py +++ b/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 diff --git a/neurodsp/tests/spectral/test_power.py b/neurodsp/tests/spectral/test_power.py index 637d0ed8..5c8dd961 100644 --- a/neurodsp/tests/spectral/test_power.py +++ b/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 - diff --git a/neurodsp/tests/spectral/test_utils.py b/neurodsp/tests/spectral/test_utils.py index 16da0420..8585dec6 100644 --- a/neurodsp/tests/spectral/test_utils.py +++ b/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