eegdash 0.3.9.dev182388821__py3-none-any.whl → 0.4.0__py3-none-any.whl

eegdash/features/serialization.py

@@ -1,8 +1,13 @@
-"""Convenience functions for storing and loading of features datasets.
+"""Convenience functions for storing and loading features datasets.
+
+See Also
+--------
+https://github.com/braindecode/braindecode//blob/master/braindecode/datautil/serialization.py#L165-L229
 
-see also: https://github.com/braindecode/braindecode//blob/master/braindecode/datautil/serialization.py#L165-L229
 """
 
+from __future__ import annotations
+
 from pathlib import Path
 
 import pandas as pd
@@ -14,32 +19,40 @@ from braindecode.datautil.serialization import _load_kwargs_json
 from .datasets import FeaturesConcatDataset, FeaturesDataset
 
 
-def load_features_concat_dataset(path, ids_to_load=None, n_jobs=1):
-    """Load a stored FeaturesConcatDataset of FeaturesDatasets from files.
+def load_features_concat_dataset(
+    path: str | Path, ids_to_load: list[int] | None = None, n_jobs: int = 1
+) -> FeaturesConcatDataset:
+    """Load a stored `FeaturesConcatDataset` from a directory.
+
+    This function reconstructs a :class:`FeaturesConcatDataset` by loading
+    individual :class:`FeaturesDataset` instances from subdirectories within
+    the given path. It uses joblib for parallel loading.
 
     Parameters
     ----------
-    path: str | pathlib.Path
-        Path to the directory of the .fif / -epo.fif and .json files.
-    ids_to_load: list of int | None
-        Ids of specific files to load.
-    n_jobs: int
-        Number of jobs to be used to read files in parallel.
+    path : str or pathlib.Path
+        The path to the directory where the dataset was saved. This directory
+        should contain subdirectories (e.g., "0", "1", "2", ...) for each
+        individual dataset.
+    ids_to_load : list of int, optional
+        A list of specific dataset IDs (subdirectory names) to load. If None,
+        all subdirectories in the path will be loaded.
+    n_jobs : int, default 1
+        The number of jobs to use for parallel loading. -1 means using all
+        processors.
 
     Returns
     -------
-    concat_dataset: FeaturesConcatDataset of FeaturesDatasets
+    eegdash.features.datasets.FeaturesConcatDataset
+        A concatenated dataset containing the loaded `FeaturesDataset` instances.
 
     """
     # Make sure we always work with a pathlib.Path
     path = Path(path)
 
-    # else we have a dataset saved in the new way with subdirectories in path
-    # for every dataset with description.json and -feat.parquet,
-    # target_name.json, raw_preproc_kwargs.json, window_kwargs.json,
-    # window_preproc_kwargs.json, features_kwargs.json
     if ids_to_load is None:
-        ids_to_load = [p.name for p in path.iterdir()]
+        # Get all subdirectories and sort them numerically
+        ids_to_load = [p.name for p in path.iterdir() if p.is_dir()]
         ids_to_load = sorted(ids_to_load, key=lambda i: int(i))
     ids_to_load = [str(i) for i in ids_to_load]
 
@@ -47,7 +60,26 @@ def load_features_concat_dataset(path, ids_to_load=None, n_jobs=1):
     return FeaturesConcatDataset(datasets)
 
 
-def _load_parallel(path, i):
+def _load_parallel(path: Path, i: str) -> FeaturesDataset:
+    """Load a single `FeaturesDataset` from its subdirectory.
+
+    This is a helper function for `load_features_concat_dataset` that handles
+    the loading of one dataset's files (features, metadata, descriptions, etc.).
+
+    Parameters
+    ----------
+    path : pathlib.Path
+        The root directory of the saved `FeaturesConcatDataset`.
+    i : str
+        The identifier of the dataset to load, corresponding to its
+        subdirectory name.
+
+    Returns
+    -------
+    eegdash.features.datasets.FeaturesDataset
+        The loaded dataset instance.
+
+    """
     sub_dir = path / i
 
     parquet_name_pattern = "{}-feat.parquet"

eegdash/features/utils.py

@@ -22,7 +22,28 @@ def _extract_features_from_windowsdataset(
     win_ds: EEGWindowsDataset | WindowsDataset,
     feature_extractor: FeatureExtractor,
     batch_size: int = 512,
-):
+) -> FeaturesDataset:
+    """Extract features from a single `WindowsDataset`.
+
+    This is a helper function that iterates through a `WindowsDataset` in
+    batches, applies a `FeatureExtractor`, and returns the results as a
+    `FeaturesDataset`.
+
+    Parameters
+    ----------
+    win_ds : EEGWindowsDataset or WindowsDataset
+        The windowed dataset to extract features from.
+    feature_extractor : FeatureExtractor
+        The feature extractor instance to apply.
+    batch_size : int, default 512
+        The number of windows to process in each batch.
+
+    Returns
+    -------
+    FeaturesDataset
+        A new dataset containing the extracted features and associated metadata.
+
+    """
     metadata = win_ds.metadata
     if not win_ds.targets_from == "metadata":
         metadata = copy.deepcopy(metadata)
@@ -51,18 +72,16 @@ def _extract_features_from_windowsdataset(
             features_dict[k].extend(v)
     features_df = pd.DataFrame(features_dict)
     if not win_ds.targets_from == "metadata":
-        metadata.set_index("orig_index", drop=False, inplace=True)
         metadata.reset_index(drop=True, inplace=True)
-        metadata.drop("orig_index", axis=1, inplace=True)
+        metadata.drop("orig_index", axis=1, inplace=True, errors="ignore")
 
-    # FUTURE: truly support WindowsDataset objects
     return FeaturesDataset(
         features_df,
         metadata=metadata,
         description=win_ds.description,
         raw_info=win_ds.raw.info,
-        raw_preproc_kwargs=win_ds.raw_preproc_kwargs,
-        window_kwargs=win_ds.window_kwargs,
+        raw_preproc_kwargs=getattr(win_ds, "raw_preproc_kwargs", None),
+        window_kwargs=getattr(win_ds, "window_kwargs", None),
         features_kwargs=feature_extractor.features_kwargs,
     )
 
@@ -73,7 +92,34 @@ def extract_features(
     *,
     batch_size: int = 512,
     n_jobs: int = 1,
-):
+) -> FeaturesConcatDataset:
+    """Extract features from a concatenated dataset of windows.
+
+    This function applies a feature extractor to each `WindowsDataset` within a
+    `BaseConcatDataset` in parallel and returns a `FeaturesConcatDataset`
+    with the results.
+
+    Parameters
+    ----------
+    concat_dataset : BaseConcatDataset
+        A concatenated dataset of `WindowsDataset` or `EEGWindowsDataset`
+        instances.
+    features : FeatureExtractor or dict or list
+        The feature extractor(s) to apply. Can be a `FeatureExtractor`
+        instance, a dictionary of named feature functions, or a list of
+        feature functions.
+    batch_size : int, default 512
+        The size of batches to use for feature extraction.
+    n_jobs : int, default 1
+        The number of parallel jobs to use for extracting features from the
+        datasets.
+
+    Returns
+    -------
+    FeaturesConcatDataset
+        A new concatenated dataset containing the extracted features.
+
+    """
     if isinstance(features, list):
         features = dict(enumerate(features))
     if not isinstance(features, FeatureExtractor):
@@ -97,7 +143,28 @@ def fit_feature_extractors(
     concat_dataset: BaseConcatDataset,
     features: FeatureExtractor | Dict[str, Callable] | List[Callable],
     batch_size: int = 8192,
-):
+) -> FeatureExtractor:
+    """Fit trainable feature extractors on a dataset.
+
+    If the provided feature extractor (or any of its sub-extractors) is
+    trainable (i.e., subclasses `TrainableFeature`), this function iterates
+    through the dataset to fit it.
+
+    Parameters
+    ----------
+    concat_dataset : BaseConcatDataset
+        The dataset to use for fitting the feature extractors.
+    features : FeatureExtractor or dict or list
+        The feature extractor(s) to fit.
+    batch_size : int, default 8192
+        The batch size to use when iterating through the dataset for fitting.
+
+    Returns
+    -------
+    FeatureExtractor
+        The fitted feature extractor.
+
+    """
     if isinstance(features, list):
         features = dict(enumerate(features))
     if not isinstance(features, FeatureExtractor):

eegdash/hbn/__init__.py

@@ -1,3 +1,14 @@
+# Authors: The EEGDash contributors.
+# License: GNU General Public License
+# Copyright the EEGDash contributors.
+
+"""Healthy Brain Network (HBN) specific utilities and preprocessing.
+
+This module provides specialized functions for working with the Healthy Brain Network
+dataset, including preprocessing pipelines, annotation handling, and windowing utilities
+tailored for HBN EEG data analysis.
+"""
+
 from .preprocessing import hbn_ec_ec_reannotation
 from .windows import (
     add_aux_anchors,

eegdash/hbn/preprocessing.py

@@ -1,35 +1,64 @@
-import logging
+# Authors: The EEGDash contributors.
+# License: GNU General Public License
+# Copyright the EEGDash contributors.
+
+"""Preprocessing utilities specific to the Healthy Brain Network dataset.
+
+This module contains preprocessing classes and functions designed specifically for
+HBN EEG data, including specialized annotation handling for eyes-open/eyes-closed
+paradigms and other HBN-specific preprocessing steps.
+"""
 
 import mne
 import numpy as np
 
 from braindecode.preprocessing import Preprocessor
 
-logger = logging.getLogger("eegdash")
+from ..logging import logger
 
 
 class hbn_ec_ec_reannotation(Preprocessor):
-    """Preprocessor to reannotate the raw data for eyes open and eyes closed events.
+    """Preprocessor to reannotate HBN data for eyes-open/eyes-closed events.
+
+    This preprocessor is specifically designed for Healthy Brain Network (HBN)
+    datasets. It identifies existing annotations for "instructed_toCloseEyes"
+    and "instructed_toOpenEyes" and creates new, regularly spaced annotations
+    for "eyes_closed" and "eyes_open" segments, respectively.
 
-    This processor is designed for HBN datasets.
+    This is useful for creating windowed datasets based on these new, more
+    precise event markers.
+
+    Notes
+    -----
+    This class inherits from :class:`braindecode.preprocessing.Preprocessor`
+    and is intended to be used within a braindecode preprocessing pipeline.
 
     """
 
     def __init__(self):
         super().__init__(fn=self.transform, apply_on_array=False)
 
-    def transform(self, raw):
-        """Reannotate the raw data to create new events for eyes open and eyes closed
+    def transform(self, raw: mne.io.Raw) -> mne.io.Raw:
+        """Create new annotations for eyes-open and eyes-closed periods.
+
+        This function finds the original "instructed_to..." annotations and
+        generates new annotations every 2 seconds within specific time ranges
+        relative to the original markers:
+        - "eyes_closed": 15s to 29s after "instructed_toCloseEyes"
+        - "eyes_open": 5s to 19s after "instructed_toOpenEyes"
 
-        This function modifies the raw MNE object by creating new events based on
-        the existing annotations for "instructed_toCloseEyes" and "instructed_toOpenEyes".
-        It generates new events every 2 seconds within specified time ranges after
-        the original events, and replaces the existing annotations with these new events.
+        The original annotations in the `mne.io.Raw` object are replaced by
+        this new set of annotations.
 
         Parameters
         ----------
         raw : mne.io.Raw
-            The raw MNE object containing EEG data and annotations.
+            The raw MNE object containing the HBN data and original annotations.
+
+        Returns
+        -------
+        mne.io.Raw
+            The raw MNE object with the modified annotations.
 
         """
         events, event_id = mne.events_from_annotations(raw)
@@ -39,15 +68,27 @@ class hbn_ec_ec_reannotation(Preprocessor):
         # Create new events array for 2-second segments
         new_events = []
         sfreq = raw.info["sfreq"]
-        for event in events[events[:, 2] == event_id["instructed_toCloseEyes"]]:
-            # For each original event, create events every 2 seconds from 15s to 29s after
-            start_times = event[0] + np.arange(15, 29, 2) * sfreq
-            new_events.extend([[int(t), 0, 1] for t in start_times])
 
-        for event in events[events[:, 2] == event_id["instructed_toOpenEyes"]]:
-            # For each original event, create events every 2 seconds from 5s to 19s after
-            start_times = event[0] + np.arange(5, 19, 2) * sfreq
-            new_events.extend([[int(t), 0, 2] for t in start_times])
+        close_event_id = event_id.get("instructed_toCloseEyes")
+        if close_event_id:
+            for event in events[events[:, 2] == close_event_id]:
+                # For each original event, create events every 2s from 15s to 29s after
+                start_times = event[0] + np.arange(15, 29, 2) * sfreq
+                new_events.extend([[int(t), 0, 1] for t in start_times])
+
+        open_event_id = event_id.get("instructed_toOpenEyes")
+        if open_event_id:
+            for event in events[events[:, 2] == open_event_id]:
+                # For each original event, create events every 2s from 5s to 19s after
+                start_times = event[0] + np.arange(5, 19, 2) * sfreq
+                new_events.extend([[int(t), 0, 2] for t in start_times])
+
+        if not new_events:
+            logger.warning(
+                "Could not find 'instructed_toCloseEyes' or 'instructed_toOpenEyes' "
+                "annotations. No new events created."
+            )
+            return raw
 
         # replace events in raw
         new_events = np.array(new_events)
@@ -56,6 +97,7 @@ class hbn_ec_ec_reannotation(Preprocessor):
             events=new_events,
             event_desc={1: "eyes_closed", 2: "eyes_open"},
             sfreq=raw.info["sfreq"],
+            orig_time=raw.info.get("meas_date"),
         )
 
         raw.set_annotations(annot_from_events)

eegdash/hbn/windows.py

@@ -1,3 +1,15 @@
+# Authors: The EEGDash contributors.
+# License: GNU General Public License
+# Copyright the EEGDash contributors.
+
+"""Windowing and trial processing utilities for HBN datasets.
+
+This module provides functions for building trial tables, adding auxiliary anchors,
+annotating trials with targets, and filtering recordings based on various criteria.
+These utilities are specifically designed for working with HBN EEG data structures
+and experimental paradigms.
+"""
+
 import logging
 
 import mne
@@ -7,11 +19,27 @@ from mne_bids import get_bids_path_from_fname
 
 from braindecode.datasets.base import BaseConcatDataset
 
-logger = logging.getLogger("eegdash")
-
 
 def build_trial_table(events_df: pd.DataFrame) -> pd.DataFrame:
-    """One row per contrast trial with stimulus/response metrics."""
+    """Build a table of contrast trials from an events DataFrame.
+
+    This function processes a DataFrame of events (typically from a BIDS
+    `events.tsv` file) to identify contrast trials and extract relevant
+    metrics like stimulus onset, response onset, and reaction times.
+
+    Parameters
+    ----------
+    events_df : pandas.DataFrame
+        A DataFrame containing event information, with at least "onset" and
+        "value" columns.
+
+    Returns
+    -------
+    pandas.DataFrame
+        A DataFrame where each row represents a single contrast trial, with
+        columns for onsets, reaction times, and response correctness.
+
+    """
     events_df = events_df.copy()
     events_df["onset"] = pd.to_numeric(events_df["onset"], errors="raise")
     events_df = events_df.sort_values("onset", kind="mergesort").reset_index(drop=True)
@@ -82,12 +110,13 @@ def build_trial_table(events_df: pd.DataFrame) -> pd.DataFrame:
     return pd.DataFrame(rows)
 
 
-# Aux functions to inject the annot
 def _to_float_or_none(x):
+    """Safely convert a value to float or None."""
     return None if pd.isna(x) else float(x)
 
 
 def _to_int_or_none(x):
+    """Safely convert a value to int or None."""
     if pd.isna(x):
         return None
     if isinstance(x, (bool, np.bool_)):
@@ -96,22 +125,55 @@ def _to_int_or_none(x):
         return int(x)
     try:
         return int(x)
-    except Exception:
+    except (ValueError, TypeError):
         return None
 
 
 def _to_str_or_none(x):
+    """Safely convert a value to string or None."""
     return None if (x is None or (isinstance(x, float) and np.isnan(x))) else str(x)
 
 
 def annotate_trials_with_target(
-    raw,
-    target_field="rt_from_stimulus",
-    epoch_length=2.0,
-    require_stimulus=True,
-    require_response=True,
-):
-    """Create 'contrast_trial_start' annotations with float target in extras."""
+    raw: mne.io.Raw,
+    target_field: str = "rt_from_stimulus",
+    epoch_length: float = 2.0,
+    require_stimulus: bool = True,
+    require_response: bool = True,
+) -> mne.io.Raw:
+    """Create trial annotations with a specified target value.
+
+    This function reads the BIDS events file associated with the `raw` object,
+    builds a trial table, and creates new MNE annotations for each trial.
+    The annotations are labeled "contrast_trial_start" and their `extras`
+    dictionary is populated with trial metrics, including a "target" key.
+
+    Parameters
+    ----------
+    raw : mne.io.Raw
+        The raw data object. Must have a single associated file name from
+        which the BIDS path can be derived.
+    target_field : str, default "rt_from_stimulus"
+        The column from the trial table to use as the "target" value in the
+        annotation extras.
+    epoch_length : float, default 2.0
+        The duration to set for each new annotation.
+    require_stimulus : bool, default True
+        If True, only include trials that have a recorded stimulus event.
+    require_response : bool, default True
+        If True, only include trials that have a recorded response event.
+
+    Returns
+    -------
+    mne.io.Raw
+        The `raw` object with the new annotations set.
+
+    Raises
+    ------
+    KeyError
+        If `target_field` is not a valid column in the built trial table.
+
+    """
     fnames = raw.filenames
     assert len(fnames) == 1, "Expected a single filename"
     bids_path = get_bids_path_from_fname(fnames[0])
@@ -142,7 +204,6 @@ def annotate_trials_with_target(
     extras = []
     for i, v in enumerate(targets):
         row = trials.iloc[i]
-
         extras.append(
             {
                 "target": _to_float_or_none(v),
@@ -159,14 +220,39 @@ def annotate_trials_with_target(
         onset=onsets,
         duration=durations,
         description=descs,
-        orig_time=raw.info["meas_date"],
+        orig_time=raw.info.get("meas_date"),
         extras=extras,
     )
     raw.set_annotations(new_ann, verbose=False)
     return raw
 
 
-def add_aux_anchors(raw, stim_desc="stimulus_anchor", resp_desc="response_anchor"):
+def add_aux_anchors(
+    raw: mne.io.Raw,
+    stim_desc: str = "stimulus_anchor",
+    resp_desc: str = "response_anchor",
+) -> mne.io.Raw:
+    """Add auxiliary annotations for stimulus and response onsets.
+
+    This function inspects existing "contrast_trial_start" annotations and
+    adds new, zero-duration "anchor" annotations at the precise onsets of
+    stimuli and responses for each trial.
+
+    Parameters
+    ----------
+    raw : mne.io.Raw
+        The raw data object with "contrast_trial_start" annotations.
+    stim_desc : str, default "stimulus_anchor"
+        The description for the new stimulus annotations.
+    resp_desc : str, default "response_anchor"
+        The description for the new response annotations.
+
+    Returns
+    -------
+    mne.io.Raw
+        The `raw` object with the auxiliary annotations added.
+
+    """
     ann = raw.annotations
     mask = ann.description == "contrast_trial_start"
     if not np.any(mask):
@@ -179,28 +265,24 @@ def add_aux_anchors(raw, stim_desc="stimulus_anchor", resp_desc="response_anchor
         ex = ann.extras[idx] if ann.extras is not None else {}
         t0 = float(ann.onset[idx])
 
-        stim_t = ex["stimulus_onset"]
-        resp_t = ex["response_onset"]
+        stim_t = ex.get("stimulus_onset")
+        resp_t = ex.get("response_onset")
 
         if stim_t is None or (isinstance(stim_t, float) and np.isnan(stim_t)):
-            rtt = ex["rt_from_trialstart"]
-            rts = ex["rt_from_stimulus"]
+            rtt = ex.get("rt_from_trialstart")
+            rts = ex.get("rt_from_stimulus")
             if rtt is not None and rts is not None:
                 stim_t = t0 + float(rtt) - float(rts)
 
         if resp_t is None or (isinstance(resp_t, float) and np.isnan(resp_t)):
-            rtt = ex["rt_from_trialstart"]
+            rtt = ex.get("rt_from_trialstart")
             if rtt is not None:
                 resp_t = t0 + float(rtt)
 
-        if (stim_t is not None) and not (
-            isinstance(stim_t, float) and np.isnan(stim_t)
-        ):
+        if stim_t is not None and not (isinstance(stim_t, float) and np.isnan(stim_t)):
             stim_onsets.append(float(stim_t))
             stim_extras.append(dict(ex, anchor="stimulus"))
-        if (resp_t is not None) and not (
-            isinstance(resp_t, float) and np.isnan(resp_t)
-        ):
+        if resp_t is not None and not (isinstance(resp_t, float) and np.isnan(resp_t)):
             resp_onsets.append(float(resp_t))
             resp_extras.append(dict(ex, anchor="response"))
 
@@ -210,7 +292,7 @@ def add_aux_anchors(raw, stim_desc="stimulus_anchor", resp_desc="response_anchor
             onset=new_onsets,
             duration=np.zeros_like(new_onsets, dtype=float),
             description=[stim_desc] * len(stim_onsets) + [resp_desc] * len(resp_onsets),
-            orig_time=raw.info["meas_date"],
+            orig_time=raw.info.get("meas_date"),
             extras=stim_extras + resp_extras,
         )
         raw.set_annotations(ann + aux, verbose=False)
@@ -218,10 +300,10 @@ def add_aux_anchors(raw, stim_desc="stimulus_anchor", resp_desc="response_anchor
 
 
 def add_extras_columns(
-    windows_concat_ds,
-    original_concat_ds,
-    desc="contrast_trial_start",
-    keys=(
+    windows_concat_ds: BaseConcatDataset,
+    original_concat_ds: BaseConcatDataset,
+    desc: str = "contrast_trial_start",
+    keys: tuple = (
         "target",
         "rt_from_stimulus",
         "rt_from_trialstart",
@@ -230,7 +312,31 @@ def add_extras_columns(
         "correct",
         "response_type",
     ),
-):
+) -> BaseConcatDataset:
+    """Add columns from annotation extras to a windowed dataset's metadata.
+
+    This function propagates trial-level information stored in the `extras`
+    of annotations to the `metadata` DataFrame of a `WindowsDataset`.
+
+    Parameters
+    ----------
+    windows_concat_ds : BaseConcatDataset
+        The windowed dataset whose metadata will be updated.
+    original_concat_ds : BaseConcatDataset
+        The original (non-windowed) dataset containing the raw data and
+        annotations with the `extras` to be added.
+    desc : str, default "contrast_trial_start"
+        The description of the annotations to source the extras from.
+    keys : tuple, default (...)
+        The keys to extract from each annotation's `extras` dictionary and
+        add as columns to the metadata.
+
+    Returns
+    -------
+    BaseConcatDataset
+        The `windows_concat_ds` with updated metadata.
+
+    """
     float_cols = {
         "target",
         "rt_from_stimulus",
@@ -282,7 +388,6 @@ def add_extras_columns(
             else:  # response_type
                 ser = pd.Series(vals, index=md.index, dtype="string")
 
-            # Replace the whole column to avoid dtype conflicts
             md[k] = ser
 
         win_ds.metadata = md.reset_index(drop=True)
@@ -293,7 +398,25 @@ def add_extras_columns(
     return windows_concat_ds
 
 
-def keep_only_recordings_with(desc, concat_ds):
+def keep_only_recordings_with(
+    desc: str, concat_ds: BaseConcatDataset
+) -> BaseConcatDataset:
+    """Filter a concatenated dataset to keep only recordings with a specific annotation.
+
+    Parameters
+    ----------
+    desc : str
+        The description of the annotation that must be present in a recording
+        for it to be kept.
+    concat_ds : BaseConcatDataset
+        The concatenated dataset to filter.
+
+    Returns
+    -------
+    BaseConcatDataset
+        A new concatenated dataset containing only the filtered recordings.
+
+    """
     kept = []
     for ds in concat_ds.datasets:
         if np.any(ds.raw.annotations.description == desc):