braindecode 1.6.1.dev1073__tar.gz → 1.6.1.dev1079__tar.gz

{braindecode-1.6.1.dev1073/braindecode.egg-info → braindecode-1.6.1.dev1079}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.6.1.dev1073
+Version: 1.6.1.dev1079
 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.6.1.dev1073 → braindecode-1.6.1.dev1079}/braindecode/datasets/__init__.py

@@ -10,6 +10,7 @@ from .base import (
 from .bcicomp import BCICompetitionIVDataset4
 from .bids import BIDSDataset, BIDSEpochsDataset
 from .chb_mit import CHBMIT
+from .collate import pad_channels_collate
 from .mne import create_from_mne_epochs, create_from_mne_raw
 from .moabb import BNCI2014_001, HGD, MOABBDataset
 from .nmt import NMT
@@ -42,4 +43,5 @@ __all__ = [
     "SleepPhysionetChallenge2018",
     "create_from_X_y",
     "BCICompetitionIVDataset4",
+    "pad_channels_collate",
 ]

{braindecode-1.6.1.dev1073 → braindecode-1.6.1.dev1079}/braindecode/datasets/base.py

@@ -192,6 +192,31 @@ def _channel_info(mne_obj):
     return n_ch, type_str, sfreq
 
 
+def _compute_ch_pos(info, targets_from):
+    """Return ``(n_ch, 3)`` float32 electrode positions aligned to ``X``'s rows.
+
+    Positions are the MNE head-frame coordinates ``info["chs"][i]["loc"][:3]``
+    (meters) -- the same representation braindecode models already consume (see
+    e.g. ``signal_jepa``/``REVE``). When ``targets_from != "metadata"`` the
+    window's misc channels are dropped from ``X`` in ``__getitem__``, so the
+    matching misc rows are dropped here too. Emits a one-time ``UserWarning`` if
+    no channel carries a real position (no montage set).
+    """
+    pos = np.asarray([ch["loc"][:3] for ch in info["chs"]], dtype="float32")
+    if targets_from != "metadata":
+        misc_mask = np.asarray(info.get_channel_types()) == "misc"
+        pos = pos[~misc_mask]
+    has_pos = np.isfinite(pos).all(axis=1) & np.any(pos != 0, axis=1)
+    if pos.size and not has_pos.any():
+        warnings.warn(
+            "Returning channel positions but no montage is set: all positions "
+            "are zero/non-finite. Call ``raw.set_montage(...)`` (or set a "
+            "montage during preprocessing) before enabling ``return_ch_pos``.",
+            UserWarning,
+        )
+    return pos
+
+
 def _window_info(crop_inds, sfreq):
     """Extract window size from crop indices.
 
@@ -448,6 +473,10 @@ class RecordDataset(Dataset[tuple[np.ndarray, int | str, tuple[int, int, int]]])
     ):
         self._description = _create_description(description)
         self.transform = transform
+        # Opt-in: when True, ``__getitem__`` appends channel positions (n_ch, 3).
+        # Default off keeps the public ``(X, y, crop_inds)`` contract intact.
+        self.return_ch_pos = False
+        self._ch_pos = None
 
     @abstractmethod
     def __len__(self) -> int:
@@ -750,6 +779,20 @@ class EEGWindowsDataset(_ZarrMixin, RecordDataset):
     @raw.setter
     def raw(self, value):
         self._raw = value
+        self._ch_pos = None  # invalidate cached positions
+
+    @property
+    def ch_pos(self):
+        """Electrode positions ``(n_ch, 3)`` (x, y, z) aligned to ``X``'s channels.
+
+        Cached MNE head-frame coordinates in meters. Rows match the channels
+        returned by :meth:`__getitem__` (misc target channels excluded when
+        ``targets_from="channels"``). All-zero/NaN when no montage is set.
+        """
+        if self._ch_pos is None:
+            info = self._raw.info if self._raw is not None else self._make_mne_info()
+            self._ch_pos = _compute_ch_pos(info, self.targets_from)
+        return self._ch_pos
 
     def __getitem__(self, index: int):
         """Get a window and its target.
@@ -793,6 +836,8 @@ class EEGWindowsDataset(_ZarrMixin, RecordDataset):
                 y = X[misc_mask, :]
             y = y.copy()
             X = X[~misc_mask, :]
+        if self.return_ch_pos:
+            return X, y, crop_inds, self.ch_pos
         return X, y, crop_inds
 
     def __len__(self):
@@ -1016,6 +1061,24 @@ class WindowsDataset(_ZarrMixin, RecordDataset):
     @windows.setter
     def windows(self, value):
         self._windows = value
+        self._ch_pos = None  # invalidate cached positions
+
+    @property
+    def ch_pos(self):
+        """Electrode positions ``(n_ch, 3)`` (x, y, z) aligned to ``X``'s channels.
+
+        Cached MNE head-frame coordinates in meters. Rows match the channels
+        returned by :meth:`__getitem__` (misc target channels excluded when
+        ``targets_from="channels"``). All-zero/NaN when no montage is set.
+        """
+        if self._ch_pos is None:
+            info = (
+                self._windows.info
+                if self._windows is not None
+                else self._make_mne_info()
+            )
+            self._ch_pos = _compute_ch_pos(info, self.targets_from)
+        return self._ch_pos
 
     @staticmethod
     def _can_use_fast_get_epoch_from_raw(epochs: mne.BaseEpochs) -> bool:
@@ -1068,6 +1131,8 @@ class WindowsDataset(_ZarrMixin, RecordDataset):
         # necessary to cast as list to get list of three tensors from batch,
         # otherwise get single 2d-tensor...
         crop_inds = self.crop_inds[index].tolist()
+        if self.return_ch_pos:
+            return X, y, crop_inds, self.ch_pos
         return X, y, crop_inds
 
     def __len__(self) -> int:
@@ -1458,6 +1523,50 @@ class BaseConcatDataset(ConcatDataset, HubDatasetMixin, Generic[T]):
                 )
         return self
 
+    def set_return_ch_pos(self, value: bool = True) -> "BaseConcatDataset":
+        """Toggle returning electrode positions from every subdataset.
+
+        When enabled, each windowed subdataset's ``__getitem__`` appends a
+        ``(n_ch, 3)`` array of electrode positions (x, y, z), so a batch yields
+        ``(X, y, crop_inds, ch_pos)``. Heterogeneous (variable-channel)
+        collections additionally need
+        :func:`braindecode.datasets.pad_channels_collate` as the DataLoader
+        ``collate_fn``. Default-off: leaves the ``(X, y, crop_inds)`` contract
+        untouched.
+
+        Parameters
+        ----------
+        value : bool
+            Whether to return channel positions (default ``True``).
+
+        Returns
+        -------
+        self : BaseConcatDataset
+        """
+        for ds in self.datasets:
+            ds.return_ch_pos = value
+        return self
+
+    @property
+    def ch_pos(self):
+        """Electrode positions ``(n_ch, 3)`` from the first recording.
+
+        Convenience accessor mirroring how the repr reports channel info "from
+        first recording". For per-sample positions in a heterogeneous
+        collection, enable :meth:`set_return_ch_pos` and read them from the
+        batch instead.
+        """
+        if not self.datasets:
+            raise ValueError("Empty BaseConcatDataset has no channel positions.")
+        first = self.datasets[0]
+        if not hasattr(first, "ch_pos"):
+            raise AttributeError(
+                f"{type(first).__name__} does not expose channel positions; "
+                "ch_pos is available on windowed datasets (EEGWindowsDataset / "
+                "WindowsDataset)."
+            )
+        return first.ch_pos
+
     def _outdated_save(self, path, overwrite=False):
         """This is a copy of the old saving function, that had inconsistent.
 

braindecode-1.6.1.dev1079/braindecode/datasets/collate.py

@@ -0,0 +1,105 @@
+"""Collate functions for braindecode datasets.
+
+:func:`pad_channels_collate` makes a :class:`~braindecode.datasets.BaseConcatDataset`
+that mixes recordings with *different channel sets* (heterogeneous montages)
+batchable: it pads every sample's channels to the batch maximum and emits a
+boolean channel mask. Pair it with ``set_return_ch_pos(True)`` so each sample
+carries its electrode positions -- channel identity is then encoded by the
+positions, not by row index, which is exactly what position-aware models
+(e.g. ``REVE``) consume.
+"""
+
+# Authors: The braindecode developers
+#
+# License: BSD (3-clause)
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+from torch.utils.data._utils.collate import default_collate
+
+__all__ = ["pad_channels_collate"]
+
+
+def pad_channels_collate(batch, pad_value: float = 0.0, pos_pad_value: float = 0.0):
+    """Collate variable-channel windows by padding channels to the batch maximum.
+
+    Each item of ``batch`` is the tuple returned by a windowed dataset's
+    ``__getitem__``: either ``(X, y, crop_inds)`` or, when
+    ``return_ch_pos=True`` (see
+    :meth:`braindecode.datasets.BaseConcatDataset.set_return_ch_pos`),
+    ``(X, y, crop_inds, ch_pos)``. Signals with different numbers of channels
+    are zero-padded along the channel axis to ``max_ch`` (the largest channel
+    count in the batch) so they can be stacked into a single tensor, and a
+    boolean ``ch_mask`` marks the real (non-padded) channels.
+
+    The time dimension is assumed uniform (fixed-length windows); a
+    :class:`ValueError` is raised otherwise.
+
+    This makes a heterogeneous collection **batchable**. *Consuming* the result
+    (so a model ignores padded channels) requires a position/channel-aware
+    model; fixed-geometry models (e.g. ``ShallowFBCSPNet``) structurally need a
+    constant channel count and cannot use variable batches.
+
+    Parameters
+    ----------
+    batch : list of tuple
+        Samples from a windowed dataset (see above).
+    pad_value : float
+        Value used to pad the signal ``X`` channel rows (default ``0.0``).
+    pos_pad_value : float
+        Value used to pad the position ``ch_pos`` rows (default ``0.0``).
+
+    Returns
+    -------
+    X : torch.Tensor
+        Padded signals, shape ``(batch, max_ch, n_times)``.
+    y : torch.Tensor
+        Targets, collated with ``default_collate``.
+    crop_inds : list
+        Crop indices, collated with ``default_collate``.
+    ch_pos : torch.Tensor
+        Only when the samples include positions. Padded positions, shape
+        ``(batch, max_ch, 3)``.
+    ch_mask : torch.Tensor
+        Boolean mask, shape ``(batch, max_ch)``; ``True`` for real channels,
+        ``False`` for padding.
+    """
+    if len(batch) == 0:
+        raise ValueError("pad_channels_collate received an empty batch.")
+
+    has_pos = len(batch[0]) >= 4
+    Xs = [np.asarray(item[0]) for item in batch]
+
+    n_times = Xs[0].shape[1]
+    if any(x.shape[1] != n_times for x in Xs):
+        raise ValueError(
+            "pad_channels_collate only pads the channel axis; all windows must "
+            f"share the same number of time samples, got {sorted({x.shape[1] for x in Xs})}."
+        )
+
+    bsz = len(batch)
+    max_ch = max(x.shape[0] for x in Xs)
+
+    X_out = torch.full((bsz, max_ch, n_times), pad_value, dtype=torch.float32)
+    ch_mask = torch.zeros((bsz, max_ch), dtype=torch.bool)
+    pos_out = (
+        torch.full((bsz, max_ch, 3), pos_pad_value, dtype=torch.float32)
+        if has_pos
+        else None
+    )
+    for i, (item, x) in enumerate(zip(batch, Xs)):
+        c = x.shape[0]
+        X_out[i, :c] = torch.as_tensor(x, dtype=torch.float32)
+        ch_mask[i, :c] = True
+        if pos_out is not None:
+            pos_out[i, :c] = torch.as_tensor(np.asarray(item[3]), dtype=torch.float32)
+
+    # y and crop_inds carry no channel axis -> standard collation.
+    y = default_collate([item[1] for item in batch])
+    crop_inds = default_collate([item[2] for item in batch])
+
+    if pos_out is not None:
+        return X_out, y, crop_inds, pos_out, ch_mask
+    return X_out, y, crop_inds, ch_mask

{braindecode-1.6.1.dev1073 → braindecode-1.6.1.dev1079}/braindecode/eegneuralnet.py

@@ -134,7 +134,14 @@ class _EEGNeuralNet(NeuralNet, abc.ABC):
         i_window_in_trials = []
         i_window_stops = []
         window_ys = []
-        for X, y, i in self.get_iterator(dataset, drop_index=False):
+        for batch in self.get_iterator(dataset, drop_index=False):
+            if len(batch) > 3:
+                raise ValueError(
+                    "Cropped prediction does not support channel positions; "
+                    "disable return_ch_pos (set_return_ch_pos(False)) for cropped "
+                    "evaluation."
+                )
+            X, y, i = batch[0], batch[1], batch[2]
             i_window_in_trials.append(i[0].cpu().numpy())
             i_window_stops.append(i[2].cpu().numpy())
             with torch.no_grad():

{braindecode-1.6.1.dev1073 → braindecode-1.6.1.dev1079}/braindecode/functional/__init__.py

@@ -5,6 +5,7 @@ from .functions import (
     identity,
     plv_time,
     safe_log,
+    sinusoidal_positional_encoding,
     square,
 )
 from .initialization import glorot_weight_zero_bias, rescale_parameter
@@ -16,6 +17,7 @@ __all__ = [
     "identity",
     "plv_time",
     "safe_log",
+    "sinusoidal_positional_encoding",
     "square",
     "glorot_weight_zero_bias",
     "rescale_parameter",

{braindecode-1.6.1.dev1073 → braindecode-1.6.1.dev1079}/braindecode/functional/functions.py

@@ -2,6 +2,8 @@
 #
 # License: BSD (3-clause)
 
+import math
+
 import torch
 import torch.nn.functional as F
 
@@ -249,3 +251,40 @@ def plv_time(x, forward_fourier=True, epsilon: float = 1e-6):
     )
 
     return plv_matrix
+
+
+def sinusoidal_positional_encoding(n_positions: int, dim: int) -> torch.Tensor:
+    r"""Fixed sine/cosine positional-encoding table of shape ``(n_positions, dim)``.
+
+    The standard encoding of Vaswani et al. (2017): for position :math:`p` and
+    channel :math:`i`, :math:`pe[p, 2i] = \sin(p / 10000^{2i/d})` and
+    :math:`pe[p, 2i+1] = \cos(p / 10000^{2i/d})`. PyTorch ships no sinusoidal
+    encoding (``nn.Embedding`` is a *learned* lookup), so braindecode models that
+    need one share this primitive instead of re-deriving it. An odd ``dim`` is
+    computed on the next even width and truncated, reproducing the per-model
+    wrappers (e.g. :class:`~braindecode.models.medformer` odd-``d_model``).
+
+    Parameters
+    ----------
+    n_positions : int
+        Number of positions (sequence length) to encode.
+    dim : int
+        Embedding dimension of each position.
+
+    Returns
+    -------
+    torch.Tensor
+        ``(n_positions, dim)`` float table; callers add their own batch axis,
+        dropout, or offset.
+    """
+    dim_even = dim + (dim % 2)
+    position = torch.arange(n_positions).unsqueeze(1).float()
+    div_term = torch.exp(
+        torch.arange(0, dim_even, 2).float() * (-math.log(10000.0) / dim_even)
+    )
+    pe = torch.zeros(n_positions, dim_even)
+    pe[:, 0::2] = torch.sin(position * div_term)
+    pe[:, 1::2] = torch.cos(position * div_term)
+    # ``.contiguous()`` so an odd-``dim`` truncation owns tight storage -- a
+    # non-contiguous view over the padded width breaks safetensors buffer saving.
+    return pe[:, :dim].contiguous()

{braindecode-1.6.1.dev1073 → braindecode-1.6.1.dev1079}/braindecode/models/__init__.py

@@ -66,6 +66,7 @@ from .util import (
     build_model_config,
     extract_channel_locations_from_chs_info,
     models_mandatory_parameters,
+    positions_from_chs_info,
 )
 
 # Call this last in order to make sure the dataset list is populated with
@@ -115,6 +116,7 @@ __all__ = [
     "Labram",
     "LUNA",
     "extract_channel_locations_from_chs_info",
+    "positions_from_chs_info",
     "MEDFormer",
     "MSVTNet",
     "PBT",

{braindecode-1.6.1.dev1073 → braindecode-1.6.1.dev1079}/braindecode/models/biot.py

@@ -1,4 +1,3 @@
-import math
 from warnings import warn
 
 import numpy as np
@@ -6,6 +5,7 @@ import torch
 import torch.nn as nn
 from linear_attention_transformer import LinearAttentionTransformer
 
+from braindecode.functional import sinusoidal_positional_encoding
 from braindecode.models.base import EEGModuleMixin
 
 # -----------------------------------------------------------------------------
@@ -388,15 +388,8 @@ class _PositionalEncoding(nn.Module):
     def __init__(self, emb_size: int, drop_prob: float = 0.1, max_len: int = 1000):
         super(_PositionalEncoding, self).__init__()
 
-        # Compute the positional encodings once in log space.
-        pe = torch.zeros(max_len, emb_size)
-        position = torch.arange(0, max_len).unsqueeze(1).float()
-        div_term = torch.exp(
-            torch.arange(0, emb_size, 2).float() * -(math.log(10000.0) / emb_size)
-        )
-        pe[:, 0::2] = torch.sin(position * div_term)
-        pe[:, 1::2] = torch.cos(position * div_term)
-        pe = pe.unsqueeze(0)
+        # Compute the positional encodings once (shared sinusoidal primitive).
+        pe = sinusoidal_positional_encoding(max_len, emb_size).unsqueeze(0)
         self.register_buffer("pe", pe)
         self.dropout = nn.Dropout(p=drop_prob)