Skip to content

Commit

Permalink
Move cVEP common functions to datasets.utils - Closes NeuroTechX#564
Browse files Browse the repository at this point in the history
  • Loading branch information
PierreGtch committed Apr 11, 2024
1 parent c8ae1fc commit 25f0e14
Show file tree
Hide file tree
Showing 4 changed files with 109 additions and 272 deletions.
89 changes: 3 additions & 86 deletions moabb/datasets/castillos2023.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,10 @@

import mne
import numpy as np
from mne import create_info
from mne.io import RawArray

from moabb.datasets import download as dl
from moabb.datasets.base import BaseDataset
from moabb.datasets.utils import add_stim_channel_epoch, add_stim_channel_trial


Castillos2023_URL = "https://zenodo.org/records/8255618"
Expand Down Expand Up @@ -41,88 +40,6 @@ def __init__(
self.n_channels = 32
self.window_size = window_size

def _add_stim_channel_trial(
self, raw, onsets, labels, offset=200, ch_name="stim_trial"
):
"""
Add a stimulus channel with trial onsets and their labels.
Parameters
----------
raw: mne.Raw
The raw object to add the stimulus channel to.
onsets: List | np.ndarray
The onsets of the trials in sample numbers.
labels: List | np.ndarray
The labels of the trials.
offset: int (default: 200)
The integer value to start markers with. For instance, if 200, then label 0 will be marker 200, label 1
will be be marker 201, etc.
ch_name: str (default: "stim_trial")
The name of the added stimulus channel.
Returns
-------
mne.Raw
The raw object with the added stimulus channel.
"""
stim_chan = np.zeros((1, len(raw)))
for onset, label in zip(onsets, labels):
stim_chan[0, onset] = offset + label
info = create_info(
ch_names=[ch_name],
ch_types=["stim"],
sfreq=raw.info["sfreq"],
verbose=False,
)
raw = raw.add_channels([RawArray(data=stim_chan, info=info, verbose=False)])
return raw

def _add_stim_channel_epoch(
self,
raw,
onsets,
labels,
offset=100,
ch_name="stim_epoch",
):
"""
Add a stimulus channel with epoch onsets and their labels, which are the values of the presented code for each
of the trials.
Parameters
----------
raw: mne.Raw
The raw object to add the stimulus channel to.
onsets: List | np.ndarray
The onsets of the trials in sample numbers.
labels: List | np.ndarray
The labels of the trials.
codes: np.ndarray
The codebook containing each presented code of shape (nb_bits, nb_codes), sampled at the presentation rate.
offset: int (default: 100)
The integer value to start markers with. For instance, if 100, then label 0 will be marker 100, label 1
will be be marker 101, etc.
ch_name: str (default: "stim_epoch")
The name of the added stimulus channel.
Returns
-------
mne.Raw
The raw object with the added stimulus channel.
"""
stim_chan = np.zeros((1, len(raw)))
for onset, label in zip(onsets, labels):
stim_chan[0, int(onset * self.sfreq)] = offset + label
info = create_info(
ch_names=[ch_name],
ch_types=["stim"],
sfreq=raw.info["sfreq"],
verbose=False,
)
raw = raw.add_channels([RawArray(data=stim_chan, info=info, verbose=False)])
return raw

def _get_single_subject_data(self, subject):
"""Return the data of a single subject."""
file_path_list = self.data_path(subject, self.paradigm_type)
Expand Down Expand Up @@ -183,10 +100,10 @@ def _get_single_subject_data(self, subject):

# Create stim channel with trial information (i.e., symbols)
# Specifically: 200 = symbol-0, 201 = symbol-1, 202 = symbol-2, etc.
raw = self._add_stim_channel_trial(raw, onset_code, labels, offset=200)
raw = add_stim_channel_trial(raw, onset_code, labels, offset=200)
# Create stim channel with epoch information (i.e., 1 / 0, or on / off)
# Specifically: 100 = "0", 101 = "1"
raw = self._add_stim_channel_epoch(
raw = add_stim_channel_epoch(
raw,
np.concatenate([onset, onset_0]),
np.concatenate([np.ones(onset.shape), np.zeros(onset_0.shape)]),
Expand Down
96 changes: 3 additions & 93 deletions moabb/datasets/thielen2015.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,10 @@
import mne
import numpy as np
from mne import create_info
from mne.io import RawArray
from scipy.io import loadmat

from moabb.datasets import download as dl
from moabb.datasets.base import BaseDataset
from moabb.datasets.utils import add_stim_channel_epoch, add_stim_channel_trial


Thielen2015_URL = "https://public.data.ru.nl/dcc/DSC_2018.00047_553_v3"
Expand Down Expand Up @@ -96,93 +95,6 @@ def __init__(self):
doi="10.34973/1ecz-1232",
)

def _add_stim_channel_trial(
self, raw, onsets, labels, offset=200, ch_name="stim_trial"
):
"""
Add a stimulus channel with trial onsets and their labels.
Parameters
----------
raw: mne.Raw
The raw object to add the stimulus channel to.
onsets: List | np.ndarray
The onsets of the trials in sample numbers.
labels: List | np.ndarray
The labels of the trials.
offset: int (default: 200)
The integer value to start markers with. For instance, if 200, then label 0 will be marker 200, label 1
will be be marker 201, etc.
ch_name: str (default: "stim_trial")
The name of the added stimulus channel.
Returns
-------
mne.Raw
The raw object with the added stimulus channel.
"""
stim_chan = np.zeros((1, len(raw)))
for onset, label in zip(onsets, labels):
stim_chan[0, onset] = offset + label
info = create_info(
ch_names=[ch_name],
ch_types=["stim"],
sfreq=raw.info["sfreq"],
verbose=False,
)
raw = raw.add_channels([RawArray(data=stim_chan, info=info, verbose=False)])
return raw

def _add_stim_channel_epoch(
self,
raw,
onsets,
labels,
codes,
presentation_rate=60,
offset=100,
ch_name="stim_epoch",
):
"""
Add a stimulus channel with epoch onsets and their labels, which are the values of the presented code for each
of the trials.
Parameters
----------
raw: mne.Raw
The raw object to add the stimulus channel to.
onsets: List | np.ndarray
The onsets of the trials in sample numbers.
labels: List | np.ndarray
The labels of the trials.
codes: np.ndarray
The codebook containing each presented code of shape (nr_bits, nr_codes), sampled at the presentation rate.
presentation_rate: int (default: 60):
The presentation rate (e.g., frame rate) at which the codes were presented in Hz.
offset: int (default: 100)
The integer value to start markers with. For instance, if 100, then label 0 will be marker 100, label 1
will be be marker 101, etc.
ch_name: str (default: "stim_epoch")
The name of the added stimulus channel.
Returns
-------
mne.Raw
The raw object with the added stimulus channel.
"""
stim_chan = np.zeros((1, len(raw)))
for onset, label in zip(onsets, labels):
idx = np.round(
onset + np.arange(codes.shape[0]) / presentation_rate * raw.info["sfreq"]
).astype("int")
stim_chan[0, idx] = offset + codes[:, label]
info = create_info(
ch_names=[ch_name],
ch_types=["stim"],
sfreq=raw.info["sfreq"],
verbose=False,
)
raw = raw.add_channels([RawArray(data=stim_chan, info=info, verbose=False)])
return raw

def _get_single_subject_data(self, subject):
"""Return the data of a single subject."""
file_path_list = self.data_path(subject)
Expand Down Expand Up @@ -231,13 +143,11 @@ def _get_single_subject_data(self, subject):

# Create stim channel with trial information (i.e., symbols)
# Specifically: 200 = symbol-0, 201 = symbol-1, 202 = symbol-2, etc.
raw = self._add_stim_channel_trial(
raw, trial_onsets, trial_labels, offset=200
)
raw = add_stim_channel_trial(raw, trial_onsets, trial_labels, offset=200)

# Create stim channel with epoch information (i.e., 1 / 0, or on / off)
# Specifically: 100 = "0", 101 = "1"
raw = self._add_stim_channel_epoch(
raw = add_stim_channel_epoch(
raw, trial_onsets, trial_labels, codes, PRESENTATION_RATE, offset=100
)

Expand Down
96 changes: 3 additions & 93 deletions moabb/datasets/thielen2021.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,11 @@
import h5py
import mne
import numpy as np
from mne import create_info
from mne.io import RawArray
from scipy.io import loadmat

from moabb.datasets import download as dl
from moabb.datasets.base import BaseDataset
from moabb.datasets.utils import add_stim_channel_epoch, add_stim_channel_trial


Thielen2021_URL = "https://public.data.ru.nl/dcc/DSC_2018.00122_448_v3"
Expand Down Expand Up @@ -140,93 +139,6 @@ def __init__(self):
doi="10.34973/9txv-z787",
)

def _add_stim_channel_trial(
self, raw, onsets, labels, offset=200, ch_name="stim_trial"
):
"""
Add a stimulus channel with trial onsets and their labels.
Parameters
----------
raw: mne.Raw
The raw object to add the stimulus channel to.
onsets: List | np.ndarray
The onsets of the trials in sample numbers.
labels: List | np.ndarray
The labels of the trials.
offset: int (default: 200)
The integer value to start markers with. For instance, if 200, then label 0 will be marker 200, label 1
will be be marker 201, etc.
ch_name: str (default: "stim_trial")
The name of the added stimulus channel.
Returns
-------
mne.Raw
The raw object with the added stimulus channel.
"""
stim_chan = np.zeros((1, len(raw)))
for onset, label in zip(onsets, labels):
stim_chan[0, onset] = offset + label
info = create_info(
ch_names=[ch_name],
ch_types=["stim"],
sfreq=raw.info["sfreq"],
verbose=False,
)
raw = raw.add_channels([RawArray(data=stim_chan, info=info, verbose=False)])
return raw

def _add_stim_channel_epoch(
self,
raw,
onsets,
labels,
codes,
presentation_rate=60,
offset=100,
ch_name="stim_epoch",
):
"""
Add a stimulus channel with epoch onsets and their labels, which are the values of the presented code for each
of the trials.
Parameters
----------
raw: mne.Raw
The raw object to add the stimulus channel to.
onsets: List | np.ndarray
The onsets of the trials in sample numbers.
labels: List | np.ndarray
The labels of the trials.
codes: np.ndarray
The codebook containing each presented code of shape (nr_bits, nr_codes), sampled at the presentation rate.
presentation_rate: int (default: 60):
The presentation rate (e.g., frame rate) at which the codes were presented in Hz.
offset: int (default: 100)
The integer value to start markers with. For instance, if 100, then label 0 will be marker 100, label 1
will be be marker 101, etc.
ch_name: str (default: "stim_epoch")
The name of the added stimulus channel.
Returns
-------
mne.Raw
The raw object with the added stimulus channel.
"""
stim_chan = np.zeros((1, len(raw)))
for onset, label in zip(onsets, labels):
idx = np.round(
onset + np.arange(codes.shape[0]) / presentation_rate * raw.info["sfreq"]
).astype("int")
stim_chan[0, idx] = offset + codes[:, label]
info = create_info(
ch_names=[ch_name],
ch_types=["stim"],
sfreq=raw.info["sfreq"],
verbose=False,
)
raw = raw.add_channels([RawArray(data=stim_chan, info=info, verbose=False)])
return raw

def _get_single_subject_data(self, subject):
"""Return the data of a single subject."""
file_path_list = self.data_path(subject)
Expand Down Expand Up @@ -277,13 +189,11 @@ def _get_single_subject_data(self, subject):

# Create stim channel with trial information (i.e., symbols)
# Specifically: 200 = symbol-0, 201 = symbol-1, 202 = symbol-2, etc.
raw = self._add_stim_channel_trial(
raw, trial_onsets, trial_labels, offset=200
)
raw = add_stim_channel_trial(raw, trial_onsets, trial_labels, offset=200)

# Create stim channel with epoch information (i.e., 1 / 0, or on / off)
# Specifically: 100 = "0", 101 = "1"
raw = self._add_stim_channel_epoch(
raw = add_stim_channel_epoch(
raw, trial_onsets, trial_labels, codes, PRESENTATION_RATE, offset=100
)

Expand Down
Loading

0 comments on commit 25f0e14

Please sign in to comment.