braindecode 1.2.0.dev168908090__py3-none-any.whl → 1.2.0.dev175267687__py3-none-any.whl

braindecode/datasets/experimental.py

@@ -0,0 +1,218 @@
+from __future__ import annotations
+
+import random
+from pathlib import Path
+from typing import Callable, Sequence
+
+import mne_bids
+from torch.utils.data import IterableDataset, get_worker_info
+
+
+class BIDSIterableDataset(IterableDataset):
+    """Dataset for loading BIDS.
+
+    .. warning::
+        This class is experimental and may change in the future.
+
+    .. warning::
+        This dataset is not consistent with the Braindecode API.
+
+    This class has the same parameters as the :func:`mne_bids.find_matching_paths` function
+    as it will be used to find the files to load. The default ``extensions`` parameter was changed.
+
+    More information on BIDS (Brain Imaging Data Structure)
+    can be found at https://bids.neuroimaging.io
+
+    Examples
+    --------
+    >>> from braindecode.datasets import BaseDataset, BaseConcatDataset
+    >>> from braindecode.datasets.bids import BIDSIterableDataset, _description_from_bids_path
+    >>> from braindecode.preprocessing import create_fixed_length_windows
+    >>>
+    >>> def my_reader_fn(path):
+    ...     raw = mne_bids.read_raw_bids(path)
+    ...     desc = _description_from_bids_path(path)
+    ...     ds = BaseDataset(raw, description=desc)
+    ...     windows_ds = create_fixed_length_windows(
+    ...         BaseConcatDataset([ds]),
+    ...         window_size_samples=400,
+    ...         window_stride_samples=200,
+    ...     )
+    ...     return windows_ds
+    >>>
+    >>> dataset = BIDSIterableDataset(
+    ...     reader_fn=my_reader_fn,
+    ...     root="root/of/my/bids/dataset/",
+    ... )
+
+    Parameters
+    ----------
+    reader_fn : Callable[[mne_bids.BIDSPath], Sequence]
+        A function that takes a BIDSPath and returns a dataset.
+    pool_size : int
+        The number of recordings to read and sample from.
+    bids_paths : list[mne_bids.BIDSPath] | None
+        A list of BIDSPaths to load. If None, will use the paths found by
+        :func:`mne_bids.find_matching_paths` and the arguments below.
+    root : pathlib.Path | str
+        The root of the BIDS path.
+    subjects : str | array-like of str | None
+        The subject ID. Corresponds to "sub".
+    sessions : str | array-like of str | None
+        The acquisition session. Corresponds to "ses".
+    tasks : str | array-like of str | None
+        The experimental task. Corresponds to "task".
+    acquisitions: str | array-like of str | None
+        The acquisition parameters. Corresponds to "acq".
+    runs : str | array-like of str | None
+        The run number. Corresponds to "run".
+    processings : str | array-like of str | None
+        The processing label. Corresponds to "proc".
+    recordings : str | array-like of str | None
+        The recording name. Corresponds to "rec".
+    spaces : str | array-like of str | None
+        The coordinate space for anatomical and sensor location
+        files (e.g., ``*_electrodes.tsv``, ``*_markers.mrk``).
+        Corresponds to "space".
+        Note that valid values for ``space`` must come from a list
+        of BIDS keywords as described in the BIDS specification.
+    splits : str | array-like of str | None
+        The split of the continuous recording file for ``.fif`` data.
+        Corresponds to "split".
+    descriptions : str | array-like of str | None
+        This corresponds to the BIDS entity ``desc``. It is used to provide
+        additional information for derivative data, e.g., preprocessed data
+        may be assigned ``description='cleaned'``.
+    suffixes : str | array-like of str | None
+        The filename suffix. This is the entity after the
+        last ``_`` before the extension. E.g., ``'channels'``.
+        The following filename suffix's are accepted:
+        'meg', 'markers', 'eeg', 'ieeg', 'T1w',
+        'participants', 'scans', 'electrodes', 'coordsystem',
+        'channels', 'events', 'headshape', 'digitizer',
+        'beh', 'physio', 'stim'
+    extensions : str | array-like of str | None
+        The extension of the filename. E.g., ``'.json'``.
+        By default, uses the ones accepted by :func:`mne_bids.read_raw_bids`.
+    datatypes : str | array-like of str | None
+        The BIDS data type, e.g., ``'anat'``, ``'func'``, ``'eeg'``, ``'meg'``,
+        ``'ieeg'``.
+    check : bool
+        If ``True``, only returns paths that conform to BIDS. If ``False``
+        (default), the ``.check`` attribute of the returned
+        :class:`mne_bids.BIDSPath` object will be set to ``True`` for paths that
+        do conform to BIDS, and to ``False`` for those that don't.
+    preload : bool
+        If True, preload the data. Defaults to False.
+    n_jobs : int
+        Number of jobs to run in parallel. Defaults to 1.
+    """
+
+    def __init__(
+        self,
+        reader_fn: Callable[[mne_bids.BIDSPath], Sequence],
+        pool_size: int = 4,
+        bids_paths: list[mne_bids.BIDSPath] | None = None,
+        root: Path | str | None = None,
+        subjects: str | list[str] | None = None,
+        sessions: str | list[str] | None = None,
+        tasks: str | list[str] | None = None,
+        acquisitions: str | list[str] | None = None,
+        runs: str | list[str] | None = None,
+        processings: str | list[str] | None = None,
+        recordings: str | list[str] | None = None,
+        spaces: str | list[str] | None = None,
+        splits: str | list[str] | None = None,
+        descriptions: str | list[str] | None = None,
+        suffixes: str | list[str] | None = None,
+        extensions: str | list[str] | None = [
+            ".con",
+            ".sqd",
+            ".pdf",
+            ".fif",
+            ".ds",
+            ".vhdr",
+            ".set",
+            ".edf",
+            ".bdf",
+            ".EDF",
+            ".snirf",
+            ".cdt",
+            ".mef",
+            ".nwb",
+        ],
+        datatypes: str | list[str] | None = None,
+        check: bool = False,
+    ):
+        if bids_paths is None:
+            bids_paths = mne_bids.find_matching_paths(
+                root=root,
+                subjects=subjects,
+                sessions=sessions,
+                tasks=tasks,
+                acquisitions=acquisitions,
+                runs=runs,
+                processings=processings,
+                recordings=recordings,
+                spaces=spaces,
+                splits=splits,
+                descriptions=descriptions,
+                suffixes=suffixes,
+                extensions=extensions,
+                datatypes=datatypes,
+                check=check,
+                ignore_json=True,
+            )
+            # Filter out _epo.fif files:
+            bids_paths = [
+                bids_path
+                for bids_path in bids_paths
+                if not (bids_path.suffix == "epo" and bids_path.extension == ".fif")
+            ]
+        self.bids_paths = bids_paths
+        self.reader_fn = reader_fn
+        self.pool_size = pool_size
+
+    def __add__(self, other):
+        assert isinstance(other, BIDSIterableDataset)
+        return BIDSIterableDataset(
+            reader_fn=self.reader_fn,
+            bids_paths=self.bids_paths + other.bids_paths,
+            pool_size=self.pool_size,
+        )
+
+    def __iadd__(self, other):
+        assert isinstance(other, BIDSIterableDataset)
+        self.bids_paths += other.bids_paths
+        return self
+
+    def __iter__(self):
+        worker_info = get_worker_info()
+        if worker_info is None:  # single-process data loading, return the full iterator
+            bids_paths = self.bids_paths
+        else:  # in a worker process
+            # split workload
+            bids_paths = self.bids_paths[worker_info.id :: worker_info.num_workers]
+
+        pool = []
+        end = False
+        paths_it = iter(random.sample(bids_paths, k=len(bids_paths)))
+        while not (end and len(pool) == 0):
+            while not end and len(pool) < self.pool_size:
+                try:
+                    bids_path = next(paths_it)
+                    ds = self.reader_fn(bids_path)
+                    if ds is None:
+                        print(f"Skipping {bids_path} as it is too short.")
+                        continue
+                    idx = iter(random.sample(range(len(ds)), k=len(ds)))
+                    pool.append((ds, idx))
+                except StopIteration:
+                    end = True
+            i_pool = random.randint(0, len(pool) - 1)
+            ds, idx = pool[i_pool]
+            try:
+                i_ds = next(idx)
+                yield ds[i_ds]
+            except StopIteration:
+                pool.pop(i_pool)

braindecode/models/__init__.py

@@ -4,6 +4,7 @@ Some predefined network architectures for EEG decoding.
 
 from .atcnet import ATCNet
 from .attentionbasenet import AttentionBaseNet
+from .attn_sleep import AttnSleep
 from .base import EEGModuleMixin
 from .biot import BIOT
 from .contrawr import ContraWR
@@ -15,9 +16,8 @@ from .eeginception_erp import EEGInceptionERP
 from .eeginception_mi import EEGInceptionMI
 from .eegitnet import EEGITNet
 from .eegminer import EEGMiner
-from .eegnet import EEGNetv1, EEGNetv4
+from .eegnet import EEGNet, EEGNetv4
 from .eegnex import EEGNeX
-from .eegresnet import EEGResNet
 from .eegsimpleconv import EEGSimpleConv
 from .eegtcnet import EEGTCNet
 from .fbcnet import FBCNet
@@ -38,12 +38,11 @@ from .signal_jepa import (
 from .sinc_shallow import SincShallowNet
 from .sleep_stager_blanco_2020 import SleepStagerBlanco2020
 from .sleep_stager_chambon_2018 import SleepStagerChambon2018
-from .sleep_stager_eldele_2021 import SleepStagerEldele2021
 from .sparcnet import SPARCNet
 from .syncnet import SyncNet
 from .tcn import BDTCN, TCN
 from .tidnet import TIDNet
-from .tsinception import TSceptionV1
+from .tsinception import TSception
 from .usleep import USleep
 from .util import _init_models_dict, models_mandatory_parameters
 
@@ -53,6 +52,7 @@ _init_models_dict()
 
 __all__ = [
     "ATCNet",
+    "AttnSleep",
     "AttentionBaseNet",
     "EEGModuleMixin",
     "BIOT",
@@ -65,10 +65,9 @@ __all__ = [
     "EEGInceptionMI",
     "EEGITNet",
     "EEGMiner",
-    "EEGNetv1",
+    "EEGNet",
     "EEGNetv4",
     "EEGNeX",
-    "EEGResNet",
     "EEGSimpleConv",
     "EEGTCNet",
     "FBCNet",
@@ -87,13 +86,12 @@ __all__ = [
     "SincShallowNet",
     "SleepStagerBlanco2020",
     "SleepStagerChambon2018",
-    "SleepStagerEldele2021",
     "SPARCNet",
     "SyncNet",
     "BDTCN",
     "TCN",
     "TIDNet",
-    "TSceptionV1",
+    "TSception",
     "USleep",
     "_init_models_dict",
     "models_mandatory_parameters",

braindecode/models/atcnet.py

@@ -138,7 +138,7 @@ class ATCNet(EEGModuleMixin, nn.Module):
     .. rubric:: Usage and Configuration
 
         - ``conv_block_n_filters (F1)``, ``conv_block_depth_mult (D)`` → capacity of the stem
-        (with ``F2 = F1·D`` feeding attention/TCN), dimensions aligned to ``F2``, like `EEGNetv4`.
+        (with ``F2 = F1·D`` feeding attention/TCN), dimensions aligned to ``F2``, like :class:`EEGNet`.
         - Pool sizes ``P1,P2`` trade temporal resolution for stability/compute; they set
         ``T_c = T/(P1·P2)`` and thus window width ``T_w``.
         - ``n_windows`` controls the ensemble over shifts (compute ∝ windows).

braindecode/models/attentionbasenet.py

@@ -26,25 +26,150 @@ from braindecode.modules.attention import (
 class AttentionBaseNet(EEGModuleMixin, nn.Module):
     """AttentionBaseNet from Wimpff M et al. (2023) [Martin2023]_.
 
+    :bdg-success:`Convolution` :bdg-info:`Small Attention`
+
     .. figure:: https://content.cld.iop.org/journals/1741-2552/21/3/036020/revision2/jnead48b9f2_hr.jpg
-       :align: center
-       :alt: Attention Base Net
+        :align: center
+        :alt: AttentionBaseNet Architecture
+        :width: 640px
+
+
+    .. rubric:: Architectural Overview
+
+    AttentionBaseNet is a *convolution-first* network with a *channel-attention* stage.
+    The end-to-end flow is:
+
+    - (i) :class:`_FeatureExtractor` learns a temporal filter bank and per-filter spatial
+      projections (depthwise across electrodes), then condenses time by pooling;
+    - (ii) **Channel Expansion** uses a ``1x1`` convolution to set the feature width;
+    - (iii) :class:`_ChannelAttentionBlock` refines features via depthwise–pointwise temporal
+      convs and an optional channel-attention module (SE/CBAM/ECA/…);
+    - (iv) **Classifier** flattens the sequence and applies a linear readout.
+
+    This design mirrors shallow CNN pipelines (EEGNet-style stem) but inserts a pluggable
+    attention unit that *re-weights channels* (and optionally temporal positions) before
+    classification.
+
+
+    .. rubric:: Macro Components
+
+    - :class:`_FeatureExtractor` **(Shallow conv stem → condensed feature map)**
+
+        - *Operations.*
+        - **Temporal conv** (:class:`torch.nn.Conv2d`) with kernel ``(1, L_t)`` creates a learned
+          FIR-like filter bank with ``n_temporal_filters`` maps.
+        - **Depthwise spatial conv** (:class:`torch.nn.Conv2d`, ``groups=n_temporal_filters``)
+          with kernel ``(n_chans, 1)`` learns per-filter spatial projections over the full montage.
+        - **BatchNorm → ELU → AvgPool → Dropout** stabilize and downsample time.
+        - Output shape: ``(B, F2, 1, T₁)`` with ``F2 = n_temporal_filters x spatial_expansion``.
+
+    *Interpretability/robustness.* Temporal kernels behave as analyzable FIR filters; the
+    depthwise spatial step yields rhythm-specific topographies. Pooling acts as a local
+    integrator that reduces variance on short EEG windows.
+
+    - **Channel Expansion**
+
+        - *Operations.*
+        - A ``1x1`` conv → BN → activation maps ``F2 → ch_dim`` without changing
+          the temporal length ``T₁`` (shape: ``(B, ch_dim, 1, T₁)``).
+          This sets the embedding width for the attention block.
+
+    - :class:`_ChannelAttentionBlock` **(temporal refinement + channel attention)**
+
+        - *Operations.*
+        - **Depthwise temporal conv** ``(1, L_a)`` (groups=``ch_dim``) + **pointwise ``1x1``**,
+          BN and activation → preserves shape ``(B, ch_dim, 1, T₁)`` while refining timing.
+        - **Optional attention module** (see *Additional Mechanisms*) applies channel reweighting
+          (some variants also apply temporal gating).
+        - **AvgPool (1, P₂)** with stride ``(1, S₂)`` and **Dropout** → outputs
+          ``(B, ch_dim, 1, T₂)``.
+
+    *Role.* Emphasizes informative channels (and, in certain modes, salient time steps)
+    before the classifier; complements the convolutional priors with adaptive re-weighting.
+
+    - **Classifier (aggregation + readout)**
+
+    *Operations.* :class:`torch.nn.Flatten` → :class:`torch.nn.Linear` from
+    ``(B, ch_dim·T₂)`` to classes.
+
+
+    .. rubric:: Convolutional Details
+
+    - **Temporal (where time-domain patterns are learned).**
+        Wide kernels in the stem (``(1, L_t)``) act as a learned filter bank for oscillatory
+        bands/transients; the attention block’s depthwise temporal conv (``(1, L_a)``) sharpens
+        short-term dynamics after downsampling. Pool sizes/strides (``P₁,S₁`` then ``P₂,S₂``)
+        set the token rate and effective temporal resolution.
+
+    - **Spatial (how electrodes are processed).**
+        A depthwise spatial conv with kernel ``(n_chans, 1)`` spans the full montage to
+        learn *per-temporal-filter* spatial projections (no cross-filter mixing at this step),
+        mirroring the interpretable spatial stage in shallow CNNs.
 
-    Neural Network from the paper: EEG motor imagery decoding:
-    A framework for comparative analysis with channel attention
-    mechanisms
+    - **Spectral (how frequency content is captured).**
+        No explicit Fourier/wavelet transform is used in the stem—spectral selectivity
+        emerges from learned temporal kernels. When ``attention_mode="fca"``, a frequency
+        channel attention (DCT-based) summarizes frequencies to drive channel weights.
 
-    The paper and original code with more details about the methodological
-    choices are available at the [Martin2023]_ and [MartinCode]_.
 
-    The AttentionBaseNet architecture is composed of four modules:
-    - Input Block that performs a temporal convolution and a spatial
-    convolution.
-    - Channel Expansion that modifies the number of channels.
-    - An attention block that performs channel attention with several
-    options
-    - ClassificationHead
+    .. rubric:: Attention / Sequential Modules
 
+    - **Type.** Channel attention chosen by ``attention_mode`` (SE, ECA, CBAM, CAT, GSoP,
+        EncNet, GE, GCT, SRM, CATLite). Most operate purely on channels; CBAM/CAT additionally
+        include temporal attention.
+
+    - **Shapes.** Input/Output around attention: ``(B, ch_dim, 1, T₁)``. Re-arrangements
+        (if any) are internal to the module; the block returns the same shape before pooling.
+
+    - **Role.** Re-weights channels (and optionally time) to highlight informative sources
+        and suppress distractors, improving SNR ahead of the linear head.
+
+
+    .. rubric:: Additional Mechanisms
+
+        - **Attention variants at a glance.**
+        - ``"se"``: Squeeze-and-Excitation (global pooling → bottleneck → gates).
+        - ``"gsop"``: Global second-order pooling (covariance-aware channel weights).
+        - ``"fca"``: Frequency Channel Attention (DCT summary; uses ``seq_len`` and ``freq_idx``).
+        - ``"encnet"``: EncNet with learned codewords (uses ``n_codewords``).
+        - ``"eca"``: Efficient Channel Attention (local 1-D conv over channel descriptor; uses ``kernel_size``).
+        - ``"ge"``: Gather–Excite (context pooling with optional MLP; can use ``extra_params``).
+        - ``"gct"``: Gated Channel Transformation (global context normalization + gating).
+        - ``"srm"``: Style-based recalibration (mean–std descriptors; optional MLP).
+        - ``"cbam"``: Channel then temporal attention (uses ``kernel_size``).
+        - ``"cat"`` / ``"catlite"``: Collaborative (channel ± temporal) attention; *lite* omits temporal.
+        - **Auto-compatibility on short inputs.**
+
+    If the input duration is too short for the configured kernels/pools, the implementation
+    **automatically rescales** temporal lengths/strides downward (with a warning) to keep
+    shapes valid and preserve the pipeline semantics.
+
+
+    .. rubric:: Usage and Configuration
+
+    - ``n_temporal_filters``, ``temporal_filter_length`` and ``spatial_expansion``:
+        control the capacity and the number of spatial projections in the stem.
+    - ``pool_length_inp``, ``pool_stride_inp`` then ``pool_length``, ``pool_stride``:
+        trade temporal resolution for compute; they determine the final sequence length ``T₂``.
+    - ``ch_dim``: width after the ``1x1`` expansion and the effective embedding size for attention.
+    - ``attention_mode`` + its specific hyperparameters (``reduction_rate``,
+        ``kernel_size``, ``seq_len``, ``freq_idx``, ``n_codewords``, ``use_mlp``):
+        select and tune the reweighting mechanism.
+    - ``drop_prob_inp`` and ``drop_prob_attn``: regularize stem and attention stages.
+    - **Training tips.**
+
+    Start with moderate pooling (e.g., ``P₁=75,S₁=15``) and ELU activations; enable attention
+    only after the stem learns stable filters. For small datasets, prefer simpler modes
+    (``"se"``, ``"eca"``) before heavier ones (``"gsop"``, ``"encnet"``).
+
+    Notes
+    -----
+    - Sequence length after each stage is computed internally; the final classifier expects
+      a flattened ``ch_dim x T₂`` vector.
+    - Attention operates on *channel* dimension by design; temporal gating exists only in
+      specific variants (CBAM/CAT).
+    - The paper and original code with more details about the methodological
+      choices are available at the [Martin2023]_ and [MartinCode]_.
     .. versionadded:: 0.9
 
     Parameters
@@ -73,18 +198,18 @@ class AttentionBaseNet(EEGModuleMixin, nn.Module):
         the depth of the network after the initial layer. Default is 16.
     attention_mode : str, optional
         The type of attention mechanism to apply. If `None`, no attention is applied.
-        - "se" for Squeeze-and-excitation network
-        - "gsop" for Global Second-Order Pooling
-        - "fca" for Frequency Channel Attention Network
-        - "encnet" for context encoding module
-        - "eca" for Efficient channel attention for deep convolutional neural networks
-        - "ge" for Gather-Excite
-        - "gct" for Gated Channel Transformation
-        - "srm" for Style-based Recalibration Module
-        - "cbam" for Convolutional Block Attention Module
-        - "cat" for Learning to collaborate channel and temporal attention
-        from multi-information fusion
-        - "catlite" for Learning to collaborate channel attention
+            - "se" for Squeeze-and-excitation network
+            - "gsop" for Global Second-Order Pooling
+            - "fca" for Frequency Channel Attention Network
+            - "encnet" for context encoding module
+            - "eca" for Efficient channel attention for deep convolutional neural networks
+            - "ge" for Gather-Excite
+            - "gct" for Gated Channel Transformation
+            - "srm" for Style-based Recalibration Module
+            - "cbam" for Convolutional Block Attention Module
+            - "cat" for Learning to collaborate channel and temporal attention
+            from multi-information fusion
+            - "catlite" for Learning to collaborate channel attention
         from multi-information fusion (lite version, cat w/o temporal attention)
     pool_length : int, default=8
         The length of the window for the average pooling operation.

braindecode/models/{sleep_stager_eldele_2021.py → attn_sleep.py}

@@ -8,18 +8,19 @@ from copy import deepcopy
 
 import torch
 import torch.nn.functional as F
+from mne.utils import deprecated
 from torch import nn
 
 from braindecode.models.base import EEGModuleMixin
 from braindecode.modules import CausalConv1d
 
 
-class SleepStagerEldele2021(EEGModuleMixin, nn.Module):
+class AttnSleep(EEGModuleMixin, nn.Module):
     """Sleep Staging Architecture from Eldele et al. (2021) [Eldele2021]_.
 
     .. figure:: https://raw.githubusercontent.com/emadeldeen24/AttnSleep/refs/heads/main/imgs/AttnSleep.png
         :align: center
-        :alt: SleepStagerEldele2021 Architecture
+        :alt: AttnSleep Architecture
 
     Attention based Neural Net for sleep staging as described in [Eldele2021]_.
     The code for the paper and this model is also available at [1]_.
@@ -533,3 +534,12 @@ class _PositionwiseFeedForward(nn.Module):
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Implements FFN equation."""
         return self.w_2(self.dropout(self.activate(self.w_1(x))))
+
+
+@deprecated(
+    "`SleepStagerEldele2021` was renamed to `AttnSleep` in v1.12 to follow original author's name; this alias will be removed in v1.14."
+)
+class SleepStagerEldele2021(AttnSleep):
+    """Deprecated alias for SleepStagerEldele2021."""
+
+    pass

braindecode/models/ctnet.py

@@ -39,7 +39,7 @@ class CTNet(EEGModuleMixin, nn.Module):
     The architecture consists of three main components:
 
     1. **Convolutional Module**:
-        - Apply EEGNetV4 to perform some feature extraction, denoted here as
+        - Apply :class:`EEGNet` to perform some feature extraction, denoted here as
         _PatchEmbeddingEEGNet module.
 
     2. **Transformer Encoder Module**:

braindecode/models/deep4.py

@@ -19,9 +19,13 @@ from braindecode.modules import (
 class Deep4Net(EEGModuleMixin, nn.Sequential):
     """Deep ConvNet model from Schirrmeister et al (2017) [Schirrmeister2017]_.
 
-     .. figure:: https://onlinelibrary.wiley.com/cms/asset/fc200ccc-d8c4-45b4-8577-56ce4d15999a/hbm23730-fig-0001-m.jpg
+    :bdg-success:`Convolution`
+
+    .. figure:: https://onlinelibrary.wiley.com/cms/asset/fc200ccc-d8c4-45b4-8577-56ce4d15999a/hbm23730-fig-0001-m.jpg
         :align: center
-        :alt: CTNet Architecture
+        :alt: Deep4Net Architecture
+        :width: 600px
+
 
     Model described in [Schirrmeister2017]_.