# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
"""Neuraltrain custom configuration for REVE.
Includes the following adaptations over the raw braindecode REVE model:
* **Channel name remapping** -- an explicit ``channel_mapping`` dict maps
dataset channel names to names recognised by REVE's position bank.
Channels whose names already appear in the bank need no entry.
* **Per-batch channel positions** -- the wrapper takes the dataset's
``channel_positions`` as REVE's ``pos`` for any montage the build could not
resolve from the position bank, so one instance serves montages it was not
built for.
* **Pretrained loading** -- REVE's ``__init__`` requires ``n_times`` and
``n_outputs`` to build its ``final_layer``, but
:class:`BaseBrainDecodeModel.build` blocks ``n_times`` for pretrained
models. ``NtReve.build()`` calls ``from_pretrained`` directly.
* **Encoder-only support** -- when ``n_outputs`` is not provided (i.e. when
a ``DownstreamWrapperModel`` handles the classification head), the wrapper
calls ``forward(return_output=True)`` and extracts the final transformer
layer's output, bypassing REVE's ``final_layer`` entirely.
"""
import logging
import re
import typing as tp
import torch
import torch.nn as nn
from .base import BaseBrainDecodeModel, RequiredBuildField
from .common import INVALID_POS_VALUE, resolve_with_anode_fallback
logger = logging.getLogger(__name__)
class _ReveWrapper(nn.Module):
"""Thin wrapper around REVE for channel handling and encoder-only output.
Combines three adaptations in a single module:
* **Channel subsetting** -- when *channel_indices* is not ``None``, the EEG
tensor and the channel positions are sliced along the channel dimension
before forwarding to REVE. No-op when ``None``.
* **Channel positions** -- when the build named every channel,
*bank_positions* wins: those are the coordinates REVE was pretrained
with, and using them keeps a result independent of how the montage was
digitised. Otherwise the batch's ``channel_positions`` become REVE's
``pos``, and channels the montage could not place (they arrive as
``INVALID_POS_VALUE``) are dropped. A bank only lines up with the
montage it was built from, so an instance that must span montages of the
same width should be built without ``chs_info``.
* **Encoder-only output** -- when *encoder_only* is ``True``, the
forward call passes ``return_output=True`` to REVE and returns the
final transformer layer output (index ``-1``), bypassing REVE's
``final_layer``. When ``False``, forwards normally.
"""
channel_indices: torch.Tensor | None
bank_positions: torch.Tensor | None
def __init__(
self,
model: nn.Module,
channel_indices: list[int] | None = None,
encoder_only: bool = False,
bank_positions: torch.Tensor | None = None,
):
super().__init__()
self.model = model
self.encoder_only = encoder_only
self.register_buffer("bank_positions", bank_positions)
if channel_indices is not None:
self.register_buffer(
"channel_indices",
torch.tensor(channel_indices, dtype=torch.long),
)
else:
self.register_buffer("channel_indices", None)
def forward(
self,
eeg: torch.Tensor,
channel_positions: torch.Tensor | None = None,
**kwargs: tp.Any,
) -> torch.Tensor:
if self.channel_indices is not None:
eeg = eeg[:, self.channel_indices]
if channel_positions is not None:
channel_positions = channel_positions[:, self.channel_indices]
bank = self.bank_positions
if bank is not None and eeg.shape[1] == bank.shape[0]:
channel_positions = bank.to(eeg).expand(eeg.shape[0], -1, -1)
elif channel_positions is not None:
unplaced = (channel_positions == INVALID_POS_VALUE).all(dim=-1)
keep = ~unplaced.any(dim=0)
eeg, channel_positions = eeg[:, keep], channel_positions[:, keep]
if self.encoder_only:
return self.model(eeg, pos=channel_positions, return_output=True, **kwargs)[
-1
]
return self.model(eeg, pos=channel_positions, **kwargs)
[docs]
class NtReve(BaseBrainDecodeModel):
"""Config for the braindecode REVE model with channel-mapping support.
Extends :class:`BaseBrainDecodeModel` with REVE-specific logic:
1. **Channel remapping** -- an explicit ``channel_mapping`` dict maps
dataset channel names to REVE position-bank names. Channels whose
names already appear in the bank (case-insensitively) need no entry.
2. **Pretrained loading** -- bypasses the base-class restriction on
``n_times`` for pretrained models, since REVE needs it to size its
``final_layer``.
3. **Encoder-only output** -- when ``n_outputs`` is ``None`` (downstream
wrapper handles the head), the model is wrapped to call
``forward(return_output=True)`` and return the final transformer
layer output, bypassing REVE's ``final_layer``.
Parameters
----------
channel_mapping : dict or None
Explicit mapping from dataset channel names to REVE position-bank
names. Useful for EEG systems whose naming convention is absent
from the bank (e.g. Neuromag ``"EEG 005"`` or easycap-M10 numeric
``"2"``).
"""
_MODEL_CLASS_PATH: tp.ClassVar[str] = "braindecode.models.REVE"
required_fields: tp.ClassVar[list[RequiredBuildField]] = ["ch_names", "n_times"]
channel_mapping: dict[str, str] | None = None
def _remap_chs_info(
self,
chs_info: list[dict[str, tp.Any]],
) -> list[dict[str, tp.Any]]:
"""Apply ``channel_mapping`` to *chs_info*, returning a new list.
MEG sensor names are also tried in either spelling of the mapping:
Neuromag / KIT whatever the separator (``MEG0113``, ``MEG-0113`` ->
``MEG 0113`` or ``MEG0113``) and CTF without the serial suffix
(``MLC11-4504`` -> ``MLC11``). Other names must match exactly.
"""
if not self.channel_mapping:
return chs_info
mapping = self.channel_mapping
def remap(name: str) -> str:
candidates = [name]
if neuromag := re.fullmatch(r"MEG[ -]?(\d+)", name):
candidates += [f"MEG {neuromag.group(1)}", f"MEG{neuromag.group(1)}"]
elif ctf := re.fullmatch(r"(M[LRZ][FCTPO]\d{2})-\d+", name):
candidates.append(ctf.group(1))
for candidate in candidates:
if candidate in mapping:
return mapping[candidate]
return name
return [{**ch, "ch_name": remap(ch["ch_name"])} for ch in chs_info]
@staticmethod
def _bank_lookup(bank: tp.Any) -> tp.Callable[[str], torch.Tensor | None]:
"""Build a case-insensitive position lookup into *bank*.
REVE's bank mixes casings (``Fp1``, ``AFz``) and duplicates only some of
them in upper case, so an exact-match lookup drops channels that arrive
in all caps. Case variants share coordinates, so folding case cannot
select a wrong position. Returns ``None`` for an unknown name.
"""
by_upper = {name.upper(): name for name in bank.mapping}
def lookup(name: str) -> torch.Tensor | None:
canonical = by_upper.get(name.upper())
if canonical is None:
return None
return bank.embedding[bank.mapping[canonical]]
return lookup
[docs]
def build(
self,
n_spatial_locations: int,
n_temporal_samples: int,
n_outputs: int | None = None,
chs_info: list[dict[str, tp.Any]] | None = None,
frequency: float | None = None,
) -> nn.Module:
if chs_info is not None:
chs_info = self._remap_chs_info(chs_info)
# Spatial size the model receives (already the adapter target when a
# channel adapter is present). Reduced below if some channels cannot
# be resolved.
n_chans = n_spatial_locations
channel_indices: list[int] | None = None
bank_positions: torch.Tensor | None = None
n_derived = 0
exact_bank_match = True # REVE's own chs_info lookup is case-sensitive
if chs_info is not None:
from braindecode.models.reve import RevePositionBank
bank = RevePositionBank()
n_original = len(chs_info)
lookup = self._bank_lookup(bank)
valid_indices: list[int] = []
valid_chs: list[dict[str, tp.Any]] = []
positions: list[torch.Tensor] = []
dropped: list[str] = []
for i, ch in enumerate(chs_info):
name = ch["ch_name"]
position, used_anode = resolve_with_anode_fallback(name, lookup)
if position is None:
dropped.append(name)
continue
valid_indices.append(i)
valid_chs.append(ch)
positions.append(position)
if used_anode:
n_derived += 1
if dropped:
logger.warning(
"Dropping %d channel(s) not resolvable from REVE position bank: %s",
len(dropped),
dropped,
)
if n_derived:
logger.info(
"Mapped %d bipolar channel(s) to REVE via anode fallback.",
n_derived,
)
if len(valid_chs) < n_original:
channel_indices = valid_indices
chs_info = valid_chs
n_chans = len(chs_info)
if not valid_chs:
raise ValueError(
"No dataset channels match the REVE position bank "
"(directly or via anode fallback). "
"Consider adding a `channel_mapping`."
)
logger.info(
"[REVE_CHANNELS] n_dataset=%d n_resolved=%d (n_derived=%d)",
n_original,
len(valid_chs),
n_derived,
)
bank_positions = torch.stack(positions)
exact_bank_match = all(ch["ch_name"] in bank.mapping for ch in valid_chs)
build_kwargs: dict[str, tp.Any] = {
"n_chans": n_chans,
"n_times": n_temporal_samples,
}
if chs_info is not None and exact_bank_match:
build_kwargs["chs_info"] = chs_info
encoder_only = n_outputs is None
if self.from_pretrained_name is not None:
build_kwargs["n_outputs"] = n_outputs if n_outputs is not None else 2
elif n_outputs is not None:
build_kwargs["n_outputs"] = n_outputs
model = self._construct(**build_kwargs)
if bank_positions is not None and not exact_bank_match:
model.default_pos = bank_positions
return _ReveWrapper(model, channel_indices, encoder_only, bank_positions)