From cd17a5a3eb5b66767314eb7ba1626d8c8761b5c8 Mon Sep 17 00:00:00 2001 From: Milo Sobral Date: Wed, 14 Jan 2026 16:20:47 -0500 Subject: [PATCH 01/13] Added basic pipeline for MOABB and physionet with new pipeline --- brainsets/moabb_pipeline.py | 193 ++++++++++++ brainsets/taxonomy/task.py | 3 + brainsets_pipelines/physionet_mi/pipeline.py | 282 ++++++++++++++++++ .../schalk_wolpaw_physionet_2009/pipeline.py | 282 ++++++++++++++++++ 4 files changed, 760 insertions(+) create mode 100644 brainsets/moabb_pipeline.py create mode 100644 brainsets_pipelines/physionet_mi/pipeline.py create mode 100644 brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py diff --git a/brainsets/moabb_pipeline.py b/brainsets/moabb_pipeline.py new file mode 100644 index 00000000..2d0a1733 --- /dev/null +++ b/brainsets/moabb_pipeline.py @@ -0,0 +1,193 @@ +"""Base pipeline class for MOABB (Mother of All BCI Benchmarks) datasets. + +This module provides a reusable base class for integrating MOABB datasets into +the brainsets pipeline framework. Subclasses define the dataset and paradigm +classes, and implement paradigm-specific processing logic. +""" + +from abc import abstractmethod +from typing import Dict, Any, Optional, Type +from pathlib import Path +import pandas as pd +import numpy as np + +from moabb.datasets.base import BaseDataset +from moabb.paradigms.base import BaseParadigm +from moabb.utils import set_download_dir +from brainsets.pipeline import BrainsetPipeline + + +class MOABBPipeline(BrainsetPipeline): + """Base class for MOABB dataset pipelines. + + Subclasses must define: + - brainset_id: str + - dataset_class: Type[BaseDataset] + - paradigm_class: Type[BaseParadigm] + - dataset_kwargs: Dict[str, Any] (optional, defaults to {}) + - paradigm_kwargs: Dict[str, Any] (optional, defaults to {}) + + Subclasses must implement: + - process(): Transform MOABB data to brainsets format + + The base class handles: + - Manifest generation from dataset metadata + - Data download via MOABB paradigm.get_data() + - Session filtering + - MNE download directory setup + """ + + dataset_class: Type[BaseDataset] + paradigm_class: Type[BaseParadigm] + dataset_kwargs: Dict[str, Any] = {} + paradigm_kwargs: Dict[str, Any] = {} + + @classmethod + def get_dataset(cls) -> BaseDataset: + """Instantiate the MOABB dataset with configured kwargs.""" + return cls.dataset_class(**cls.dataset_kwargs) + + @classmethod + def get_paradigm(cls) -> BaseParadigm: + """Instantiate the MOABB paradigm with configured kwargs.""" + kwargs = {k: v for k, v in cls.paradigm_kwargs.items() if v is not None} + return cls.paradigm_class(**kwargs) + + @classmethod + def get_manifest(cls, raw_dir: Path, args) -> pd.DataFrame: + """Auto-generate manifest from MOABB dataset metadata. + + Creates a manifest with one row per subject-session combination. + The manifest index is 'session_id' formatted as 'subj-{subject:03d}_sess-{session}'. + + Parameters + ---------- + raw_dir : Path + Raw data directory (unused but required by base class) + args : Namespace + Pipeline arguments (unused but required by base class) + + Returns + ------- + pd.DataFrame + Manifest with columns: subject, session, session_id (as index) + """ + dataset = cls.get_dataset() + manifest_list = [] + + for subject in dataset.subject_list: + for session in range(dataset.n_sessions): + session_id = f"subj-{subject:03d}_sess-{session}" + manifest_list.append( + { + "subject": subject, + "session": session, + "session_id": session_id, + } + ) + + return pd.DataFrame(manifest_list).set_index("session_id") + + def download(self, manifest_item) -> Dict[str, Any]: + """Download and extract data using MOABB paradigm. + + This method: + 1. Sets up MNE download directory + 2. Calls paradigm.get_data() to get epoched arrays + 3. Filters results to the specific session from manifest_item + + Parameters + ---------- + manifest_item : NamedTuple + Row from manifest containing subject and session info + + Returns + ------- + dict + Dictionary containing: + - X: np.ndarray of shape (n_epochs, n_channels, n_samples) + - labels: np.ndarray of shape (n_epochs,) with event names + - meta: pd.DataFrame with columns: subject, session, run + - info: MNE Info object from first epoch (for channel info) + """ + self.update_status("DOWNLOADING") + + set_download_dir(str(self.raw_dir)) + + dataset = self.get_dataset() + paradigm = self.get_paradigm() + + X, labels, meta = paradigm.get_data( + dataset=dataset, + subjects=[manifest_item.subject.item()], + return_epochs=False, + ) + + if len(X) == 0: + raise ValueError( + f"No epochs found for subject {manifest_item.subject}, " + f"session {manifest_item.session}" + ) + + session_values = sorted(meta["session"].unique()) + if isinstance(manifest_item.session, int): + if manifest_item.session < len(session_values): + session_key = session_values[manifest_item.session] + else: + raise ValueError( + f"Session index {manifest_item.session} out of range for subject {manifest_item.subject}. " + f"Available {len(session_values)} sessions: {list(session_values)}" + ) + else: + session_key = str(manifest_item.session.item()) + if session_key not in session_values: + raise ValueError( + f"Session {session_key} not found for subject {manifest_item.subject}. " + f"Available sessions: {list(session_values)}" + ) + + session_mask = meta["session"] == session_key + if not session_mask.any(): + raise ValueError( + f"No epochs found for subject {manifest_item.subject}, " + f"session {session_key}" + ) + + X_filtered = X[session_mask] + labels_filtered = labels[session_mask] + meta_filtered = meta[session_mask].reset_index(drop=True) + + epochs, labels_epochs, meta_epochs = paradigm.get_data( + dataset=dataset, + subjects=[manifest_item.subject.item()], + return_epochs=True, + ) + + session_mask_epochs = meta_epochs["session"] == session_key + epochs_filtered = epochs[session_mask_epochs] + + info = epochs_filtered[0].info if len(epochs_filtered) > 0 else None + + return { + "X": X_filtered, + "labels": labels_filtered, + "meta": meta_filtered, + "info": info, + "epochs": epochs_filtered, + } + + @abstractmethod + def process(self, download_output: Dict[str, Any]) -> None: + """Transform MOABB data to brainsets format. + + Subclasses implement paradigm-specific processing: + - Motor Imagery: extract trials with movement labels + - P300: extract target/non-target epochs + - SSVEP: extract frequency-tagged responses + + Parameters + ---------- + download_output : dict + Dictionary returned by download() containing X, labels, meta, info + """ + ... diff --git a/brainsets/taxonomy/task.py b/brainsets/taxonomy/task.py index afb7be6c..6312ceab 100644 --- a/brainsets/taxonomy/task.py +++ b/brainsets/taxonomy/task.py @@ -23,6 +23,9 @@ class Task(StringIntEnum): # Full sentence speaking CONTINUOUS_SPEAKING_SENTENCE = 7 + # Motor imagery tasks + MOTOR_IMAGERY = 8 + class Stimulus(StringIntEnum): """Stimuli can variously act like inputs (for conditioning) or like outputs.""" diff --git a/brainsets_pipelines/physionet_mi/pipeline.py b/brainsets_pipelines/physionet_mi/pipeline.py new file mode 100644 index 00000000..7831ea56 --- /dev/null +++ b/brainsets_pipelines/physionet_mi/pipeline.py @@ -0,0 +1,282 @@ +# /// brainset-pipeline +# python-version = "3.11" +# dependencies = [ +# "mne==1.11.0", +# "moabb==1.4.3", +# "scikit-learn==1.8.0", +# ] +# /// + +"""Pipeline for PhysionetMI Motor Imagery dataset using MOABB. + +This pipeline downloads and processes EEG motor imagery data from the PhysioNet +dataset using the MOABB dataset loader. The dataset consists of over 1500 one- +and two-minute EEG recordings obtained from 109 volunteers performing motor +imagery tasks. +""" + +from argparse import ArgumentParser +from typing import NamedTuple +import logging +import datetime + +import h5py +import numpy as np + +from moabb.datasets import PhysionetMI +from moabb.paradigms import MotorImagery + +from temporaldata import Data, RegularTimeSeries, Interval, ArrayDict +from brainsets import serialize_fn_map +from brainsets.descriptions import ( + BrainsetDescription, + SessionDescription, + SubjectDescription, + DeviceDescription, +) +from brainsets.taxonomy import Species, Task +from brainsets.moabb_pipeline import MOABBPipeline +from brainsets.utils.split import generate_stratified_folds + + +logging.basicConfig(level=logging.INFO) + +parser = ArgumentParser() +parser.add_argument("--redownload", action="store_true") +parser.add_argument("--reprocess", action="store_true") + + +MOVEMENT_ID_MAP = { + "left_hand": 0, + "right_hand": 1, +} + + +class Pipeline(MOABBPipeline): + brainset_id = "physionet_mi" + parser = parser + + dataset_class = PhysionetMI + paradigm_class = MotorImagery + dataset_kwargs = {"imagined": True, "executed": False} + + def process(self, download_output): + """Process downloaded MOABB data into standardized brainsets format. + + Parameters + ---------- + download_output : dict + Dictionary containing X, labels, meta, info, epochs from download() + """ + X = download_output["X"] + labels = download_output["labels"] + meta = download_output["meta"] + info = download_output["info"] + epochs = download_output["epochs"] + + self.update_status("PROCESSING") + self.processed_dir.mkdir(exist_ok=True, parents=True) + + subject_id = f"S{meta.iloc[0]['subject']:03d}" + session_id = f"{subject_id}_sess-{meta.iloc[0]['session']}" + + store_path = self.processed_dir / f"{session_id}.h5" + if store_path.exists() and not self.args.reprocess: + self.update_status("Skipped Processing") + return + + self.update_status("Creating Descriptions") + brainset_description = BrainsetDescription( + id="physionet_mi", + origin_version="unknown", + derived_version="1.0.0", + source="https://moabb.neurotechx.com/docs/generated/moabb.datasets.PhysionetMI.html", + description="PhysioNet Motor Imagery dataset: over 1500 EEG recordings " + "from 109 volunteers performing motor imagery tasks.", + ) + + subject_description = SubjectDescription( + id=subject_id, + species=Species.HOMO_SAPIENS, + ) + + if info is None: + raise ValueError("No MNE Info object available from epochs") + + recording_date = info.get("meas_date") + if recording_date is None: + recording_date = datetime.datetime.now() + + session_description = SessionDescription( + id=session_id, + recording_date=recording_date, + task=Task.MOTOR_IMAGERY, + ) + + device_description = DeviceDescription( + id=f"{subject_id}_{recording_date.strftime('%Y%m%d')}", + ) + + self.update_status("Extracting EEG") + eeg, units, epoch_intervals = self._extract_eeg_data(X, info, epochs) + + self.update_status("Extracting Trials") + trials = self._extract_motor_imagery_trials(X, labels, info) + + self.update_status("Generating Splits") + folds = generate_stratified_folds( + trials, + stratify_by="movements", + n_folds=5, + val_ratio=0.2, + seed=42, + ) + + folds_dict = {f"fold_{i}": fold for i, fold in enumerate(folds)} + splits = Data(**folds_dict, domain=trials) + + self.update_status("Creating Data Object") + data = Data( + brainset=brainset_description, + subject=subject_description, + session=session_description, + device=device_description, + eeg=eeg, + units=units, + motor_imagery_trials=trials, + splits=splits, + domain=eeg.domain, + ) + + self.update_status("Storing") + with h5py.File(store_path, "w") as file: + data.to_hdf5(file, serialize_fn_map=serialize_fn_map) + + logging.info(f"Saved processed data to: {store_path}") + + def _extract_eeg_data(self, X, info, epochs): + """Extract EEG data from epoched arrays. + + Parameters + ---------- + X : np.ndarray + Array of shape (n_epochs, n_channels, n_samples) + info : mne.Info + MNE Info object with channel and sampling rate information + epochs : mne.Epochs + MNE Epochs object (used to get channel types) + + Returns + ------- + eeg : RegularTimeSeries + Concatenated EEG signals + units : ArrayDict + Channel IDs and types + epoch_intervals : Interval + Time intervals for each epoch + """ + sfreq = info["sfreq"] + n_epochs, n_channels, n_samples = X.shape + + eeg_signals = np.concatenate([X[i].T for i in range(n_epochs)], axis=0) + + epoch_starts = [] + epoch_ends = [] + current_time = 0.0 + + for i in range(n_epochs): + epoch_duration = n_samples / sfreq + epoch_starts.append(current_time) + epoch_ends.append(current_time + epoch_duration) + current_time += epoch_duration + + eeg = RegularTimeSeries( + signal=eeg_signals, + sampling_rate=sfreq, + domain=Interval( + start=np.array([0.0]), + end=np.array([(len(eeg_signals) - 1) / sfreq]), + ), + ) + + ch_names = info["ch_names"] + if len(epochs) > 0: + ch_types = epochs[0].get_channel_types() + else: + ch_types = [] + for ch_name in ch_names: + ch_idx = info["ch_names"].index(ch_name) + ch_kind = info["chs"][ch_idx]["kind"] + ch_type_map = { + 2: "EEG", + 3: "EOG", + 4: "EMG", + 5: "ECG", + 301: "MISC", + } + ch_types.append(ch_type_map.get(ch_kind, "MISC")) + + units = ArrayDict( + id=np.array(ch_names, dtype="U"), + types=np.array(ch_types, dtype="U"), + ) + + epoch_intervals = Interval( + start=np.array(epoch_starts), + end=np.array(epoch_ends), + ) + + return eeg, units, epoch_intervals + + def _extract_motor_imagery_trials(self, X, labels, info): + """Extract motor imagery trial intervals with movement labels. + + Parameters + ---------- + X : np.ndarray + Array of shape (n_epochs, n_channels, n_samples) + labels : np.ndarray + Array of shape (n_epochs,) with movement labels + info : mne.Info + MNE Info object with sampling rate information + + Returns + ------- + trials : Interval + Interval object with start/end times and movement labels + """ + sfreq = info["sfreq"] + n_epochs, _, n_samples = X.shape + + start_times = [] + end_times = [] + movements = [] + movement_ids = [] + + current_time = 0.0 + + for i in range(n_epochs): + epoch_duration = n_samples / sfreq + + start_times.append(current_time) + end_times.append(current_time + epoch_duration) + + label = labels[i] + movements.append(label) + movement_ids.append(MOVEMENT_ID_MAP.get(label, -1)) + + current_time += epoch_duration + + trials = Interval( + start=np.array(start_times), + end=np.array(end_times), + timestamps=(np.array(start_times) + np.array(end_times)) / 2, + movements=np.array(movements), + movement_ids=np.array(movement_ids), + timekeys=["start", "end", "timestamps"], + ) + + if not trials.is_disjoint(): + raise ValueError("Found overlapping trials") + + return trials diff --git a/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py b/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py new file mode 100644 index 00000000..7831ea56 --- /dev/null +++ b/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py @@ -0,0 +1,282 @@ +# /// brainset-pipeline +# python-version = "3.11" +# dependencies = [ +# "mne==1.11.0", +# "moabb==1.4.3", +# "scikit-learn==1.8.0", +# ] +# /// + +"""Pipeline for PhysionetMI Motor Imagery dataset using MOABB. + +This pipeline downloads and processes EEG motor imagery data from the PhysioNet +dataset using the MOABB dataset loader. The dataset consists of over 1500 one- +and two-minute EEG recordings obtained from 109 volunteers performing motor +imagery tasks. +""" + +from argparse import ArgumentParser +from typing import NamedTuple +import logging +import datetime + +import h5py +import numpy as np + +from moabb.datasets import PhysionetMI +from moabb.paradigms import MotorImagery + +from temporaldata import Data, RegularTimeSeries, Interval, ArrayDict +from brainsets import serialize_fn_map +from brainsets.descriptions import ( + BrainsetDescription, + SessionDescription, + SubjectDescription, + DeviceDescription, +) +from brainsets.taxonomy import Species, Task +from brainsets.moabb_pipeline import MOABBPipeline +from brainsets.utils.split import generate_stratified_folds + + +logging.basicConfig(level=logging.INFO) + +parser = ArgumentParser() +parser.add_argument("--redownload", action="store_true") +parser.add_argument("--reprocess", action="store_true") + + +MOVEMENT_ID_MAP = { + "left_hand": 0, + "right_hand": 1, +} + + +class Pipeline(MOABBPipeline): + brainset_id = "physionet_mi" + parser = parser + + dataset_class = PhysionetMI + paradigm_class = MotorImagery + dataset_kwargs = {"imagined": True, "executed": False} + + def process(self, download_output): + """Process downloaded MOABB data into standardized brainsets format. + + Parameters + ---------- + download_output : dict + Dictionary containing X, labels, meta, info, epochs from download() + """ + X = download_output["X"] + labels = download_output["labels"] + meta = download_output["meta"] + info = download_output["info"] + epochs = download_output["epochs"] + + self.update_status("PROCESSING") + self.processed_dir.mkdir(exist_ok=True, parents=True) + + subject_id = f"S{meta.iloc[0]['subject']:03d}" + session_id = f"{subject_id}_sess-{meta.iloc[0]['session']}" + + store_path = self.processed_dir / f"{session_id}.h5" + if store_path.exists() and not self.args.reprocess: + self.update_status("Skipped Processing") + return + + self.update_status("Creating Descriptions") + brainset_description = BrainsetDescription( + id="physionet_mi", + origin_version="unknown", + derived_version="1.0.0", + source="https://moabb.neurotechx.com/docs/generated/moabb.datasets.PhysionetMI.html", + description="PhysioNet Motor Imagery dataset: over 1500 EEG recordings " + "from 109 volunteers performing motor imagery tasks.", + ) + + subject_description = SubjectDescription( + id=subject_id, + species=Species.HOMO_SAPIENS, + ) + + if info is None: + raise ValueError("No MNE Info object available from epochs") + + recording_date = info.get("meas_date") + if recording_date is None: + recording_date = datetime.datetime.now() + + session_description = SessionDescription( + id=session_id, + recording_date=recording_date, + task=Task.MOTOR_IMAGERY, + ) + + device_description = DeviceDescription( + id=f"{subject_id}_{recording_date.strftime('%Y%m%d')}", + ) + + self.update_status("Extracting EEG") + eeg, units, epoch_intervals = self._extract_eeg_data(X, info, epochs) + + self.update_status("Extracting Trials") + trials = self._extract_motor_imagery_trials(X, labels, info) + + self.update_status("Generating Splits") + folds = generate_stratified_folds( + trials, + stratify_by="movements", + n_folds=5, + val_ratio=0.2, + seed=42, + ) + + folds_dict = {f"fold_{i}": fold for i, fold in enumerate(folds)} + splits = Data(**folds_dict, domain=trials) + + self.update_status("Creating Data Object") + data = Data( + brainset=brainset_description, + subject=subject_description, + session=session_description, + device=device_description, + eeg=eeg, + units=units, + motor_imagery_trials=trials, + splits=splits, + domain=eeg.domain, + ) + + self.update_status("Storing") + with h5py.File(store_path, "w") as file: + data.to_hdf5(file, serialize_fn_map=serialize_fn_map) + + logging.info(f"Saved processed data to: {store_path}") + + def _extract_eeg_data(self, X, info, epochs): + """Extract EEG data from epoched arrays. + + Parameters + ---------- + X : np.ndarray + Array of shape (n_epochs, n_channels, n_samples) + info : mne.Info + MNE Info object with channel and sampling rate information + epochs : mne.Epochs + MNE Epochs object (used to get channel types) + + Returns + ------- + eeg : RegularTimeSeries + Concatenated EEG signals + units : ArrayDict + Channel IDs and types + epoch_intervals : Interval + Time intervals for each epoch + """ + sfreq = info["sfreq"] + n_epochs, n_channels, n_samples = X.shape + + eeg_signals = np.concatenate([X[i].T for i in range(n_epochs)], axis=0) + + epoch_starts = [] + epoch_ends = [] + current_time = 0.0 + + for i in range(n_epochs): + epoch_duration = n_samples / sfreq + epoch_starts.append(current_time) + epoch_ends.append(current_time + epoch_duration) + current_time += epoch_duration + + eeg = RegularTimeSeries( + signal=eeg_signals, + sampling_rate=sfreq, + domain=Interval( + start=np.array([0.0]), + end=np.array([(len(eeg_signals) - 1) / sfreq]), + ), + ) + + ch_names = info["ch_names"] + if len(epochs) > 0: + ch_types = epochs[0].get_channel_types() + else: + ch_types = [] + for ch_name in ch_names: + ch_idx = info["ch_names"].index(ch_name) + ch_kind = info["chs"][ch_idx]["kind"] + ch_type_map = { + 2: "EEG", + 3: "EOG", + 4: "EMG", + 5: "ECG", + 301: "MISC", + } + ch_types.append(ch_type_map.get(ch_kind, "MISC")) + + units = ArrayDict( + id=np.array(ch_names, dtype="U"), + types=np.array(ch_types, dtype="U"), + ) + + epoch_intervals = Interval( + start=np.array(epoch_starts), + end=np.array(epoch_ends), + ) + + return eeg, units, epoch_intervals + + def _extract_motor_imagery_trials(self, X, labels, info): + """Extract motor imagery trial intervals with movement labels. + + Parameters + ---------- + X : np.ndarray + Array of shape (n_epochs, n_channels, n_samples) + labels : np.ndarray + Array of shape (n_epochs,) with movement labels + info : mne.Info + MNE Info object with sampling rate information + + Returns + ------- + trials : Interval + Interval object with start/end times and movement labels + """ + sfreq = info["sfreq"] + n_epochs, _, n_samples = X.shape + + start_times = [] + end_times = [] + movements = [] + movement_ids = [] + + current_time = 0.0 + + for i in range(n_epochs): + epoch_duration = n_samples / sfreq + + start_times.append(current_time) + end_times.append(current_time + epoch_duration) + + label = labels[i] + movements.append(label) + movement_ids.append(MOVEMENT_ID_MAP.get(label, -1)) + + current_time += epoch_duration + + trials = Interval( + start=np.array(start_times), + end=np.array(end_times), + timestamps=(np.array(start_times) + np.array(end_times)) / 2, + movements=np.array(movements), + movement_ids=np.array(movement_ids), + timekeys=["start", "end", "timestamps"], + ) + + if not trials.is_disjoint(): + raise ValueError("Found overlapping trials") + + return trials From 873a031d9d7e1da3f1ab2e9bf00132d4c46557f2 Mon Sep 17 00:00:00 2001 From: Milo Sobral Date: Wed, 14 Jan 2026 16:56:10 -0500 Subject: [PATCH 02/13] Renaming the pipelines --- brainsets/moabb_pipeline.py | 333 +++++++++++++++++- brainsets/taxonomy/task.py | 3 + .../pipeline.py | 63 ++++ brainsets_pipelines/physionet_mi/pipeline.py | 282 --------------- .../schalk_wolpaw_physionet_2009/pipeline.py | 255 +------------- 5 files changed, 410 insertions(+), 526 deletions(-) create mode 100644 brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py delete mode 100644 brainsets_pipelines/physionet_mi/pipeline.py diff --git a/brainsets/moabb_pipeline.py b/brainsets/moabb_pipeline.py index 2d0a1733..e55847fc 100644 --- a/brainsets/moabb_pipeline.py +++ b/brainsets/moabb_pipeline.py @@ -6,15 +6,28 @@ """ from abc import abstractmethod -from typing import Dict, Any, Optional, Type +from typing import Dict, Any, Type from pathlib import Path import pandas as pd import numpy as np +import datetime +import logging +import h5py from moabb.datasets.base import BaseDataset from moabb.paradigms.base import BaseParadigm from moabb.utils import set_download_dir +from temporaldata import Data, RegularTimeSeries, Interval, ArrayDict from brainsets.pipeline import BrainsetPipeline +from brainsets import serialize_fn_map +from brainsets.descriptions import ( + BrainsetDescription, + SessionDescription, + SubjectDescription, + DeviceDescription, +) +from brainsets.taxonomy import Species, Task +from brainsets.utils.split import generate_stratified_folds class MOABBPipeline(BrainsetPipeline): @@ -26,15 +39,22 @@ class MOABBPipeline(BrainsetPipeline): - paradigm_class: Type[BaseParadigm] - dataset_kwargs: Dict[str, Any] (optional, defaults to {}) - paradigm_kwargs: Dict[str, Any] (optional, defaults to {}) + - task: Task (e.g., Task.MOTOR_IMAGERY, Task.P300) + - trial_key: str (e.g., "motor_imagery_trials", "p300_trials") + - label_field: str (e.g., "movements", "targets") + - id_field: str (e.g., "movement_ids", "target_ids") + - stratify_field: str (e.g., "movements", "targets") + - label_map: Dict[str, int] (mapping from label strings to integer IDs) Subclasses must implement: - - process(): Transform MOABB data to brainsets format + - get_brainset_description(): Return dataset-specific BrainsetDescription The base class handles: - Manifest generation from dataset metadata - Data download via MOABB paradigm.get_data() - Session filtering - MNE download directory setup + - Default process() workflow (EEG extraction, trial extraction, splits, storage) """ dataset_class: Type[BaseDataset] @@ -42,6 +62,13 @@ class MOABBPipeline(BrainsetPipeline): dataset_kwargs: Dict[str, Any] = {} paradigm_kwargs: Dict[str, Any] = {} + task: Task + trial_key: str + label_field: str + id_field: str + stratify_field: str + label_map: Dict[str, int] + @classmethod def get_dataset(cls) -> BaseDataset: """Instantiate the MOABB dataset with configured kwargs.""" @@ -176,18 +203,308 @@ def download(self, manifest_item) -> Dict[str, Any]: "epochs": epochs_filtered, } + def _extract_eeg_data(self, X, info, epochs): + """Extract EEG data from epoched arrays. + + Parameters + ---------- + X : np.ndarray + Array of shape (n_epochs, n_channels, n_samples) + info : mne.Info + MNE Info object with channel and sampling rate information + epochs : mne.Epochs + MNE Epochs object (used to get channel types) + + Returns + ------- + eeg : RegularTimeSeries + Concatenated EEG signals + units : ArrayDict + Channel IDs and types + epoch_intervals : Interval + Time intervals for each epoch + """ + sfreq = info["sfreq"] + n_epochs, n_channels, n_samples = X.shape + + eeg_signals = np.concatenate([X[i].T for i in range(n_epochs)], axis=0) + + epoch_starts = [] + epoch_ends = [] + current_time = 0.0 + + for i in range(n_epochs): + epoch_duration = n_samples / sfreq + epoch_starts.append(current_time) + epoch_ends.append(current_time + epoch_duration) + current_time += epoch_duration + + eeg = RegularTimeSeries( + signal=eeg_signals, + sampling_rate=sfreq, + domain=Interval( + start=np.array([0.0]), + end=np.array([(len(eeg_signals) - 1) / sfreq]), + ), + ) + + ch_names = info["ch_names"] + if len(epochs) > 0: + ch_types = epochs[0].get_channel_types() + else: + ch_types = [] + for ch_name in ch_names: + ch_idx = info["ch_names"].index(ch_name) + ch_kind = info["chs"][ch_idx]["kind"] + ch_type_map = { + 2: "EEG", + 3: "EOG", + 4: "EMG", + 5: "ECG", + 301: "MISC", + } + ch_types.append(ch_type_map.get(ch_kind, "MISC")) + + units = ArrayDict( + id=np.array(ch_names, dtype="U"), + types=np.array(ch_types, dtype="U"), + ) + + epoch_intervals = Interval( + start=np.array(epoch_starts), + end=np.array(epoch_ends), + ) + + return eeg, units, epoch_intervals + @abstractmethod + def get_brainset_description(self) -> BrainsetDescription: + """Return dataset-specific BrainsetDescription. + + Returns + ------- + BrainsetDescription + Description object with dataset metadata + """ + ... + + def _get_subject_description(self, subject_id: str) -> SubjectDescription: + """Create subject description from subject ID. + + Parameters + ---------- + subject_id : str + Subject identifier + + Returns + ------- + SubjectDescription + Subject description object + """ + return SubjectDescription( + id=subject_id, + species=Species.HOMO_SAPIENS, + ) + + def _get_session_description( + self, session_id: str, info: Any + ) -> SessionDescription: + """Create session description from session ID and MNE info. + + Parameters + ---------- + session_id : str + Session identifier + info : mne.Info + MNE Info object with recording date + + Returns + ------- + SessionDescription + Session description object + """ + if info is None: + raise ValueError("No MNE Info object available from epochs") + + recording_date = info.get("meas_date") + if recording_date is None: + recording_date = datetime.datetime.now() + + return SessionDescription( + id=session_id, + recording_date=recording_date, + task=self.task, + ) + + def _get_device_description(self, subject_id: str, info: Any) -> DeviceDescription: + """Create device description from subject ID and MNE info. + + Parameters + ---------- + subject_id : str + Subject identifier + info : mne.Info + MNE Info object with recording date + + Returns + ------- + DeviceDescription + Device description object + """ + if info is None: + raise ValueError("No MNE Info object available from epochs") + + recording_date = info.get("meas_date") + if recording_date is None: + recording_date = datetime.datetime.now() + + return DeviceDescription( + id=f"{subject_id}_{recording_date.strftime('%Y%m%d')}", + ) + + def _extract_trials(self, X, labels, info): + """Extract trial intervals with labels using configured field names. + + Parameters + ---------- + X : np.ndarray + Array of shape (n_epochs, n_channels, n_samples) + labels : np.ndarray + Array of shape (n_epochs,) with event labels + info : mne.Info + MNE Info object with sampling rate information + + Returns + ------- + trials : Interval + Interval object with start/end times and configured label fields + """ + sfreq = info["sfreq"] + n_epochs, _, n_samples = X.shape + + start_times = [] + end_times = [] + label_values = [] + id_values = [] + + current_time = 0.0 + + for i in range(n_epochs): + epoch_duration = n_samples / sfreq + + start_times.append(current_time) + end_times.append(current_time + epoch_duration) + + label = labels[i] + label_values.append(label) + id_values.append(self.label_map.get(label, -1)) + + current_time += epoch_duration + + trial_kwargs = { + "start": np.array(start_times), + "end": np.array(end_times), + "timestamps": (np.array(start_times) + np.array(end_times)) / 2, + self.label_field: np.array(label_values), + self.id_field: np.array(id_values), + "timekeys": ["start", "end", "timestamps"], + } + + trials = Interval(**trial_kwargs) + + if not trials.is_disjoint(): + raise ValueError("Found overlapping trials") + + return trials + + def _generate_splits(self, trials): + """Generate stratified folds for trials. + + Parameters + ---------- + trials : Interval + Trial intervals with label fields + + Returns + ------- + splits : Data + Data object containing fold splits + """ + folds = generate_stratified_folds( + trials, + stratify_by=self.stratify_field, + n_folds=5, + val_ratio=0.2, + seed=42, + ) + + folds_dict = {f"fold_{i}": fold for i, fold in enumerate(folds)} + return Data(**folds_dict, domain=trials) + def process(self, download_output: Dict[str, Any]) -> None: """Transform MOABB data to brainsets format. - Subclasses implement paradigm-specific processing: - - Motor Imagery: extract trials with movement labels - - P300: extract target/non-target epochs - - SSVEP: extract frequency-tagged responses + This default implementation handles the common workflow: + 1. Extract EEG data and channel information + 2. Extract trial intervals with labels + 3. Generate stratified splits + 4. Create and store Data object + + Subclasses can override for custom processing, but typically only + need to implement get_brainset_description(). Parameters ---------- download_output : dict - Dictionary returned by download() containing X, labels, meta, info + Dictionary returned by download() containing X, labels, meta, info, epochs """ - ... + X = download_output["X"] + labels = download_output["labels"] + meta = download_output["meta"] + info = download_output["info"] + epochs = download_output["epochs"] + + self.update_status("PROCESSING") + self.processed_dir.mkdir(exist_ok=True, parents=True) + + subject_id = f"S{meta.iloc[0]['subject']:03d}" + session_id = f"{subject_id}_sess-{meta.iloc[0]['session']}" + + store_path = self.processed_dir / f"{session_id}.h5" + if store_path.exists() and not self.args.reprocess: + self.update_status("Skipped Processing") + return + + self.update_status("Creating Descriptions") + brainset_description = self.get_brainset_description() + subject_description = self._get_subject_description(subject_id) + session_description = self._get_session_description(session_id, info) + device_description = self._get_device_description(subject_id, info) + + self.update_status("Extracting EEG") + eeg, units, _ = self._extract_eeg_data(X, info, epochs) + + self.update_status("Extracting Trials") + trials = self._extract_trials(X, labels, info) + + self.update_status("Generating Splits") + splits = self._generate_splits(trials) + + self.update_status("Creating Data Object") + data = Data( + brainset=brainset_description, + subject=subject_description, + session=session_description, + device=device_description, + eeg=eeg, + units=units, + **{self.trial_key: trials}, + splits=splits, + domain=eeg.domain, + ) + + self.update_status("Storing") + with h5py.File(store_path, "w") as file: + data.to_hdf5(file, serialize_fn_map=serialize_fn_map) + + logging.info(f"Saved processed data to: {store_path}") diff --git a/brainsets/taxonomy/task.py b/brainsets/taxonomy/task.py index 6312ceab..dfe4dcf4 100644 --- a/brainsets/taxonomy/task.py +++ b/brainsets/taxonomy/task.py @@ -26,6 +26,9 @@ class Task(StringIntEnum): # Motor imagery tasks MOTOR_IMAGERY = 8 + # P300 event-related potential tasks + P300 = 9 + class Stimulus(StringIntEnum): """Stimuli can variously act like inputs (for conditioning) or like outputs.""" diff --git a/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py b/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py new file mode 100644 index 00000000..34d3bbfc --- /dev/null +++ b/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py @@ -0,0 +1,63 @@ +# /// brainset-pipeline +# python-version = "3.11" +# dependencies = [ +# "mne==1.11.0", +# "moabb==1.4.3", +# "scikit-learn==1.8.0", +# ] +# /// + +"""Pipeline for BI2014a P300 dataset using MOABB. + +This pipeline downloads and processes EEG P300 data from the Brain Invaders 2014a +dataset using the MOABB dataset loader. The dataset consists of EEG recordings +from 71 subjects performing a visual P300 Brain-Computer Interface task using +16 active dry electrodes across up to 3 sessions. +""" + +from argparse import ArgumentParser +import logging + +from moabb.datasets import BI2014a +from moabb.paradigms import P300 + +from brainsets.descriptions import BrainsetDescription +from brainsets.taxonomy import Task +from brainsets.moabb_pipeline import MOABBPipeline + + +logging.basicConfig(level=logging.INFO) + +parser = ArgumentParser() +parser.add_argument("--redownload", action="store_true") +parser.add_argument("--reprocess", action="store_true") + + +class Pipeline(MOABBPipeline): + brainset_id = "korczowski_brain_invaders_2014a" + parser = parser + + dataset_class = BI2014a + paradigm_class = P300 + dataset_kwargs = {} + + task = Task.P300 + trial_key = "p300_trials" + label_field = "targets" + id_field = "target_ids" + stratify_field = "targets" + label_map = { + "Target": 1, + "NonTarget": 0, + } + + def get_brainset_description(self): + return BrainsetDescription( + id="korczowski_brain_invaders_2014a", + origin_version="unknown", + derived_version="1.0.0", + source="https://moabb.neurotechx.com/docs/generated/moabb.datasets.BI2014a.html", + description="Brain Invaders 2014a P300 dataset: EEG recordings from " + "71 subjects performing a visual P300 Brain-Computer Interface task " + "using 16 active dry electrodes.", + ) diff --git a/brainsets_pipelines/physionet_mi/pipeline.py b/brainsets_pipelines/physionet_mi/pipeline.py deleted file mode 100644 index 7831ea56..00000000 --- a/brainsets_pipelines/physionet_mi/pipeline.py +++ /dev/null @@ -1,282 +0,0 @@ -# /// brainset-pipeline -# python-version = "3.11" -# dependencies = [ -# "mne==1.11.0", -# "moabb==1.4.3", -# "scikit-learn==1.8.0", -# ] -# /// - -"""Pipeline for PhysionetMI Motor Imagery dataset using MOABB. - -This pipeline downloads and processes EEG motor imagery data from the PhysioNet -dataset using the MOABB dataset loader. The dataset consists of over 1500 one- -and two-minute EEG recordings obtained from 109 volunteers performing motor -imagery tasks. -""" - -from argparse import ArgumentParser -from typing import NamedTuple -import logging -import datetime - -import h5py -import numpy as np - -from moabb.datasets import PhysionetMI -from moabb.paradigms import MotorImagery - -from temporaldata import Data, RegularTimeSeries, Interval, ArrayDict -from brainsets import serialize_fn_map -from brainsets.descriptions import ( - BrainsetDescription, - SessionDescription, - SubjectDescription, - DeviceDescription, -) -from brainsets.taxonomy import Species, Task -from brainsets.moabb_pipeline import MOABBPipeline -from brainsets.utils.split import generate_stratified_folds - - -logging.basicConfig(level=logging.INFO) - -parser = ArgumentParser() -parser.add_argument("--redownload", action="store_true") -parser.add_argument("--reprocess", action="store_true") - - -MOVEMENT_ID_MAP = { - "left_hand": 0, - "right_hand": 1, -} - - -class Pipeline(MOABBPipeline): - brainset_id = "physionet_mi" - parser = parser - - dataset_class = PhysionetMI - paradigm_class = MotorImagery - dataset_kwargs = {"imagined": True, "executed": False} - - def process(self, download_output): - """Process downloaded MOABB data into standardized brainsets format. - - Parameters - ---------- - download_output : dict - Dictionary containing X, labels, meta, info, epochs from download() - """ - X = download_output["X"] - labels = download_output["labels"] - meta = download_output["meta"] - info = download_output["info"] - epochs = download_output["epochs"] - - self.update_status("PROCESSING") - self.processed_dir.mkdir(exist_ok=True, parents=True) - - subject_id = f"S{meta.iloc[0]['subject']:03d}" - session_id = f"{subject_id}_sess-{meta.iloc[0]['session']}" - - store_path = self.processed_dir / f"{session_id}.h5" - if store_path.exists() and not self.args.reprocess: - self.update_status("Skipped Processing") - return - - self.update_status("Creating Descriptions") - brainset_description = BrainsetDescription( - id="physionet_mi", - origin_version="unknown", - derived_version="1.0.0", - source="https://moabb.neurotechx.com/docs/generated/moabb.datasets.PhysionetMI.html", - description="PhysioNet Motor Imagery dataset: over 1500 EEG recordings " - "from 109 volunteers performing motor imagery tasks.", - ) - - subject_description = SubjectDescription( - id=subject_id, - species=Species.HOMO_SAPIENS, - ) - - if info is None: - raise ValueError("No MNE Info object available from epochs") - - recording_date = info.get("meas_date") - if recording_date is None: - recording_date = datetime.datetime.now() - - session_description = SessionDescription( - id=session_id, - recording_date=recording_date, - task=Task.MOTOR_IMAGERY, - ) - - device_description = DeviceDescription( - id=f"{subject_id}_{recording_date.strftime('%Y%m%d')}", - ) - - self.update_status("Extracting EEG") - eeg, units, epoch_intervals = self._extract_eeg_data(X, info, epochs) - - self.update_status("Extracting Trials") - trials = self._extract_motor_imagery_trials(X, labels, info) - - self.update_status("Generating Splits") - folds = generate_stratified_folds( - trials, - stratify_by="movements", - n_folds=5, - val_ratio=0.2, - seed=42, - ) - - folds_dict = {f"fold_{i}": fold for i, fold in enumerate(folds)} - splits = Data(**folds_dict, domain=trials) - - self.update_status("Creating Data Object") - data = Data( - brainset=brainset_description, - subject=subject_description, - session=session_description, - device=device_description, - eeg=eeg, - units=units, - motor_imagery_trials=trials, - splits=splits, - domain=eeg.domain, - ) - - self.update_status("Storing") - with h5py.File(store_path, "w") as file: - data.to_hdf5(file, serialize_fn_map=serialize_fn_map) - - logging.info(f"Saved processed data to: {store_path}") - - def _extract_eeg_data(self, X, info, epochs): - """Extract EEG data from epoched arrays. - - Parameters - ---------- - X : np.ndarray - Array of shape (n_epochs, n_channels, n_samples) - info : mne.Info - MNE Info object with channel and sampling rate information - epochs : mne.Epochs - MNE Epochs object (used to get channel types) - - Returns - ------- - eeg : RegularTimeSeries - Concatenated EEG signals - units : ArrayDict - Channel IDs and types - epoch_intervals : Interval - Time intervals for each epoch - """ - sfreq = info["sfreq"] - n_epochs, n_channels, n_samples = X.shape - - eeg_signals = np.concatenate([X[i].T for i in range(n_epochs)], axis=0) - - epoch_starts = [] - epoch_ends = [] - current_time = 0.0 - - for i in range(n_epochs): - epoch_duration = n_samples / sfreq - epoch_starts.append(current_time) - epoch_ends.append(current_time + epoch_duration) - current_time += epoch_duration - - eeg = RegularTimeSeries( - signal=eeg_signals, - sampling_rate=sfreq, - domain=Interval( - start=np.array([0.0]), - end=np.array([(len(eeg_signals) - 1) / sfreq]), - ), - ) - - ch_names = info["ch_names"] - if len(epochs) > 0: - ch_types = epochs[0].get_channel_types() - else: - ch_types = [] - for ch_name in ch_names: - ch_idx = info["ch_names"].index(ch_name) - ch_kind = info["chs"][ch_idx]["kind"] - ch_type_map = { - 2: "EEG", - 3: "EOG", - 4: "EMG", - 5: "ECG", - 301: "MISC", - } - ch_types.append(ch_type_map.get(ch_kind, "MISC")) - - units = ArrayDict( - id=np.array(ch_names, dtype="U"), - types=np.array(ch_types, dtype="U"), - ) - - epoch_intervals = Interval( - start=np.array(epoch_starts), - end=np.array(epoch_ends), - ) - - return eeg, units, epoch_intervals - - def _extract_motor_imagery_trials(self, X, labels, info): - """Extract motor imagery trial intervals with movement labels. - - Parameters - ---------- - X : np.ndarray - Array of shape (n_epochs, n_channels, n_samples) - labels : np.ndarray - Array of shape (n_epochs,) with movement labels - info : mne.Info - MNE Info object with sampling rate information - - Returns - ------- - trials : Interval - Interval object with start/end times and movement labels - """ - sfreq = info["sfreq"] - n_epochs, _, n_samples = X.shape - - start_times = [] - end_times = [] - movements = [] - movement_ids = [] - - current_time = 0.0 - - for i in range(n_epochs): - epoch_duration = n_samples / sfreq - - start_times.append(current_time) - end_times.append(current_time + epoch_duration) - - label = labels[i] - movements.append(label) - movement_ids.append(MOVEMENT_ID_MAP.get(label, -1)) - - current_time += epoch_duration - - trials = Interval( - start=np.array(start_times), - end=np.array(end_times), - timestamps=(np.array(start_times) + np.array(end_times)) / 2, - movements=np.array(movements), - movement_ids=np.array(movement_ids), - timekeys=["start", "end", "timestamps"], - ) - - if not trials.is_disjoint(): - raise ValueError("Found overlapping trials") - - return trials diff --git a/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py b/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py index 7831ea56..58569705 100644 --- a/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py +++ b/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py @@ -16,27 +16,14 @@ """ from argparse import ArgumentParser -from typing import NamedTuple import logging -import datetime - -import h5py -import numpy as np from moabb.datasets import PhysionetMI from moabb.paradigms import MotorImagery -from temporaldata import Data, RegularTimeSeries, Interval, ArrayDict -from brainsets import serialize_fn_map -from brainsets.descriptions import ( - BrainsetDescription, - SessionDescription, - SubjectDescription, - DeviceDescription, -) -from brainsets.taxonomy import Species, Task +from brainsets.descriptions import BrainsetDescription +from brainsets.taxonomy import Task from brainsets.moabb_pipeline import MOABBPipeline -from brainsets.utils.split import generate_stratified_folds logging.basicConfig(level=logging.INFO) @@ -46,237 +33,33 @@ parser.add_argument("--reprocess", action="store_true") -MOVEMENT_ID_MAP = { - "left_hand": 0, - "right_hand": 1, -} - - class Pipeline(MOABBPipeline): - brainset_id = "physionet_mi" + brainset_id = "schalk_wolpaw_physionet_2009" parser = parser dataset_class = PhysionetMI paradigm_class = MotorImagery dataset_kwargs = {"imagined": True, "executed": False} - def process(self, download_output): - """Process downloaded MOABB data into standardized brainsets format. - - Parameters - ---------- - download_output : dict - Dictionary containing X, labels, meta, info, epochs from download() - """ - X = download_output["X"] - labels = download_output["labels"] - meta = download_output["meta"] - info = download_output["info"] - epochs = download_output["epochs"] - - self.update_status("PROCESSING") - self.processed_dir.mkdir(exist_ok=True, parents=True) - - subject_id = f"S{meta.iloc[0]['subject']:03d}" - session_id = f"{subject_id}_sess-{meta.iloc[0]['session']}" - - store_path = self.processed_dir / f"{session_id}.h5" - if store_path.exists() and not self.args.reprocess: - self.update_status("Skipped Processing") - return - - self.update_status("Creating Descriptions") - brainset_description = BrainsetDescription( - id="physionet_mi", + task = Task.MOTOR_IMAGERY + trial_key = "motor_imagery_trials" + label_field = "movements" + id_field = "movement_ids" + stratify_field = "movements" + label_map = { + "left_hand": 0, + "right_hand": 1, + "hands": 2, + "feet": 3, + "rest": 4, + } + + def get_brainset_description(self): + return BrainsetDescription( + id="schalk_wolpaw_physionet_2009", origin_version="unknown", derived_version="1.0.0", source="https://moabb.neurotechx.com/docs/generated/moabb.datasets.PhysionetMI.html", description="PhysioNet Motor Imagery dataset: over 1500 EEG recordings " "from 109 volunteers performing motor imagery tasks.", ) - - subject_description = SubjectDescription( - id=subject_id, - species=Species.HOMO_SAPIENS, - ) - - if info is None: - raise ValueError("No MNE Info object available from epochs") - - recording_date = info.get("meas_date") - if recording_date is None: - recording_date = datetime.datetime.now() - - session_description = SessionDescription( - id=session_id, - recording_date=recording_date, - task=Task.MOTOR_IMAGERY, - ) - - device_description = DeviceDescription( - id=f"{subject_id}_{recording_date.strftime('%Y%m%d')}", - ) - - self.update_status("Extracting EEG") - eeg, units, epoch_intervals = self._extract_eeg_data(X, info, epochs) - - self.update_status("Extracting Trials") - trials = self._extract_motor_imagery_trials(X, labels, info) - - self.update_status("Generating Splits") - folds = generate_stratified_folds( - trials, - stratify_by="movements", - n_folds=5, - val_ratio=0.2, - seed=42, - ) - - folds_dict = {f"fold_{i}": fold for i, fold in enumerate(folds)} - splits = Data(**folds_dict, domain=trials) - - self.update_status("Creating Data Object") - data = Data( - brainset=brainset_description, - subject=subject_description, - session=session_description, - device=device_description, - eeg=eeg, - units=units, - motor_imagery_trials=trials, - splits=splits, - domain=eeg.domain, - ) - - self.update_status("Storing") - with h5py.File(store_path, "w") as file: - data.to_hdf5(file, serialize_fn_map=serialize_fn_map) - - logging.info(f"Saved processed data to: {store_path}") - - def _extract_eeg_data(self, X, info, epochs): - """Extract EEG data from epoched arrays. - - Parameters - ---------- - X : np.ndarray - Array of shape (n_epochs, n_channels, n_samples) - info : mne.Info - MNE Info object with channel and sampling rate information - epochs : mne.Epochs - MNE Epochs object (used to get channel types) - - Returns - ------- - eeg : RegularTimeSeries - Concatenated EEG signals - units : ArrayDict - Channel IDs and types - epoch_intervals : Interval - Time intervals for each epoch - """ - sfreq = info["sfreq"] - n_epochs, n_channels, n_samples = X.shape - - eeg_signals = np.concatenate([X[i].T for i in range(n_epochs)], axis=0) - - epoch_starts = [] - epoch_ends = [] - current_time = 0.0 - - for i in range(n_epochs): - epoch_duration = n_samples / sfreq - epoch_starts.append(current_time) - epoch_ends.append(current_time + epoch_duration) - current_time += epoch_duration - - eeg = RegularTimeSeries( - signal=eeg_signals, - sampling_rate=sfreq, - domain=Interval( - start=np.array([0.0]), - end=np.array([(len(eeg_signals) - 1) / sfreq]), - ), - ) - - ch_names = info["ch_names"] - if len(epochs) > 0: - ch_types = epochs[0].get_channel_types() - else: - ch_types = [] - for ch_name in ch_names: - ch_idx = info["ch_names"].index(ch_name) - ch_kind = info["chs"][ch_idx]["kind"] - ch_type_map = { - 2: "EEG", - 3: "EOG", - 4: "EMG", - 5: "ECG", - 301: "MISC", - } - ch_types.append(ch_type_map.get(ch_kind, "MISC")) - - units = ArrayDict( - id=np.array(ch_names, dtype="U"), - types=np.array(ch_types, dtype="U"), - ) - - epoch_intervals = Interval( - start=np.array(epoch_starts), - end=np.array(epoch_ends), - ) - - return eeg, units, epoch_intervals - - def _extract_motor_imagery_trials(self, X, labels, info): - """Extract motor imagery trial intervals with movement labels. - - Parameters - ---------- - X : np.ndarray - Array of shape (n_epochs, n_channels, n_samples) - labels : np.ndarray - Array of shape (n_epochs,) with movement labels - info : mne.Info - MNE Info object with sampling rate information - - Returns - ------- - trials : Interval - Interval object with start/end times and movement labels - """ - sfreq = info["sfreq"] - n_epochs, _, n_samples = X.shape - - start_times = [] - end_times = [] - movements = [] - movement_ids = [] - - current_time = 0.0 - - for i in range(n_epochs): - epoch_duration = n_samples / sfreq - - start_times.append(current_time) - end_times.append(current_time + epoch_duration) - - label = labels[i] - movements.append(label) - movement_ids.append(MOVEMENT_ID_MAP.get(label, -1)) - - current_time += epoch_duration - - trials = Interval( - start=np.array(start_times), - end=np.array(end_times), - timestamps=(np.array(start_times) + np.array(end_times)) / 2, - movements=np.array(movements), - movement_ids=np.array(movement_ids), - timekeys=["start", "end", "timestamps"], - ) - - if not trials.is_disjoint(): - raise ValueError("Found overlapping trials") - - return trials From 82bdf7a6d6952f7775b146604cae2563feb4b6d4 Mon Sep 17 00:00:00 2001 From: Milo Sobral Date: Thu, 15 Jan 2026 10:46:55 -0500 Subject: [PATCH 03/13] Fixed bug with manifest --- brainsets/moabb_pipeline.py | 25 +++++++++++++------------ 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/brainsets/moabb_pipeline.py b/brainsets/moabb_pipeline.py index e55847fc..205d239f 100644 --- a/brainsets/moabb_pipeline.py +++ b/brainsets/moabb_pipeline.py @@ -104,6 +104,7 @@ def get_manifest(cls, raw_dir: Path, args) -> pd.DataFrame: for subject in dataset.subject_list: for session in range(dataset.n_sessions): + # Make sure session is an integer session_id = f"subj-{subject:03d}_sess-{session}" manifest_list.append( { @@ -143,41 +144,41 @@ def download(self, manifest_item) -> Dict[str, Any]: dataset = self.get_dataset() paradigm = self.get_paradigm() + subject = int(manifest_item.subject) + session = int(manifest_item.session) X, labels, meta = paradigm.get_data( dataset=dataset, - subjects=[manifest_item.subject.item()], + subjects=[subject], return_epochs=False, ) if len(X) == 0: raise ValueError( - f"No epochs found for subject {manifest_item.subject}, " - f"session {manifest_item.session}" + f"No epochs found for subject {subject}, " f"session {session}" ) session_values = sorted(meta["session"].unique()) - if isinstance(manifest_item.session, int): - if manifest_item.session < len(session_values): - session_key = session_values[manifest_item.session] + if isinstance(session, int): + if session < len(session_values): + session_key = session_values[session] else: raise ValueError( - f"Session index {manifest_item.session} out of range for subject {manifest_item.subject}. " + f"Session index {session} out of range for subject {subject}. " f"Available {len(session_values)} sessions: {list(session_values)}" ) else: - session_key = str(manifest_item.session.item()) + session_key = str(session) if session_key not in session_values: raise ValueError( - f"Session {session_key} not found for subject {manifest_item.subject}. " + f"Session {session_key} not found for subject {subject}. " f"Available sessions: {list(session_values)}" ) session_mask = meta["session"] == session_key if not session_mask.any(): raise ValueError( - f"No epochs found for subject {manifest_item.subject}, " - f"session {session_key}" + f"No epochs found for subject {subject}, " f"session {session_key}" ) X_filtered = X[session_mask] @@ -186,7 +187,7 @@ def download(self, manifest_item) -> Dict[str, Any]: epochs, labels_epochs, meta_epochs = paradigm.get_data( dataset=dataset, - subjects=[manifest_item.subject.item()], + subjects=[subject], return_epochs=True, ) From 0777a4776e9b18a0dfb260764f87827e73ce1f4f Mon Sep 17 00:00:00 2001 From: Milo Sobral Date: Thu, 15 Jan 2026 11:26:00 -0500 Subject: [PATCH 04/13] Fixes and optimization for the pipeline --- brainsets/moabb_pipeline.py | 163 ++++++++++++++---------------------- 1 file changed, 63 insertions(+), 100 deletions(-) diff --git a/brainsets/moabb_pipeline.py b/brainsets/moabb_pipeline.py index 205d239f..edac23e0 100644 --- a/brainsets/moabb_pipeline.py +++ b/brainsets/moabb_pipeline.py @@ -185,7 +185,7 @@ def download(self, manifest_item) -> Dict[str, Any]: labels_filtered = labels[session_mask] meta_filtered = meta[session_mask].reset_index(drop=True) - epochs, labels_epochs, meta_epochs = paradigm.get_data( + epochs, _, meta_epochs = paradigm.get_data( dataset=dataset, subjects=[subject], return_epochs=True, @@ -204,8 +204,38 @@ def download(self, manifest_item) -> Dict[str, Any]: "epochs": epochs_filtered, } - def _extract_eeg_data(self, X, info, epochs): - """Extract EEG data from epoched arrays. + def _get_channel_types(self, info, epochs): + """Extract channel types from MNE info and epochs objects. + + Parameters + ---------- + info : mne.Info + MNE Info object with channel information + epochs : mne.Epochs + MNE Epochs object + + Returns + ------- + list[str] + List of channel type strings (e.g., "EEG", "EOG", "EMG") + """ + if len(epochs) > 0: + return epochs[0].get_channel_types() + + ch_type_map = { + 2: "EEG", + 3: "EOG", + 4: "EMG", + 5: "ECG", + 301: "MISC", + } + return [ + ch_type_map.get(info["chs"][info["ch_names"].index(ch)]["kind"], "MISC") + for ch in info["ch_names"] + ] + + def _extract_eeg_data(self, X, info, ch_types, labels): + """Extract EEG data and trial intervals from epoched arrays. Parameters ---------- @@ -213,8 +243,10 @@ def _extract_eeg_data(self, X, info, epochs): Array of shape (n_epochs, n_channels, n_samples) info : mne.Info MNE Info object with channel and sampling rate information - epochs : mne.Epochs - MNE Epochs object (used to get channel types) + ch_types : list[str] + List of channel type strings for each channel + labels : np.ndarray + Array of shape (n_epochs,) with event labels Returns ------- @@ -222,23 +254,13 @@ def _extract_eeg_data(self, X, info, epochs): Concatenated EEG signals units : ArrayDict Channel IDs and types - epoch_intervals : Interval - Time intervals for each epoch + trials : Interval + Trial intervals with start/end times and label fields """ sfreq = info["sfreq"] n_epochs, n_channels, n_samples = X.shape - eeg_signals = np.concatenate([X[i].T for i in range(n_epochs)], axis=0) - - epoch_starts = [] - epoch_ends = [] - current_time = 0.0 - - for i in range(n_epochs): - epoch_duration = n_samples / sfreq - epoch_starts.append(current_time) - epoch_ends.append(current_time + epoch_duration) - current_time += epoch_duration + eeg_signals = X.transpose(0, 2, 1).reshape(-1, n_channels) eeg = RegularTimeSeries( signal=eeg_signals, @@ -249,34 +271,32 @@ def _extract_eeg_data(self, X, info, epochs): ), ) - ch_names = info["ch_names"] - if len(epochs) > 0: - ch_types = epochs[0].get_channel_types() - else: - ch_types = [] - for ch_name in ch_names: - ch_idx = info["ch_names"].index(ch_name) - ch_kind = info["chs"][ch_idx]["kind"] - ch_type_map = { - 2: "EEG", - 3: "EOG", - 4: "EMG", - 5: "ECG", - 301: "MISC", - } - ch_types.append(ch_type_map.get(ch_kind, "MISC")) - units = ArrayDict( - id=np.array(ch_names, dtype="U"), + id=np.array(info["ch_names"], dtype="U"), types=np.array(ch_types, dtype="U"), ) - epoch_intervals = Interval( - start=np.array(epoch_starts), - end=np.array(epoch_ends), + sample_boundaries = np.arange(n_epochs + 1) * n_samples + time_boundaries = sample_boundaries / sfreq + start_times = time_boundaries[:-1] + end_times = time_boundaries[1:] + id_values = np.array([self.label_map.get(label, -1) for label in labels]) + + trials = Interval( + start=start_times, + end=end_times, + timestamps=(start_times + end_times) / 2, + timekeys=["start", "end", "timestamps"], + **{ + self.label_field: np.asarray(labels), + self.id_field: id_values, + }, ) - return eeg, units, epoch_intervals + if not trials.is_disjoint(): + raise ValueError("Found overlapping trials") + + return eeg, units, trials @abstractmethod def get_brainset_description(self) -> BrainsetDescription: @@ -363,61 +383,6 @@ def _get_device_description(self, subject_id: str, info: Any) -> DeviceDescripti id=f"{subject_id}_{recording_date.strftime('%Y%m%d')}", ) - def _extract_trials(self, X, labels, info): - """Extract trial intervals with labels using configured field names. - - Parameters - ---------- - X : np.ndarray - Array of shape (n_epochs, n_channels, n_samples) - labels : np.ndarray - Array of shape (n_epochs,) with event labels - info : mne.Info - MNE Info object with sampling rate information - - Returns - ------- - trials : Interval - Interval object with start/end times and configured label fields - """ - sfreq = info["sfreq"] - n_epochs, _, n_samples = X.shape - - start_times = [] - end_times = [] - label_values = [] - id_values = [] - - current_time = 0.0 - - for i in range(n_epochs): - epoch_duration = n_samples / sfreq - - start_times.append(current_time) - end_times.append(current_time + epoch_duration) - - label = labels[i] - label_values.append(label) - id_values.append(self.label_map.get(label, -1)) - - current_time += epoch_duration - - trial_kwargs = { - "start": np.array(start_times), - "end": np.array(end_times), - "timestamps": (np.array(start_times) + np.array(end_times)) / 2, - self.label_field: np.array(label_values), - self.id_field: np.array(id_values), - "timekeys": ["start", "end", "timestamps"], - } - - trials = Interval(**trial_kwargs) - - if not trials.is_disjoint(): - raise ValueError("Found overlapping trials") - - return trials - def _generate_splits(self, trials): """Generate stratified folds for trials. @@ -482,11 +447,9 @@ def process(self, download_output: Dict[str, Any]) -> None: session_description = self._get_session_description(session_id, info) device_description = self._get_device_description(subject_id, info) - self.update_status("Extracting EEG") - eeg, units, _ = self._extract_eeg_data(X, info, epochs) - - self.update_status("Extracting Trials") - trials = self._extract_trials(X, labels, info) + self.update_status("Extracting EEG and Trials") + ch_types = self._get_channel_types(info, epochs) + eeg, units, trials = self._extract_eeg_data(X, info, ch_types, labels) self.update_status("Generating Splits") splits = self._generate_splits(trials) From 20f0f1f8613b2bc756e35755ed794692ff8ae543 Mon Sep 17 00:00:00 2001 From: Milo Sobral Date: Thu, 15 Jan 2026 14:26:51 -0500 Subject: [PATCH 05/13] Better split functionality --- brainsets/moabb_pipeline.py | 16 +- brainsets/utils/split.py | 140 +++++++++++++++++- .../pipeline.py | 32 ++++ .../schalk_wolpaw_physionet_2009/pipeline.py | 48 ++++++ 4 files changed, 228 insertions(+), 8 deletions(-) diff --git a/brainsets/moabb_pipeline.py b/brainsets/moabb_pipeline.py index edac23e0..36255bf9 100644 --- a/brainsets/moabb_pipeline.py +++ b/brainsets/moabb_pipeline.py @@ -252,7 +252,7 @@ def _extract_eeg_data(self, X, info, ch_types, labels): ------- eeg : RegularTimeSeries Concatenated EEG signals - units : ArrayDict + channels : ArrayDict Channel IDs and types trials : Interval Trial intervals with start/end times and label fields @@ -271,7 +271,7 @@ def _extract_eeg_data(self, X, info, ch_types, labels): ), ) - units = ArrayDict( + channels = ArrayDict( id=np.array(info["ch_names"], dtype="U"), types=np.array(ch_types, dtype="U"), ) @@ -296,7 +296,7 @@ def _extract_eeg_data(self, X, info, ch_types, labels): if not trials.is_disjoint(): raise ValueError("Found overlapping trials") - return eeg, units, trials + return eeg, channels, trials @abstractmethod def get_brainset_description(self) -> BrainsetDescription: @@ -383,13 +383,15 @@ def _get_device_description(self, subject_id: str, info: Any) -> DeviceDescripti id=f"{subject_id}_{recording_date.strftime('%Y%m%d')}", ) - def _generate_splits(self, trials): + def _generate_splits(self, trials, subject_id: str = None): """Generate stratified folds for trials. Parameters ---------- trials : Interval Trial intervals with label fields + subject_id : str, optional + Subject identifier for subject-level splits. Default is None. Returns ------- @@ -449,10 +451,10 @@ def process(self, download_output: Dict[str, Any]) -> None: self.update_status("Extracting EEG and Trials") ch_types = self._get_channel_types(info, epochs) - eeg, units, trials = self._extract_eeg_data(X, info, ch_types, labels) + eeg, channels, trials = self._extract_eeg_data(X, info, ch_types, labels) self.update_status("Generating Splits") - splits = self._generate_splits(trials) + splits = self._generate_splits(trials, subject_id=subject_id) self.update_status("Creating Data Object") data = Data( @@ -461,7 +463,7 @@ def process(self, download_output: Dict[str, Any]) -> None: session=session_description, device=device_description, eeg=eeg, - units=units, + channels=channels, **{self.trial_key: trials}, splits=splits, domain=eeg.domain, diff --git a/brainsets/utils/split.py b/brainsets/utils/split.py index 881686a6..65e5c2e0 100644 --- a/brainsets/utils/split.py +++ b/brainsets/utils/split.py @@ -1,6 +1,8 @@ import logging +import hashlib import numpy as np -from typing import List +from typing import List, Dict +from collections import Counter from temporaldata import Interval, Data @@ -325,3 +327,139 @@ def generate_stratified_folds( folds.append(fold_data) return folds + + +def generate_task_kfold_splits( + trials: Interval, + task_configs: Dict[str, List[str]], + label_field: str, + n_folds: int = 5, + val_ratio: float = 0.2, + seed: int = 42, +) -> Dict[str, Interval]: + """ + Generate k-fold cross-validation train/valid/test splits for multiple tasks. + + For each task, creates n_folds stratified train/valid/test splits using the + provided trials Interval object. The splits are returned as Interval objects + with names formatted as "{task_name}_fold{k}_train", "{task_name}_fold{k}_valid", + and "{task_name}_fold{k}_test" for each fold. + + Args + ---- + trials : Interval + An Interval object containing trial information, including labels. + task_configs : Dict[str, List[str]] + Dictionary mapping task names to lists of labels to include for that task. + Example: {"MotorImagery": ["left_hand", "right_hand", "feet"], ...} + label_field : str + The attribute name in trials that contains the labels. + n_folds : int + Number of folds for cross-validation. Default is 5. + val_ratio : float + Ratio of validation set relative to train+valid combined. Default is 0.2. + seed : int + Random seed for reproducibility. Default is 42. + + Returns + ------- + Dict[str, Interval] + Dictionary mapping split names to Interval objects. + """ + if not hasattr(trials, label_field): + raise ValueError( + f"Trials must have a '{label_field}' attribute for task filtering." + ) + + all_labels = getattr(trials, label_field) + splits_dict = {} + + for task_name, include_labels in task_configs.items(): + logging.info(f"\nGenerating {task_name} k-fold train/valid/test splits") + + task_mask = np.isin(all_labels, include_labels) + task_trials = trials.select_by_mask(task_mask) + + if len(task_trials) < n_folds: + logging.warning( + f"Task {task_name} has only {len(task_trials)} trials, " + f"skipping (need at least {n_folds})" + ) + continue + + folds = generate_stratified_folds( + task_trials, + stratify_by=label_field, + n_folds=n_folds, + val_ratio=val_ratio, + seed=seed, + ) + + for k, fold_data in enumerate(folds): + task_labels_train = getattr(fold_data.train, label_field) + task_labels_valid = getattr(fold_data.valid, label_field) + task_labels_test = getattr(fold_data.test, label_field) + + logging.info(f"Fold {k}:") + logging.info(f" Train label counts: {dict(Counter(task_labels_train))}") + logging.info(f" Valid label counts: {dict(Counter(task_labels_valid))}") + logging.info(f" Test label counts: {dict(Counter(task_labels_test))}") + + splits_dict[f"{task_name}_fold{k}_train"] = fold_data.train + splits_dict[f"{task_name}_fold{k}_valid"] = fold_data.valid + splits_dict[f"{task_name}_fold{k}_test"] = fold_data.test + + return splits_dict + + +def compute_subject_kfold_assignments( + subject_id: str, n_folds: int = 5, val_ratio: float = 0.2, seed: int = 42 +) -> Dict[str, str]: + """ + Compute deterministic subject-level k-fold train/valid/test assignments. + + Uses hash-based assignment to deterministically assign subjects to folds. + For each fold k: + - Subjects in bucket k are assigned to test + - Remaining subjects are split into train/valid based on val_ratio + + Args + ---- + subject_id : str + Subject identifier (e.g., "S001", "subj-001") + n_folds : int + Number of folds for cross-validation. Default is 5. + val_ratio : float + Ratio of validation set relative to train+valid combined. Default is 0.2. + seed : int + Random seed for reproducibility. Default is 42. + + Returns + ------- + Dict[str, str] + Dictionary mapping "SubjectSplit_fold{k}" to "train", "valid", or "test" + for each fold k in range(n_folds). + """ + subject_str = f"{subject_id}_{seed}" + subject_bytes = subject_str.encode("utf-8") + hash_obj = hashlib.md5(subject_bytes) + hash_int = int(hash_obj.hexdigest(), 16) + bucket = hash_int % n_folds + + assignments = {} + + for k in range(n_folds): + if bucket == k: + assignments[f"SubjectSplit_fold{k}"] = "test" + else: + fold_str = f"{subject_id}_{seed}_{k}" + fold_bytes = fold_str.encode("utf-8") + fold_hash_obj = hashlib.md5(fold_bytes) + fold_hash_int = int(fold_hash_obj.hexdigest(), 16) + normalized_hash = (fold_hash_int % 10000) / 10000.0 + if normalized_hash < val_ratio: + assignments[f"SubjectSplit_fold{k}"] = "valid" + else: + assignments[f"SubjectSplit_fold{k}"] = "train" + + return assignments diff --git a/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py b/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py index 34d3bbfc..7ac6094c 100644 --- a/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py +++ b/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py @@ -21,9 +21,11 @@ from moabb.datasets import BI2014a from moabb.paradigms import P300 +from temporaldata import Data from brainsets.descriptions import BrainsetDescription from brainsets.taxonomy import Task from brainsets.moabb_pipeline import MOABBPipeline +from brainsets.utils.split import compute_subject_kfold_assignments logging.basicConfig(level=logging.INFO) @@ -51,6 +53,36 @@ class Pipeline(MOABBPipeline): "NonTarget": 0, } + def _generate_splits(self, trials, subject_id: str = None): + """Generate stratified folds and subject-level k-fold assignments. + + Generates: + 1. Standard stratified k-fold splits (from parent class) + 2. Subject-level k-fold assignments (train/valid/test per fold) + + Parameters + ---------- + trials : Interval + Trial intervals with label fields + subject_id : str + Subject identifier for subject-level splits + + Returns + ------- + splits : Data + Data object containing all split masks + """ + splits = super()._generate_splits(trials, subject_id=subject_id) + + if subject_id is not None: + subject_assignments = compute_subject_kfold_assignments( + subject_id, n_folds=5, val_ratio=0.2, seed=42 + ) + for key, value in subject_assignments.items(): + setattr(splits, key, value) + + return splits + def get_brainset_description(self): return BrainsetDescription( id="korczowski_brain_invaders_2014a", diff --git a/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py b/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py index 58569705..3afa3db0 100644 --- a/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py +++ b/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py @@ -21,9 +21,14 @@ from moabb.datasets import PhysionetMI from moabb.paradigms import MotorImagery +from temporaldata import Data from brainsets.descriptions import BrainsetDescription from brainsets.taxonomy import Task from brainsets.moabb_pipeline import MOABBPipeline +from brainsets.utils.split import ( + generate_task_kfold_splits, + compute_subject_kfold_assignments, +) logging.basicConfig(level=logging.INFO) @@ -54,6 +59,49 @@ class Pipeline(MOABBPipeline): "rest": 4, } + TASK_CONFIGS = { + "MotorImagery": ["left_hand", "right_hand", "hands", "feet", "rest"], + "LeftRightImagery": ["left_hand", "right_hand"], + "RightHandFeetImagery": ["right_hand", "feet"], + } + + def _generate_splits(self, trials, subject_id: str = None): + """Generate task-specific and subject-level k-fold splits. + + Generates: + 1. Task-specific within-session stratified k-fold splits for MotorImagery, + LeftRightImagery, and RightHandFeetImagery tasks + 2. Subject-level k-fold assignments (train/valid/test per fold) + + Parameters + ---------- + trials : Interval + Trial intervals with label fields + subject_id : str + Subject identifier for subject-level splits + + Returns + ------- + splits : Data + Data object containing all split masks + """ + task_splits = generate_task_kfold_splits( + trials, + task_configs=self.TASK_CONFIGS, + label_field=self.label_field, + n_folds=5, + val_ratio=0.2, + seed=42, + ) + + if subject_id is not None: + subject_assignments = compute_subject_kfold_assignments( + subject_id, n_folds=5, val_ratio=0.2, seed=42 + ) + return Data(**task_splits, **subject_assignments, domain=trials) + else: + return Data(**task_splits, domain=trials) + def get_brainset_description(self): return BrainsetDescription( id="schalk_wolpaw_physionet_2009", From 7eda26dbee94101e3b6ccf2de64fa2f6f880d9fc Mon Sep 17 00:00:00 2001 From: Milo Sobral Date: Thu, 15 Jan 2026 14:59:09 -0500 Subject: [PATCH 06/13] small fixes based on initial review plus new tests for split stuff --- brainsets/moabb_pipeline.py | 48 +++++---- tests/test_split_utils.py | 204 ++++++++++++++++++++++++++++++++++++ 2 files changed, 229 insertions(+), 23 deletions(-) diff --git a/brainsets/moabb_pipeline.py b/brainsets/moabb_pipeline.py index 36255bf9..3cafad4f 100644 --- a/brainsets/moabb_pipeline.py +++ b/brainsets/moabb_pipeline.py @@ -6,7 +6,7 @@ """ from abc import abstractmethod -from typing import Dict, Any, Type +from typing import Dict, Any, Type, Optional from pathlib import Path import pandas as pd import numpy as np @@ -147,13 +147,13 @@ def download(self, manifest_item) -> Dict[str, Any]: subject = int(manifest_item.subject) session = int(manifest_item.session) - X, labels, meta = paradigm.get_data( + epochs, labels, meta = paradigm.get_data( dataset=dataset, subjects=[subject], - return_epochs=False, + return_epochs=True, ) - if len(X) == 0: + if len(epochs) == 0: raise ValueError( f"No epochs found for subject {subject}, " f"session {session}" ) @@ -181,19 +181,11 @@ def download(self, manifest_item) -> Dict[str, Any]: f"No epochs found for subject {subject}, " f"session {session_key}" ) - X_filtered = X[session_mask] + epochs_filtered = epochs[session_mask] labels_filtered = labels[session_mask] meta_filtered = meta[session_mask].reset_index(drop=True) - epochs, _, meta_epochs = paradigm.get_data( - dataset=dataset, - subjects=[subject], - return_epochs=True, - ) - - session_mask_epochs = meta_epochs["session"] == session_key - epochs_filtered = epochs[session_mask_epochs] - + X_filtered = np.concatenate([ep.get_data() for ep in epochs_filtered], axis=0) info = epochs_filtered[0].info if len(epochs_filtered) > 0 else None return { @@ -222,12 +214,13 @@ def _get_channel_types(self, info, epochs): if len(epochs) > 0: return epochs[0].get_channel_types() + # MNE channel kind constants (mne.io.constants.FIFF) ch_type_map = { - 2: "EEG", - 3: "EOG", - 4: "EMG", - 5: "ECG", - 301: "MISC", + 2: "EEG", # FIFFV_EEG_CH + 3: "EOG", # FIFFV_EOG_CH + 4: "EMG", # FIFFV_EMG_CH + 5: "ECG", # FIFFV_ECG_CH + 301: "MISC", # FIFFV_MISC_CH } return [ ch_type_map.get(info["chs"][info["ch_names"].index(ch)]["kind"], "MISC") @@ -349,7 +342,13 @@ def _get_session_description( recording_date = info.get("meas_date") if recording_date is None: - recording_date = datetime.datetime.now() + recording_date = datetime.datetime( + 2026, 1, 15, tzinfo=datetime.timezone.utc + ) + logging.warning( + f"Missing meas_date for session '{session_id}' (task={self.task}); " + "using sentinel date 2026-01-15" + ) return SessionDescription( id=session_id, @@ -377,13 +376,15 @@ def _get_device_description(self, subject_id: str, info: Any) -> DeviceDescripti recording_date = info.get("meas_date") if recording_date is None: - recording_date = datetime.datetime.now() + recording_date = datetime.datetime( + 2026, 1, 15, tzinfo=datetime.timezone.utc + ) return DeviceDescription( id=f"{subject_id}_{recording_date.strftime('%Y%m%d')}", ) - def _generate_splits(self, trials, subject_id: str = None): + def _generate_splits(self, trials, subject_id: Optional[str] = None): """Generate stratified folds for trials. Parameters @@ -439,7 +440,8 @@ def process(self, download_output: Dict[str, Any]) -> None: session_id = f"{subject_id}_sess-{meta.iloc[0]['session']}" store_path = self.processed_dir / f"{session_id}.h5" - if store_path.exists() and not self.args.reprocess: + safe_reprocess = getattr(getattr(self, "args", None), "reprocess", False) + if store_path.exists() and not safe_reprocess: self.update_status("Skipped Processing") return diff --git a/tests/test_split_utils.py b/tests/test_split_utils.py index 1a7b6946..cfa51cd0 100644 --- a/tests/test_split_utils.py +++ b/tests/test_split_utils.py @@ -4,6 +4,8 @@ from brainsets.utils.split import ( chop_intervals, generate_stratified_folds, + generate_task_kfold_splits, + compute_subject_kfold_assignments, ) @@ -199,3 +201,205 @@ def test_generate_stratified_folds_missing_attribute(): with pytest.raises(ValueError, match="must have a 'label' attribute"): generate_stratified_folds(intervals, stratify_by="label", n_folds=5) + + +class TestGenerateTaskKfoldSplits: + def test_basic_functionality(self): + n_samples = 100 + start = np.arange(n_samples, dtype=float) + end = start + 1.0 + labels = np.array( + ["left_hand"] * 30 + ["right_hand"] * 30 + ["feet"] * 20 + ["rest"] * 20 + ) + + rng = np.random.default_rng(42) + perm = rng.permutation(n_samples) + labels = labels[perm] + start = start[perm] + end = end[perm] + + trials = Interval(start=start, end=end, label=labels) + + task_configs = { + "MotorImagery": ["left_hand", "right_hand", "feet"], + "BinaryTask": ["left_hand", "rest"], + } + + splits_dict = generate_task_kfold_splits( + trials, + task_configs=task_configs, + label_field="label", + n_folds=5, + val_ratio=0.2, + seed=42, + ) + + assert isinstance(splits_dict, dict) + + for k in range(5): + assert f"MotorImagery_fold{k}_train" in splits_dict + assert f"MotorImagery_fold{k}_valid" in splits_dict + assert f"MotorImagery_fold{k}_test" in splits_dict + + assert f"BinaryTask_fold{k}_train" in splits_dict + assert f"BinaryTask_fold{k}_valid" in splits_dict + assert f"BinaryTask_fold{k}_test" in splits_dict + + def test_labels_preserved_per_task(self): + n_samples = 60 + start = np.arange(n_samples, dtype=float) + end = start + 1.0 + labels = np.array(["A"] * 20 + ["B"] * 20 + ["C"] * 20) + + trials = Interval(start=start, end=end, label=labels) + + task_configs = {"TaskAB": ["A", "B"]} + + splits_dict = generate_task_kfold_splits( + trials, + task_configs=task_configs, + label_field="label", + n_folds=5, + val_ratio=0.2, + seed=42, + ) + + for k in range(5): + train = splits_dict[f"TaskAB_fold{k}_train"] + valid = splits_dict[f"TaskAB_fold{k}_valid"] + test = splits_dict[f"TaskAB_fold{k}_test"] + + all_labels_in_fold = np.concatenate([train.label, valid.label, test.label]) + assert set(all_labels_in_fold) == {"A", "B"} + assert "C" not in all_labels_in_fold + + def test_missing_label_field_raises(self): + start = np.arange(10, dtype=float) + end = start + 1.0 + trials = Interval(start=start, end=end) + + with pytest.raises(ValueError, match="must have a 'label' attribute"): + generate_task_kfold_splits( + trials, + task_configs={"Task": ["A"]}, + label_field="label", + ) + + def test_skips_task_with_insufficient_trials(self, caplog): + start = np.arange(14, dtype=float) + end = start + 1.0 + labels = np.array(["A"] * 4 + ["B"] * 10) + trials = Interval(start=start, end=end, label=labels) + + task_configs = { + "SmallTask": ["A"], + "LargerTask": ["A", "B"], + } + + import logging + + with caplog.at_level(logging.WARNING): + splits_dict = generate_task_kfold_splits( + trials, + task_configs=task_configs, + label_field="label", + n_folds=5, + val_ratio=0.2, + seed=42, + ) + + assert "SmallTask_fold0_train" not in splits_dict + assert "LargerTask_fold0_train" in splits_dict + + +class TestComputeSubjectKfoldAssignments: + def test_basic_output_structure(self): + assignments = compute_subject_kfold_assignments("S001", n_folds=5) + + assert isinstance(assignments, dict) + assert len(assignments) == 5 + + for k in range(5): + assert f"SubjectSplit_fold{k}" in assignments + assert assignments[f"SubjectSplit_fold{k}"] in ["train", "valid", "test"] + + def test_exactly_one_test_assignment(self): + assignments = compute_subject_kfold_assignments("S001", n_folds=5) + + test_count = sum(1 for v in assignments.values() if v == "test") + assert test_count == 1 + + def test_deterministic_assignment(self): + assignments1 = compute_subject_kfold_assignments("S001", n_folds=5, seed=42) + assignments2 = compute_subject_kfold_assignments("S001", n_folds=5, seed=42) + + assert assignments1 == assignments2 + + def test_different_subjects_different_assignments(self): + assignments_s1 = compute_subject_kfold_assignments("S001", n_folds=5, seed=42) + assignments_s2 = compute_subject_kfold_assignments("S002", n_folds=5, seed=42) + + assert assignments_s1 != assignments_s2 + + def test_different_seeds_different_assignments(self): + assignments_seed1 = compute_subject_kfold_assignments( + "S001", n_folds=5, seed=42 + ) + assignments_seed2 = compute_subject_kfold_assignments( + "S001", n_folds=5, seed=99 + ) + + assert assignments_seed1 != assignments_seed2 + + def test_val_ratio_affects_valid_proportion(self): + n_subjects = 100 + n_folds = 5 + + valid_counts_low = 0 + valid_counts_high = 0 + + for i in range(n_subjects): + assignments_low = compute_subject_kfold_assignments( + f"S{i:03d}", n_folds=n_folds, val_ratio=0.1, seed=42 + ) + assignments_high = compute_subject_kfold_assignments( + f"S{i:03d}", n_folds=n_folds, val_ratio=0.4, seed=42 + ) + + valid_counts_low += sum(1 for v in assignments_low.values() if v == "valid") + valid_counts_high += sum( + 1 for v in assignments_high.values() if v == "valid" + ) + + assert valid_counts_high > valid_counts_low + + def test_distribution_across_many_subjects(self): + n_subjects = 1000 + n_folds = 5 + val_ratio = 0.2 + + test_fold_counts = {k: 0 for k in range(n_folds)} + valid_count = 0 + train_count = 0 + + for i in range(n_subjects): + assignments = compute_subject_kfold_assignments( + f"subject_{i}", n_folds=n_folds, val_ratio=val_ratio, seed=42 + ) + + for k in range(n_folds): + assignment = assignments[f"SubjectSplit_fold{k}"] + if assignment == "test": + test_fold_counts[k] += 1 + elif assignment == "valid": + valid_count += 1 + else: + train_count += 1 + + expected_per_fold = n_subjects / n_folds + for k, count in test_fold_counts.items(): + assert abs(count - expected_per_fold) < expected_per_fold * 0.3 + + non_test_total = valid_count + train_count + actual_val_ratio = valid_count / non_test_total + assert abs(actual_val_ratio - val_ratio) < 0.05 From f5bef8f7d18eb72a2836f157c74bc81efd36b266 Mon Sep 17 00:00:00 2001 From: Milo Sobral Date: Thu, 15 Jan 2026 15:46:28 -0500 Subject: [PATCH 07/13] Small fix for issue wih download location --- brainsets/moabb_pipeline.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/brainsets/moabb_pipeline.py b/brainsets/moabb_pipeline.py index 3cafad4f..7091c15e 100644 --- a/brainsets/moabb_pipeline.py +++ b/brainsets/moabb_pipeline.py @@ -8,6 +8,7 @@ from abc import abstractmethod from typing import Dict, Any, Type, Optional from pathlib import Path +import os import pandas as pd import numpy as np import datetime @@ -16,7 +17,6 @@ from moabb.datasets.base import BaseDataset from moabb.paradigms.base import BaseParadigm -from moabb.utils import set_download_dir from temporaldata import Data, RegularTimeSeries, Interval, ArrayDict from brainsets.pipeline import BrainsetPipeline from brainsets import serialize_fn_map @@ -140,7 +140,8 @@ def download(self, manifest_item) -> Dict[str, Any]: """ self.update_status("DOWNLOADING") - set_download_dir(str(self.raw_dir)) + self.raw_dir.mkdir(exist_ok=True, parents=True) + os.environ["MNE_DATA"] = str(self.raw_dir.resolve()) dataset = self.get_dataset() paradigm = self.get_paradigm() @@ -185,7 +186,7 @@ def download(self, manifest_item) -> Dict[str, Any]: labels_filtered = labels[session_mask] meta_filtered = meta[session_mask].reset_index(drop=True) - X_filtered = np.concatenate([ep.get_data() for ep in epochs_filtered], axis=0) + X_filtered = np.concatenate(epochs_filtered, axis=0) info = epochs_filtered[0].info if len(epochs_filtered) > 0 else None return { From 6b2fbe627908fd325450a15e66e3713b869768c1 Mon Sep 17 00:00:00 2001 From: Milo Sobral Date: Thu, 22 Jan 2026 11:51:12 -0500 Subject: [PATCH 08/13] Minor reorganization, changed to three splits, and improved docs and naming --- .../moabb/pipeline.py} | 12 +- brainsets/utils/split.py | 188 +++++++++++++----- .../kemp_sleep_edf_2013/pipeline.py | 4 +- .../pipeline.py | 8 +- .../schalk_wolpaw_physionet_2009/pipeline.py | 17 +- tests/test_split_utils.py | 52 ++--- 6 files changed, 188 insertions(+), 93 deletions(-) rename brainsets/{moabb_pipeline.py => utils/moabb/pipeline.py} (98%) diff --git a/brainsets/moabb_pipeline.py b/brainsets/utils/moabb/pipeline.py similarity index 98% rename from brainsets/moabb_pipeline.py rename to brainsets/utils/moabb/pipeline.py index 7091c15e..ab69f972 100644 --- a/brainsets/moabb_pipeline.py +++ b/brainsets/utils/moabb/pipeline.py @@ -27,7 +27,7 @@ DeviceDescription, ) from brainsets.taxonomy import Species, Task -from brainsets.utils.split import generate_stratified_folds +from brainsets.utils.split import generate_trial_folds class MOABBPipeline(BrainsetPipeline): @@ -218,10 +218,10 @@ def _get_channel_types(self, info, epochs): # MNE channel kind constants (mne.io.constants.FIFF) ch_type_map = { 2: "EEG", # FIFFV_EEG_CH - 3: "EOG", # FIFFV_EOG_CH - 4: "EMG", # FIFFV_EMG_CH - 5: "ECG", # FIFFV_ECG_CH - 301: "MISC", # FIFFV_MISC_CH + 202: "EOG", # FIFFV_EOG_CH + 302: "EMG", # FIFFV_EMG_CH + 402: "ECG", # FIFFV_ECG_CH + 502: "MISC", # FIFFV_MISC_CH } return [ ch_type_map.get(info["chs"][info["ch_names"].index(ch)]["kind"], "MISC") @@ -400,7 +400,7 @@ def _generate_splits(self, trials, subject_id: Optional[str] = None): splits : Data Data object containing fold splits """ - folds = generate_stratified_folds( + folds = generate_trial_folds( trials, stratify_by=self.stratify_field, n_folds=5, diff --git a/brainsets/utils/split.py b/brainsets/utils/split.py index 65e5c2e0..d2f563c0 100644 --- a/brainsets/utils/split.py +++ b/brainsets/utils/split.py @@ -239,39 +239,74 @@ def _create_interval_split(intervals: Interval, indices: np.ndarray) -> Interval return split -def generate_stratified_folds( - intervals: Interval, +def generate_trial_folds( + trials: Interval, stratify_by: str, n_folds: int = 5, val_ratio: float = 0.2, seed: int = 42, ) -> List[Data]: """ - Generates stratified train/valid/test splits using a two-stage splitting process. + Generate stratified k-fold train/valid/test splits at the trial level. - The splitting is performed in two stages: - 1. Outer split (StratifiedKFold): The intervals are divided into n_folds, - where each fold uses one partition as the test set and the remaining - partitions as train+valid. Stratification ensures each fold maintains - the class distribution of the original data. - 2. Inner split (StratifiedShuffleSplit): The train+valid portion of each fold - is further split into train and valid sets using val_ratio, while preserving - the class distribution. + This function performs **intra-session splitting** at the trial level. Individual + trials are distributed across folds while maintaining the class distribution + (stratification). Use this when you want to evaluate generalization across trials + within the same session. - Args: - intervals: The intervals to split. - n_folds: Number of folds for cross-validation. - val_ratio: Ratio of validation set relative to train+valid combined. - seed: Random seed. - stratify_by: The attribute name to use for stratification (e.g., "id", "label", - "class"). The intervals must have this attribute. + For **cross-subject splitting** (where entire subjects are held out), use + :func:`generate_subject_kfold_assignment` instead. - Returns: - List of Data objects, one for each fold. + The splitting is performed in two stages: + 1. Outer split (StratifiedKFold): Trials are divided into n_folds, where each + fold uses one partition as test and the rest as train+valid. Stratification + ensures each fold maintains the class distribution. + 2. Inner split (StratifiedShuffleSplit): The train+valid portion is further + split into train and valid sets using val_ratio, preserving class distribution. - Raises: - ValueError: If the intervals don't have the specified stratify_by attribute. - ValueError: If there are fewer samples than n_folds. + Args + ---- + trials : Interval + The trials to split. Must have an attribute matching `stratify_by`. + stratify_by : str + The attribute name to use for stratification (e.g., "label", "class"). + n_folds : int + Number of folds for cross-validation. Default is 5. + val_ratio : float + Ratio of validation set relative to train+valid combined. Default is 0.2. + seed : int + Random seed for reproducibility. Default is 42. + + Returns + ------- + List[Data] + List of Data objects, one per fold. Each Data object contains: + - train: Interval with training trials + - valid: Interval with validation trials + - test: Interval with test trials + - domain: Combined interval of all trials in the fold + + Raises + ------ + ValueError + If trials don't have the specified stratify_by attribute. + ValueError + If there are fewer trials than n_folds. + + Examples + -------- + >>> folds = generate_trial_folds( + ... trials=session.trials, + ... stratify_by="label", + ... n_folds=5, + ... ) + >>> for k, fold in enumerate(folds): + ... print(f"Fold {k}: train={len(fold.train)}, valid={len(fold.valid)}, test={len(fold.test)}") + + See Also + -------- + generate_trial_folds_by_task : Higher-level function for multi-task trial splitting. + generate_subject_kfold_assignment : For cross-subject (leave-subject-out) splitting. """ try: from sklearn.model_selection import StratifiedKFold, StratifiedShuffleSplit @@ -281,25 +316,25 @@ def generate_stratified_folds( "`pip install scikit-learn`" ) - if not hasattr(intervals, stratify_by): + if not hasattr(trials, stratify_by): raise ValueError( - f"Intervals must have a '{stratify_by}' attribute for stratification." + f"Trials must have a '{stratify_by}' attribute for stratification." ) - class_labels = getattr(intervals, stratify_by) + class_labels = getattr(trials, stratify_by) if len(class_labels) < n_folds: raise ValueError( - f"Not enough samples ({len(class_labels)}) for {n_folds} folds." + f"Not enough trials ({len(class_labels)}) for {n_folds} folds." ) outer_splitter = StratifiedKFold(n_splits=n_folds, shuffle=True, random_state=seed) folds = [] - sample_indices = np.arange(len(intervals)) + sample_indices = np.arange(len(trials)) for fold_idx, (train_val_indices, test_indices) in enumerate( outer_splitter.split(sample_indices, class_labels) ): - test_split = _create_interval_split(intervals, test_indices) + test_split = _create_interval_split(trials, test_indices) train_val_labels = class_labels[train_val_indices] inner_splitter = StratifiedShuffleSplit( @@ -312,8 +347,8 @@ def generate_stratified_folds( train_original_indices = train_val_indices[train_indices] val_original_indices = train_val_indices[val_indices] - train_split = _create_interval_split(intervals, train_original_indices) - val_split = _create_interval_split(intervals, val_original_indices) + train_split = _create_interval_split(trials, train_original_indices) + val_split = _create_interval_split(trials, val_original_indices) combined_domain = train_split | val_split | test_split @@ -329,7 +364,7 @@ def generate_stratified_folds( return folds -def generate_task_kfold_splits( +def generate_trial_folds_by_task( trials: Interval, task_configs: Dict[str, List[str]], label_field: str, @@ -338,22 +373,30 @@ def generate_task_kfold_splits( seed: int = 42, ) -> Dict[str, Interval]: """ - Generate k-fold cross-validation train/valid/test splits for multiple tasks. + Generate trial-level k-fold cross-validation splits within a single session. + + This function performs **intra-session splitting** at the trial level. It takes + trials from a single recording session and creates stratified k-fold splits, + ensuring trials from the same session can appear in different folds. Use this + when you want to evaluate generalization across trials within the same session. - For each task, creates n_folds stratified train/valid/test splits using the - provided trials Interval object. The splits are returned as Interval objects - with names formatted as "{task_name}_fold{k}_train", "{task_name}_fold{k}_valid", - and "{task_name}_fold{k}_test" for each fold. + For **cross-subject splitting** (where entire subjects are held out), use + :func:`generate_subject_kfold_assignment` instead. + + For each task configuration, creates n_folds stratified train/valid/test splits. + The splits are returned as Interval objects with names formatted as + "{task_name}_fold{k}_train", "{task_name}_fold{k}_valid", and + "{task_name}_fold{k}_test". Args ---- trials : Interval - An Interval object containing trial information, including labels. + An Interval object containing trial information from a single session, + including labels specified by `label_field`. task_configs : Dict[str, List[str]] Dictionary mapping task names to lists of labels to include for that task. - Example: {"MotorImagery": ["left_hand", "right_hand", "feet"], ...} label_field : str - The attribute name in trials that contains the labels. + The attribute name in trials that contains the labels for stratification. n_folds : int Number of folds for cross-validation. Default is 5. val_ratio : float @@ -365,6 +408,26 @@ def generate_task_kfold_splits( ------- Dict[str, Interval] Dictionary mapping split names to Interval objects. + + Examples + -------- + >>> task_configs = { + ... "MotorImagery": ["left_hand", "right_hand", "feet"], + ... "RestVsActive": ["rest", "left_hand", "right_hand"], + ... } + >>> splits = generate_trial_folds_by_task( + ... trials=session.trials, + ... task_configs=task_configs, + ... label_field="label", + ... n_folds=5, + ... ) + >>> # Returns keys like: "MotorImagery_fold0_train", "MotorImagery_fold0_valid", ... + >>> train_trials = splits["MotorImagery_fold0_train"] + + See Also + -------- + generate_subject_kfold_assignment : For cross-subject (leave-subject-out) splitting. + generate_trial_folds : Lower-level function for stratified fold generation. """ if not hasattr(trials, label_field): raise ValueError( @@ -387,7 +450,7 @@ def generate_task_kfold_splits( ) continue - folds = generate_stratified_folds( + folds = generate_trial_folds( task_trials, stratify_by=label_field, n_folds=n_folds, @@ -412,21 +475,33 @@ def generate_task_kfold_splits( return splits_dict -def compute_subject_kfold_assignments( +def generate_subject_kfold_assignment( subject_id: str, n_folds: int = 5, val_ratio: float = 0.2, seed: int = 42 ) -> Dict[str, str]: """ - Compute deterministic subject-level k-fold train/valid/test assignments. + Generate cross-subject k-fold train/valid/test assignments for a single subject. + + This function performs **cross-subject splitting** (also known as leave-subject-out + or between-subject splitting). It deterministically assigns a subject to train, + valid, or test for each fold using hash-based assignment. Use this when you want + to evaluate generalization across different subjects, ensuring no data leakage + between subjects in different splits. + + For **intra-session splitting** (trial-level splits within a session), use + :func:`generate_trial_folds_by_task` instead. + + The assignment is deterministic: the same subject_id with the same seed will + always receive the same assignments. This allows processing subjects independently + (e.g., in parallel) while ensuring consistent fold assignments. - Uses hash-based assignment to deterministically assign subjects to folds. For each fold k: - - Subjects in bucket k are assigned to test - - Remaining subjects are split into train/valid based on val_ratio + - Subjects hashed to bucket k are assigned to "test" + - Remaining subjects are assigned to "train" or "valid" based on val_ratio Args ---- subject_id : str - Subject identifier (e.g., "S001", "subj-001") + Subject identifier (e.g., "S001", "sub-01"). n_folds : int Number of folds for cross-validation. Default is 5. val_ratio : float @@ -439,6 +514,25 @@ def compute_subject_kfold_assignments( Dict[str, str] Dictionary mapping "SubjectSplit_fold{k}" to "train", "valid", or "test" for each fold k in range(n_folds). + + Examples + -------- + >>> assignments = generate_subject_kfold_assignment("sub-01", n_folds=5) + >>> assignments + {'SubjectSplit_fold0': 'train', 'SubjectSplit_fold1': 'test', + 'SubjectSplit_fold2': 'train', 'SubjectSplit_fold3': 'valid', + 'SubjectSplit_fold4': 'train'} + + >>> # Use in a pipeline to tag sessions with their fold assignments + >>> for subject_id in subject_ids: + ... assignments = generate_subject_kfold_assignment(subject_id) + ... session = load_session(subject_id) + ... for key, value in assignments.items(): + ... setattr(session, key, value) + + See Also + -------- + generate_trial_folds_by_task : For intra-session (trial-level) splitting. """ subject_str = f"{subject_id}_{seed}" subject_bytes = subject_str.encode("utf-8") diff --git a/brainsets_pipelines/kemp_sleep_edf_2013/pipeline.py b/brainsets_pipelines/kemp_sleep_edf_2013/pipeline.py index 1359ab0f..1bf60fc4 100644 --- a/brainsets_pipelines/kemp_sleep_edf_2013/pipeline.py +++ b/brainsets_pipelines/kemp_sleep_edf_2013/pipeline.py @@ -28,7 +28,7 @@ from brainsets.pipeline import BrainsetPipeline from brainsets.utils.split import ( chop_intervals, - generate_stratified_folds, + generate_trial_folds, ) from brainsets.utils.s3_utils import get_s3_client_for_download from temporaldata import Data, Interval, RegularTimeSeries, ArrayDict @@ -396,7 +396,7 @@ def create_splits( if len(filtered) == 0: raise ValueError("No valid epochs remaining after filtering") - folds = generate_stratified_folds( + folds = generate_trial_folds( filtered, stratify_by="id", n_folds=n_folds, diff --git a/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py b/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py index 7ac6094c..0ba74d8e 100644 --- a/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py +++ b/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py @@ -24,8 +24,8 @@ from temporaldata import Data from brainsets.descriptions import BrainsetDescription from brainsets.taxonomy import Task -from brainsets.moabb_pipeline import MOABBPipeline -from brainsets.utils.split import compute_subject_kfold_assignments +from brainsets.utils.moabb.pipeline import MOABBPipeline +from brainsets.utils.split import generate_subject_kfold_assignment logging.basicConfig(level=logging.INFO) @@ -75,8 +75,8 @@ def _generate_splits(self, trials, subject_id: str = None): splits = super()._generate_splits(trials, subject_id=subject_id) if subject_id is not None: - subject_assignments = compute_subject_kfold_assignments( - subject_id, n_folds=5, val_ratio=0.2, seed=42 + subject_assignments = generate_subject_kfold_assignment( + subject_id, n_folds=3, val_ratio=0.2, seed=42 ) for key, value in subject_assignments.items(): setattr(splits, key, value) diff --git a/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py b/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py index 3afa3db0..f414a5df 100644 --- a/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py +++ b/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py @@ -24,10 +24,10 @@ from temporaldata import Data from brainsets.descriptions import BrainsetDescription from brainsets.taxonomy import Task -from brainsets.moabb_pipeline import MOABBPipeline +from brainsets.utils.moabb.pipeline import MOABBPipeline from brainsets.utils.split import ( - generate_task_kfold_splits, - compute_subject_kfold_assignments, + generate_trial_folds_by_task, + generate_subject_kfold_assignment, ) @@ -60,7 +60,8 @@ class Pipeline(MOABBPipeline): } TASK_CONFIGS = { - "MotorImagery": ["left_hand", "right_hand", "hands", "feet", "rest"], + "MotorImagery_all": ["left_hand", "right_hand", "hands", "feet", "rest"], + "MotorImagery_norest": ["left_hand", "right_hand", "hands", "feet"], "LeftRightImagery": ["left_hand", "right_hand"], "RightHandFeetImagery": ["right_hand", "feet"], } @@ -85,18 +86,18 @@ def _generate_splits(self, trials, subject_id: str = None): splits : Data Data object containing all split masks """ - task_splits = generate_task_kfold_splits( + task_splits = generate_trial_folds_by_task( trials, task_configs=self.TASK_CONFIGS, label_field=self.label_field, - n_folds=5, + n_folds=3, val_ratio=0.2, seed=42, ) if subject_id is not None: - subject_assignments = compute_subject_kfold_assignments( - subject_id, n_folds=5, val_ratio=0.2, seed=42 + subject_assignments = generate_subject_kfold_assignment( + subject_id, n_folds=3, val_ratio=0.2, seed=42 ) return Data(**task_splits, **subject_assignments, domain=trials) else: diff --git a/tests/test_split_utils.py b/tests/test_split_utils.py index cfa51cd0..f0763f67 100644 --- a/tests/test_split_utils.py +++ b/tests/test_split_utils.py @@ -3,9 +3,9 @@ from temporaldata import Data, Interval from brainsets.utils.split import ( chop_intervals, - generate_stratified_folds, - generate_task_kfold_splits, - compute_subject_kfold_assignments, + generate_trial_folds, + generate_trial_folds_by_task, + generate_subject_kfold_assignment, ) @@ -86,7 +86,7 @@ def test_chop_intervals_overlapping_no_check(): assert len(chopped) == 20 -def test_generate_stratified_folds(): +def test_generate_trial_folds(): n_samples = 100 start = np.arange(n_samples, dtype=float) end = start + 1.0 @@ -105,7 +105,7 @@ def test_generate_stratified_folds(): n_folds = 5 val_ratio = 0.25 - folds = generate_stratified_folds( + folds = generate_trial_folds( intervals, stratify_by="id", n_folds=n_folds, val_ratio=val_ratio, seed=42 ) @@ -164,7 +164,7 @@ def test_generate_stratified_folds(): assert np.allclose(all_test_starts_sorted, original_starts_sorted) -def test_generate_stratified_folds_custom_attribute(): +def test_generate_trial_folds_custom_attribute(): n_samples = 50 start = np.arange(n_samples, dtype=float) end = start + 1.0 @@ -179,7 +179,7 @@ def test_generate_stratified_folds_custom_attribute(): intervals = Interval(start=start, end=end, label=labels) - folds = generate_stratified_folds( + folds = generate_trial_folds( intervals, stratify_by="label", n_folds=5, val_ratio=0.25, seed=42 ) @@ -194,16 +194,16 @@ def test_generate_stratified_folds_custom_attribute(): assert all(c == 5 for c in counts) -def test_generate_stratified_folds_missing_attribute(): +def test_generate_trial_folds_missing_attribute(): start = np.arange(10, dtype=float) end = start + 1.0 intervals = Interval(start=start, end=end) with pytest.raises(ValueError, match="must have a 'label' attribute"): - generate_stratified_folds(intervals, stratify_by="label", n_folds=5) + generate_trial_folds(intervals, stratify_by="label", n_folds=5) -class TestGenerateTaskKfoldSplits: +class TestGenerateTrialFoldsByTask: def test_basic_functionality(self): n_samples = 100 start = np.arange(n_samples, dtype=float) @@ -225,7 +225,7 @@ def test_basic_functionality(self): "BinaryTask": ["left_hand", "rest"], } - splits_dict = generate_task_kfold_splits( + splits_dict = generate_trial_folds_by_task( trials, task_configs=task_configs, label_field="label", @@ -255,7 +255,7 @@ def test_labels_preserved_per_task(self): task_configs = {"TaskAB": ["A", "B"]} - splits_dict = generate_task_kfold_splits( + splits_dict = generate_trial_folds_by_task( trials, task_configs=task_configs, label_field="label", @@ -279,7 +279,7 @@ def test_missing_label_field_raises(self): trials = Interval(start=start, end=end) with pytest.raises(ValueError, match="must have a 'label' attribute"): - generate_task_kfold_splits( + generate_trial_folds_by_task( trials, task_configs={"Task": ["A"]}, label_field="label", @@ -299,7 +299,7 @@ def test_skips_task_with_insufficient_trials(self, caplog): import logging with caplog.at_level(logging.WARNING): - splits_dict = generate_task_kfold_splits( + splits_dict = generate_trial_folds_by_task( trials, task_configs=task_configs, label_field="label", @@ -312,9 +312,9 @@ def test_skips_task_with_insufficient_trials(self, caplog): assert "LargerTask_fold0_train" in splits_dict -class TestComputeSubjectKfoldAssignments: +class TestGenerateSubjectKfoldAssignment: def test_basic_output_structure(self): - assignments = compute_subject_kfold_assignments("S001", n_folds=5) + assignments = generate_subject_kfold_assignment("S001", n_folds=5) assert isinstance(assignments, dict) assert len(assignments) == 5 @@ -324,28 +324,28 @@ def test_basic_output_structure(self): assert assignments[f"SubjectSplit_fold{k}"] in ["train", "valid", "test"] def test_exactly_one_test_assignment(self): - assignments = compute_subject_kfold_assignments("S001", n_folds=5) + assignments = generate_subject_kfold_assignment("S001", n_folds=5) test_count = sum(1 for v in assignments.values() if v == "test") assert test_count == 1 def test_deterministic_assignment(self): - assignments1 = compute_subject_kfold_assignments("S001", n_folds=5, seed=42) - assignments2 = compute_subject_kfold_assignments("S001", n_folds=5, seed=42) + assignments1 = generate_subject_kfold_assignment("S001", n_folds=5, seed=42) + assignments2 = generate_subject_kfold_assignment("S001", n_folds=5, seed=42) assert assignments1 == assignments2 def test_different_subjects_different_assignments(self): - assignments_s1 = compute_subject_kfold_assignments("S001", n_folds=5, seed=42) - assignments_s2 = compute_subject_kfold_assignments("S002", n_folds=5, seed=42) + assignments_s1 = generate_subject_kfold_assignment("S001", n_folds=5, seed=42) + assignments_s2 = generate_subject_kfold_assignment("S002", n_folds=5, seed=42) assert assignments_s1 != assignments_s2 def test_different_seeds_different_assignments(self): - assignments_seed1 = compute_subject_kfold_assignments( + assignments_seed1 = generate_subject_kfold_assignment( "S001", n_folds=5, seed=42 ) - assignments_seed2 = compute_subject_kfold_assignments( + assignments_seed2 = generate_subject_kfold_assignment( "S001", n_folds=5, seed=99 ) @@ -359,10 +359,10 @@ def test_val_ratio_affects_valid_proportion(self): valid_counts_high = 0 for i in range(n_subjects): - assignments_low = compute_subject_kfold_assignments( + assignments_low = generate_subject_kfold_assignment( f"S{i:03d}", n_folds=n_folds, val_ratio=0.1, seed=42 ) - assignments_high = compute_subject_kfold_assignments( + assignments_high = generate_subject_kfold_assignment( f"S{i:03d}", n_folds=n_folds, val_ratio=0.4, seed=42 ) @@ -383,7 +383,7 @@ def test_distribution_across_many_subjects(self): train_count = 0 for i in range(n_subjects): - assignments = compute_subject_kfold_assignments( + assignments = generate_subject_kfold_assignment( f"subject_{i}", n_folds=n_folds, val_ratio=val_ratio, seed=42 ) From cb788e0f834e5bfd8f66464529804e1314ca4848 Mon Sep 17 00:00:00 2001 From: Milo Sobral Date: Wed, 28 Jan 2026 15:09:38 -0500 Subject: [PATCH 09/13] Minor reformatting and changed to only 3 folds instead of 5 --- brainsets/utils/moabb/pipeline.py | 6 +++--- .../korczowski_brain_invaders_2014a/pipeline.py | 1 - 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/brainsets/utils/moabb/pipeline.py b/brainsets/utils/moabb/pipeline.py index ab69f972..1abf37b5 100644 --- a/brainsets/utils/moabb/pipeline.py +++ b/brainsets/utils/moabb/pipeline.py @@ -156,7 +156,7 @@ def download(self, manifest_item) -> Dict[str, Any]: if len(epochs) == 0: raise ValueError( - f"No epochs found for subject {subject}, " f"session {session}" + f"No epochs found for subject {subject}, session {session}" ) session_values = sorted(meta["session"].unique()) @@ -179,7 +179,7 @@ def download(self, manifest_item) -> Dict[str, Any]: session_mask = meta["session"] == session_key if not session_mask.any(): raise ValueError( - f"No epochs found for subject {subject}, " f"session {session_key}" + f"No epochs found for subject {subject}, session {session_key}" ) epochs_filtered = epochs[session_mask] @@ -403,7 +403,7 @@ def _generate_splits(self, trials, subject_id: Optional[str] = None): folds = generate_trial_folds( trials, stratify_by=self.stratify_field, - n_folds=5, + n_folds=3, val_ratio=0.2, seed=42, ) diff --git a/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py b/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py index 0ba74d8e..16ce2900 100644 --- a/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py +++ b/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py @@ -21,7 +21,6 @@ from moabb.datasets import BI2014a from moabb.paradigms import P300 -from temporaldata import Data from brainsets.descriptions import BrainsetDescription from brainsets.taxonomy import Task from brainsets.utils.moabb.pipeline import MOABBPipeline From c4bdeadde39099d18e117deba8def7df42ed9170 Mon Sep 17 00:00:00 2001 From: Milo Sobral Date: Tue, 3 Feb 2026 11:41:44 -0500 Subject: [PATCH 10/13] Add bandpass filtering and resampling options to MOABBPipeline, and refactored trial slicing to keep all data. - Introduced command-line arguments for bandpass filter settings and resampling rate. - Implemented validation for bandpass parameters against Nyquist frequency. - Updated data retrieval methods to support new filtering options. - Refactored related pipelines to utilize the new base argument parser. --- brainsets/utils/moabb/pipeline.py | 285 +++++++++++++----- .../pipeline.py | 6 +- .../schalk_wolpaw_physionet_2009/pipeline.py | 6 +- 3 files changed, 209 insertions(+), 88 deletions(-) diff --git a/brainsets/utils/moabb/pipeline.py b/brainsets/utils/moabb/pipeline.py index 1abf37b5..ba1554c0 100644 --- a/brainsets/utils/moabb/pipeline.py +++ b/brainsets/utils/moabb/pipeline.py @@ -8,12 +8,14 @@ from abc import abstractmethod from typing import Dict, Any, Type, Optional from pathlib import Path +from argparse import ArgumentParser import os import pandas as pd import numpy as np import datetime import logging import h5py +import mne from moabb.datasets.base import BaseDataset from moabb.paradigms.base import BaseParadigm @@ -30,6 +32,33 @@ from brainsets.utils.split import generate_trial_folds +_base_parser = ArgumentParser(add_help=False) +_base_parser.add_argument( + "--redownload", action="store_true", help="Force redownload of raw data" +) +_base_parser.add_argument( + "--reprocess", action="store_true", help="Force reprocessing of data" +) +_base_parser.add_argument( + "--bandpass-low", + type=float, + default=None, + help="Low cutoff frequency for bandpass filter in Hz (default: None, no filtering)", +) +_base_parser.add_argument( + "--bandpass-high", + type=float, + default=None, + help="High cutoff frequency for bandpass filter in Hz (default: None, no filtering)", +) +_base_parser.add_argument( + "--resample", + type=float, + default=None, + help="Target sampling rate in Hz (default: no resampling)", +) + + class MOABBPipeline(BrainsetPipeline): """Base class for MOABB dataset pipelines. @@ -75,9 +104,29 @@ def get_dataset(cls) -> BaseDataset: return cls.dataset_class(**cls.dataset_kwargs) @classmethod - def get_paradigm(cls) -> BaseParadigm: - """Instantiate the MOABB paradigm with configured kwargs.""" + def get_paradigm(cls, args=None) -> BaseParadigm: + """Instantiate the MOABB paradigm with configured kwargs. + + Parameters + ---------- + args : Namespace, optional + CLI arguments containing bandpass_low, bandpass_high, resample params + + Returns + ------- + BaseParadigm + Paradigm instance with filtering parameters from args + """ kwargs = {k: v for k, v in cls.paradigm_kwargs.items() if v is not None} + + if args is not None: + if args.bandpass_low is not None: + kwargs["fmin"] = args.bandpass_low + if args.bandpass_high is not None: + kwargs["fmax"] = args.bandpass_high + if args.resample is not None: + kwargs["resample"] = args.resample + return cls.paradigm_class(**kwargs) @classmethod @@ -116,13 +165,56 @@ def get_manifest(cls, raw_dir: Path, args) -> pd.DataFrame: return pd.DataFrame(manifest_list).set_index("session_id") + def _validate_bandpass_params(self, dataset: BaseDataset, subject: int) -> None: + """Validate bandpass parameters against Nyquist frequency. + + Loads raw data briefly to check sampling rate and raises an error if + bandpass_high exceeds the Nyquist frequency. + + Parameters + ---------- + dataset : BaseDataset + MOABB dataset instance + subject : int + Subject ID to check + + Raises + ------ + ValueError + If bandpass_high >= Nyquist frequency + """ + fmax = self.args.bandpass_high + if fmax is None: + return + + data_path = dataset.data_path(subject) + if isinstance(data_path, dict): + first_path = next(iter(data_path.values())) + if isinstance(first_path, list): + first_path = first_path[0] + elif isinstance(data_path, list): + first_path = data_path[0] + else: + first_path = data_path + + raw_check = mne.io.read_raw(first_path, preload=False, verbose=False) + sfreq = raw_check.info["sfreq"] + nyquist = sfreq / 2.0 + + if fmax >= nyquist: + raise ValueError( + f"Bandpass high ({fmax} Hz) exceeds Nyquist frequency ({nyquist} Hz). " + f"Use a value less than {nyquist} Hz or omit --bandpass-high for no filtering." + ) + def download(self, manifest_item) -> Dict[str, Any]: """Download and extract data using MOABB paradigm. This method: 1. Sets up MNE download directory - 2. Calls paradigm.get_data() to get epoched arrays - 3. Filters results to the specific session from manifest_item + 2. Validates bandpass parameters against Nyquist frequency + 3. Calls paradigm.get_data(return_raws=True) to get filtered Raw objects + 4. Filters results to the specific session from manifest_item Parameters ---------- @@ -133,10 +225,9 @@ def download(self, manifest_item) -> Dict[str, Any]: ------- dict Dictionary containing: - - X: np.ndarray of shape (n_epochs, n_channels, n_samples) - - labels: np.ndarray of shape (n_epochs,) with event names + - raws: list of mne.io.Raw objects (filtered, continuous) - meta: pd.DataFrame with columns: subject, session, run - - info: MNE Info object from first epoch (for channel info) + - dataset: BaseDataset instance """ self.update_status("DOWNLOADING") @@ -144,20 +235,20 @@ def download(self, manifest_item) -> Dict[str, Any]: os.environ["MNE_DATA"] = str(self.raw_dir.resolve()) dataset = self.get_dataset() - paradigm = self.get_paradigm() subject = int(manifest_item.subject) session = int(manifest_item.session) - epochs, labels, meta = paradigm.get_data( + self._validate_bandpass_params(dataset, subject) + paradigm = self.get_paradigm(self.args) + + raws, labels, meta = paradigm.get_data( dataset=dataset, subjects=[subject], - return_epochs=True, + return_raws=True, ) - if len(epochs) == 0: - raise ValueError( - f"No epochs found for subject {subject}, session {session}" - ) + if len(raws) == 0: + raise ValueError(f"No data found for subject {subject}, session {session}") session_values = sorted(meta["session"].unique()) if isinstance(session, int): @@ -179,42 +270,32 @@ def download(self, manifest_item) -> Dict[str, Any]: session_mask = meta["session"] == session_key if not session_mask.any(): raise ValueError( - f"No epochs found for subject {subject}, session {session_key}" + f"No data found for subject {subject}, session {session_key}" ) - epochs_filtered = epochs[session_mask] - labels_filtered = labels[session_mask] meta_filtered = meta[session_mask].reset_index(drop=True) - X_filtered = np.concatenate(epochs_filtered, axis=0) - info = epochs_filtered[0].info if len(epochs_filtered) > 0 else None + raws_filtered = [raws[i] for i in range(len(raws)) if session_mask.iloc[i]] return { - "X": X_filtered, - "labels": labels_filtered, + "raws": raws_filtered, "meta": meta_filtered, - "info": info, - "epochs": epochs_filtered, + "dataset": dataset, } - def _get_channel_types(self, info, epochs): - """Extract channel types from MNE info and epochs objects. + def _get_channel_types(self, info): + """Extract channel types from MNE info object. Parameters ---------- info : mne.Info MNE Info object with channel information - epochs : mne.Epochs - MNE Epochs object Returns ------- list[str] List of channel type strings (e.g., "EEG", "EOG", "EMG") """ - if len(epochs) > 0: - return epochs[0].get_channel_types() - # MNE channel kind constants (mne.io.constants.FIFF) ch_type_map = { 2: "EEG", # FIFFV_EEG_CH @@ -228,61 +309,61 @@ def _get_channel_types(self, info, epochs): for ch in info["ch_names"] ] - def _extract_eeg_data(self, X, info, ch_types, labels): - """Extract EEG data and trial intervals from epoched arrays. + def _extract_trials_from_raw(self, raw, dataset) -> Interval: + """Extract trial intervals from Raw annotations using actual trial durations. + + Instead of using a fixed duration, this method calculates trial duration + by finding the time between consecutive events. Each trial starts at an + event onset and ends when the next event begins. Parameters ---------- - X : np.ndarray - Array of shape (n_epochs, n_channels, n_samples) - info : mne.Info - MNE Info object with channel and sampling rate information - ch_types : list[str] - List of channel type strings for each channel - labels : np.ndarray - Array of shape (n_epochs,) with event labels + raw : mne.io.Raw + Raw MNE object with annotations + dataset : BaseDataset + MOABB dataset instance Returns ------- - eeg : RegularTimeSeries - Concatenated EEG signals - channels : ArrayDict - Channel IDs and types - trials : Interval + Interval Trial intervals with start/end times and label fields """ - sfreq = info["sfreq"] - n_epochs, n_channels, n_samples = X.shape + events, _ = mne.events_from_annotations( + raw, event_id=dataset.event_id, verbose=False + ) - eeg_signals = X.transpose(0, 2, 1).reshape(-1, n_channels) + if len(events) == 0: + raise ValueError("No events found in Raw annotations") - eeg = RegularTimeSeries( - signal=eeg_signals, - sampling_rate=sfreq, - domain=Interval( - start=np.array([0.0]), - end=np.array([(len(eeg_signals) - 1) / sfreq]), - ), - ) + sfreq = raw.info["sfreq"] + event_id_to_name = {v: k for k, v in dataset.event_id.items()} - channels = ArrayDict( - id=np.array(info["ch_names"], dtype="U"), - types=np.array(ch_types, dtype="U"), + starts = events[:, 0] / sfreq + + ends = np.zeros(len(events)) + for i in range(len(events) - 1): + ends[i] = events[i + 1, 0] / sfreq + + last_sample = raw.n_times - 1 + ends[-1] = last_sample / sfreq + + labels = np.array( + [ + event_id_to_name.get(event_code, "unknown") + for event_code in events[:, 2] + ], + dtype="U", ) - sample_boundaries = np.arange(n_epochs + 1) * n_samples - time_boundaries = sample_boundaries / sfreq - start_times = time_boundaries[:-1] - end_times = time_boundaries[1:] id_values = np.array([self.label_map.get(label, -1) for label in labels]) trials = Interval( - start=start_times, - end=end_times, - timestamps=(start_times + end_times) / 2, + start=starts, + end=ends, + timestamps=(starts + ends) / 2, timekeys=["start", "end", "timestamps"], **{ - self.label_field: np.asarray(labels), + self.label_field: labels, self.id_field: id_values, }, ) @@ -290,7 +371,42 @@ def _extract_eeg_data(self, X, info, ch_types, labels): if not trials.is_disjoint(): raise ValueError("Found overlapping trials") - return eeg, channels, trials + return trials + + def _extract_continuous_eeg(self, raw): + """Extract continuous EEG signal from Raw object. + + Parameters + ---------- + raw : mne.io.Raw + Raw MNE object with continuous data + + Returns + ------- + eeg : RegularTimeSeries + Continuous EEG signals + channels : ArrayDict + Channel IDs and types + """ + data, times = raw.get_data(return_times=True) + info = raw.info + + eeg = RegularTimeSeries( + signal=data.T, # (n_samples, n_channels) + sampling_rate=info["sfreq"], + domain=Interval( + start=np.array([times[0]]), + end=np.array([times[-1]]), + ), + ) + + ch_types = self._get_channel_types(info) + channels = ArrayDict( + id=np.array(info["ch_names"], dtype="U"), + types=np.array(ch_types, dtype="U"), + ) + + return eeg, channels @abstractmethod def get_brainset_description(self) -> BrainsetDescription: @@ -415,10 +531,11 @@ def process(self, download_output: Dict[str, Any]) -> None: """Transform MOABB data to brainsets format. This default implementation handles the common workflow: - 1. Extract EEG data and channel information - 2. Extract trial intervals with labels - 3. Generate stratified splits - 4. Create and store Data object + 1. Concatenate runs into continuous recording + 2. Extract continuous EEG data and channel information + 3. Extract trial intervals from annotations + 4. Generate stratified splits + 5. Create and store Data object Subclasses can override for custom processing, but typically only need to implement get_brainset_description(). @@ -426,13 +543,11 @@ def process(self, download_output: Dict[str, Any]) -> None: Parameters ---------- download_output : dict - Dictionary returned by download() containing X, labels, meta, info, epochs + Dictionary returned by download() containing raws, meta, dataset """ - X = download_output["X"] - labels = download_output["labels"] + raws = download_output["raws"] meta = download_output["meta"] - info = download_output["info"] - epochs = download_output["epochs"] + dataset = download_output["dataset"] self.update_status("PROCESSING") self.processed_dir.mkdir(exist_ok=True, parents=True) @@ -446,15 +561,25 @@ def process(self, download_output: Dict[str, Any]) -> None: self.update_status("Skipped Processing") return + self.update_status("Concatenating Runs") + if len(raws) > 1: + raw = mne.concatenate_raws(raws, verbose=False) + else: + raw = raws[0] + + info = raw.info + self.update_status("Creating Descriptions") brainset_description = self.get_brainset_description() subject_description = self._get_subject_description(subject_id) session_description = self._get_session_description(session_id, info) device_description = self._get_device_description(subject_id, info) - self.update_status("Extracting EEG and Trials") - ch_types = self._get_channel_types(info, epochs) - eeg, channels, trials = self._extract_eeg_data(X, info, ch_types, labels) + self.update_status("Extracting Continuous EEG") + eeg, channels = self._extract_continuous_eeg(raw) + + self.update_status("Extracting Trial Intervals") + trials = self._extract_trials_from_raw(raw, dataset) self.update_status("Generating Splits") splits = self._generate_splits(trials, subject_id=subject_id) diff --git a/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py b/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py index 16ce2900..cb991d26 100644 --- a/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py +++ b/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py @@ -23,15 +23,13 @@ from brainsets.descriptions import BrainsetDescription from brainsets.taxonomy import Task -from brainsets.utils.moabb.pipeline import MOABBPipeline +from brainsets.utils.moabb.pipeline import MOABBPipeline, _base_parser from brainsets.utils.split import generate_subject_kfold_assignment logging.basicConfig(level=logging.INFO) -parser = ArgumentParser() -parser.add_argument("--redownload", action="store_true") -parser.add_argument("--reprocess", action="store_true") +parser = ArgumentParser(parents=[_base_parser]) class Pipeline(MOABBPipeline): diff --git a/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py b/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py index f414a5df..3f12adb3 100644 --- a/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py +++ b/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py @@ -24,7 +24,7 @@ from temporaldata import Data from brainsets.descriptions import BrainsetDescription from brainsets.taxonomy import Task -from brainsets.utils.moabb.pipeline import MOABBPipeline +from brainsets.utils.moabb.pipeline import MOABBPipeline, _base_parser from brainsets.utils.split import ( generate_trial_folds_by_task, generate_subject_kfold_assignment, @@ -33,9 +33,7 @@ logging.basicConfig(level=logging.INFO) -parser = ArgumentParser() -parser.add_argument("--redownload", action="store_true") -parser.add_argument("--reprocess", action="store_true") +parser = ArgumentParser(parents=[_base_parser]) class Pipeline(MOABBPipeline): From 4babb2161b7b9ff56084252d2519c0132f94b058 Mon Sep 17 00:00:00 2001 From: Milo Sobral Date: Tue, 3 Feb 2026 13:52:42 -0500 Subject: [PATCH 11/13] Fixed error with processing arguments --- brainsets/utils/moabb/pipeline.py | 45 ++++++++++++++++++++++++++----- 1 file changed, 39 insertions(+), 6 deletions(-) diff --git a/brainsets/utils/moabb/pipeline.py b/brainsets/utils/moabb/pipeline.py index ba1554c0..2654150b 100644 --- a/brainsets/utils/moabb/pipeline.py +++ b/brainsets/utils/moabb/pipeline.py @@ -19,6 +19,11 @@ from moabb.datasets.base import BaseDataset from moabb.paradigms.base import BaseParadigm +from moabb.datasets.preprocessing import ( + get_filter_pipeline, + get_resample_pipeline, + make_fixed_pipeline, +) from temporaldata import Data, RegularTimeSeries, Interval, ArrayDict from brainsets.pipeline import BrainsetPipeline from brainsets import serialize_fn_map @@ -168,8 +173,7 @@ def get_manifest(cls, raw_dir: Path, args) -> pd.DataFrame: def _validate_bandpass_params(self, dataset: BaseDataset, subject: int) -> None: """Validate bandpass parameters against Nyquist frequency. - Loads raw data briefly to check sampling rate and raises an error if - bandpass_high exceeds the Nyquist frequency. + Raises an error if bandpass_high exceeds the Nyquist frequency. Parameters ---------- @@ -213,8 +217,9 @@ def download(self, manifest_item) -> Dict[str, Any]: This method: 1. Sets up MNE download directory 2. Validates bandpass parameters against Nyquist frequency - 3. Calls paradigm.get_data(return_raws=True) to get filtered Raw objects - 4. Filters results to the specific session from manifest_item + 3. If no filtering requested: modifies paradigm to skip filtering step + 4. Calls paradigm.get_data(return_raws=True) to get Raw objects + 5. Filters results to the specific session from manifest_item Parameters ---------- @@ -225,7 +230,7 @@ def download(self, manifest_item) -> Dict[str, Any]: ------- dict Dictionary containing: - - raws: list of mne.io.Raw objects (filtered, continuous) + - raws: list of mne.io.Raw objects (filtered or unfiltered) - meta: pd.DataFrame with columns: subject, session, run - dataset: BaseDataset instance """ @@ -241,7 +246,35 @@ def download(self, manifest_item) -> Dict[str, Any]: self._validate_bandpass_params(dataset, subject) paradigm = self.get_paradigm(self.args) - raws, labels, meta = paradigm.get_data( + no_filtering = ( + self.args.bandpass_low is None and self.args.bandpass_high is None + ) + needs_resample = self.args.resample is not None + + if no_filtering: + if needs_resample: + resample_rate = self.args.resample + paradigm._get_raw_pipelines = lambda: [ + get_resample_pipeline(resample_rate) + ] + else: + paradigm._get_raw_pipelines = lambda: [None] + elif needs_resample: + resample_rate = self.args.resample + + def make_filter_resample_pipelines(): + pipelines = [] + for fmin, fmax in paradigm.filters: + combined = make_fixed_pipeline( + get_filter_pipeline(fmin, fmax), + get_resample_pipeline(resample_rate), + ) + pipelines.append(combined) + return pipelines + + paradigm._get_raw_pipelines = make_filter_resample_pipelines + + raws, _, meta = paradigm.get_data( dataset=dataset, subjects=[subject], return_raws=True, From bdc5a485aafaefd9af06ae96136779d92a585a03 Mon Sep 17 00:00:00 2001 From: Milo Sobral Date: Wed, 4 Feb 2026 11:08:21 -0500 Subject: [PATCH 12/13] Add dataset_sign attribute to MOABBPipeline and its subclasses for better dataset identification Updated the Korczowski and Schalk pipelines to include their respective dataset_sign values ("BRAININVADERS2014A" and "EEGBCI"). --- brainsets/utils/moabb/pipeline.py | 17 ++++++++++++++--- .../korczowski_brain_invaders_2014a/pipeline.py | 1 + .../schalk_wolpaw_physionet_2009/pipeline.py | 1 + 3 files changed, 16 insertions(+), 3 deletions(-) diff --git a/brainsets/utils/moabb/pipeline.py b/brainsets/utils/moabb/pipeline.py index 2654150b..bad4816e 100644 --- a/brainsets/utils/moabb/pipeline.py +++ b/brainsets/utils/moabb/pipeline.py @@ -9,7 +9,6 @@ from typing import Dict, Any, Type, Optional from pathlib import Path from argparse import ArgumentParser -import os import pandas as pd import numpy as np import datetime @@ -71,6 +70,7 @@ class MOABBPipeline(BrainsetPipeline): - brainset_id: str - dataset_class: Type[BaseDataset] - paradigm_class: Type[BaseParadigm] + - dataset_sign: str (e.g., "EEGBCI", "BRAININVADERS2014A") - dataset_kwargs: Dict[str, Any] (optional, defaults to {}) - paradigm_kwargs: Dict[str, Any] (optional, defaults to {}) - task: Task (e.g., Task.MOTOR_IMAGERY, Task.P300) @@ -87,12 +87,13 @@ class MOABBPipeline(BrainsetPipeline): - Manifest generation from dataset metadata - Data download via MOABB paradigm.get_data() - Session filtering - - MNE download directory setup + - MNE download directory setup (via mne.set_config) - Default process() workflow (EEG extraction, trial extraction, splits, storage) """ dataset_class: Type[BaseDataset] paradigm_class: Type[BaseParadigm] + dataset_sign: str dataset_kwargs: Dict[str, Any] = {} paradigm_kwargs: Dict[str, Any] = {} @@ -237,7 +238,17 @@ def download(self, manifest_item) -> Dict[str, Any]: self.update_status("DOWNLOADING") self.raw_dir.mkdir(exist_ok=True, parents=True) - os.environ["MNE_DATA"] = str(self.raw_dir.resolve()) + + # Set dataset-specific MNE path to override any stale config entries. + # This prevents failures when users have deleted directories referenced + # in ~/.mne/mne-python.json or have misconfigured MNE_DATA. The + # dataset-specific key (MNE_DATASETS_{SIGN}_PATH) takes precedence + # over the generic MNE_DATA setting. + mne.set_config( + f"MNE_DATASETS_{self.dataset_sign}_PATH", + str(self.raw_dir.resolve()), + set_env=True, + ) dataset = self.get_dataset() subject = int(manifest_item.subject) diff --git a/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py b/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py index cb991d26..48b0cfa4 100644 --- a/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py +++ b/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py @@ -38,6 +38,7 @@ class Pipeline(MOABBPipeline): dataset_class = BI2014a paradigm_class = P300 + dataset_sign = "BRAININVADERS2014A" dataset_kwargs = {} task = Task.P300 diff --git a/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py b/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py index 3f12adb3..6f46faeb 100644 --- a/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py +++ b/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py @@ -42,6 +42,7 @@ class Pipeline(MOABBPipeline): dataset_class = PhysionetMI paradigm_class = MotorImagery + dataset_sign = "EEGBCI" dataset_kwargs = {"imagined": True, "executed": False} task = Task.MOTOR_IMAGERY From 94d71644de2228aeb6f50d85926cb17dc1ba4e51 Mon Sep 17 00:00:00 2001 From: Milo Sobral Date: Fri, 20 Feb 2026 14:21:14 -0500 Subject: [PATCH 13/13] Fixes to both pipelines --- brainsets/utils/moabb/pipeline.py | 173 ++++++++---------- brainsets/utils/split.py | 4 +- .../pipeline.py | 28 ++- .../schalk_wolpaw_physionet_2009/pipeline.py | 24 ++- 4 files changed, 108 insertions(+), 121 deletions(-) diff --git a/brainsets/utils/moabb/pipeline.py b/brainsets/utils/moabb/pipeline.py index bad4816e..1281dd20 100644 --- a/brainsets/utils/moabb/pipeline.py +++ b/brainsets/utils/moabb/pipeline.py @@ -5,7 +5,6 @@ classes, and implement paradigm-specific processing logic. """ -from abc import abstractmethod from typing import Dict, Any, Type, Optional from pathlib import Path from argparse import ArgumentParser @@ -18,11 +17,6 @@ from moabb.datasets.base import BaseDataset from moabb.paradigms.base import BaseParadigm -from moabb.datasets.preprocessing import ( - get_filter_pipeline, - get_resample_pipeline, - make_fixed_pipeline, -) from temporaldata import Data, RegularTimeSeries, Interval, ArrayDict from brainsets.pipeline import BrainsetPipeline from brainsets import serialize_fn_map @@ -36,26 +30,26 @@ from brainsets.utils.split import generate_trial_folds -_base_parser = ArgumentParser(add_help=False) -_base_parser.add_argument( +base_parser = ArgumentParser(add_help=False) +base_parser.add_argument( "--redownload", action="store_true", help="Force redownload of raw data" ) -_base_parser.add_argument( +base_parser.add_argument( "--reprocess", action="store_true", help="Force reprocessing of data" ) -_base_parser.add_argument( +base_parser.add_argument( "--bandpass-low", type=float, default=None, help="Low cutoff frequency for bandpass filter in Hz (default: None, no filtering)", ) -_base_parser.add_argument( +base_parser.add_argument( "--bandpass-high", type=float, default=None, help="High cutoff frequency for bandpass filter in Hz (default: None, no filtering)", ) -_base_parser.add_argument( +base_parser.add_argument( "--resample", type=float, default=None, @@ -68,6 +62,7 @@ class MOABBPipeline(BrainsetPipeline): Subclasses must define: - brainset_id: str + - brainset_description: BrainsetDescription - dataset_class: Type[BaseDataset] - paradigm_class: Type[BaseParadigm] - dataset_sign: str (e.g., "EEGBCI", "BRAININVADERS2014A") @@ -80,9 +75,6 @@ class MOABBPipeline(BrainsetPipeline): - stratify_field: str (e.g., "movements", "targets") - label_map: Dict[str, int] (mapping from label strings to integer IDs) - Subclasses must implement: - - get_brainset_description(): Return dataset-specific BrainsetDescription - The base class handles: - Manifest generation from dataset metadata - Data download via MOABB paradigm.get_data() @@ -94,6 +86,7 @@ class MOABBPipeline(BrainsetPipeline): dataset_class: Type[BaseDataset] paradigm_class: Type[BaseParadigm] dataset_sign: str + brainset_description: BrainsetDescription dataset_kwargs: Dict[str, Any] = {} paradigm_kwargs: Dict[str, Any] = {} @@ -213,13 +206,18 @@ def _validate_bandpass_params(self, dataset: BaseDataset, subject: int) -> None: ) def download(self, manifest_item) -> Dict[str, Any]: - """Download and extract data using MOABB paradigm. + """Download and extract data using MOABB dataset. + + Uses dataset._get_single_subject_data() instead of paradigm.get_data() + to avoid MOABB's SetRawAnnotations transformer, which assigns a fixed + duration to all annotations and silently drops the last event in each + run when that duration exceeds the recording length. This method: 1. Sets up MNE download directory 2. Validates bandpass parameters against Nyquist frequency - 3. If no filtering requested: modifies paradigm to skip filtering step - 4. Calls paradigm.get_data(return_raws=True) to get Raw objects + 3. Loads Raw objects directly via dataset._get_single_subject_data() + 4. Applies bandpass filtering and/or resampling if requested 5. Filters results to the specific session from manifest_item Parameters @@ -255,75 +253,57 @@ def download(self, manifest_item) -> Dict[str, Any]: session = int(manifest_item.session) self._validate_bandpass_params(dataset, subject) - paradigm = self.get_paradigm(self.args) no_filtering = ( self.args.bandpass_low is None and self.args.bandpass_high is None ) needs_resample = self.args.resample is not None - if no_filtering: - if needs_resample: - resample_rate = self.args.resample - paradigm._get_raw_pipelines = lambda: [ - get_resample_pipeline(resample_rate) - ] - else: - paradigm._get_raw_pipelines = lambda: [None] - elif needs_resample: - resample_rate = self.args.resample - - def make_filter_resample_pipelines(): - pipelines = [] - for fmin, fmax in paradigm.filters: - combined = make_fixed_pipeline( - get_filter_pipeline(fmin, fmax), - get_resample_pipeline(resample_rate), - ) - pipelines.append(combined) - return pipelines - - paradigm._get_raw_pipelines = make_filter_resample_pipelines - - raws, _, meta = paradigm.get_data( - dataset=dataset, - subjects=[subject], - return_raws=True, - ) - - if len(raws) == 0: - raise ValueError(f"No data found for subject {subject}, session {session}") + # Load data directly from the dataset, bypassing paradigm.get_data() + # and its SetRawAnnotations transformer which drops events. + recording_data = dataset._get_single_subject_data(subject) - session_values = sorted(meta["session"].unique()) - if isinstance(session, int): - if session < len(session_values): - session_key = session_values[session] - else: - raise ValueError( - f"Session index {session} out of range for subject {subject}. " - f"Available {len(session_values)} sessions: {list(session_values)}" - ) + session_values = sorted(recording_data.keys()) + if session < len(session_values): + session_key = session_values[session] else: - session_key = str(session) - if session_key not in session_values: - raise ValueError( - f"Session {session_key} not found for subject {subject}. " - f"Available sessions: {list(session_values)}" - ) + raise ValueError( + f"Session index {session} out of range for subject {subject}. " + f"Available {len(session_values)} sessions: {list(session_values)}" + ) + + session_data = recording_data[session_key] - session_mask = meta["session"] == session_key - if not session_mask.any(): + if len(session_data) == 0: raise ValueError( f"No data found for subject {subject}, session {session_key}" ) - meta_filtered = meta[session_mask].reset_index(drop=True) + raws = list(session_data.values()) + run_keys = list(session_data.keys()) - raws_filtered = [raws[i] for i in range(len(raws)) if session_mask.iloc[i]] + # Apply filtering and/or resampling if requested + for raw in raws: + if not no_filtering: + raw.filter( + self.args.bandpass_low, + self.args.bandpass_high, + verbose=False, + ) + if needs_resample: + raw.resample(self.args.resample, verbose=False) + + meta = pd.DataFrame( + { + "subject": [subject] * len(raws), + "session": [session_key] * len(raws), + "run": run_keys, + } + ) return { - "raws": raws_filtered, - "meta": meta_filtered, + "raws": raws, + "meta": meta, "dataset": dataset, } @@ -372,15 +352,36 @@ def _extract_trials_from_raw(self, raw, dataset) -> Interval: Interval Trial intervals with start/end times and label fields """ - events, _ = mne.events_from_annotations( - raw, event_id=dataset.event_id, verbose=False - ) + stim_channels = mne.utils._get_stim_channel(None, raw.info, raise_error=False) + if len(stim_channels) > 0: + events = mne.find_events(raw, shortest_event=0, verbose=False) + + valid_codes = set() + for v in dataset.event_id.values(): + if isinstance(v, list): + valid_codes.update(v) + else: + valid_codes.add(v) + + mask = np.isin(events[:, 2], list(valid_codes)) + events = events[mask] + else: + events, _ = mne.events_from_annotations( + raw, event_id=dataset.event_id, verbose=False + ) if len(events) == 0: - raise ValueError("No events found in Raw annotations") + raise ValueError("No events found in Raw annotations or stim channels") sfreq = raw.info["sfreq"] - event_id_to_name = {v: k for k, v in dataset.event_id.items()} + + event_id_to_name = {} + for name, codes in dataset.event_id.items(): + if isinstance(codes, list): + for code in codes: + event_id_to_name[code] = name + else: + event_id_to_name[codes] = name starts = events[:, 0] / sfreq @@ -452,17 +453,6 @@ def _extract_continuous_eeg(self, raw): return eeg, channels - @abstractmethod - def get_brainset_description(self) -> BrainsetDescription: - """Return dataset-specific BrainsetDescription. - - Returns - ------- - BrainsetDescription - Description object with dataset metadata - """ - ... - def _get_subject_description(self, subject_id: str) -> SubjectDescription: """Create subject description from subject ID. @@ -545,7 +535,7 @@ def _get_device_description(self, subject_id: str, info: Any) -> DeviceDescripti id=f"{subject_id}_{recording_date.strftime('%Y%m%d')}", ) - def _generate_splits(self, trials, subject_id: Optional[str] = None): + def generate_splits(self, trials, subject_id: Optional[str] = None): """Generate stratified folds for trials. Parameters @@ -581,8 +571,7 @@ def process(self, download_output: Dict[str, Any]) -> None: 4. Generate stratified splits 5. Create and store Data object - Subclasses can override for custom processing, but typically only - need to implement get_brainset_description(). + Subclasses can override for custom processing when needed. Parameters ---------- @@ -614,7 +603,7 @@ def process(self, download_output: Dict[str, Any]) -> None: info = raw.info self.update_status("Creating Descriptions") - brainset_description = self.get_brainset_description() + brainset_description = self.brainset_description subject_description = self._get_subject_description(subject_id) session_description = self._get_session_description(session_id, info) device_description = self._get_device_description(subject_id, info) @@ -626,7 +615,7 @@ def process(self, download_output: Dict[str, Any]) -> None: trials = self._extract_trials_from_raw(raw, dataset) self.update_status("Generating Splits") - splits = self._generate_splits(trials, subject_id=subject_id) + splits = self.generate_splits(trials, subject_id=subject_id) self.update_status("Creating Data Object") data = Data( diff --git a/brainsets/utils/split.py b/brainsets/utils/split.py index d2f563c0..a2f00d36 100644 --- a/brainsets/utils/split.py +++ b/brainsets/utils/split.py @@ -151,7 +151,9 @@ def generate_train_valid_test_splits(epoch_dict, grid): for name, epoch in epoch_dict.items(): if name == "invalid_presentation_epochs": - logging.warn(f"Found invalid presentation epochs, which will be excluded.") + logging.warning( + "Found invalid presentation epochs, which will be excluded." + ) continue if len(epoch) == 1: train, valid, test = split_one_epoch(epoch, grid) diff --git a/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py b/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py index 48b0cfa4..d2591dae 100644 --- a/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py +++ b/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py @@ -23,17 +23,26 @@ from brainsets.descriptions import BrainsetDescription from brainsets.taxonomy import Task -from brainsets.utils.moabb.pipeline import MOABBPipeline, _base_parser +from brainsets.utils.moabb.pipeline import MOABBPipeline, base_parser from brainsets.utils.split import generate_subject_kfold_assignment logging.basicConfig(level=logging.INFO) -parser = ArgumentParser(parents=[_base_parser]) +parser = ArgumentParser(parents=[base_parser]) class Pipeline(MOABBPipeline): brainset_id = "korczowski_brain_invaders_2014a" + brainset_description = BrainsetDescription( + id="korczowski_brain_invaders_2014a", + origin_version="unknown", + derived_version="1.0.0", + source="https://moabb.neurotechx.com/docs/generated/moabb.datasets.BI2014a.html", + description="Brain Invaders 2014a P300 dataset: EEG recordings from " + "71 subjects performing a visual P300 Brain-Computer Interface task " + "using 16 active dry electrodes.", + ) parser = parser dataset_class = BI2014a @@ -51,7 +60,7 @@ class Pipeline(MOABBPipeline): "NonTarget": 0, } - def _generate_splits(self, trials, subject_id: str = None): + def generate_splits(self, trials, subject_id: str = None): """Generate stratified folds and subject-level k-fold assignments. Generates: @@ -70,7 +79,7 @@ def _generate_splits(self, trials, subject_id: str = None): splits : Data Data object containing all split masks """ - splits = super()._generate_splits(trials, subject_id=subject_id) + splits = super().generate_splits(trials, subject_id=subject_id) if subject_id is not None: subject_assignments = generate_subject_kfold_assignment( @@ -80,14 +89,3 @@ def _generate_splits(self, trials, subject_id: str = None): setattr(splits, key, value) return splits - - def get_brainset_description(self): - return BrainsetDescription( - id="korczowski_brain_invaders_2014a", - origin_version="unknown", - derived_version="1.0.0", - source="https://moabb.neurotechx.com/docs/generated/moabb.datasets.BI2014a.html", - description="Brain Invaders 2014a P300 dataset: EEG recordings from " - "71 subjects performing a visual P300 Brain-Computer Interface task " - "using 16 active dry electrodes.", - ) diff --git a/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py b/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py index 6f46faeb..3c7164b3 100644 --- a/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py +++ b/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py @@ -24,7 +24,7 @@ from temporaldata import Data from brainsets.descriptions import BrainsetDescription from brainsets.taxonomy import Task -from brainsets.utils.moabb.pipeline import MOABBPipeline, _base_parser +from brainsets.utils.moabb.pipeline import MOABBPipeline, base_parser from brainsets.utils.split import ( generate_trial_folds_by_task, generate_subject_kfold_assignment, @@ -33,11 +33,19 @@ logging.basicConfig(level=logging.INFO) -parser = ArgumentParser(parents=[_base_parser]) +parser = ArgumentParser(parents=[base_parser]) class Pipeline(MOABBPipeline): brainset_id = "schalk_wolpaw_physionet_2009" + brainset_description = BrainsetDescription( + id="schalk_wolpaw_physionet_2009", + origin_version="unknown", + derived_version="1.0.0", + source="https://moabb.neurotechx.com/docs/generated/moabb.datasets.PhysionetMI.html", + description="PhysioNet Motor Imagery dataset: over 1500 EEG recordings " + "from 109 volunteers performing motor imagery tasks.", + ) parser = parser dataset_class = PhysionetMI @@ -65,7 +73,7 @@ class Pipeline(MOABBPipeline): "RightHandFeetImagery": ["right_hand", "feet"], } - def _generate_splits(self, trials, subject_id: str = None): + def generate_splits(self, trials, subject_id: str = None): """Generate task-specific and subject-level k-fold splits. Generates: @@ -101,13 +109,3 @@ def _generate_splits(self, trials, subject_id: str = None): return Data(**task_splits, **subject_assignments, domain=trials) else: return Data(**task_splits, domain=trials) - - def get_brainset_description(self): - return BrainsetDescription( - id="schalk_wolpaw_physionet_2009", - origin_version="unknown", - derived_version="1.0.0", - source="https://moabb.neurotechx.com/docs/generated/moabb.datasets.PhysionetMI.html", - description="PhysioNet Motor Imagery dataset: over 1500 EEG recordings " - "from 109 volunteers performing motor imagery tasks.", - )