diff --git a/brainsets/taxonomy/task.py b/brainsets/taxonomy/task.py index afb7be6c..dfe4dcf4 100644 --- a/brainsets/taxonomy/task.py +++ b/brainsets/taxonomy/task.py @@ -23,6 +23,12 @@ class Task(StringIntEnum): # Full sentence speaking CONTINUOUS_SPEAKING_SENTENCE = 7 + # 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/utils/moabb/pipeline.py b/brainsets/utils/moabb/pipeline.py new file mode 100644 index 00000000..1281dd20 --- /dev/null +++ b/brainsets/utils/moabb/pipeline.py @@ -0,0 +1,637 @@ +"""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 typing import Dict, Any, Type, Optional +from pathlib import Path +from argparse import ArgumentParser +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 +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_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. + + Subclasses must define: + - brainset_id: str + - brainset_description: BrainsetDescription + - 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) + - 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) + + The base class handles: + - Manifest generation from dataset metadata + - Data download via MOABB paradigm.get_data() + - Session filtering + - 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 + brainset_description: BrainsetDescription + 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.""" + return cls.dataset_class(**cls.dataset_kwargs) + + @classmethod + 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 + 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): + # Make sure session is an integer + 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 _validate_bandpass_params(self, dataset: BaseDataset, subject: int) -> None: + """Validate bandpass parameters against Nyquist frequency. + + 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 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. 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 + ---------- + manifest_item : NamedTuple + Row from manifest containing subject and session info + + Returns + ------- + dict + Dictionary containing: + - raws: list of mne.io.Raw objects (filtered or unfiltered) + - meta: pd.DataFrame with columns: subject, session, run + - dataset: BaseDataset instance + """ + self.update_status("DOWNLOADING") + + self.raw_dir.mkdir(exist_ok=True, parents=True) + + # 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) + session = int(manifest_item.session) + + self._validate_bandpass_params(dataset, subject) + + no_filtering = ( + self.args.bandpass_low is None and self.args.bandpass_high is None + ) + needs_resample = self.args.resample is not None + + # 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(recording_data.keys()) + 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_data = recording_data[session_key] + + if len(session_data) == 0: + raise ValueError( + f"No data found for subject {subject}, session {session_key}" + ) + + raws = list(session_data.values()) + run_keys = list(session_data.keys()) + + # 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, + "meta": meta, + "dataset": dataset, + } + + def _get_channel_types(self, info): + """Extract channel types from MNE info object. + + Parameters + ---------- + info : mne.Info + MNE Info object with channel information + + Returns + ------- + list[str] + List of channel type strings (e.g., "EEG", "EOG", "EMG") + """ + # MNE channel kind constants (mne.io.constants.FIFF) + ch_type_map = { + 2: "EEG", # FIFFV_EEG_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") + for ch in info["ch_names"] + ] + + 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 + ---------- + raw : mne.io.Raw + Raw MNE object with annotations + dataset : BaseDataset + MOABB dataset instance + + Returns + ------- + Interval + Trial intervals with start/end times and label fields + """ + 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 or stim channels") + + sfreq = raw.info["sfreq"] + + 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 + + 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", + ) + + id_values = np.array([self.label_map.get(label, -1) for label in labels]) + + trials = Interval( + start=starts, + end=ends, + timestamps=(starts + ends) / 2, + timekeys=["start", "end", "timestamps"], + **{ + self.label_field: labels, + self.id_field: id_values, + }, + ) + + if not trials.is_disjoint(): + raise ValueError("Found overlapping 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 + + 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( + 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, + 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( + 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: Optional[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 + ------- + splits : Data + Data object containing fold splits + """ + folds = generate_trial_folds( + trials, + stratify_by=self.stratify_field, + n_folds=3, + 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. + + This default implementation handles the common workflow: + 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 when needed. + + Parameters + ---------- + download_output : dict + Dictionary returned by download() containing raws, meta, dataset + """ + raws = download_output["raws"] + meta = download_output["meta"] + dataset = download_output["dataset"] + + 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" + safe_reprocess = getattr(getattr(self, "args", None), "reprocess", False) + if store_path.exists() and not safe_reprocess: + 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.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 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) + + self.update_status("Creating Data Object") + data = Data( + brainset=brainset_description, + subject=subject_description, + session=session_description, + device=device_description, + eeg=eeg, + channels=channels, + **{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/utils/split.py b/brainsets/utils/split.py index 881686a6..a2f00d36 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 @@ -149,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) @@ -237,39 +241,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. - - Returns: - List of Data objects, one for each fold. + For **cross-subject splitting** (where entire subjects are held out), use + :func:`generate_subject_kfold_assignment` instead. - Raises: - ValueError: If the intervals don't have the specified stratify_by attribute. - ValueError: If there are fewer samples than n_folds. + 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. + + 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 @@ -279,25 +318,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( @@ -310,8 +349,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 @@ -325,3 +364,198 @@ def generate_stratified_folds( folds.append(fold_data) return folds + + +def generate_trial_folds_by_task( + 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 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 **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 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. + label_field : str + 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 + 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. + + 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( + 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_trial_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 generate_subject_kfold_assignment( + subject_id: str, n_folds: int = 5, val_ratio: float = 0.2, seed: int = 42 +) -> Dict[str, str]: + """ + 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. + + For each fold k: + - 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", "sub-01"). + 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). + + 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") + 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/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 new file mode 100644 index 00000000..d2591dae --- /dev/null +++ b/brainsets_pipelines/korczowski_brain_invaders_2014a/pipeline.py @@ -0,0 +1,91 @@ +# /// 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.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]) + + +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 + paradigm_class = P300 + dataset_sign = "BRAININVADERS2014A" + 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 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 = 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) + + return splits 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..3c7164b3 --- /dev/null +++ b/brainsets_pipelines/schalk_wolpaw_physionet_2009/pipeline.py @@ -0,0 +1,111 @@ +# /// 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 +import logging + +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.utils.moabb.pipeline import MOABBPipeline, base_parser +from brainsets.utils.split import ( + generate_trial_folds_by_task, + generate_subject_kfold_assignment, +) + + +logging.basicConfig(level=logging.INFO) + +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 + paradigm_class = MotorImagery + dataset_sign = "EEGBCI" + dataset_kwargs = {"imagined": True, "executed": False} + + 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, + } + + TASK_CONFIGS = { + "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"], + } + + 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_trial_folds_by_task( + trials, + task_configs=self.TASK_CONFIGS, + label_field=self.label_field, + n_folds=3, + val_ratio=0.2, + seed=42, + ) + + if subject_id is not None: + 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: + return Data(**task_splits, domain=trials) diff --git a/tests/test_split_utils.py b/tests/test_split_utils.py index 1a7b6946..f0763f67 100644 --- a/tests/test_split_utils.py +++ b/tests/test_split_utils.py @@ -3,7 +3,9 @@ from temporaldata import Data, Interval from brainsets.utils.split import ( chop_intervals, - generate_stratified_folds, + generate_trial_folds, + generate_trial_folds_by_task, + generate_subject_kfold_assignment, ) @@ -84,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 @@ -103,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 ) @@ -162,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 @@ -177,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 ) @@ -192,10 +194,212 @@ 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 TestGenerateTrialFoldsByTask: + 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_trial_folds_by_task( + 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_trial_folds_by_task( + 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_trial_folds_by_task( + 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_trial_folds_by_task( + 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 TestGenerateSubjectKfoldAssignment: + def test_basic_output_structure(self): + assignments = generate_subject_kfold_assignment("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 = 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 = 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 = 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 = generate_subject_kfold_assignment( + "S001", n_folds=5, seed=42 + ) + assignments_seed2 = generate_subject_kfold_assignment( + "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 = generate_subject_kfold_assignment( + f"S{i:03d}", n_folds=n_folds, val_ratio=0.1, seed=42 + ) + assignments_high = generate_subject_kfold_assignment( + 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 = generate_subject_kfold_assignment( + 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