braindecode 1.5.0.dev985__tar.gz → 1.5.0.dev172791986__tar.gz

{braindecode-1.5.0.dev985/braindecode.egg-info → braindecode-1.5.0.dev172791986}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.5.0.dev985
+Version: 1.5.0.dev172791986
 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.5.0.dev985 → braindecode-1.5.0.dev172791986}/braindecode/models/__init__.py

@@ -4,8 +4,8 @@ from .atcnet import ATCNet
 from .attentionbasenet import AttentionBaseNet
 from .attn_sleep import AttnSleep
 from .base import EEGModuleMixin
-from .bendr import BENDR
-from .biot import BIOT
+from .bendr import BENDR, InterpolatedBENDR
+from .biot import BIOT, InterpolatedBIOT
 from .brainmodule import BrainModule
 from .cbramod import CBraMod
 from .codebrain import CodeBrain
@@ -30,7 +30,8 @@ from .fblightconvnet import FBLightConvNet
 from .fbmsnet import FBMSNet
 from .hybrid import HybridNet
 from .ifnet import IFNet
-from .labram import Labram
+from .interpolated import InterpolatedModel
+from .labram import InterpolatedLaBraM, Labram
 from .luna import LUNA
 from .medformer import MEDFormer
 from .msvtnet import MSVTNet
@@ -39,6 +40,7 @@ from .reve import REVE
 from .sccnet import SCCNet
 from .shallow_fbcsp import ShallowFBCSPNet
 from .signal_jepa import (
+    InterpolatedSignalJEPA,
     SignalJEPA,
     SignalJEPA_Contextual,
     SignalJEPA_PostLocal,
@@ -97,6 +99,11 @@ __all__ = [
     "FBMSNet",
     "HybridNet",
     "IFNet",
+    "InterpolatedBENDR",
+    "InterpolatedBIOT",
+    "InterpolatedLaBraM",
+    "InterpolatedModel",
+    "InterpolatedSignalJEPA",
     "Labram",
     "LUNA",
     "extract_channel_locations_from_chs_info",

{braindecode-1.5.0.dev985 → braindecode-1.5.0.dev172791986}/braindecode/models/bendr.py

@@ -1,12 +1,60 @@
 import copy
 
+import numpy as np
 import torch
 from einops.layers.torch import Rearrange
 from torch import nn
-from torch.nn.utils.parametrize import register_parametrization
 
 from braindecode.models.base import EEGModuleMixin
-from braindecode.modules.parametrization import MaxNormParametrize
+
+# The 20 channels used to pre-train BENDR, in the order expected by the
+# `braindecode/braindecode-bendr` checkpoint. The first 19 entries are the
+# EEG channels taken verbatim from `dn3.transforms.instance.To1020.EEG_20_div`
+# (https://github.com/SPOClab-ca/dn3/blob/master/dn3/transforms/instance.py).
+# Their positions come from MNE's ``standard_1005`` montage (T5/T6 are
+# legacy names that share positions with P7/P8 there).
+#
+# The 20th entry is ``SCALE``, a relative-amplitude statistic (not an
+# electrode) appended by ``To1020(include_scale_ch=True)`` during
+# pre-training. Since it has no physical position, the ``loc`` below is
+# the centroid of the 19 EEG positions — purely a placeholder so that
+# :class:`~braindecode.modules.ChannelInterpolationLayer` (used by
+# :class:`InterpolatedBENDR`) can build a valid spline interpolation
+# matrix. It is NOT the SCALE the pre-training pipeline computes
+# (which is an RMS-like amplitude via ``dn3.MappingDeep1010``); users
+# who need a faithful SCALE must compute it themselves and feed 20
+# channels to :class:`BENDR` directly.
+_BENDR_TARGET_CHS_TUPLES: list[tuple[str, tuple[float, float, float]]] = [
+    ("FP1", (-0.0294367, +0.0839171, -0.0069900)),  # standard_1005
+    ("FP2", (+0.0298723, +0.0848959, -0.0070800)),  # standard_1005
+    ("F7", (-0.0702629, +0.0424743, -0.0114200)),  # standard_1005
+    ("F3", (-0.0502438, +0.0531112, +0.0421920)),  # standard_1005
+    ("FZ", (+0.0003122, +0.0585120, +0.0664620)),  # standard_1005
+    ("F4", (+0.0518362, +0.0543048, +0.0408140)),  # standard_1005
+    ("F8", (+0.0730431, +0.0444217, -0.0120000)),  # standard_1005
+    ("T7", (-0.0841611, -0.0160187, -0.0093460)),  # standard_1005
+    ("C3", (-0.0653581, -0.0116317, +0.0643580)),  # standard_1005
+    ("CZ", (+0.0004009, -0.0091670, +0.1002440)),  # standard_1005
+    ("C4", (+0.0671179, -0.0109003, +0.0635800)),  # standard_1005
+    ("T8", (+0.0850799, -0.0150203, -0.0094900)),  # standard_1005
+    ("T5", (-0.0724343, -0.0734527, -0.0024870)),  # standard_1005 (= P7)
+    ("P3", (-0.0530073, -0.0787878, +0.0559400)),  # standard_1005
+    ("PZ", (+0.0003247, -0.0811150, +0.0826150)),  # standard_1005
+    ("P4", (+0.0556667, -0.0785602, +0.0565610)),  # standard_1005
+    ("T6", (+0.0730557, -0.0730683, -0.0025400)),  # standard_1005 (= P8)
+    ("O1", (-0.0294134, -0.1124490, +0.0088390)),  # standard_1005
+    ("O2", (+0.0298426, -0.1121560, +0.0088000)),  # standard_1005
+    (
+        "SCALE",
+        (+0.0006439, -0.0131942, +0.0278448),
+    ),  # centroid of the 19 EEG positions (placeholder; see comment above)
+]
+
+_BENDR_TARGET_CHS_INFO: list[dict] = [
+    {"ch_name": ch, "kind": "eeg", "loc": np.asarray(loc, dtype=float)}
+    for ch, loc in _BENDR_TARGET_CHS_TUPLES
+]
+BENDR_CHANNEL_ORDER: list[str] = [ch for ch, _ in _BENDR_TARGET_CHS_TUPLES]
 
 
 class BENDR(EEGModuleMixin, nn.Module):
@@ -195,15 +243,6 @@ class BENDR(EEGModuleMixin, nn.Module):
         The contextualizer is still created (to allow loading pretrained weights) but is not
         used in the forward pass. Requires input length of at least
         ``4 * product(enc_downsample)`` samples (384 with default downsampling of 96x).
-    n_chans_pretrained : int or None, default=None
-        Number of input channels the pretrained weights expect (20 for the official BENDR
-        checkpoint). When set and ``n_chans != n_chans_pretrained``, a 1x1 Conv1d with
-        max-norm constraint projects from ``n_chans`` to ``n_chans_pretrained`` before the
-        encoder. This allows fine-tuning pretrained BENDR on datasets with arbitrary channel
-        counts. When using ``from_pretrained``, pass ``strict=False`` since the checkpoint
-        will not contain ``channel_projection`` weights.
-    chan_proj_max_norm : float, default=1.0
-        Max-norm constraint value for the channel projection weights.
     """
 
     def __init__(
@@ -231,8 +270,6 @@ class BENDR(EEGModuleMixin, nn.Module):
         start_token=-5,  # Value for start token embedding
         final_layer=True,  # Whether to include the final linear layer
         encoder_only=False,  # If True, bypass contextualizer and use 4-chunk pooling
-        n_chans_pretrained=None,  # Expected input channels of pretrained weights
-        chan_proj_max_norm=1.0,  # Max-norm for channel projection weights
     ):
         super().__init__(
             n_outputs=n_outputs,
@@ -246,25 +283,34 @@ class BENDR(EEGModuleMixin, nn.Module):
         # Keep these parameters if needed later, otherwise they are captured by the mixin
         del n_outputs, n_chans, chs_info, n_times, input_window_seconds, sfreq
 
+        # If the user supplies chs_info, require it to match BENDR_CHANNEL_ORDER
+        # exactly (case-insensitive). Arbitrary channel sets should go through
+        # InterpolatedBENDR — same pattern as Labram / InterpolatedLaBraM.
+        # When chs_info is absent (the usual n_chans=20 path, incl.
+        # from_pretrained), no check is performed.
+        try:
+            _chs_info = self.chs_info
+        except ValueError:
+            _chs_info = None
+        if _chs_info is not None:
+            user_names = [ch["ch_name"] for ch in _chs_info]  # type: ignore[index]
+            canonical = BENDR_CHANNEL_ORDER
+            if [n.lower() for n in user_names] != [n.lower() for n in canonical]:
+                raise ValueError(
+                    f"BENDR requires chs_info to match BENDR_CHANNEL_ORDER exactly "
+                    f"({len(canonical)} channels, specific order; last is 'SCALE'). "
+                    f"Got {len(user_names)} channel(s). For arbitrary channel sets, "
+                    f"use InterpolatedBENDR "
+                    f"(from braindecode.models import InterpolatedBENDR)."
+                )
+
         self.encoder_h = encoder_h
         self.contextualizer_hidden = contextualizer_hidden
         self.include_final_layer = final_layer
         self.encoder_only = encoder_only
 
-        # Channel projection for pretrained weight compatibility
-        encoder_n_chans = self.n_chans
-        if n_chans_pretrained is not None and self.n_chans != n_chans_pretrained:
-            conv = nn.Conv1d(self.n_chans, n_chans_pretrained, 1, bias=False)
-            register_parametrization(
-                conv, "weight", MaxNormParametrize(chan_proj_max_norm)
-            )
-            self.channel_projection = conv
-            encoder_n_chans = n_chans_pretrained
-        else:
-            self.channel_projection = None
-
         self.encoder = _ConvEncoderBENDR(
-            in_features=encoder_n_chans,
+            in_features=self.n_chans,
             encoder_h=encoder_h,
             dropout=drop_prob,
             projection_head=projection_head,
@@ -308,8 +354,6 @@ class BENDR(EEGModuleMixin, nn.Module):
         self._build_head(n_outputs)
 
     def forward(self, x, return_features=False):
-        if self.channel_projection is not None:
-            x = self.channel_projection(x)
         encoded = self.encoder(x)
         # encoded: [batch_size, encoder_h, n_encoded_times]
 
@@ -552,3 +596,23 @@ class _BENDRContextualizer(nn.Module):
         # x: [batch_size, in_features, seq_len + 1]
 
         return x
+
+
+# -----------------------------------------------------------------------------
+# InterpolatedBENDR — experimental channel-interpolation variant of BENDR
+# -----------------------------------------------------------------------------
+# Wraps :class:`BENDR` with an MNE-backed channel-interpolation layer that
+# projects arbitrary user ``chs_info`` to the canonical 20-channel BENDR
+# input (:data:`_BENDR_TARGET_CHS_INFO` — the 19 pre-training EEG channels
+# plus a ``SCALE`` placeholder at the centroid of those 19 positions).
+# Frozen by default; set ``trainable=True`` to fine-tune the projection.
+#
+# NOTE: the ``SCALE`` target has no physical position, so the row of the
+# interpolation matrix that produces it is a spatial spline of the user's
+# EEG channels — *not* the dn3 ``MappingDeep1010`` RMS statistic the
+# checkpoint saw during pre-training. Expect degraded zero-shot transfer
+# from the SCALE channel; downstream fine-tuning should still work.
+
+from braindecode.models.interpolated import InterpolatedModel  # noqa: E402
+
+InterpolatedBENDR = InterpolatedModel(BENDR, _BENDR_TARGET_CHS_INFO)

{braindecode-1.5.0.dev985 → braindecode-1.5.0.dev172791986}/braindecode/models/biot.py

@@ -1,12 +1,57 @@
 import math
 from warnings import warn
 
+import numpy as np
 import torch
 import torch.nn as nn
 from linear_attention_transformer import LinearAttentionTransformer
 
 from braindecode.models.base import EEGModuleMixin
 
+# -----------------------------------------------------------------------------
+# Canonical channel order for InterpolatedBIOT — the 18-channel TCP bipolar
+# montage used by BIOT's shhs-prest and six-datasets pretrained checkpoints.
+# Source: https://github.com/ycq091044/BIOT (README + datasets/TUAB/process.py
+# + datasets/SHHS/process.py). Indices 0-15 are the TCP 16-channel bipolar
+# derivations; indices 16-17 are SHHS differential channels.
+#
+# The `loc` values are only used to build an MNE interpolation matrix for
+# InterpolatedBIOT. All entries are bipolar / differential derivations.
+# TODO: positions are stored as the midpoint of the two constituent
+# electrodes. This is a simplification — a bipolar signal V(A)-V(B) cannot
+# be faithfully recovered by spatial interpolation at the midpoint. Revisit
+# in a follow-up PR (e.g. a dedicated BipolarDerivationLayer).
+# -----------------------------------------------------------------------------
+
+# fmt: off
+_BIOT_TARGET_CHS_TUPLES: list[tuple[str, tuple[float, float, float]]] = [
+    ("FP1-F7", (-0.04984980, 0.06319570, -0.00920500)),
+    ("F7-T7", (-0.07721200, 0.01322780, -0.01038300)),
+    ("T7-P7", (-0.07829770, -0.04473570, -0.00591650)),
+    ("P7-O1", (-0.05092385, -0.09295085, 0.00317600)),
+    ("FP2-F8", (0.05145770, 0.06465880, -0.00954000)),
+    ("F8-T8", (0.07906150, 0.01470070, -0.01074500)),
+    ("T8-P8", (0.07906780, -0.04404430, -0.00601500)),
+    ("P8-O2", (0.05144915, -0.09261215, 0.00313000)),
+    ("FP1-F3", (-0.03984025, 0.06851415, 0.01760100)),
+    ("F3-C3", (-0.05780095, 0.02073975, 0.05327500)),
+    ("C3-P3", (-0.05918270, -0.04520975, 0.06014900)),
+    ("P3-O1", (-0.04121035, -0.09561840, 0.03238950)),
+    ("FP2-F4", (0.04085425, 0.06960035, 0.01686700)),
+    ("F4-C4", (0.05947705, 0.02170225, 0.05219700)),
+    ("C4-P4", (0.06139230, -0.04473025, 0.06007050)),
+    ("P4-O2", (0.04275465, -0.09535810, 0.03268050)),
+    ("C3-A2", (0.01021790, -0.01832050, -0.00183650)),
+    ("C4-A1", (-0.00947910, -0.01794500, -0.00220300)),
+]
+# fmt: on
+
+_BIOT_TARGET_CHS_INFO = [
+    {"ch_name": ch, "kind": "eeg", "loc": np.asarray(loc, dtype=float)}
+    for ch, loc in _BIOT_TARGET_CHS_TUPLES
+]
+BIOT_CHANNEL_ORDER = [ch for ch, _ in _BIOT_TARGET_CHS_TUPLES]
+
 
 class BIOT(EEGModuleMixin, nn.Module):
     r"""BIOT from Yang et al (2023) [Yang2023]_
@@ -439,7 +484,9 @@ class _BIOTEncoder(nn.Module):
         self.channel_tokens = nn.Embedding(
             num_embeddings=n_chans, embedding_dim=emb_size
         )
-        self.register_buffer("index", torch.arange(n_chans, dtype=torch.long))
+        self.register_buffer(
+            "index", torch.arange(n_chans, dtype=torch.long), persistent=False
+        )
 
     def stft(self, sample):
         """
@@ -553,3 +600,16 @@ class _BIOTEncoder(nn.Module):
         # (batch_size, emb)
         emb = self.transformer(emb).mean(dim=1)
         return emb
+
+
+# -----------------------------------------------------------------------------
+# InterpolatedBIOT — experimental channel-interpolation variant of BIOT
+# -----------------------------------------------------------------------------
+# Wraps :class:`BIOT` with an MNE-backed channel-interpolation layer that
+# projects arbitrary user ``chs_info`` to the canonical 18-channel BIOT
+# montage (:data:`_BIOT_TARGET_CHS_INFO`). Frozen by default; set
+# ``trainable=True`` to fine-tune the projection matrix.
+
+from braindecode.models.interpolated import InterpolatedModel  # noqa: E402
+
+InterpolatedBIOT = InterpolatedModel(BIOT, _BIOT_TARGET_CHS_INFO)

braindecode-1.5.0.dev172791986/braindecode/models/interpolated.py

@@ -0,0 +1,182 @@
+# Authors: Pierre Guetschel
+#
+# License: BSD (3-clause)
+from __future__ import annotations
+
+from typing import Literal, Optional, Type
+
+import numpy as np
+import torch.nn as nn
+
+from braindecode.modules.interpolation import ChannelInterpolationLayer
+
+
+def InterpolatedModel(
+    model_cls: Type,
+    target_chs_info: list[dict],
+    name: Optional[str] = None,
+) -> Type:
+    """Return a subclass of ``model_cls`` that interpolates channels to ``target_chs_info``.
+
+    .. warning:: Experimental. Public API may change without a deprecation cycle.
+
+    Parameters
+    ----------
+    model_cls : type
+        A braindecode model class (subclass of
+        :class:`~braindecode.models.base.EEGModuleMixin`).
+    target_chs_info : list of dict
+        The canonical channel set the backbone expects internally. Every
+        instance of the returned class projects its input channels to
+        this set via :class:`~braindecode.modules.ChannelInterpolationLayer`.
+    name : str, optional
+        ``__name__`` to assign to the returned class. Defaults to
+        ``f"Interpolated{model_cls.__name__}"``.
+
+    Returns
+    -------
+    type
+        A new subclass of ``model_cls`` whose ``__init__`` accepts
+        arbitrary user ``chs_info`` and automatically inserts a frozen
+        (by default) channel-interpolation layer before the backbone.
+    """
+    _is_sequential = issubclass(model_cls, nn.Sequential)
+
+    class _Interpolated(model_cls):
+        _TARGET_CHS_INFO = target_chs_info
+
+        def __init__(
+            self,
+            chs_info,
+            n_outputs=None,
+            n_times=None,
+            input_window_seconds=None,
+            sfreq=None,
+            n_chans=None,
+            interpolation_method: str = "spline",
+            interpolation_mode: Literal["always", "name_match"] = "name_match",
+            trainable: bool = False,
+            **kwargs,
+        ):
+            # Signal-related params are declared EXPLICITLY here so that
+            # skorch's ``EEGClassifier._set_signal_args`` (which inspects
+            # the ``__init__`` signature via ``inspect.signature``) can
+            # auto-forward them from the training dataset. They would not
+            # be discoverable if they were collapsed into ``**kwargs``.
+            # ``n_chans`` is declared for the same discoverability reason
+            # but intentionally ignored: the backbone must see
+            # ``len(target_chs_info)`` derived from ``chs_info``.
+            del n_chans
+            # Backbone init uses the target channels. During this call,
+            # some backbones run a dummy forward (e.g. to size the head);
+            # ``self.interpolation_layer`` does not exist yet — the
+            # ``forward`` override below falls back to pass-through when
+            # the attribute is absent. Assigning an ``nn.Identity()``
+            # before ``super()`` is impossible: ``nn.Module.__setattr__``
+            # requires ``nn.Module.__init__`` to have run, and the chain
+            # would wipe ``self._modules`` when it reaches it again.
+            super().__init__(
+                chs_info=target_chs_info,
+                n_outputs=n_outputs,
+                n_times=n_times,
+                input_window_seconds=input_window_seconds,
+                sfreq=sfreq,
+                **kwargs,
+            )
+
+            layer = ChannelInterpolationLayer(
+                src_chs_info=chs_info,
+                tgt_chs_info=target_chs_info,
+                mode=interpolation_mode,
+                method=interpolation_method,
+                trainable=trainable,
+            )
+            if _is_sequential:
+                # For nn.Sequential subclasses, prepend the interpolation
+                # layer so that nn.Sequential.forward runs it first.
+                # Registering via attribute assignment appends to _modules;
+                # instead we rebuild _modules with the layer first.
+                old_modules = list(self._modules.items())  # type: ignore[has-type]
+                self._modules.clear()  # type: ignore[has-type]
+                self._modules["interpolation_layer"] = layer  # type: ignore[index]
+                for k, v in old_modules:
+                    self._modules[k] = v  # type: ignore[index]
+            else:
+                self.interpolation_layer = layer
+
+            # Rebind private attrs so the user-facing view (.chs_info,
+            # .n_chans, .input_shape, build_model_config) reflects the
+            # user's channels. Properties are NOT overridden — we mutate
+            # the private attrs the base-class properties read from.
+            self._chs_info = chs_info
+            self._n_chans = len(chs_info)
+
+        if not _is_sequential:
+
+            def forward(self, x, *args, **kwargs):
+                # During super().__init__() the interpolation_layer attr
+                # does not exist yet; any dummy forward call (e.g. from
+                # get_output_shape) must pass through unchanged so the
+                # backbone sees its expected target-shape input.
+                # Forward *args / **kwargs so backbone-specific flags
+                # like ``return_features`` keep working through the
+                # wrapper.
+                interp = getattr(self, "interpolation_layer", None)
+                if interp is not None:
+                    x = interp(x)
+                return super().forward(x, *args, **kwargs)
+
+    _Interpolated.__name__ = name or f"Interpolated{model_cls.__name__}"
+    _Interpolated.__qualname__ = _Interpolated.__name__
+    # Propagate the backbone docstring so Sphinx and the categorization tests
+    # can read the class badges.  Prepend a short header so the class shows
+    # up clearly in documentation as distinct from the backbone.
+    backbone_doc = model_cls.__doc__ or ""
+    _Interpolated.__doc__ = (
+        f"Channel-interpolating wrapper around :class:`{model_cls.__name__}`.\n\n"
+        ":bdg-dark-line:`Channel`\n\n"
+        f"Accepts arbitrary user ``chs_info`` and projects them to the\n"
+        f"backbone's canonical channel set via\n"
+        f":class:`~braindecode.modules.ChannelInterpolationLayer`.\n\n"
+        f"For all other parameters and behavior see the backbone\n"
+        f"documentation reproduced below.\n\n" + backbone_doc
+    )
+    return _Interpolated
+
+
+def _build_chs_info_from_montage(names: list[str], montage: str) -> list[dict]:
+    """Build a ``list[dict]`` ``chs_info`` from channel names + an MNE montage.
+
+    Each returned dict has ``ch_name``, ``kind="eeg"``, and ``loc`` (shape
+    ``(3,)``). Used by braindecode's shipped ``Interpolated*`` variants to
+    turn a bare list of canonical channel names into the dict form
+    ``ChannelInterpolationLayer`` expects.
+
+    Parameters
+    ----------
+    names : list of str
+        Channel names in the desired order.
+    montage : str
+        Name of an MNE standard montage (e.g. ``"standard_1005"``).
+
+    Returns
+    -------
+    list of dict
+
+    Raises
+    ------
+    ValueError
+        If a name is not found in the montage.
+    """
+    import mne
+
+    mtg = mne.channels.make_standard_montage(montage)
+    ch_pos = mtg.get_positions()["ch_pos"]
+    out = []
+    for n in names:
+        if n not in ch_pos:
+            raise ValueError(f"Channel {n!r} not found in montage {montage!r}.")
+        out.append(
+            {"ch_name": n, "kind": "eeg", "loc": np.asarray(ch_pos[n], dtype=float)}
+        )
+    return out