Conn input conversion function

Author: EtienneCmb

examples/conn/plot_dfc.py

@@ -72,7 +72,7 @@
 # each of the temporal window
 
 # compute DFC
-dfc, pairs, roi_p = conn_dfc(x, times, roi, win_sample, n_jobs=1)
+dfc, pairs, roi_p = conn_dfc(x, win_sample, times=times, roi=roi, n_jobs=1)
 
 plt.figure(figsize=(10, 8))
 plt.plot(times_p, dfc.mean('trials').T)

examples/tutorials/plot_stim_spec_network.py

@@ -148,7 +148,8 @@ def plot_mi(mi, pv):
 # compute the DFC for each subject
 dfc = []
 for n_s in range(n_subjects):
-    _dfc = conn_dfc(x[n_s].data, times, roi, win_sample, verbose=False)[0]
+    _dfc = conn_dfc(x[n_s].data, win_sample, times=times, roi=roi,
+                    verbose=False)[0]
     # reset trials dimension
     _dfc['trials'] = x[n_s]['trials'].data
     dfc += [_dfc]

frites/conn/__init__.py

@@ -1,6 +1,7 @@
 """Information-theritical measures of connectivity."""
 from .conn_covgc import conn_covgc  # noqa
 from .conn_dfc import conn_dfc  # noqa
-from .conn_transfer_entropy import conn_transfer_entropy  # noqa
 from .conn_fit import conn_fit  # noqa
+from .conn_io import conn_io  # noqa
+from .conn_transfer_entropy import conn_transfer_entropy  # noqa
 from .conn_utils import (conn_reshape_undirected, conn_reshape_directed)  # noqa

frites/conn/conn_dfc.py

@@ -7,11 +7,14 @@
 from frites.config import CONFIG
 from frites.utils import parallel_func
 
+from frites.conn.conn_io import conn_io
+
 from mne.utils import ProgressBar
 
 
 
-def conn_dfc(data, times, roi, win_sample, n_jobs=1, gcrn=True, verbose=None):
+def conn_dfc(data, win_sample, times=None, roi=None, n_jobs=1, gcrn=True,
+             verbose=None):
     """Compute the Dynamic Functional Connectivity using the GCMI.
 
     This function computes the Dynamic Functional Connectivity (DFC) using the
@@ -24,14 +27,14 @@ def conn_dfc(data, times, roi, win_sample, n_jobs=1, gcrn=True, verbose=None):
     data : array_like
         Electrophysiological data array of a single subject organized as
         (n_epochs, n_roi, n_times)
-    times : array_like
-        Time vector array of shape (n_times,)
-    roi : array_like
-        ROI names of a single subject
     win_sample : array_like
         Array of shape (n_windows, 2) describing where each window start and
         finish. You can use the function :func:`frites.utils.define_windows`
         to define either manually either sliding windows.
+    times : array_like | None
+        Time vector array of shape (n_times,)
+    roi : array_like | None
+        ROI names of a single subject
     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.
@@ -54,9 +57,12 @@ def conn_dfc(data, times, roi, win_sample, n_jobs=1, gcrn=True, verbose=None):
     define_windows, conn_covgc
     """
     set_log_level(verbose)
+    # -------------------------------------------------------------------------
+    # inputs conversion
+    da, roi, times = conn_io(data, roi=roi, times=times, verbose=verbose)
+
     # -------------------------------------------------------------------------
     # data checking
-    assert isinstance(data, np.ndarray) and (data.ndim == 3)
     n_epochs, n_roi, n_pts = data.shape
     assert (len(roi) == n_roi) and (len(times) == n_pts)
     assert isinstance(win_sample, np.ndarray) and (win_sample.ndim == 2)

frites/conn/conn_io.py

@@ -0,0 +1,49 @@
+"""Manage I/O for connectivity functions."""
+import numpy as np
+import xarray as xr
+
+
+from frites.io import set_log_level, logger
+
+
+def conn_io(da, roi=None, times=None, verbose=None):
+    """I/O conversion for connectivity functions.
+
+    Parameters
+    ----------
+    da : array_like
+        Array of electrophysiological data of shape (n_trials, n_roi, n_times)
+    roi : array_like | None
+        List of roi names or string corresponding to the dimension name in a
+        DataArray
+    times : array_like | None
+        Time vector or string corresponding to the dimension name in a
+        DataArray
+    """
+    set_log_level(verbose)
+    assert isinstance(da, np.ndarray) or isinstance(da, xr.DataArray)
+    assert da.ndim == 3
+    n_trials, n_roi, n_times = da.shape
+    logger.info(f"Inputs conversion (n_trials={n_trials}, n_roi={n_roi}, "
+                "n_times={n_times})")
+
+    # _____________________________ Empty inputs ______________________________
+    if roi is None:
+        roi = [f"roi_{k}" for k in range(n_roi)]
+    if times is None:
+        times = np.arange(n_times)
+
+    # _______________________________ Xarray case _____________________________
+    if isinstance(da, xr.DataArray):
+        # get roi and times
+        if isinstance(roi, str):
+            roi = da[roi].data
+        if isinstance(times, str):
+            times = da[times].data
+        da = da.data
+
+    # _______________________________ Final check _____________________________
+    assert isinstance(da, np.ndarray)
+    assert da.shape == (n_trials, len(roi), len(times))
+
+    return da, roi, times

frites/conn/tests/test_conn.py

@@ -44,9 +44,9 @@ def test_conn_dfc(self):
         roi = [f"roi_{k}" for k in range(n_roi)]
         x = np.random.rand(n_epochs, n_roi, n_times)
 
-        dfc = conn_dfc(x, times, roi, win_sample)[0]
+        dfc = conn_dfc(x, win_sample, times=times, roi=roi)[0]
         assert dfc.shape == (n_epochs, 3, 2)
-        dfc = conn_dfc(x, times, roi, win_sample)[0]
+        dfc = conn_dfc(x, win_sample, times=times, roi=roi)[0]
         assert isinstance(dfc, xr.DataArray)
 
     def test_conn_covgc(self):
@@ -75,7 +75,7 @@ def test_conn_reshape_undirected(self):
         roi = [f"roi_{k}" for k in range(n_roi)]
         order = ['roi_2', 'roi_1']
         x = np.random.rand(n_epochs, n_roi, n_times)
-        dfc = conn_dfc(x, times, roi, win_sample)[0].mean('trials')
+        dfc = conn_dfc(x, win_sample, times=times, roi=roi)[0].mean('trials')
         # reshape it without the time dimension
         dfc_mean = conn_reshape_undirected(dfc.mean('times'))
         assert dfc_mean.shape == (n_roi, n_roi, 1)

frites/dataset/ds_ephy_io.py

@@ -3,11 +3,9 @@
 
 import numpy as np
 
-from frites.io import set_log_level
+from frites.io import set_log_level, logger
 from frites.config import CONFIG
 
-logger = logging.getLogger("frites")
-
 
 ###############################################################################
 ###############################################################################