[ENH] - Add autocorrelation related functionality

doc/api.rst

@@ -207,6 +207,9 @@ Auto-correlation Analyses
     :toctree: generated/
 
     compute_autocorr
+    compute_decay_time
+    fit_autocorr
+    compute_ac_fit
 
 Fluctuation Analyses
 ~~~~~~~~~~~~~~~~~~~~

neurodsp/aperiodic/autocorr.py

@@ -1,6 +1,8 @@
-"""Autocorrelation analyses of time series."""
+"""Autocorrelation related analyses of time series."""
 
 import numpy as np
+from scipy.optimize import curve_fit
+
 from neurodsp.utils.decorators import multidim
 
 ###################################################################################################
@@ -26,7 +28,7 @@
     timepoints : 1d array
         Time points, in samples, at which autocorrelations are computed.
     autocorrs : array
-        Autocorrelation values, for across time lags.
+        Autocorrelation values, across time lags.
 
     Examples
     --------
@@ -47,3 +49,161 @@
     timepoints = np.arange(0, max_lag+1, lag_step)
 
     return timepoints, autocorrs
+
+
+def compute_decay_time(timepoints, autocorrs, fs, level=0):
+    """Compute autocorrelation decay time, from precomputed autocorrelation.
+
+    Parameters
+    ----------
+    timepoints : 1d array
+        Timepoints for the computed autocorrelations.
+    autocorrs : 1d array
+        Autocorrelation values.
+    fs : int
+        Sampling rate of the signal.
+    level : float
+        Autocorrelation decay threshold.
+
+    Returns
+    -------
+    result : float
+        Autocorrelation decay time.
+        If decay time value not found, returns nan.
+
+    Notes
+    -----
+    The autocorrelation decay time is computed as the time delay for the
+    autocorrelation to drop to (or below) the decay time threshold.
+    """
+
+    val_checks = autocorrs <= level
+
+    if np.any(val_checks):
+        # Get the first value to cross the threshold, and convert to time value
+        result = timepoints[np.argmax(val_checks)] / fs
+    else:
+        result = np.nan
+
+    return result
+
+
+def fit_autocorr(timepoints, autocorrs, fit_function='single_exp', bounds=None):
+    """Fit autocorrelation function, returning timescale estimate.
+
+    Parameters
+    ----------
+    timepoints : 1d array
+        Timepoints for the computed autocorrelations, in samples or seconds.
+    autocorrs : 1d array
+        Autocorrelation values.
+    fs : int, optional
+        Sampling rate of the signal.
+        If provided, timepoints are converted to time values.
+    fit_func : {'single_exp', 'double_exp'}
+        Which fitting function to use to fit the autocorrelation results.
+    bounds : tuple of list
+        Parameter bounds for fitting.
+        Organized as ([min_p1, min_p1, ...], [max_p1, max_p2, ...]).
+
+    Returns
+    -------
+    popts
+        Fit parameters. Parameters depend on the fitting function.
+        If `fit_func` is 'single_exp', fit parameters are: tau, scale, offset
+        If `fit_func` is 'douple_exp', fit parameters are: tau1, tau2, scale1, scale2, offset
+        See fit function for more details.
+
+    Notes
+    -----
+    The values / units of the returned parameters are dependent on the units of samples.
+    For example, if the timepoints input is in samples, the fit tau value is too.
+    If providing parameter bounds, these also need to match the unit of timepoints.
+    """
+
+    if not bounds:
+        if fit_function == 'single_exp':
+            bounds = ([0, 0, 0], [np.inf, np.inf, np.inf])
+        elif fit_function == 'double_exp':
+            bounds = ([0, 0, 0, 0, 0], [np.inf, np.inf, np.inf, np.inf, np.inf])
+
+    popts, _ = curve_fit(AC_FIT_FUNCS[fit_function], timepoints, autocorrs, bounds=bounds)
+
+    return popts
+
+
+## AC FITTING
+
+def exp_decay_func(timepoints, tau, scale, offset):
+    """Exponential decay fit function.
+
+    Parameters
+    ----------
+    timepoints : 1d array
+        Time values.
+    tau : float
+        Timescale value.
+    scale : float
+        Scaling factor, which captures the start value of the function.
+    offset : float
+        Offset factor, which captures the end value of the function.
+
+    Returns
+    -------
+    ac_fit : 1d array
+        Result of fitting the function to the autocorrelation.
+    """
+
+    return scale * (np.exp(-timepoints / tau) + offset)
+
+
+def double_exp_decay_func(timepoints, tau1, tau2, scale1, scale2, offset):
+    """Exponential decay fit function with two timescales.
+
+    Parameters
+    ----------
+    timepoints : 1d array
+        Time values.
+    tau1, tau2 : float
+        Timescale values, for the 1st and 2nd timescale.
+    scale1, scale2 : float
+        Scaling factors, for the 1st and 2nd timescale.
+    offset : float
+        Offset factor.
+
+    Returns
+    -------
+    ac_fit : 1d array
+        Result of fitting the function to the autocorrelation.
+    """
+
+    return scale1 * np.exp(-timepoints / tau1) + scale2 * np.exp(-timepoints / tau2) + offset
+
+
+AC_FIT_FUNCS = {
+    'single_exp' : exp_decay_func,
+    'double_exp' : double_exp_decay_func,
+}
+
+
+def compute_ac_fit(timepoints, *popts, fit_function='single_exp'):
+    """Regenerate values of the exponential decay fit.
+
+    Parameters
+    ----------
+    timepoints : 1d array
+        Time values, in samples or seconds.
+    *popts
+        Fit parameters.
+    fit_func : {'single_exp', 'double_exp'}
+        Which fit function to use to fit the autocorrelation results.
+
+    Returns
+    -------
+    fit_values : 1d array
+        Values of the fit to the autocorrelation values.
+    """
+
+    fit_func = AC_FIT_FUNCS[fit_function]
+
+    return fit_func(timepoints, *popts)

neurodsp/tests/aperiodic/test_autocorr.py

@@ -1,5 +1,7 @@
 """Tests for neurodsp.aperiodic.autocorr."""
 
+from neurodsp.tests.settings import FS
+
 from neurodsp.aperiodic.autocorr import *
 
 ###################################################################################################
@@ -10,3 +12,46 @@ def test_compute_autocorr(tsig):
     max_lag = 500
     timepoints, autocorrs = compute_autocorr(tsig, max_lag=max_lag)
     assert len(timepoints) == len(autocorrs) == max_lag + 1
+
+def test_compute_decay_time(tsig):
+
+    timepoints, autocorrs = compute_autocorr(tsig, max_lag=500)
+    decay_time = compute_decay_time(timepoints, autocorrs, FS)
+    assert isinstance(decay_time, float)
+
+def test_fit_autocorr(tsig):
+    # This is a smoke test - check it runs with no accuracy checking
+
+    timepoints, autocorrs = compute_autocorr(tsig, max_lag=500)
+
+    popts1 = fit_autocorr(timepoints, autocorrs, fit_function='single_exp')
+    fit_vals1 = compute_ac_fit(timepoints, *popts1, fit_function='single_exp')
+    assert np.all(fit_vals1)
+
+    # Test with bounds passed in
+    bounds = ([0, 0, 0], [10, np.inf, np.inf])
+    popts1 = fit_autocorr(timepoints, autocorrs, 'single_exp', bounds)
+    fit_vals1 = compute_ac_fit(timepoints, *popts1, fit_function='single_exp')
+    assert np.all(fit_vals1)
+
+    popts2 = fit_autocorr(timepoints, autocorrs, fit_function='double_exp')
+    fit_vals2 = compute_ac_fit(timepoints, *popts2, fit_function='double_exp')
+    assert np.all(fit_vals2)
+
+def test_fit_autocorr_acc():
+    # This test includes some basic accuracy checking
+
+    fs = 100
+    tau = 0.015
+    lags = np.linspace(0, 100, 1000)
+
+    # Test can fit and recompute the tau value
+    corrs1 = exp_decay_func(lags, tau, 1, 0)
+    params1 = fit_autocorr(lags, corrs1)
+    assert np.isclose(params1[0], tau, 0.001)
+
+    # Test can fit and recompute the tau value - timepoints as time values
+    lags = lags / fs
+    corrs2 = exp_decay_func(lags, tau, 1, 0)
+    params2 = fit_autocorr(lags, corrs2)
+    assert np.isclose(params2[0], tau, 0.001)