[WIP] - Add / use / update compute_spectrogram

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

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))

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]