braindecode 1.7.0.dev182642410__tar.gz → 1.7.0.dev183981614__tar.gz

{braindecode-1.7.0.dev182642410/braindecode.egg-info → braindecode-1.7.0.dev183981614}/PKG-INFO

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

{braindecode-1.7.0.dev182642410 → braindecode-1.7.0.dev183981614}/braindecode/classifier.py

@@ -3,6 +3,7 @@
 #          Lukas Gemein <l.gemein@gmail.com>
 #          Bruno Aristimunha <b.aristimunha@gmail.com>
 #          Pierre Guetschel <pierre.guetschel@gmail.com>
+#          Sarthak Tayal <sarthaktayal2@gmail.com>
 #
 # License: BSD (3-clause)
 
@@ -195,9 +196,7 @@ class EEGClassifier(_EEGNeuralNet, NeuralNetClassifier):
         return self.predict_proba(X).argmax(1)
 
     def predict_trials(self, X, return_targets=True):
-        """Create trialwise predictions and optionally also return trialwise.
-
-        labels from cropped dataset.
+        """Create trialwise predictions from a cropped dataset.
 
         Parameters
         ----------
@@ -212,10 +211,11 @@ class EEGClassifier(_EEGNeuralNet, NeuralNetClassifier):
             3-dimensional array (n_trials x n_classes x n_predictions), where
             the number of predictions depend on the chosen window size and the
             receptive field of the network.
-        trial_labels : np.ndarray
-            2-dimensional array (n_trials x n_targets) where the number of
-            targets depends on the decoding paradigm and can be either a single
-            value, multiple values, or a sequence.
+        trial_targets : np.ndarray
+            Ground-truth targets from the dataset in a 2-dimensional array
+            (n_trials x n_targets). Only returned when ``return_targets=True``.
+            The number of targets depends on the decoding paradigm and can be
+            either a single value, multiple values, or a sequence.
         """
         if not self.cropped:
             warnings.warn(

{braindecode-1.7.0.dev182642410 → braindecode-1.7.0.dev183981614}/braindecode/functional/__init__.py

@@ -5,6 +5,8 @@ from .functions import (
     dwt_max_level,
     hilbert_freq,
     identity,
+    iou_1d,
+    pairwise_iou_1d,
     plv_time,
     safe_log,
     sinusoidal_positional_encoding,
@@ -20,6 +22,8 @@ __all__ = [
     "dwt_max_level",
     "hilbert_freq",
     "identity",
+    "iou_1d",
+    "pairwise_iou_1d",
     "plv_time",
     "safe_log",
     "sinusoidal_positional_encoding",

{braindecode-1.7.0.dev182642410 → braindecode-1.7.0.dev183981614}/braindecode/functional/functions.py

@@ -1,5 +1,6 @@
 # Authors: Robin Schirrmeister <robintibor@gmail.com>
 #          Bruno Aristimunha <b.aristimunha@gmail.com>
+#          Sarthak Tayal <sarthaktayal2@gmail.com>
 #
 # License: BSD (3-clause)
 
@@ -109,7 +110,7 @@ def _get_gaussian_kernel1d(kernel_size: int, sigma: float) -> torch.Tensor:
     return kernel1d
 
 
-def hilbert_freq(x, forward_fourier=True):
+def hilbert_freq(x: torch.Tensor, forward_fourier: bool = True) -> torch.Tensor:
     r"""
     Compute the Hilbert transform using PyTorch, separating the real and
     imaginary parts.
@@ -178,7 +179,9 @@ def hilbert_freq(x, forward_fourier=True):
     return x
 
 
-def plv_time(x, forward_fourier=True, epsilon: float = 1e-6):
+def plv_time(
+    x: torch.Tensor, forward_fourier: bool = True, epsilon: float = 1e-6
+) -> torch.Tensor:
     """Compute the Phase Locking Value (PLV) metric in the time domain.
 
     The Phase Locking Value (PLV) is a measure of the synchronization between
@@ -257,6 +260,40 @@ def plv_time(x, forward_fourier=True, epsilon: float = 1e-6):
     return plv_matrix
 
 
+# -----------------------------------------------------------------------------
+# DANCE functional helpers
+#
+# Authors: Bruno Aristimunha <b.aristimunha@gmail.com>
+#
+# License: MIT
+#
+# Ported from the DANCE event-detection model (facebookresearch/dance, MIT).
+# -----------------------------------------------------------------------------
+
+
+def iou_1d(s1, e1, s2, e2, eps: float = 1e-7):
+    """ELEMENTWISE 1-D temporal IoU. All four inputs share the same shape S
+    (e.g. ``(B, Q)``); returns IoU of shape S. Used by ``DanceLoss`` on the
+    matched ``(B, Q)`` spans. For (Q,)x(T,) pairwise IoU use ``pairwise_iou_1d``.
+    """
+    inter = (torch.minimum(e1, e2) - torch.maximum(s1, s2)).clamp(min=0)
+    union = (e1 - s1) + (e2 - s2) - inter
+    return inter / (union + eps)
+
+
+def pairwise_iou_1d(s1, e1, s2, e2, eps: float = 1e-7):
+    """PAIRWISE 1-D temporal IoU. ``s1,e1`` shape ``(Q,)``, ``s2,e2`` shape
+    ``(T,)``; returns ``(Q, T)`` (broadcast ``[:, None]`` x ``[None, :]``).
+    Transcribed from ``dance/matcher.py:33-39`` (``_pairwise_iou``). Used by
+    ``HungarianMatcher`` to build the ``(Q, n_targets)`` cost matrix.
+    """
+    inter_start = torch.maximum(s1[:, None], s2[None, :])
+    inter_end = torch.minimum(e1[:, None], e2[None, :])
+    inter = (inter_end - inter_start).clamp(min=0)
+    union = (e1 - s1)[:, None] + (e2 - s2)[None, :] - inter
+    return inter / (union + eps)
+
+
 def daubechies_filters(n_vanishing: int) -> torch.Tensor:
     r"""Daubechies ``db<n_vanishing>`` wavelet decomposition filters.
 

{braindecode-1.7.0.dev182642410 → braindecode-1.7.0.dev183981614}/braindecode/models/__init__.py

@@ -11,6 +11,7 @@ from .cbramod import CBraMod
 from .codebrain import CodeBrain
 from .contrawr import ContraWR
 from .ctnet import CTNet
+from .dance import DANCE
 from .deep4 import Deep4Net
 from .deepsleepnet import DeepSleepNet
 from .dgcnn import DGCNN
@@ -86,6 +87,7 @@ __all__ = [
     "CodeBrain",
     "ContraWR",
     "CTNet",
+    "DANCE",
     "Deep4Net",
     "DeepSleepNet",
     "BrainModule",

{braindecode-1.7.0.dev182642410 → braindecode-1.7.0.dev183981614}/braindecode/models/bendr.py

@@ -1,3 +1,8 @@
+# Authors: Bruno Aristimunha <b.aristimunha@gmail.com>
+#          Sarthak Tayal <sarthaktayal2@gmail.com>
+#
+# License: BSD (3-clause)
+
 import copy
 
 import numpy as np
@@ -435,7 +440,9 @@ class _ConvEncoderBENDR(nn.Module):
                         padding=width
                         // 2,  # Correct padding for 'same' output length before stride
                     ),
-                    nn.Dropout2d(dropout),  # 2D dropout (matches paper specification)
+                    # channel-wise dropout (matches paper specification), on
+                    # (batch, channels, times) activations
+                    nn.Dropout1d(dropout),
                     nn.GroupNorm(
                         encoder_h // 2, encoder_h
                     ),  # Consider making num_groups configurable or ensure encoder_h is divisible by 2

braindecode-1.7.0.dev183981614/braindecode/models/dance.py

@@ -0,0 +1,368 @@
+# Authors: Bruno Aristimunha <b.aristimunha@gmail.com>
+#          Meta Platforms, Inc. and affiliates (original DANCE)
+#
+# License: MIT
+# Adapted from https://github.com/facebookresearch/dance (MIT).
+"""``DANCE``: detect-and-classify EEG events (DETR for 1-D EEG)."""
+
+from __future__ import annotations
+
+import warnings
+
+import torch
+from torch import nn
+
+from braindecode.models.base import EEGModuleMixin
+from braindecode.models.util import (
+    _disable_batch_norm_training_if_batch_size_one,  # decorates ``forward``
+)
+from braindecode.models.util import has_valid_locations as _has_valid_locations
+from braindecode.models.util import positions_from_chs_info as _positions_from_chs_info
+from braindecode.modules import ChannelMerger, Perceiver, SimpleConv
+from braindecode.modules.dance_modules import DanceDetrDecoder
+
+
+class DANCE(EEGModuleMixin, nn.Module):
+    r"""DANCE from Lévy et al (2026) [dance]_.
+
+    :bdg-success:`Convolution` :bdg-info:`Attention/Transformer` :bdg-dark-line:`Channel`
+
+    DANCE frames EEG decoding as event *set prediction*: a long, unaligned
+    window is mapped to a set of events ``(t_start, t_end, class)`` with
+    normalized ``[0, 1]`` spans -- DETR for 1-D EEG.
+
+    .. rubric:: Architecture Overview
+
+    ``ChannelMerger`` (spatial Fourier attention over electrode positions) ->
+    ``SimpleConv`` (dilated conv stack) -> ``Perceiver`` (cross-attention to a
+    fixed 256-latent grid) -> a dense head (``forward`` output) and a DETR
+    cross-attention decoder (``detect`` output).
+
+    .. rubric:: Macro Components
+
+    ``DANCE.conv.merger`` (Spatial ChannelMerger)
+        **Operations.** :class:`~braindecode.modules.ChannelMerger` Fourier-embeds
+        each electrode's ``(x, y)`` position (:class:`~braindecode.modules.FourierEmb`)
+        and computes ``n_virtual_channels`` softmax-attention combinations of the
+        input channels, mapping ``(B, n_chans, T) -> (B, n_virtual_channels, T)``.
+        Nested inside ``self.conv`` (matching upstream); disabled (``None``) when
+        ``chs_info`` has no usable locations. Controlled by ``merger_drop_prob``.
+        **Role.** Makes the model montage-agnostic by projecting any electrode
+        layout onto a fixed virtual-channel basis.
+
+    ``DANCE.conv`` (SimpleConv dilated front-end)
+        **Operations.** :class:`~braindecode.modules.SimpleConv` runs the nested
+        merger, a ``1x1`` ``initial_linear`` projection, then ``conv_depth``
+        residual dilated ``Conv1d`` blocks (``nn.ReLU``, dilation accumulating as
+        ``int(dilation * conv_dilation_growth)`` per block, same-padding preserves
+        ``T``), mapping ``-> (B, embed_dim, T)``.
+        **Role.** Builds the temporal feature representation (Défossez lineage).
+
+    ``DANCE.perceiver`` (Perceiver bottleneck)
+        **Operations.** :class:`~braindecode.modules.Perceiver` Fourier-encodes the
+        time axis, then cross-attends a fixed ``num_latents``-token learnable grid to
+        the conv features over ``perceiver_depth`` blocks, mapping
+        ``(B, T, embed_dim) -> (B, num_latents, embed_dim)``.
+        **Role.** Makes the model length-agnostic (any ``T`` -> ``num_latents``
+        tokens) and forms the detection time-grid (``num_latents / duration``
+        tokens/s).
+
+    ``DANCE.decoder`` (DETR cross-attention decoder)
+        **Operations.** ``DanceDetrDecoder`` projects the latents to
+        ``decoder_dim``, then ``decoder_depth`` self-/cross-attention layers update
+        ``n_queries`` learnable event queries; per-query heads emit class logits and
+        sigmoid ``start``/``end`` spans. Used only by :meth:`detect` (not
+        :meth:`forward`).
+        **Role.** Produces the event-set prediction ``{class, start, end}``.
+
+    ``DANCE.final_layer`` (dense per-token head)
+        **Operations.** ``nn.Linear(embed_dim, n_outputs)`` applied to every latent
+        token, mapping ``(B, num_latents, embed_dim) -> (B, num_latents, n_outputs)``.
+        **Role.** The :meth:`forward` output (dense per-token class logits) and the
+        ``dense`` term consumed by the consistency loss.
+
+    .. rubric:: Temporal, Spatial, and Spectral Encoding
+
+    Temporal context comes from the dilated conv stack; spatial structure from
+    the Fourier-position channel merge; the Perceiver fourier-encodes the time
+    axis before cross-attention.
+
+    .. rubric:: Additional Mechanisms
+
+    The DETR decoder uses learnable event queries and emits per-query class
+    logits plus sigmoid start/end spans. The dense head provides per-token
+    class logits used as ``forward``'s output and for the consistency loss.
+
+    .. versionadded:: 1.6.1
+
+    Parameters
+    ----------
+    n_queries : int, optional
+        Number of learnable event queries. The default is ``100``.
+    use_channel_merger : bool, optional
+        Enable the spatial Fourier ChannelMerger. Auto-disabled if ``chs_info``
+        has no usable electrode locations. The default is ``True``.
+    n_virtual_channels : int, optional
+        Merger output channels. The default is ``270``.
+    fourier_emb_dim : int, optional
+        Fourier position embedding dim. The default is ``2048``.
+    merger_drop_prob : float, optional
+        Spatial dropout of the ChannelMerger ONLY (``self.conv.merger``);
+        bans whole channels within a random radius during training. The
+        default is ``0.2``.
+    embed_dim : int, optional
+        Conv/Perceiver feature dim. The default is ``128``.
+    conv_hidden : int, optional
+        SimpleConv hidden width. The default is ``512``.
+    conv_depth : int, optional
+        Number of dilated conv blocks. The default is ``10``.
+    conv_kernel_size : int, optional
+        Conv kernel size. The default is ``9``.
+    conv_dilation_growth : float, optional
+        Per-block dilation growth. The default is ``2.5``.
+    conv_initial_linear : int, optional
+        1x1 projection width. The default is ``256``.
+    conv_initial_depth : int, optional
+        Number of 1x1 projections. The default is ``1``.
+    conv_drop_prob : float, optional
+        Conv-block dropout knob, exposed for API symmetry. NOTE: the verified
+        DANCE config has NO conv dropout, so this maps to no upstream weights
+        and does not affect parity; kept for forward-compatibility. The
+        default is ``0.2`` (the paper value).
+    num_latents : int, optional
+        Perceiver latent count (detection token grid). The default is ``256``.
+    perceiver_depth : int, optional
+        Perceiver cross-attn blocks. The default is ``6``.
+    cross_attn_heads : int, optional
+        Perceiver cross-attn heads. The default is ``2``.
+    latent_attn_heads : int, optional
+        Perceiver self-attn heads. The default is ``2``.
+    cross_dim_head : int, optional
+        Perceiver cross-attn head dim. The default is ``64``.
+    latent_dim_head : int, optional
+        Perceiver self-attn head dim. The default is ``64``.
+    max_freq : float, optional
+        Perceiver fourier max frequency. The default is ``10.0``.
+    num_freq_bands : int, optional
+        Perceiver fourier bands. The default is ``6``.
+    decoder_dim : int, optional
+        DETR decoder dim. The default is ``256``.
+    decoder_depth : int, optional
+        DETR decoder layers. The default is ``4``.
+    decoder_heads : int, optional
+        DETR decoder heads. The default is ``4``.
+    activation : type[nn.Module], optional
+        Accepted for interface symmetry but currently INERT: it is forwarded to
+        ``self.decoder`` yet the decoder feed-forward hardwires GEGLU
+        (:class:`~braindecode.modules.dance_modules._FeedForward`), and the
+        ``SimpleConv`` front-end hardcodes ``nn.ReLU`` to match upstream. No
+        submodule reads it today. The default is ``nn.GELU``.
+    drop_prob : float, optional
+        Dropout applied to the raw input (``self.input_drop``) and inside the
+        DETR decoder (``self.decoder``). Does NOT touch the merger (use
+        ``merger_drop_prob``) or the conv stack (no conv dropout upstream).
+        The default is ``0.1``.
+
+    References
+    ----------
+    .. [dance] Lévy, Banville, Rapin, King, Moreau, d'Ascoli (2026). DANCE:
+       Detect and Classify Events in EEG. arXiv:2605.10688.
+    .. [defossez2023] Défossez et al. (2023). Decoding speech from
+       non-invasive brain recordings.
+    .. [perceiver2021] Jaegle et al. (2021). Perceiver: General perception
+       with iterative attention.
+    .. [detr2020] Carion et al. (2020). End-to-end object detection with
+       transformers.
+    """
+
+    def __init__(
+        self,
+        n_outputs=None,
+        n_chans=None,
+        chs_info=None,
+        n_times=None,
+        input_window_seconds=None,
+        sfreq=None,
+        *,
+        n_queries: int = 100,
+        use_channel_merger: bool = True,
+        n_virtual_channels: int = 270,
+        fourier_emb_dim: int = 2048,
+        merger_drop_prob: float = 0.2,
+        embed_dim: int = 128,
+        conv_hidden: int = 512,
+        conv_depth: int = 10,
+        conv_kernel_size: int = 9,
+        conv_dilation_growth: float = 2.5,
+        conv_initial_linear: int = 256,
+        conv_initial_depth: int = 1,
+        conv_drop_prob: float = 0.2,
+        num_latents: int = 256,
+        perceiver_depth: int = 6,
+        cross_attn_heads: int = 2,
+        latent_attn_heads: int = 2,
+        cross_dim_head: int = 64,
+        latent_dim_head: int = 64,
+        max_freq: float = 10.0,
+        num_freq_bands: int = 6,
+        decoder_dim: int = 256,
+        decoder_depth: int = 4,
+        decoder_heads: int = 4,
+        activation: type[nn.Module] = nn.GELU,
+        drop_prob: float = 0.1,
+    ):
+        super().__init__(
+            n_outputs=n_outputs,
+            n_chans=n_chans,
+            chs_info=chs_info,
+            n_times=n_times,
+            input_window_seconds=input_window_seconds,
+            sfreq=sfreq,
+        )
+        del n_outputs, n_chans, sfreq, n_times, input_window_seconds, chs_info
+
+        self.num_latents = num_latents
+        self.n_queries = n_queries
+        self.embed_dim = embed_dim
+
+        # Decide the merger up front (fall back to no-merger if no locations).
+        if use_channel_merger and not _has_valid_locations(self.chs_info):
+            warnings.warn(
+                "DANCE: chs_info has no usable electrode locations "
+                "('loc' missing or all-zero); disabling the ChannelMerger "
+                "(use_channel_merger=False).",
+                UserWarning,
+            )
+            use_channel_merger = False
+        # Public flag mirroring the fallback decision (``self.conv.merger`` is the
+        # source of truth in forward); kept for introspection/tests.
+        self.use_channel_merger = use_channel_merger
+
+        # The merger is NESTED INSIDE SimpleConv (self.conv.merger), matching
+        # upstream SimpleConvModel's call path and state_dict layout. There is
+        # NO top-level self.channel_merger. The (n_chans, 2) positions buffer is
+        # derived from chs_info and broadcast in forward; subject_ids is always
+        # None (braindecode has no subjects; per_subject=False).
+        merger = None
+        if use_channel_merger:
+            positions = _positions_from_chs_info(self.chs_info)
+            self.register_buffer(
+                "channel_positions",
+                torch.as_tensor(positions, dtype=torch.float32),
+                persistent=False,
+            )
+            merger = ChannelMerger(
+                out_channels=n_virtual_channels,
+                pos_dim=fourier_emb_dim,
+                dropout=merger_drop_prob,
+            )
+
+        self.input_drop = nn.Dropout(drop_prob)
+        self.conv = SimpleConv(
+            in_channels=self.n_chans,
+            out_channels=embed_dim,
+            hidden=conv_hidden,
+            depth=conv_depth,
+            kernel_size=conv_kernel_size,
+            dilation_growth=conv_dilation_growth,
+            initial_linear=conv_initial_linear,
+            initial_depth=conv_initial_depth,
+            drop_prob=conv_drop_prob,
+            activation=nn.ReLU,
+            merger=merger,
+        )
+        # The conv stack is SAME-padded, so it preserves length for any T >= 1;
+        # the dilated receptive field is not an input-length requirement. The
+        # genuine minimum is one first-block kernel (dilation 1), so the guard
+        # uses ``conv_kernel_size`` rather than the full receptive field.
+        self._min_n_times = conv_kernel_size
+        self.perceiver = Perceiver(
+            input_dim=embed_dim,
+            num_latents=num_latents,
+            latent_dim=embed_dim,
+            depth=perceiver_depth,
+            cross_heads=cross_attn_heads,
+            latent_heads=latent_attn_heads,
+            cross_dim_head=cross_dim_head,
+            latent_dim_head=latent_dim_head,
+            max_freq=max_freq,
+            num_freq_bands=num_freq_bands,
+        )
+        self.decoder = DanceDetrDecoder(
+            input_dim=embed_dim,
+            dim=decoder_dim,
+            depth=decoder_depth,
+            heads=decoder_heads,
+            n_queries=n_queries,
+            n_outputs=self.n_outputs,
+            drop_prob=drop_prob,
+            activation=activation,
+        )
+        # final_layer LAST so it lands in the last two named_children(); init
+        # weights AFTER it exists so the dense head gets the custom init too.
+        self.final_layer = nn.Linear(embed_dim, self.n_outputs)
+        self.apply(self._init_weights)
+
+    @staticmethod
+    def _init_weights(module: nn.Module) -> None:
+        if isinstance(module, nn.Linear):
+            nn.init.trunc_normal_(module.weight, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.ones_(module.weight)
+            nn.init.zeros_(module.bias)
+
+    def _encode(self, x: torch.Tensor) -> torch.Tensor:
+        if x.ndim != 3 or x.shape[-2] != self.n_chans:
+            raise ValueError(
+                f"expected (batch, {self.n_chans}, T) input; got {tuple(x.shape)}."
+            )
+        if x.shape[-1] < self._min_n_times:
+            raise ValueError(
+                f"n_times={x.shape[-1]} is shorter than the minimum input "
+                f"length ({self._min_n_times} samples = one conv kernel)."
+            )
+        x = self.input_drop(x)
+        # The merger is nested inside self.conv; pass positions through it.
+        # subject_ids is always None (braindecode has no subjects).
+        if self.conv.merger is not None:
+            pos = self.channel_positions.unsqueeze(0).expand(x.size(0), -1, -1)
+            x = self.conv(x, positions=pos)  # merger -> initial_linear -> blocks
+        else:
+            x = self.conv(x)  # (B, embed_dim, T)
+        x = x.transpose(2, 1)  # (B, T, embed_dim)
+        x = self.perceiver(x)  # (B, num_latents, embed_dim)
+        return x
+
+    @_disable_batch_norm_training_if_batch_size_one
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        latents = self._encode(x)  # (B, num_latents, embed_dim)
+        return self.final_layer(latents)  # (B, num_latents, n_outputs)
+
+    @_disable_batch_norm_training_if_batch_size_one
+    def detect(self, x: torch.Tensor) -> dict:
+        latents = self._encode(x)  # (B, num_latents, embed_dim)
+        events = self.decoder(latents)  # {class, start, end}
+        events["dense"] = self.final_layer(latents)  # (B, num_latents, n_outputs)
+        return events
+
+    def reset_head(self, n_outputs: int) -> None:
+        """Replace the dense head and the DETR class head for a new ``n_outputs``."""
+        if n_outputs <= 0:
+            raise ValueError(f"n_outputs must be positive; got {n_outputs}.")
+        old = self.final_layer
+        self.final_layer = nn.Linear(old.in_features, n_outputs).to(
+            device=old.weight.device, dtype=old.weight.dtype
+        )
+        ch = self.decoder.class_head
+        self.decoder.class_head = nn.Linear(ch.in_features, n_outputs).to(
+            device=ch.weight.device, dtype=ch.weight.dtype
+        )
+        self._n_outputs = n_outputs
+        init_kwargs = getattr(self, "_braindecode_init_kwargs", None)
+        if init_kwargs is not None and "n_outputs" in init_kwargs:
+            init_kwargs["n_outputs"] = n_outputs
+        hub_config = getattr(self, "_hub_mixin_config", None)
+        if hub_config is not None and "n_outputs" in hub_config:
+            hub_config["n_outputs"] = n_outputs

{braindecode-1.7.0.dev182642410 → braindecode-1.7.0.dev183981614}/braindecode/models/eegminer.py

@@ -5,10 +5,9 @@
 * https://www.ipo.gov.uk/p-ipsum/Case/ApplicationNumber/GB2113420.0
 """
 
-from functools import partial
+# Authors: Sarthak Tayal <sarthaktayal2@gmail.com>
 
 import torch
-from einops.layers.torch import Rearrange
 from torch import nn
 
 import braindecode.functional as F
@@ -19,6 +18,36 @@ from braindecode.modules import GeneralizedGaussianFilter
 _eeg_miner_methods = ["mag", "corr", "plv"]
 
 
+class _EEGMinerFeatures(nn.Module):
+    def __init__(self, method, n_chans, n_filters, n_times):
+        super().__init__()
+        self.method = method
+        self.n_chans = n_chans
+        self.n_filters = n_filters
+        self.n_times = n_times
+        triu = torch.triu_indices(n_chans, n_chans, 1)
+        self.register_buffer("triu_row", triu[0], persistent=False)
+        self.register_buffer("triu_col", triu[1], persistent=False)
+
+    def forward(self, x):
+        if self.method == "mag":
+            return torch.sqrt((x * x).mean(dim=-1))
+        if self.method == "corr":
+            x = x.reshape(
+                x.shape[0], self.n_chans, self.n_filters, self.n_times
+            ).transpose(-3, -2)
+            x = (x - x.mean(dim=-1, keepdim=True)) / torch.sqrt(
+                x.var(dim=-1, keepdim=True) + 1e-6
+            )
+            x = torch.matmul(x, x.transpose(-2, -1)) / x.shape[-1]
+            x = x.permute(0, 2, 3, 1).abs()
+        else:
+            x = x.transpose(-4, -3)
+            x = F.plv_time(x, forward_fourier=False)
+            x = x.permute(0, 2, 3, 1)
+        return x[:, self.triu_row, self.triu_col, :]
+
+
 class EEGMiner(EEGModuleMixin, nn.Module):
     r"""EEGMiner from Ludwig et al (2024) [eegminer]_.
 
@@ -168,24 +197,18 @@ class EEGMiner(EEGModuleMixin, nn.Module):
         )
 
         # Forward method
+        self.ensure_dim = nn.Identity()
         if self.method == "mag":
-            self.method_forward = self._apply_mag_forward
             self.n_features = self.n_chans * self.n_filters
-            self.ensure_dim = nn.Identity()
         elif self.method == "corr":
-            self.method_forward = partial(
-                self._apply_corr_forward,
-                n_chans=self.n_chans,
-                n_filters=self.n_filters,
-                n_times=self.n_times,
-            )
             self.n_features = self.n_filters * self.n_chans * (self.n_chans - 1) // 2
-            self.ensure_dim = nn.Identity()
         elif self.method == "plv":
-            self.method_forward = partial(self._apply_plv, n_chans=self.n_chans)
-            self.ensure_dim = Rearrange("... d -> ... 1 d")
+            self.ensure_dim = nn.Unflatten(-1, (1, self.n_times))
             self.n_features = (self.n_filters * self.n_chans * (self.n_chans - 1)) // 2
 
+        self.feature_layer = _EEGMinerFeatures(
+            self.method, self.n_chans, self.n_filters, self.n_times
+        )
         self.flatten_layer = nn.Flatten()
         # Classifier
         self.batch_layer = nn.BatchNorm1d(self.n_features, affine=False)
@@ -195,64 +218,17 @@ class EEGMiner(EEGModuleMixin, nn.Module):
     @_disable_batch_norm_training_if_batch_size_one
     def forward(self, x):
         """x: (batch, electrodes, time)"""
-        batch = x.shape[0]
         x = self.ensure_dim(x)
         # Apply Gaussian filters in frequency domain
         # x -> (batch, electrodes * filters, time)
         x = self.filter(x)
 
-        x = self.method_forward(x=x, batch=batch)
+        x = self.feature_layer(x)
         # Classifier
         # Note that the order of dimensions before flattening the feature vector is important
         # for attributing feature weights during interpretation.
-        x = x.reshape(batch, self.n_features)
+        x = x.reshape(x.shape[0], self.n_features)
         x = self.batch_layer(x)
         x = self.final_layer(x)
 
         return x
-
-    @staticmethod
-    def _apply_mag_forward(x, batch=None):
-        # Signal magnitude
-        x = x * x
-        x = x.mean(dim=-1)
-        x = torch.sqrt(x)
-        return x
-
-    @staticmethod
-    def _apply_corr_forward(
-        x, batch, n_chans, n_filters, n_times, epilson: float = 1e-6
-    ):
-        x = x.reshape(batch, n_chans, n_filters, n_times).transpose(-3, -2)
-        x = (x - x.mean(dim=-1, keepdim=True)) / torch.sqrt(
-            x.var(dim=-1, keepdim=True) + epilson
-        )
-        x = torch.matmul(x, x.transpose(-2, -1)) / x.shape[-1]
-        # Original tensor shape: [batch, n_filters, chans, chans]
-        x = x.permute(0, 2, 3, 1)
-        # New tensor shape: [batch, chans, chans, n_filters]
-        # move filter channels to the end
-        x = x.abs()
-
-        # Get upper triu of symmetric connectivity matrix
-        triu = torch.triu_indices(n_chans, n_chans, 1)
-        x = x[:, triu[0], triu[1], :]
-
-        return x
-
-    @staticmethod
-    def _apply_plv(x, n_chans, batch=None):
-        # Compute PLV connectivity
-        # x -> (batch, electrodes, electrodes, filters)
-        x = x.transpose(-4, -3)  # swap electrodes and filters
-        # adjusting to compute the plv
-        x = F.plv_time(x, forward_fourier=False)
-        # batch, number of filters, connectivity matrix
-        # [batch, n_filters, chans, chans]
-        x = x.permute(0, 2, 3, 1)
-        # [batch, chans, chans, n_filters]
-
-        # Get upper triu of symmetric connectivity matrix
-        triu = torch.triu_indices(n_chans, n_chans, 1)
-        x = x[:, triu[0], triu[1], :]
-        return x