Add neo objects as ephy data sources

frites/dataset/tests/test_ds_ephy.py

@@ -66,6 +66,26 @@ def test_repr(self):
         ds.__repr__()
         ds._repr_html_()
 
+    def test_definition_as_neo(self):
+        """Test function definition."""
+        # test array definition
+        data, roi, time = sim_multi_suj_ephy(modality="intra", n_times=57,
+                                             n_roi=5, n_sites_per_roi=7,
+                                             n_epochs=10, n_subjects=5,
+                                             random_state=0)
+        y, _ = sim_mi_cc(data, snr=.8)
+        z = [np.random.randint(0, 3, (10,)) for _ in range(len(y))]
+        dt = DatasetEphy(data, y, roi, z=z, times=time)
+        dt.groupby('roi')
+        # test mne definition
+        data, roi, time = sim_multi_suj_ephy(modality="meeg", n_times=57,
+                                             n_roi=5, n_sites_per_roi=1,
+                                             n_epochs=7, n_subjects=5,
+                                             as_mne=False, as_neo=True,
+                                             random_state=0)
+        y, _ = sim_mi_cc(data, snr=.8)
+        ds = DatasetEphy(data, y, roi, times=time)
+
     def test_multiconditions(self):
         """Test multi-conditions remapping."""
         d_3d = self._get_data(3)

frites/simulations/sim_generate_data.py

@@ -4,6 +4,13 @@
 from scipy.signal import savgol_filter
 from itertools import product
 
+try:
+    import neo
+    import quantities as pq
+    HAS_NEO = True
+except ModuleNotFoundError:
+    HAS_NEO = False
+
 MA_NAMES = ['L_VCcm', 'L_VCl', 'L_VCs', 'L_Cu', 'L_VCrm', 'L_ITCm', 'L_ITCr',
             'L_MTCc', 'L_STCc', 'L_STCr', 'L_MTCr', 'L_ICC', 'L_IPCv',
             'L_IPCd', 'L_SPC', 'L_SPCm', 'L_PCm', 'L_PCC', 'L_Sv', 'L_Sdl',
@@ -23,7 +30,8 @@
 
 def sim_single_suj_ephy(modality="meeg", sf=512., n_times=1000, n_roi=1,
                         n_sites_per_roi=1, n_epochs=100, n_sines=100, f_min=.5,
-                        f_max=160., noise=10, as_mne=False, random_state=None):
+                        f_max=160., noise=10, as_mne=False, as_neo=False,
+                        random_state=None):
     """Simulate electrophysiological data of a single subject.
 
     This function generate some illustrative random electrophysiological data
@@ -54,6 +62,8 @@ def sim_single_suj_ephy(modality="meeg", sf=512., n_times=1000, n_roi=1,
         Noise level.
     as_mne : bool | False
         If True, data are converted to a mne.EpochsArray structure
+    as_neo : bool | False
+        If True, data are converted to a neo.Block structure
     random_state : int | None
         Fix the random state for the reproducibility.
 
@@ -103,6 +113,18 @@ def sim_single_suj_ephy(modality="meeg", sf=512., n_times=1000, n_roi=1,
         from mne import create_info, EpochsArray
         info = create_info(roi.tolist(), sf, ch_types='seeg')
         signal = EpochsArray(signal, info, tmin=float(time[0]), verbose=False)
+    elif as_neo:
+        if not HAS_NEO:
+            raise ImportError('`as_neo` requires neo to be installed.')
+        block = neo.Block()
+        for epoch_id in range(len(signal)):
+            seg = neo.Segment()
+            anasig = neo.AnalogSignal(signal[epoch_id].T*pq.dimensionless,
+                                      t_start=time[0, 0]*pq.s,
+                                      sampling_rate=sf*pq.Hz)
+            seg.analogsignals.append(anasig)
+            block.segments.append(seg)
+        signal = block
     return signal, roi, time.squeeze()
 
 

frites/simulations/sim_mi.py

@@ -75,7 +75,16 @@ def sim_mi_cc(x, snr=.9):
     # if mne types, turn into arrays
     if isinstance(x[0], CONFIG["MNE_EPOCHS_TYPE"]):
         x = [x[k].get_data() for k in range(len(x))]
+    elif 'neo.core' in str(type(x[0])):
+        subject_list = []
+        for block in x:
+            subject_data = np.stack([seg.analogsignals[0].magnitude for seg in block.segments])
+            # reorder dimensions to match (n_epochs, n_channels, n_times)
+            subject_list.append(subject_data.swapaxes(1, 2))
+        x = subject_list
+
     n_times = x[0].shape[-1]
+
     # cluster definition (20% length around central point)
     cluster = _get_cluster(n_times, location='center', perc=.2)
     # ground truth definition

frites/simulations/tests/test_sim_generate_data.py

@@ -2,6 +2,12 @@
 import numpy as np
 from mne import EpochsArray
 
+try:
+    import neo
+    HAS_NEO = True
+except ModuleNotFoundError:
+    HAS_NEO = False
+
 from frites.simulations import (sim_single_suj_ephy, sim_multi_suj_ephy)
 
 
@@ -19,6 +25,10 @@ def test_sim_single_suj_ephy(self):
         # mne type
         data, _, _ = sim_single_suj_ephy(as_mne=True)
         assert isinstance(data, EpochsArray)
+        # neo type
+        if HAS_NEO:
+            data, _, _ = sim_single_suj_ephy(as_neo=True)
+            assert isinstance(data, neo.Block)
 
     def test_sim_multi_suj_ephy(self):
         """Test function sim_multi_suj_ephy."""
@@ -34,3 +44,8 @@ def test_sim_multi_suj_ephy(self):
         # mne type
         data, _, _ = sim_multi_suj_ephy(n_subjects=5, as_mne=True)
         assert all([isinstance(k, EpochsArray) for k in data])
+        # neo type
+        if HAS_NEO:
+            data, _, _ = sim_multi_suj_ephy(n_subjects=5, as_neo=True)
+            assert all([isinstance(k, neo.Block) for k in data])
+