Implementation + test of the cross-correlation function

Author: EtienneCmb

docs/source/api/api_connectivity.rst

@@ -17,6 +17,7 @@ Connectivity metrics
 
    conn_dfc
    conn_covgc
+   conn_ccf
    conn_transfer_entropy
 
 Utility functions

frites/conn/__init__.py

@@ -14,6 +14,7 @@
 from .conn_io import conn_io  # noqa
 
 # connectivity metrics
+from .conn_ccf import conn_ccf  # noqa
 from .conn_covgc import conn_covgc  # noqa
 from .conn_dfc import conn_dfc  # noqa
 from .conn_transfer_entropy import conn_transfer_entropy  # noqa

frites/conn/conn_ccf.py

@@ -0,0 +1,132 @@
+"""Cross-correlation function."""
+import numpy as np
+import xarray as xr
+
+from frites.conn import conn_io
+from frites.io import logger
+from frites.estimator import GCMIEstimator
+from frites.utils import parallel_func
+from frites.utils.preproc import _acf
+
+
+
+def conn_ccf(data, times=None, roi=None, normalized=True, n_jobs=1,
+             verbose=None):
+    """Single trial Cross-Correlation Function.
+
+    This function computes the pairwise Cross Correlation (CCF) at the single
+    trial level. This can be particulary usefull to find whether there are
+    temporal delays between times series.
+
+    Parameters
+    ----------
+    data : array_like
+        Electrophysiological data. Several input types are supported :
+
+            * Standard NumPy arrays of shape (n_epochs, n_roi, n_times)
+            * mne.Epochs
+            * xarray.DataArray of shape (n_epochs, n_roi, n_times)
+
+    times : array_like | None
+        Time vector array of shape (n_times,). If the input is an xarray, the
+        name of the time dimension can be provided
+    roi : array_like | None
+        ROI names of a single subject. If the input is an xarray, the
+        name of the ROI dimension can be provided
+    normalized : bool | True
+        Z-score normalization of the data. By default, it set to true.
+    n_jobs : int | 1
+        Number of jobs to use for parallel computing (use -1 to use all
+        jobs). The parallel loop is set at the pair level.
+
+    Returns
+    -------
+    ccf : array_like
+        The Cross-Correlation array of shape (n_epochs, n_pairs, n_times). When
+        the peak of correlation occurs at a negative time it means that the
+        target has to be moved **toward** the source. On the contrary, if the
+        peak occurs at positive time it means that the target is moved **away**
+        of the source.
+    """
+    # ________________________________ INPUTS _________________________________
+    # inputs conversion
+    data, cfg = conn_io(
+        data, times=times, roi=roi, agg_ch=False, win_sample=None, pairs=None,
+        sort=True, name='CCF', verbose=verbose,
+    )
+
+    # extract variables
+    x, trials, attrs = data.data, data['y'].data, cfg['attrs']
+    x_s, x_t = cfg['x_s'], cfg['x_t']
+    roi_p, roi_idx = cfg['roi_p'], cfg['roi_idx']
+    times = data['times'].data
+    n_pairs = len(x_s)
+
+    # data normalization
+    if normalized:
+        x = (x - x.mean(-1, keepdims=True)) / x.std(-1, keepdims=True)
+
+    # __________________________________ CCF __________________________________
+    # function to put in parallel
+    def para_fun(xs, xt):
+        n_trials = xs.shape[0]
+        corr = np.zeros((n_trials, int(2 * len(times)) - 1))
+        for n_t in range(n_trials):
+            corr[n_t, :] = _acf(xs[n_t, :], xt[n_t, :])
+        return corr
+
+    # prepare parallel function
+    n_jobs = 1 if n_pairs == 1 else n_jobs
+    parallel, p_fun = parallel_func(para_fun, n_jobs=n_jobs, verbose=verbose,
+                                    total=n_pairs, mesg='Estimating CCF')
+
+    logger.info(f'Computing CCF between {n_pairs} pairs')
+
+    # compute ccf
+    ccf = parallel(
+        p_fun(x[:, i_s, :], x[:, i_t, :]) for i_s, i_t in zip(x_s, x_t))
+    ccf = np.stack(ccf, axis=1)
+
+    # ________________________________ OUTPUTS ________________________________
+    # dataarray conversion
+    times_n = np.arange(ccf.shape[-1]).astype(float)# / cfg['sfreq']
+    times_n -= times_n.mean()
+    ccf = xr.DataArray(ccf, dims=('trials', 'roi', 'times'), name=f'CCF',
+                       coords=(trials, roi_p, times_n))
+
+    # add the windows used in the attributes
+    ccf.attrs = {**dict(type='ccf', normalized=int(normalized)), **attrs}
+
+    return ccf
+
+if __name__ == '__main__':
+    import matplotlib.pyplot as plt
+    from frites.estimator import CorrEstimator
+
+    n_trials = 20
+    n_roi = 3
+    n_times = 1000
+    # create coordinates
+    trials = np.arange(n_trials)
+    roi = [f"roi_{k}" for k in range(n_roi)]
+    times = (np.arange(n_times) - 200) / 64.
+    # data creation
+    x = np.random.rand(n_trials, n_roi, n_times)
+    # inject relation
+    bump = np.hanning(200).reshape(1, -1)
+    x[:, 0, 200:400] += bump
+    x[:, 1, 220:420] += bump
+    x[:, 2, 260:460] += bump
+    # xarray conversion
+    x = xr.DataArray(x, dims=('trials', 'roi', 'times'),
+                     coords=(trials, roi, times))
+    plt.figure(figsize=(15, 6))
+
+    # compute delayed dfc
+    ccf = conn_ccf(x, times='times', roi='roi', n_jobs=1, verbose=False)
+
+    plt.subplot(121)
+    x.mean('trials').plot(x='times', hue='roi')
+    plt.subplot(122)
+    ccf.mean('trials').plot(x='times', hue='roi')
+    plt.show()

frites/conn/tests/test_conn.py

@@ -2,7 +2,7 @@
 import numpy as np
 import xarray as xr
 
-from frites.conn import (conn_covgc, conn_transfer_entropy, conn_dfc)
+from frites.conn import (conn_covgc, conn_transfer_entropy, conn_dfc, conn_ccf)
 
 
 class TestConn(object):
@@ -70,3 +70,41 @@ def test_conn_covgc(self):
         assert isinstance(gc, xr.DataArray)
         gc = conn_covgc(x, dt, lag, t0, n_jobs=1, method='gc',
                         conditional=True, norm=False)
+
+    def test_conn_ccf(self):
+        """Test function conn_ccf."""
+        n_trials, n_roi, n_times = 20, 3, 1000
+        # create coordinates
+        trials = np.arange(n_trials)
+        roi = [f"roi_{k}" for k in range(n_roi)]
+        times = (np.arange(n_times) - 200) / 64.
+        # data creation
+        rnd = np.random.RandomState(0)
+        x = .1 * rnd.rand(n_trials, n_roi, n_times)
+        # inject relation
+        bump = np.hanning(200).reshape(1, -1)
+        x[:, 0, 200:400] += bump
+        x[:, 1, 220:420] += bump
+        x[:, 2, 150:350] += bump
+        # xarray conversion
+        x = xr.DataArray(x, dims=('trials', 'roi', 'times'),
+                         coords=(trials, roi, times))
+        # compute delayed dfc
+        conn_ccf(x, times='times', roi='roi', n_jobs=1, verbose=False,
+                 normalized=False)
+        ccf = conn_ccf(x, times='times', roi='roi', n_jobs=1, verbose=False)
+        # shape and dimension checking
+        assert ccf.ndim == 3
+        assert ccf.dims == ('trials', 'roi', 'times')
+        assert len(ccf['trials']) == len(trials)
+        np.testing.assert_array_equal(ccf['trials'].data, trials)
+        assert len(ccf['roi']) == 3
+        # peak detection
+        ccf_m = ccf.mean('trials')
+        is_peaks = np.where(ccf_m == ccf_m.max('times'))
+        peaks = ccf['times'].data[is_peaks[1]]
+        # peak checking
+        tol = 5
+        assert -20 - tol <= peaks[0] <= -20 + tol
+        assert 50 - tol <= peaks[1] <= 50 + tol
+        assert 70 - tol <= peaks[2] <= 70 + tol

frites/utils/preproc.py

@@ -120,11 +120,15 @@ def kernel_smoothing(x, kernel, axis=-1):
     return x
 
 
-def _acf(xd):
-    """Auto-correlation on a single time-series."""
-    n = len(xd)
-    acov = np.correlate(xd, xd, "full")[n - 1:] / n
-    return acov[: n + 1] / acov[0]
+def _acf(xs, xt=None):
+    """Auto- or cross-correlation of time-series."""
+    n = len(xs)
+
+    if xt is None:  # auto-correlation
+        acov = np.correlate(xs, xs, "full")[n - 1:] / n
+        return acov[: n + 1] / acov[0]
+    else:           # cross-correlation
+        return np.correlate(xs, xt, mode="full") / len(xs)
 
 
 def acf(x, axis=-1, demean=True):