[ENH] - Add IRASA

neurodsp/aperiodic/irasa.py

@@ -0,0 +1,123 @@
+"""IRASA method."""
+
+import fractions
+
+import numpy as np
+
+from scipy import signal
+from scipy.optimize import curve_fit
+
+from neurodsp.spectral import compute_spectrum, trim_spectrum
+
+###################################################################################################
+###################################################################################################
+
+def irasa(sig, fs=None, f_range=(1, 30), hset=None, **spectrum_kwargs):
+    """Separate the aperiodic and periodic components using the IRASA method.
+
+    Parameters
+    ----------
+    sig : 1d array
+        Time series.
+    fs : float
+        The sampling frequency of sig.
+    f_range : tuple or None
+        Frequency range.
+    hset : 1d array
+        Resampling factors used in IRASA calculation.
+        If not provided, defautls to values from 1.1 to 1.9 with an increment of 0.05.
+    spectrum_kwargs : dict
+        Optional keywords arguments that are passed to `compute_spectrum`.
+
+    Returns
+    -------
+    freqs : 1d array
+        Frequency vector.
+    psd_aperiodic : 1d array
+        The aperiodic component of the power spectrum.
+    psd_periodic : 1d array
+        The periodic component of the power spectrum.
+
+    Notes
+    -----
+    Irregular-Resampling Auto-Spectral Analysis (IRASA) is described in Wen & Liu (2016).
+    Briefly, it aims to separate 1/f and periodic components by resampling time series, and
+    computing power spectra, effectively averaging away any activity that is frequency specific.
+
+    References
+    ----------
+    Wen, H., & Liu, Z. (2016). Separating Fractal and Oscillatory Components in the Power Spectrum
+    of Neurophysiological Signal. Brain Topography, 29(1), 13–26. DOI: 10.1007/s10548-015-0448-0
+    """
+
+    # Check & get hset. The rounding avoids floating precision errors
+    hset = np.arange(1.1, 1.95, 0.05) if not hset else hset
+    hset = np.round(hset, 4)
+
+    # `nperseg` needs to be set to lock in the size of the FFT's
+    if 'nperseg' not in spectrum_kwargs:
+        spectrum_kwargs['nperseg'] = int(4 * fs)
+
+    # Calculate the original spectrum across the whole signal
+    freqs, psd = compute_spectrum(sig, fs, **spectrum_kwargs)
+
+    # Do the IRASA resampling procedure
+    psds = np.zeros((len(hset), *psd.shape))
+
+    for ind, h in enumerate(hset):
+
+        # Get the upsampling/downsampling (h, 1/h) factors as integer
+        rat = fractions.Fraction(str(h))
+        up, dn = rat.numerator, rat.denominator
+
+        # Resample signal
+        sig_up = signal.resample_poly(sig, up, dn, axis=-1)
+        sig_dn = signal.resample_poly(sig, dn, up, axis=-1)
+
+        # Calculate the PSD using same params as original
+        freqs_up, psd_up = compute_spectrum(sig_up, h * fs, **spectrum_kwargs)
+        freqs_dn, psd_dn = compute_spectrum(sig_dn, fs / h, **spectrum_kwargs)
+
+        # Geometric mean of h and 1/h
+        psds[ind, :] = np.sqrt(psd_up * psd_dn)
+
+    # Now we take the median resampled spectra, as an estimate of the aperiodic component
+    psd_aperiodic = np.median(psds, axis=0)
+
+    # Subtract aperiodic from original, to get the periodic component
+    psd_periodic = psd - psd_aperiodic
+
+    # Restrict spectrum to requested range
+    if f_range:
+        psds = np.array([psd_aperiodic, psd_periodic])
+        freqs, (psd_aperiodic, psd_periodic) = trim_spectrum(freqs, psds, f_range)
+
+    return freqs, psd_aperiodic, psd_periodic
+
+
+def fit_irasa(freqs, psd_aperiodic):
+    """Fit the IRASA derived aperiodic component of the spectrum.
+
+    Parameters
+    ----------
+    freqs : 1d array
+        Frequency vector, in linear space.
+    psd_aperidic : 1d array
+        Power values, in linear space.
+
+    Returns
+    -------
+    intercept : float
+        Fit intercept value.
+    slope : float
+        Fit slope value.
+    """
+
+    popt, _ = curve_fit(fit_func, np.log(freqs), np.log(psd_aperiodic))
+    intercept, slope = popt
+
+    return intercept, slope
+
+
+def fit_func(freqs, intercept, slope):
+    return slope * freqs + intercept

neurodsp/tests/aperiodic/test_irasa.py

@@ -0,0 +1,59 @@
+"""Tests for IRASA functions."""
+
+import numpy as np
+
+from neurodsp.tests.settings import FS, N_SECONDS_LONG
+
+from neurodsp.sim import sim_combined
+from neurodsp.spectral import compute_spectrum, trim_spectrum
+from neurodsp.aperiodic.irasa import irasa, fit_irasa, fit_func
+
+###################################################################################################
+###################################################################################################
+
+
+def test_irasa():
+
+    # Simulate a signal
+    sim_components = {'sim_powerlaw': {'exponent' : -2},
+                      'sim_oscillation': {'freq' : 10}}
+    sig = sim_combined(n_seconds=N_SECONDS_LONG, fs=FS, components=sim_components)
+    _, powers = trim_spectrum(*compute_spectrum(sig, FS, nperseg=int(4*FS)), [1, 30])
+
+    # Estimate periodic and aperiodic components with IRASA
+    freqs, psd_ap, psd_pe = irasa(sig, FS, noverlap=int(2*FS))
+
+    # Compute r-squared for the full model
+    r_sq = np.corrcoef(np.array([powers, psd_ap+psd_pe]))[0][1]
+
+    assert r_sq > .95
+    assert len(freqs) == len(psd_ap) == len(psd_pe)
+
+
+def test_fit_irasa():
+
+    knee = -1
+
+    # Simulate a signal
+    sim_components = {'sim_powerlaw': {'exponent' : knee},
+                      'sim_oscillation': {'freq' : 10}}
+    sig = sim_combined(n_seconds=N_SECONDS_LONG, fs=FS, components=sim_components)
+
+    # Estimate periodic and aperiodic components with IRASA
+    freqs, psd_ap, _ = irasa(sig, FS, noverlap=int(2*FS))
+
+    # Get aperiodic coefficients
+    b0, b1 = fit_irasa(freqs, psd_ap)
+
+    assert round(b1) == knee
+    assert np.abs(b0 - np.log((psd_ap)[0])) < 1
+
+
+def test_fit_func():
+
+    freqs = np.arange(30)
+    intercept = -2
+    slope = -2
+
+    fit = fit_func(freqs, intercept, slope)
+    assert (fit == slope * freqs + intercept).all()