braindecode 1.2.0.dev176358851__py3-none-any.whl → 1.2.0.dev180217551__py3-none-any.whl

braindecode/models/atcnet.py

@@ -13,13 +13,154 @@ from braindecode.modules import CausalConv1d, Ensure4d, MaxNormLinear
 
 
 class ATCNet(EEGModuleMixin, nn.Module):
-    """ATCNet model from Altaheri et al. (2022) [1]_
+    """ATCNet from Altaheri et al. (2022) [1]_.
 
-    Pytorch implementation based on official tensorflow code [2]_.
+    :bdg-success:`Convolution` :bdg-info:`Small Attention`
 
     .. figure:: https://user-images.githubusercontent.com/25565236/185449791-e8539453-d4fa-41e1-865a-2cf7e91f60ef.png
-       :align: center
-       :alt: ATCNet Architecture
+        :align: center
+        :alt: ATCNet Architecture
+        :width: 650px
+
+    .. rubric:: Architectural Overview
+
+    ATCNet is a *convolution-first* architecture augmented with a *lightweight attention–TCN*
+    sequence module. The end-to-end flow is:
+
+    - (i) :class:`_ConvBlock` learns temporal filter-banks and spatial projections (EEGNet-style),
+      downsampling time to a compact feature map;
+
+    - (ii) Sliding Windows carve overlapping temporal windows from this map;
+
+    - (iii) for each window, :class:`_AttentionBlock` applies small multi-head self-attention
+      over time, followed by a :class:`_TCNResidualBlock` stack (causal, dilated);
+
+    - (iv) window-level features are aggregated (mean of window logits or concatenation)
+      and mapped via a max-norm–constrained linear layer.
+
+    Relative to ViT, ATCNet replaces linear patch projection with learned *temporal–spatial*
+    convolutions; it processes *parallel* window encoders (attention→TCN) instead of a deep
+    stack; and swaps the MLP head for a TCN suited to 1-D EEG sequences.
+
+    .. rubric:: Macro Components
+
+    - :class:`_ConvBlock` **(Shallow conv stem → feature map)**
+
+        - *Operations.*
+        - **Temporal conv** (:class:`torch.nn.Conv2d`) with kernel ``(L_t, 1)`` builds a
+            FIR-like filter bank (``F1`` maps).
+        - **Depthwise spatial conv** (:class:`torch.nn.Conv2d`, ``groups=F1``) with kernel
+          ``(1, n_chans)`` learns per-filter spatial projections (akin to EEGNet’s CSP-like step).
+        - **BN → ELU → AvgPool → Dropout** to stabilize and condense activations.
+        - **Refining temporal conv** (:class:`torch.nn.Conv2d`) with kernel ``(L_r, 1)`` +
+          **BN → ELU → AvgPool → Dropout**.
+
+    The output shape is ``(B, F2, T_c, 1)`` with ``F2 = F1·D`` and ``T_c = T/(P1·P2)``.
+    Temporal kernels behave as FIR filters; the depthwise-spatial conv yields frequency-specific
+    topographies. Pooling acts as a local integrator, reducing variance and imposing a
+    useful inductive bias on short EEG windows.
+
+    - **Sliding-Window Sequencer**
+
+      From the condensed time axis (length ``T_c``), ATCNet forms ``n`` overlapping windows
+      of width ``T_w = T_c - n + 1`` (one start per index). Each window produces a sequence
+      ``(B, F2, T_w)`` forwarded to its own attention–TCN branch. This creates *parallel*
+      encoders over shifted contexts and is key to robustness on nonstationary EEG.
+
+    - :class:`_AttentionBlock` **(small MHA on temporal positions)**
+
+        - *Operations.*
+        - Rearrange to ``(B, T_w, F2)``,
+        - Normalization :class:`torch.nn.LayerNorm`
+        - Custom MultiHeadAttention :class:`_MHA` (``num_heads=H``, per-head dim ``d_h``) + residual add,
+        - Dropout :class:`torch.nn.Dropout`
+        - Rearrange back to ``(B, F2, T_w)``.
+
+
+    **Note**: Attention is *local to a window* and purely temporal.
+
+    *Role.* Re-weights evidence across the window, letting the model emphasize informative
+    segments (onsets, bursts) before causal convolutions aggregate history.
+
+    - :class:`_TCNResidualBlock` **(causal dilated temporal CNN)**
+
+        - *Operations.*
+        - Two :class:`braindecode.modules.CausalConv1d` layers per block with dilation  ``1, 2, 4, …``
+        - Across blocks of `torch.nn.ELU` + `torch.nn.BatchNorm1d` + `torch.nn.Dropout`) +
+          a residual (identity or 1x1 mapping).
+        - The final feature used per window is the *last* causal step ``[..., -1]`` (forecast-style).
+
+    *Role.* Efficient long-range temporal integration with stable gradients; the dilated
+    receptive field complements attention’s soft selection.
+
+    - **Aggregation & Classifier**
+
+        - *Operations.*
+        - Either (a) map each window feature ``(B, F2)`` to logits via :class:`braindecode.modules.MaxNormLinear`
+        and **average** across windows (default, matching official code), or
+        - (b) **concatenate** all window features ``(B, n·F2)`` and apply a single :class:`MaxNormLinear`.
+        The max-norm constraint regularizes the readout.
+
+    .. rubric:: Convolutional Details
+
+    - **Temporal.** Temporal structure is learned in three places:
+        - (1) the stem’s wide ``(L_t, 1)`` conv (learned filter bank),
+        - (2) the refining ``(L_r, 1)`` conv after pooling (short-term dynamics), and
+        - (3) the TCN’s causal 1-D convolutions with exponentially increasing dilation
+          (long-range dependencies). The minimum sequence length required by the TCN stack is
+          ``(K_t - 1)·2^{L-1} + 1``; the implementation *auto-scales* kernels/pools/windows
+          when inputs are shorter to preserve feasibility.
+
+    - **Spatial.** A depthwise spatial conv spans the **full montage** (kernel ``(1, n_chans)``),
+        producing *per-temporal-filter* spatial projections (no cross-filter mixing at this step).
+        This mirrors EEGNet’s interpretability: each temporal filter has its own spatial pattern.
+
+
+    .. rubric:: Attention / Sequential Modules
+
+    - **Type.** Multi-head self-attention with ``H`` heads and per-head dim ``d_h`` implemented
+      in :class:`_MHA`, allowing ``embed_dim = H·d_h`` independent of input and output dims.
+    - **Shapes.** ``(B, F2, T_w) → (B, T_w, F2) → (B, F2, T_w)``. Attention operates along
+      the **temporal** axis within a window; channels/features stay in the embedding dim ``F2``.
+    - **Role.** Highlights salient temporal positions prior to causal convolution; small attention
+      keeps compute modest while improving context modeling over pooled features.
+
+    .. rubric:: Additional Mechanisms
+
+    - **Parallel encoders over shifted windows.** Improves montage/phase robustness by
+      ensembling nearby contexts rather than committing to a single segmentation.
+    - **Max-norm classifier.** Enforces weight norm constraints at the readout, a common
+      stabilization trick in EEG decoding.
+    - **ViT vs. ATCNet (design choices).** Convolutional *nonlinear* projection rather than
+      linear patchification; attention followed by **TCN** (not MLP); *parallel* window
+      encoders rather than stacked encoders.
+
+    .. 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`.
+        - 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).
+        - ``att_num_heads``, ``att_head_dim`` set attention capacity; keep ``H·d_h ≈ F2``.
+        - ``tcn_depth``, ``tcn_kernel_size`` govern receptive field; larger values demand
+        longer inputs (see minimum length above). The implementation warns and *rescales*
+        kernels/pools/windows if inputs are too short.
+        - **Aggregation choice.** ``concat=False`` (default, average of per-window logits) matches
+        the official code; ``concat=True`` mirrors the paper’s concatenation variant.
+
+
+    Notes
+    -----
+    - Inputs substantially shorter than the implied minimum length trigger **automatic
+      downscaling** of kernels, pools, windows, and TCN kernel size to maintain validity.
+    - The attention–TCN sequence operates **per window**; the last causal step is used as the
+      window feature, aligning the temporal semantics across windows.
+
+    .. versionadded:: 1.1
+
+        - More detailed documentation of the model.
+
 
     Parameters
     ----------
@@ -85,15 +226,13 @@ class ATCNet(EEGModuleMixin, nn.Module):
         Maximum L2-norm constraint imposed on weights of the last
         fully-connected layer. Defaults to 0.25.
 
-
     References
     ----------
-    .. [1] H. Altaheri, G. Muhammad and M. Alsulaiman,
-        Physics-informed attention temporal convolutional network for EEG-based
-        motor imagery classification in IEEE Transactions on Industrial Informatics,
-        2022, doi: 10.1109/TII.2022.3197419.
-    .. [2] EEE-ATCNet implementation.
-       https://github.com/Altaheri/EEG-ATCNet/blob/main/models.py
+    .. [1] H. Altaheri, G. Muhammad, M. Alsulaiman (2022).
+        *Physics-informed attention temporal convolutional network for EEG-based motor imagery classification.*
+        IEEE Transactions on Industrial Informatics. doi:10.1109/TII.2022.3197419.
+    .. [2] Official EEG-ATCNet implementation (TensorFlow):
+        https://github.com/Altaheri/EEG-ATCNet/blob/main/models.py
     """
 
     def __init__(
@@ -556,7 +695,8 @@ class _TCNResidualBlock(nn.Module):
         # Reshape the input for the residual connection when necessary
         if in_channels != n_filters:
             self.reshaping_conv = nn.Conv1d(
-                n_filters,
+                in_channels=in_channels,
+                out_channels=n_filters,
                 kernel_size=1,
                 padding="same",
             )

braindecode/models/attentionbasenet.py

@@ -24,26 +24,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.
+
+    - **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.
 
-    Neural Network from the paper: EEG motor imagery decoding:
-    A framework for comparative analysis with channel attention
-    mechanisms
 
-    The paper and original code with more details about the methodological
-    choices are available at the [Martin2023]_ and [MartinCode]_.
+    .. rubric:: Attention / Sequential Modules
 
-    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
+    - **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).
 
     .. versionadded:: 0.9
 
@@ -73,18 +197,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/eegconformer.py

@@ -27,22 +27,25 @@ class EEGConformer(EEGModuleMixin, nn.Module):
     EEG-Conformer is a *convolution-first* model augmented with a *lightweight transformer
     encoder*. The end-to-end flow is:
 
-    - (i) :class:`_PatchEmbedding` converts the continuous EEG into a compact sequence of tokens via a :class:`ShallowFBCSPNet` temporal–spatial conv stem and temporal pooling;
-    - (ii) :class:`_TransformerEncoder applies small multi-head self-attention to integrate longer-range temporal context across tokens;
+    - (i) :class:`_PatchEmbedding` converts the continuous EEG into a compact sequence of tokens via a
+      :class:`ShallowFBCSPNet` temporal–spatial conv stem and temporal pooling;
+    - (ii) :class:`_TransformerEncoder` applies small multi-head self-attention to integrate
+      longer-range temporal context across tokens;
     - (iii) :class:`_ClassificationHead` aggregates the sequence and performs a linear readout.
-        This preserves the strong inductive biases of shallow CNN filter banks while adding
-        just enough attention to capture dependencies beyond the pooling horizon [song2022]_.
+      This preserves the strong inductive biases of shallow CNN filter banks while adding
+      just enough attention to capture dependencies beyond the pooling horizon [song2022]_.
 
     .. rubric:: Macro Components
 
     - :class:`_PatchEmbedding` **(Shallow conv stem → tokens)**
 
-    - *Operations.*
-        - A temporal convolution (`:class:`torch.nn.Conv2d`) ``(1 x L_t)`` forms a data-driven "filter bank";
-        - A spatial convolution (`:class:`torch.nn.Conv2d`) (n_chans x 1)`` projects across electrodes, collapsing the channel axis into a virtual channel.
-        - **Normalization function** `:class:torch.nn.BatchNorm`
-        - **Activation function** `:class:torch.nn.ELU`
-        - **Average Pooling** `:class:torch.nn.AvgPool` along time (kernel ``(1, P)`` with stride ``(1, S)``)
+        - *Operations.*
+        - A temporal convolution (`:class:torch.nn.Conv2d`) ``(1 x L_t)`` forms a data-driven "filter bank";
+        - A spatial convolution (`:class:torch.nn.Conv2d`) (n_chans x 1)`` projects across electrodes,
+          collapsing the channel axis into a virtual channel.
+        - **Normalization function** :class:`torch.nn.BatchNorm`
+        - **Activation function** :class:`torch.nn.ELU`
+        - **Average Pooling** :class:`torch.nn.AvgPool` along time (kernel ``(1, P)`` with stride ``(1, S)``)
         -  final ``1x1`` :class:`torch.nn.Linear` projection.
 
     The result is rearranged to a token sequence ``(B, S_tokens, D)``, where ``D = n_filters_time``.
@@ -53,7 +56,7 @@ class EEGConformer(EEGModuleMixin, nn.Module):
 
     - :class:`_TransformerEncoder` **(context over temporal tokens)**
 
-    - *Operations.*
+        - *Operations.*
         - A stack of ``att_depth`` encoder blocks. :class:`_TransformerEncoderBlock`
         - Each block applies LayerNorm :class:`torch.nn.LayerNorm`
         - Multi-Head Self-Attention (``att_heads``) with dropout + residual :class:`MultiHeadAttention` (:class:`torch.nn.Dropout`)
@@ -67,7 +70,7 @@ class EEGConformer(EEGModuleMixin, nn.Module):
 
     - :class:`ClassificationHead` **(aggregation + readout)**
 
-    - *Operations*.
+        - *Operations*.
         - Flatten, :class:`torch.nn.Flatten` the sequence ``(B, S_tokens·D)`` -
         - MLP (:class:`torch.nn.Linear` → activation (default: :class:`torch.nn.ELU`) → :class:`torch.nn.Dropout` → :class:`torch.nn.Linear`)
         - final Linear to classes.
@@ -100,8 +103,8 @@ class EEGConformer(EEGModuleMixin, nn.Module):
     - **Type.** Standard multi-head self-attention (MHA) with ``att_heads`` heads over the token sequence.
     - **Shapes.** Input/Output: ``(B, S_tokens, D)``; attention operates along the ``S_tokens`` axis.
     - **Role.** Re-weights and integrates evidence across pooled windows, capturing dependencies
-        longer than any single token while leaving channel relationships to the convolutional stem.
-        The design is intentionally *small*—attention refines rather than replaces convolutional feature extraction.
+      longer than any single token while leaving channel relationships to the convolutional stem.
+      The design is intentionally *small*—attention refines rather than replaces convolutional feature extraction.
 
     .. rubric:: Additional Mechanisms
 
@@ -112,20 +115,20 @@ class EEGConformer(EEGModuleMixin, nn.Module):
         refine temporal context before classification.
 
     - **Tokenization knob.** ``pool_time_length`` and especially ``pool_time_stride`` set
-    the number of tokens ``S_tokens``. Smaller strides → more tokens and higher attention
-    capacity (but higher compute); larger strides → fewer tokens and stronger inductive bias.
+        the number of tokens ``S_tokens``. Smaller strides → more tokens and higher attention
+        capacity (but higher compute); larger strides → fewer tokens and stronger inductive bias.
 
     - **Embedding dimension = filters.** ``n_filters_time`` serves double duty as both the
-    number of temporal filters in the stem and the transformer’s embedding size ``D``,
-    simplifying dimensional alignment.
+        number of temporal filters in the stem and the transformer’s embedding size ``D``,
+        simplifying dimensional alignment.
 
     .. rubric:: Usage and Configuration
 
     - **Instantiation.** Choose ``n_filters_time`` (embedding size ``D``) and
-    ``filter_time_length`` to match the rhythms of interest. Tune
-    ``pool_time_length/stride`` to trade temporal resolution for sequence length.
-    Keep ``att_depth`` modest (e.g., 4–6) and set ``att_heads`` to divide ``D``.
-    ``final_fc_length="auto"`` infers the flattened size from PatchEmbedding.
+        ``filter_time_length`` to match the rhythms of interest. Tune
+        ``pool_time_length/stride`` to trade temporal resolution for sequence length.
+        Keep ``att_depth`` modest (e.g., 4–6) and set ``att_heads`` to divide ``D``.
+        ``final_fc_length="auto"`` infers the flattened size from PatchEmbedding.
 
     Notes
     -----

braindecode/models/eegnet.py

@@ -31,8 +31,7 @@ class EEGNetv4(EEGModuleMixin, nn.Sequential):
 
     .. rubric:: Architectural Overview
 
-    EEGNetv4 is a compact convolutional network designed for EEG decoding with a
-    pipeline that mirrors classical EEG processing:
+    EEGNetv4 is a compact convolutional network designed for EEG decoding with a pipeline that mirrors classical EEG processing:
     - (i) learn temporal frequency-selective filters,
     - (ii) learn spatial filters for those frequencies, and
     - (iii) condense features with depthwise–separable convolutions before a lightweight classifier.
@@ -56,16 +55,16 @@ class EEGNetv4(EEGModuleMixin, nn.Sequential):
 
     .. rubric:: Convolutional Details
 
-    **Temporal.** The initial temporal convs serve as a *learned filter bank*:
-    long 1-D kernels (implemented as 2-D with singleton spatial extent) emphasize oscillatory bands and transients.
-    Because this stage is linear prior to BN/ELU, kernels can be analyzed as FIR filters to reveal each feature’s spectrum [Lawhern2018]_.
+    - **Temporal.** The initial temporal convs serve as a *learned filter bank*:
+      long 1-D kernels (implemented as 2-D with singleton spatial extent) emphasize oscillatory bands and transients.
+      Because this stage is linear prior to BN/ELU, kernels can be analyzed as FIR filters to reveal each feature’s spectrum [Lawhern2018]_.
 
-    **Spatial.** The depthwise spatial conv spans the full channel axis (kernel height = #electrodes; temporal size = 1).
-    With ``groups = F1``, each temporal filter learns its own set of ``D`` spatial projections—akin to CSP, learned end-to-end and
-    typically regularized with max-norm.
+    - **Spatial.** The depthwise spatial conv spans the full channel axis (kernel height = #electrodes; temporal size = 1).
+      With ``groups = F1``, each temporal filter learns its own set of ``D`` spatial projections—akin to CSP, learned end-to-end and
+      typically regularized with max-norm.
 
-    **Spectral.** No explicit Fourier/wavelet transform is used. Frequency structure
-    is captured implicitly by the temporal filter bank; later depthwise temporal kernels act as short-time integrators/refiners.
+    - **Spectral.** No explicit Fourier/wavelet transform is used. Frequency structure
+      is captured implicitly by the temporal filter bank; later depthwise temporal kernels act as short-time integrators/refiners.
 
     .. rubric:: Additional Comments
 

braindecode/models/eegnex.py

@@ -16,9 +16,123 @@ from braindecode.modules import Conv2dWithConstraint, LinearWithConstraint
 class EEGNeX(EEGModuleMixin, nn.Module):
     """EEGNeX model from Chen et al. (2024) [eegnex]_.
 
+    :bdg-success:`Convolution`
+
     .. figure:: https://braindecode.org/dev/_static/model/eegnex.jpg
         :align: center
         :alt: EEGNeX Architecture
+        :width: 620px
+
+    .. rubric:: Architectural Overview
+
+    EEGNeX is a **purely convolutional** architecture that refines the EEGNet-style stem
+    and deepens the temporal stack with **dilated temporal convolutions**. The end-to-end
+    flow is:
+
+    - (i) **Block-1/2**: two temporal convolutions ``(1 x L)`` with BN refine a
+      learned FIR-like *temporal filter bank* (no pooling yet);
+    - (ii) **Block-3**: depthwise **spatial** convolution across electrodes
+      ``(n_chans x 1)`` with max-norm constraint, followed by ELU → AvgPool (time) → Dropout;
+    - (iii) **Block-4/5**: two additional **temporal** convolutions with increasing **dilation**
+      to expand the receptive field; the last block applies ELU → AvgPool → Dropout → Flatten;
+    - (iv) **Classifier**: a max-norm–constrained linear layer.
+
+    The published work positions EEGNeX as a compact, conv-only alternative that consistently
+    outperforms prior baselines across MOABB-style benchmarks, with the popular
+    “EEGNeX-8,32” shorthand denoting *8 temporal filters* and *kernel length 32*.
+
+
+    .. rubric:: Macro Components
+
+    - **Block-1 / Block-2 — Temporal filter (learned).**
+
+       - *Operations.*
+       - :class:`torch.nn.Conv2d` with kernels ``(1, L)``
+       - :class:`torch.nn.BatchNorm2d` (no nonlinearity until Block-3, mirroring a linear FIR analysis stage).
+         These layers set up frequency-selective detectors before spatial mixing.
+
+    - *Interpretability.* Kernels can be inspected as FIR filters; two stacked temporal
+      convs allow longer effective kernels without parameter blow-up.
+
+    - **Block-3 — Spatial projection + condensation.**
+
+        - *Operations.*
+        - :class:`braindecode.modules.Conv2dWithConstraint` with kernel``(n_chans, 1)``
+          and ``groups = filter_2`` (depthwise across filters)
+        - :class:`torch.nn.BatchNorm2d`
+        - :class:`torch.nn.ELU`
+        - :class:`torch.nn.AvgPool2d` (time)
+        - :class:`torch.nn.Dropout`.
+
+    *Role.* Learns per-filter spatial patterns over the **full montage** while temporal
+      pooling stabilizes and compresses features; max-norm encourages well-behaved spatial
+      weights similar to EEGNet practice.
+
+    - **Block-4 / Block-5 — Dilated temporal integration.**
+
+        - *Operations.*
+        - :class:`torch.nn.Conv2d` with kernels ``(1, k)`` and **dilations**
+        (e.g., 2 then 4);
+        - :class:`torch.nn.BatchNorm2d`
+        - :class:`torch.nn.ELU`
+        - :class:`torch.nn.AvgPool2d` (time)
+        - :class:`torch.nn.Dropout`
+        - :class:`torch.nn.Flatten`.
+
+    *Role.* Expands the temporal receptive field efficiently to capture rhythms and
+      long-range context after condensation.
+
+    - **Final Classifier — Max-norm linear.**
+
+        - *Operations.*
+        - :class:`braindecode.modules.LinearWithConstraint` maps the flattened
+          vector to the target classes; the max-norm constraint regularizes the readout.
+
+
+    .. rubric:: Convolutional Details
+
+    - **Temporal (where time-domain patterns are learned).**
+      Blocks 1-2 learn the primary filter bank (oscillations/transients), while Blocks 4-5
+      use **dilation** to integrate over longer horizons without extra pooling. The final
+      AvgPool in Block-5 sets the output token rate and helps noise suppression.
+
+    - **Spatial (how electrodes are processed).**
+      A *single* depthwise spatial conv (Block-3) spans the entire electrode set
+      (kernel ``(n_chans, 1)``), producing per-temporal-filter topographies; no cross-filter
+      mixing occurs at this stage, aiding interpretability.
+
+    - **Spectral (how frequency content is captured).**
+      Frequency selectivity emerges from the learned temporal kernels; dilation broadens effective
+      bandwidth coverage by composing multiple scales.
+
+    .. rubric:: Additional Mechanisms
+
+    - **EEGNeX-8,32 naming.** “8,32” indicates *8 temporal filters* and *kernel length 32*,
+      reflecting the paper's ablation path from EEGNet-8,2 toward thicker temporal kernels
+      and a deeper conv stack.
+    - **Max-norm constraints.** Spatial (Block-3) and final linear layers use max-norm
+      regularization—standard in EEG CNNs—to reduce overfitting and encourage stable spatial
+      patterns.
+
+    .. rubric:: Usage and Configuration
+
+    - **Kernel schedule.** Start with the canonical **EEGNeX-8,32** (``filter_1=8``,
+      ``kernel_block_1_2=32``) and keep **Block-3** depth multiplier modest (e.g., 2) to match
+      the paper's “pure conv” profile.
+    - **Pooling vs. dilation.** Use pooling in Blocks 3 and 5 to control compute and variance;
+      increase dilations (Blocks 4-5) to widen temporal context when windows are short.
+    - **Regularization.** Combine dropout (Blocks 3 & 5) with max-norm on spatial and
+      classifier layers; prefer ELU activations for stable training on small EEG datasets.
+
+
+    Notes
+    -----
+    - The braindecode implementation follows the paper's conv-only design with five blocks
+      and reproduces the depthwise spatial step and dilated temporal stack. See the class
+      reference for exact kernel sizes, dilations, and pooling defaults.
+
+    .. versionadded:: 1.1
+
 
     Parameters
     ----------

braindecode/version.py

@@ -1 +1 @@
-__version__ = "1.2.0.dev176358851"
+__version__ = "1.2.0.dev180217551"

{braindecode-1.2.0.dev176358851.dist-info → braindecode-1.2.0.dev180217551.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.2.0.dev176358851
+Version: 1.2.0.dev180217551
 Summary: Deep learning software to decode EEG, ECG or MEG signals
 Author-email: Robin Tibor Schirrmeister <robintibor@gmail.com>
 Maintainer-email: Alexandre Gramfort <agramfort@meta.com>, Bruno Aristimunha Pinto <b.aristimunha@gmail.com>, Robin Tibor Schirrmeister <robintibor@gmail.com>

{braindecode-1.2.0.dev176358851.dist-info → braindecode-1.2.0.dev180217551.dist-info}/RECORD

@@ -3,7 +3,7 @@ braindecode/classifier.py,sha256=k9vSCtfQbld0YVleDi5rrrmk6k_k5JYEPPBYcNxYjZ8,980
 braindecode/eegneuralnet.py,sha256=dz8k_-2jV7WqkaX4bQG-dmr-vRT7ZtOwJqomXyC9PTw,15287
 braindecode/regressor.py,sha256=VLfrpiXklwI4onkwue3QmzlBWcvspu0tlrLo9RT1Oiw,9375
 braindecode/util.py,sha256=J-tBcDJNlMTIFW2mfOy6Ko0nsgdP4obRoEVDeg2rFH0,12686
-braindecode/version.py,sha256=pqOP2XIkF23sSxHn98gks8lm_NZ40LR5V6vz_cuxJQo,35
+braindecode/version.py,sha256=Y04aPFEYg6CFbsSt6BtmY3lNGnXYQ3FAwT0WYnn5X2Q,35
 braindecode/augmentation/__init__.py,sha256=LG7ONqCufYAF9NZt8POIp10lYXb8iSueYkF-CWGK2Ls,1001
 braindecode/augmentation/base.py,sha256=gg7wYsVfa9jfqBddtE03B5ZrPHFFmPl2sa3LOrRnGfo,7325
 braindecode/augmentation/functional.py,sha256=ygkMNEFHaUdRQfk7meMML19FnM406Uf34h-ztKXdJwM,37978
@@ -27,21 +27,21 @@ braindecode/functional/__init__.py,sha256=JPUDFeKtfogEzfrwPaZRBmxexPjBw7AglYMlIm
 braindecode/functional/functions.py,sha256=CoEweM6YLhigx0tNmmz6yAc8iQ078sTFY2GeCjK5fFs,8622
 braindecode/functional/initialization.py,sha256=BUSC7y2TMsfShpMYBVwm3xg3ODFqWp-STH7yD4sn8zk,1388
 braindecode/models/__init__.py,sha256=xv1QPELZxocPgbc_mz-eYM5w08ZDNOsDV4pOnIFhUww,2551
-braindecode/models/atcnet.py,sha256=PhDJl6nBChButabjsmLz_heRcGFCCMKoeUt7k7neNzs,24483
-braindecode/models/attentionbasenet.py,sha256=1uwrtsdEGiBwokkO8A_2SR5zapOTQUBZd4q7hIpR0cw,23359
+braindecode/models/atcnet.py,sha256=Pn5KzQjv7YxSNDr_CY6O_Yg9K4m9XJ7btCIqyzkcPxc,32102
+braindecode/models/attentionbasenet.py,sha256=zqSzrFAjl89EoacWHuPLbjBRY4RW2awdhWElb-d9JjY,30178
 braindecode/models/base.py,sha256=9icrWNZBGbh_VLyB9m8g_K1QyK7s3mh8X-hJ29gEbWs,10802
 braindecode/models/biot.py,sha256=T4PymX3penMJcrdfb5Nq6B3P-jyP2laAIu_R9o3uCXo,17512
 braindecode/models/contrawr.py,sha256=eeR_ik4gNZ3rJLM6Mw9gJ2gTMkZ8CU8C4rN_GQMQTAE,10044
 braindecode/models/ctnet.py,sha256=-J9QtUM8kcntz_xinfuBBvwDMECHiMPMcr2MS4GDPEY,17308
 braindecode/models/deep4.py,sha256=YJQUw-0EuFUi4qjm8caJGB8wRM_aeJa5X_d8jrGaQAI,14588
 braindecode/models/deepsleepnet.py,sha256=RrciuVJtZ-fhiUl-yLPfK2FP-G29V5Wor6pPlrMHQWQ,9218
-braindecode/models/eegconformer.py,sha256=OSORbNYwMA0hvUMjuyB8wI8qBKVSiraioLHTHmt8sdQ,17376
+braindecode/models/eegconformer.py,sha256=rxMAmqErDVLq7nS77CnTtpcC3C2OR_EoZ8-jG-dKP9I,17433
 braindecode/models/eeginception_erp.py,sha256=mwh3rGSHAJVvnbOlYTuWWkKxlmFAdAXBNCrq4IPgOS4,11408
 braindecode/models/eeginception_mi.py,sha256=aKJRFuYrpbcRbmmT2xVghKbK8pnl7fzu5hrV0ybRKso,12424
 braindecode/models/eegitnet.py,sha256=feXFmPCd-Ejxt7jgWPen1Ag0-oSclDVQai0Atwu9d_A,9827
 braindecode/models/eegminer.py,sha256=ouKZah9Q7_sxT7DJJMcPObwVxNQE87sEljJg6QwiQNw,9847
-braindecode/models/eegnet.py,sha256=YeBCmU6Al9FDS4MZQTOLd0MCUfPbM6tmVlGWpb59Qzg,19256
-braindecode/models/eegnex.py,sha256=KNJIh8pFNhY087Bey2OPzDD4Uqw9pS6UkwMjnOngBzg,8497
+braindecode/models/eegnet.py,sha256=CtfQuw7iaxQh3j1dRmF_UhdjfO3uHOlObnmasHk_boM,19268
+braindecode/models/eegnex.py,sha256=xfaISCgW5ShlUh9fFBHc5ylz80OhW5C69Fi-_0hVS5U,13767
 braindecode/models/eegresnet.py,sha256=cqWOSGqfJN_dNYUU9l8nYd_S3T1N-UX5-encKQzfBlg,12057
 braindecode/models/eegsimpleconv.py,sha256=sHpK-7ZGOCMuXsdkSVuarFTd1T0jMJUP_xwXP3gxQwc,7268
 braindecode/models/eegtcnet.py,sha256=np-93Ttctp2uaEYpMrfXfH5bJmCOUZZHLjv8GJEEym4,10830
@@ -93,9 +93,9 @@ braindecode/training/scoring.py,sha256=WRkwqbitA3m_dzRnGp2ZIZPge5Nhx9gAEQhIHzeH4
 braindecode/visualization/__init__.py,sha256=4EER_xHqZIDzEvmgUEm7K1bgNKpyZAIClR9ZCkMuY4M,240
 braindecode/visualization/confusion_matrices.py,sha256=qIWMLEHow5CJ7PhGggD8mnD55Le6xhma9HSzt4R33fc,9509
 braindecode/visualization/gradients.py,sha256=KZo-GA0uwiwty2_94j2IjmCR2SKcfPb1Bi3sQq7vpTk,2170
-braindecode-1.2.0.dev176358851.dist-info/licenses/LICENSE.txt,sha256=7rg7k6hyj8m9whQ7dpKbqnCssoOEx_Mbtqb4uSOjljE,1525
-braindecode-1.2.0.dev176358851.dist-info/licenses/NOTICE.txt,sha256=sOxuTbalPxTM8H6VqtvGbXCt_BoOF7JevEYG_knqbm4,620
-braindecode-1.2.0.dev176358851.dist-info/METADATA,sha256=23uR3nYKaKV2G4EtrutqVU6r40qFe4_Oh9DuXDX5fFI,6883
-braindecode-1.2.0.dev176358851.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-braindecode-1.2.0.dev176358851.dist-info/top_level.txt,sha256=pHsWQmSy0uhIez62-HA9j0iaXKvSbUL39ifFRkFnChA,12
-braindecode-1.2.0.dev176358851.dist-info/RECORD,,
+braindecode-1.2.0.dev180217551.dist-info/licenses/LICENSE.txt,sha256=7rg7k6hyj8m9whQ7dpKbqnCssoOEx_Mbtqb4uSOjljE,1525
+braindecode-1.2.0.dev180217551.dist-info/licenses/NOTICE.txt,sha256=sOxuTbalPxTM8H6VqtvGbXCt_BoOF7JevEYG_knqbm4,620
+braindecode-1.2.0.dev180217551.dist-info/METADATA,sha256=3_WlN-hqNZ-mvQRBABpJzmU2cIb_RKSpjWtgghqT4A0,6883
+braindecode-1.2.0.dev180217551.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+braindecode-1.2.0.dev180217551.dist-info/top_level.txt,sha256=pHsWQmSy0uhIez62-HA9j0iaXKvSbUL39ifFRkFnChA,12
+braindecode-1.2.0.dev180217551.dist-info/RECORD,,