diff --git a/neurodsp/spectral/power.py b/neurodsp/spectral/power.py index d763b877..e8de1fa8 100644 --- a/neurodsp/spectral/power.py +++ b/neurodsp/spectral/power.py @@ -381,3 +381,43 @@ def compute_spectrum_multitaper(sig, fs, bandwidth=None, n_tapers=None, spectrum = spectrum[0] return freqs, spectrum + + +def compute_spectrogram(sig, fs, window='hann', nperseg=None, noverlap=0, nfft=None): + """Compute a spectrogram. + + Parameters + ---------- + sig : array + Time series. + fs : float + Sampling rate, in Hz. + 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. + nfft : int, optional + Number of samples per window. Requires nfft > nperseg. + Windows are zero-padded by the difference, nfft - nperseg. + + Returns + ------- + freqs : 1d array + Frequencies at which the spectrogram was calculated. + times : 1d array + Time values at which the spectrogram was calculated. + spg : 2d array + Power values of the spectrogram. + """ + + nperseg, noverlap = check_windowing_settings(fs, window, nperseg, noverlap) + + freqs, times, spg = spectrogram(sig, fs, window, nperseg, noverlap, nfft) + + return freqs, times, spg diff --git a/neurodsp/spectral/variance.py b/neurodsp/spectral/variance.py index 68026614..2076d6ab 100644 --- a/neurodsp/spectral/variance.py +++ b/neurodsp/spectral/variance.py @@ -1,11 +1,11 @@ """Compute spectral variance.""" import numpy as np -from scipy.signal import spectrogram from neurodsp.utils import discard_outliers from neurodsp.utils.decorators import multidim from neurodsp.utils.checks import check_param_options +from neurodsp.spectral.power import compute_spectrogram from neurodsp.spectral.utils import trim_spectrum from neurodsp.spectral.checks import check_windowing_settings @@ -56,9 +56,7 @@ def compute_scv(sig, fs, window='hann', nperseg=None, noverlap=0, outlier_pct=No >>> freqs, scv = compute_scv(sig, fs=500) """ - # Compute spectrogram of data - nperseg, noverlap = check_windowing_settings(fs, window, nperseg, noverlap) - freqs, _, spg = spectrogram(sig, fs, window, nperseg, noverlap) + freqs, _, spg = compute_spectrogram(sig, fs, window, nperseg, noverlap) if outlier_pct is not None: spg = discard_outliers(spg, outlier_pct) @@ -134,10 +132,8 @@ def compute_scv_rs(sig, fs, window='hann', nperseg=None, noverlap=0, """ check_param_options(method, 'method', ['bootstrap', 'rolling']) - nperseg, noverlap = check_windowing_settings(fs, window, nperseg, noverlap) - # Compute spectrogram of data - freqs, ts, spg = spectrogram(sig, fs, window, nperseg, noverlap) + freqs, ts, spg = compute_spectrogram(sig, fs, window, nperseg, noverlap) if method == 'bootstrap': @@ -228,9 +224,7 @@ def compute_spectral_hist(sig, fs, window='hann', nperseg=None, noverlap=None, >>> freqs, power_bins, spectral_hist = compute_spectral_hist(sig, fs=500) """ - # Compute spectrogram of data - nperseg, noverlap = check_windowing_settings(fs, window, nperseg, noverlap) - freqs, _, spg = spectrogram(sig, fs, window, nperseg, noverlap, return_onesided=True) + freqs, _, spg = compute_spectrogram(sig, fs, window, nperseg, noverlap) # Get log10 power & limit to frequency range of interest before binning ps = np.transpose(np.log10(spg)) diff --git a/neurodsp/tests/spectral/test_power.py b/neurodsp/tests/spectral/test_power.py index 5c8dd961..c76caee1 100644 --- a/neurodsp/tests/spectral/test_power.py +++ b/neurodsp/tests/spectral/test_power.py @@ -136,3 +136,9 @@ 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 + +def test_compute_spectrogram(tsig): + + freqs, times, spg = compute_spectrogram(tsig, FS) + assert freqs.shape[0] == spg.shape[0] + assert times.shape[0] == spg.shape[1]