S-CEReBrO Integration

CONTRIBUTING.md

@@ -10,15 +10,17 @@ BioFoundation is both a research repository and an onboarding codebase. Contribu
 - Preserve the full Apache 2.0 header in every Python source and Hydra YAML file.
 - Normalize task input through `biofoundation.core.batch.as_signal_batch`.
 - Keep model-family metadata in `biofoundation/model_registry.py`, including its paper and Hugging Face repository.
+- Change `biofoundation/` additively. New fields, functions, and classes are fine; existing names must not change signature or behaviour, and new defaults must leave existing callers observing no difference. A change that cannot be made additively needs a deprecation period and a major version bump of `biofoundation.__version__`.
+- Record decisions that are expensive to reverse as an ADR under [`docs/adr`](docs/adr/).
 - Use Hydra configs for reproducible settings instead of embedding experiment paths or hyperparameters in Python.
 - Mark local values that a new user must supply with `#CHANGEME` and document the expected value.
 
 ## Adding a Model
 
-1. Add the `nn.Module` implementation under `models/` and its Hydra model config under `config/model/`.
+1. Add the `nn.Module` implementation under `models/` and its Hydra model config under `config/model/`. A family may either bundle its output layer into the model, as the five original families do, or separate an encoder from a prediction head under `models/model_heads/` with a matching `config/model_head/` entry. The second shape is described by the protocols in `biofoundation/core/protocols.py` and lets one pre-trained encoder serve every downstream task.
 2. Add pre-training and fine-tuning experiments under `config/experiment/`.
 3. Use an existing task when its behavior fits. New task steps must use the shared batch adapter.
-4. Register the model in `biofoundation/model_registry.py`, including modalities, architecture, experiment names, batch requirements, paper, and Hugging Face URL.
+4. Register the model in `biofoundation/model_registry.py`, including modalities, architecture, experiment names, batch requirements, venue, paper, and Hugging Face URL. Use the case-folded display name as the key, and list any prediction heads in `head_targets`. Set `model_family` in the experiment so tasks can enforce the declared batch requirements at runtime.
 5. Add a model page under `docs/model/` with input assumptions, training details, and checkpoint usage.
 6. Extend the contract tests for any new shared behavior.
 
@@ -41,6 +43,14 @@ python -m unittest discover -s tests -p 'test_*.py' -v
 python -m compileall -q biofoundation run_train.py models tasks datasets data_module
 ```
 
+Tests that import PyTorch live under `tests/model_tests/` and run with `pytest`. That
+directory is not a package, so `unittest discover` does not traverse it and the fast
+suite stays runnable without the training dependencies installed.
+
+```bash
+pytest tests/model_tests -v
+```
+
 ## Official Checkpoint Improvements
 
 The published weights are licensed under CC BY-ND 4.0. Modified weights, including adapters, deltas, pruned variants, and quantized variants, may not be redistributed.

README.md

@@ -6,7 +6,7 @@ Authors: Thorir Mar Ingolfsson, Anna Tegon, Berkay Döner, Xiaying Wang, Matteo
 
 > **TL;DR:** Choose a model from the table below, install the training dependencies, set `DATA_PATH` and `CHECKPOINT_DIR`, then run `python -u run_train.py +experiment=<MODEL>_pretrain` or the matching fine-tuning experiment. Each model page links its Hugging Face weights and exact checkpoint command. ARES is separate and only needed for embedded deployment.
 
-BioFoundation is a research and onboarding codebase for foundation models across EEG, sEMG, ECG, and PPG. It collects the model implementations, Hydra experiments, preprocessing tools, and pretrained releases behind five model families.
+BioFoundation is a research and onboarding codebase for foundation models across EEG, sEMG, ECG, and PPG. It collects the model implementations, Hydra experiments, preprocessing tools, and pretrained releases behind six model families.
 
 The training stack is built on PyTorch Lightning and Hydra. Embedded deployment through ARES is maintained as a separate toolchain inside the repository.
 
@@ -19,6 +19,7 @@ The training stack is built on PyTorch Lightning and Hydra. Embedded deployment
 | [TinyMyo](docs/model/TinyMyo.md) | sEMG | Rotary Transformer | [Paper](https://arxiv.org/abs/2512.15729) / [Hugging Face](https://huggingface.co/PulpBio/TinyMyo) |
 | [LuMamba](docs/model/LuMamba.md) | EEG | Query-unified Mamba | [Paper](https://arxiv.org/abs/2603.19100) / [Hugging Face](https://huggingface.co/PulpBio/LuMamba) |
 | [PanLUNA](docs/model/PanLUNA.md) | EEG, ECG, PPG | Multimodal query-unified Transformer | [Paper](https://arxiv.org/abs/2604.04297) / [Hugging Face](https://huggingface.co/PulpBio/PanLUNA) |
+| [S-CEReBrO](docs/model/SCEReBrO.md) | EEG | Windowed alternating-attention Transformer | [Paper](https://arxiv.org/abs/2607.27913) / [Hugging Face](https://huggingface.co/PulpBio/S-CEReBrO) |
 
 The machine-readable [`model_registry.py`](biofoundation/model_registry.py) records the experiment names, papers, Hugging Face repositories, modalities, and batch metadata requirements for these families.
 
@@ -74,6 +75,7 @@ Choose another `+experiment` from the model registry. Before a long run, review
 | [`datasets`](datasets/) | Dataset readers and sample contracts. |
 | [`data_module`](data_module/) | Lightning data modules and loader composition. |
 | [`config`](config/) | Hydra defaults, modules, and reproducible experiments. |
+| [`docs/adr`](docs/adr/) | Architecture decision records for shared contracts. |
 | [`make_datasets`](make_datasets/) | Raw-data preprocessing and HDF5 conversion. |
 | [`criterion`](criterion/) | Training objectives. |
 | [`tests`](tests/) | Fast repository and refactoring contracts. |

biofoundation/__init__.py

@@ -17,5 +17,21 @@
 #* Author:  BioFoundation Contributors                                       *
 #*----------------------------------------------------------------------------*
 
-"""Shared infrastructure for the BioFoundation model zoo."""
+"""Shared infrastructure for the BioFoundation model zoo.
+
+This package holds the contracts that model families, tasks, and datasets agree on:
+the batch layout, the encoder and prediction-head protocols, checkpoint entry points,
+environment validation, and the model registry. Everything here is imported by code
+that must run without PyTorch installed, so this package has no heavyweight runtime
+dependencies.
+
+Changes follow an additive rule. New fields, functions, and classes may be added;
+existing names do not change signature or behaviour, and defaults are chosen so that
+an existing caller observes no difference. A change that cannot be made additively is
+a major version bump and needs a deprecation period first.
+"""
+
+__version__ = "0.2.0"
+
+__all__ = ["__version__"]
 

biofoundation/core/__init__.py

@@ -20,11 +20,17 @@
 """Stable contracts shared by model-specific implementations."""
 
 from biofoundation.core.batch import BatchRequirements, SignalBatch, as_signal_batch, require_batch_fields
+from biofoundation.core.checkpoints import SafetensorsCheckpointMixin, split_state_dict_by_prefix
+from biofoundation.core.protocols import PredictionHead, SignalEncoder
 
 __all__ = [
     "BatchRequirements",
+    "PredictionHead",
+    "SafetensorsCheckpointMixin",
     "SignalBatch",
+    "SignalEncoder",
     "as_signal_batch",
     "require_batch_fields",
+    "split_state_dict_by_prefix",
 ]
 

biofoundation/core/batch.py

@@ -29,24 +29,54 @@ class SignalBatch(TypedDict, total=False):
 
     Only ``input`` is universally required. Model adapters can require channel
     or sensor metadata without forcing simpler models to manufacture it.
+
+    Two independent electrode-geometry representations are supported, and a model
+    requires exactly one of them. They are peers rather than alternatives to convert
+    between, so a dataset declares which one it produces and the registry records
+    which one each model consumes:
+
+    ``channel_locations``
+        Shape ``(batch, channels, 3)``. One 3D coordinate per channel. For a bipolar
+        derivation this is the midpoint of the two electrodes. Consumed by LUNA,
+        LuMamba and PanLUNA.
+
+    ``channel_coords``
+        Shape ``(batch, channels, 2, 3)``. Both electrodes of every channel kept
+        separate, so a bipolar pair and a scalp-plus-reference channel stay
+        distinguishable. Consumed by S-CEReBrO.
+
+    Padding metadata describes how much of a sample is filler, which lets montages of
+    different sizes share one batch. It is only meaningful on mapping-shaped batches;
+    the tuple form accepted by :func:`as_signal_batch` cannot carry it.
     """
 
     input: Any
     label: Any
     channel_names: Any
     channel_locations: Any
+    channel_coords: Any
     sensor_type: Any
+    num_padded_channels: Any
+    num_padded_timesteps: Any
     metadata: Mapping[str, Any]
 
 
 @dataclass(frozen=True)
 class BatchRequirements:
-    """Metadata fields required by a particular model adapter."""
+    """Metadata fields required by a particular model adapter.
+
+    Every field defaults to ``False``, so a model only declares what it actually
+    reads. ``channel_locations`` and ``channel_coords`` are the two electrode-geometry
+    representations described on :class:`SignalBatch`; a model sets one of them.
+    """
 
     label: bool = False
     channel_names: bool = False
     channel_locations: bool = False
+    channel_coords: bool = False
     sensor_type: bool = False
+    num_padded_channels: bool = False
+    num_padded_timesteps: bool = False
 
 
 def as_signal_batch(batch: Any) -> SignalBatch:
@@ -80,7 +110,10 @@ def require_batch_fields(batch: SignalBatch, requirements: BatchRequirements) ->
         ("label", requirements.label),
         ("channel_names", requirements.channel_names),
         ("channel_locations", requirements.channel_locations),
+        ("channel_coords", requirements.channel_coords),
         ("sensor_type", requirements.sensor_type),
+        ("num_padded_channels", requirements.num_padded_channels),
+        ("num_padded_timesteps", requirements.num_padded_timesteps),
     )
     missing = [name for name, required in required_fields if required and name not in batch]
     if missing:

biofoundation/core/checkpoints.py

@@ -0,0 +1,113 @@
+#*----------------------------------------------------------------------------*
+#* Copyright (C) 2026 ETH Zurich, Switzerland                                 *
+#* SPDX-License-Identifier: Apache-2.0                                        *
+#*                                                                            *
+#* Licensed under the Apache License, Version 2.0 (the "License");            *
+#* you may not use this file except in compliance with the License.           *
+#* You may obtain a copy of the License at                                    *
+#*                                                                            *
+#* http://www.apache.org/licenses/LICENSE-2.0                                 *
+#*                                                                            *
+#* Unless required by applicable law or agreed to in writing, software        *
+#* distributed under the License is distributed on an "AS IS" BASIS,          *
+#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.   *
+#* See the License for the specific language governing permissions and        *
+#* limitations under the License.                                             *
+#*                                                                            *
+#* Author:  BioFoundation Contributors                                       *
+#*----------------------------------------------------------------------------*
+
+"""Checkpoint entry points shared by tasks that separate encoder from head.
+
+``run_train.py`` loads pre-trained weights by calling ``load_pretrained_checkpoint``
+for a Lightning ``.ckpt`` and ``load_safetensors_checkpoint`` for a ``.safetensors``
+file. Tasks that implement Lightning's own ``load_from_checkpoint`` instead can mix
+this class in to expose both names without duplicating either loader.
+
+The mixin is deliberately thin: it maps names and formats, and delegates the actual
+tensor matching to the task. Loading policy, including which shape mismatches are
+tolerated, stays with the task that owns the model.
+"""
+
+from typing import Any, Dict, Optional, Protocol
+
+
+class _CheckpointLoadable(Protocol):  # pragma: no cover - typing only
+    """The single method a task must provide for the mixin to delegate to."""
+
+    def load_from_checkpoint(self, checkpoint_path: str, **kwargs: Any) -> Any: ...
+
+
+class SafetensorsCheckpointMixin:
+    """Expose BioFoundation's two checkpoint entry points over ``load_from_checkpoint``.
+
+    Mix in ahead of :class:`pytorch_lightning.LightningModule` on tasks whose loading
+    logic already lives in ``load_from_checkpoint``:
+
+    .. code-block:: python
+
+        class MyTask(SafetensorsCheckpointMixin, pl.LightningModule):
+            def load_from_checkpoint(self, checkpoint_path, **kwargs): ...
+
+    ``load_safetensors_checkpoint`` converts the flat ``.safetensors`` mapping into the
+    ``{"state_dict": ...}`` layout ``load_from_checkpoint`` expects, writes it to a
+    temporary file, and delegates. Keys are given a ``model.`` prefix when they lack
+    one, matching how the pre-training task saves an encoder.
+    """
+
+    def load_pretrained_checkpoint(self, model_ckpt: str, **kwargs: Any) -> Any:
+        """Load a Lightning ``.ckpt`` by delegating to ``load_from_checkpoint``."""
+
+        return self.load_from_checkpoint(checkpoint_path=model_ckpt, **kwargs)
+
+    def load_safetensors_checkpoint(self, model_ckpt: str, **kwargs: Any) -> Any:
+        """Load a ``.safetensors`` file through the same path as a Lightning checkpoint."""
+
+        import tempfile
+        from pathlib import Path
+
+        import torch
+        from safetensors.torch import load_file
+
+        state_dict = {
+            key if key.startswith(("model.", "model_head.")) else f"model.{key}": value
+            for key, value in load_file(model_ckpt).items()
+        }
+
+        with tempfile.TemporaryDirectory() as directory:
+            converted = Path(directory) / "converted.ckpt"
+            torch.save({"state_dict": state_dict}, converted)
+            return self.load_from_checkpoint(checkpoint_path=str(converted), **kwargs)
+
+
+def split_state_dict_by_prefix(
+    state_dict: Dict[str, Any],
+    prefixes: tuple[str, ...] = ("model_head.", "model."),
+) -> Dict[str, Dict[str, Any]]:
+    """Group a flat Lightning ``state_dict`` by top-level module prefix.
+
+    Args:
+        state_dict: Mapping of parameter name to tensor, as stored by Lightning.
+        prefixes: Prefixes to split on, longest first so that ``model_head.`` is
+            matched before ``model.``.
+
+    Returns:
+        Mapping of prefix-without-the-dot to a prefix-stripped state dict. Keys that
+        match no prefix are collected under ``""``.
+    """
+
+    grouped: Dict[str, Dict[str, Any]] = {prefix.rstrip("."): {} for prefix in prefixes}
+    grouped[""] = {}
+
+    for key, value in state_dict.items():
+        for prefix in prefixes:
+            if key.startswith(prefix):
+                grouped[prefix.rstrip(".")][key[len(prefix):]] = value
+                break
+        else:
+            grouped[""][key] = value
+
+    return grouped
+
+
+__all__ = ["SafetensorsCheckpointMixin", "split_state_dict_by_prefix"]

biofoundation/core/protocols.py

@@ -0,0 +1,81 @@
+#*----------------------------------------------------------------------------*
+#* Copyright (C) 2026 ETH Zurich, Switzerland                                 *
+#* SPDX-License-Identifier: Apache-2.0                                        *
+#*                                                                            *
+#* Licensed under the Apache License, Version 2.0 (the "License");            *
+#* you may not use this file except in compliance with the License.           *
+#* You may obtain a copy of the License at                                    *
+#*                                                                            *
+#* http://www.apache.org/licenses/LICENSE-2.0                                 *
+#*                                                                            *
+#* Unless required by applicable law or agreed to in writing, software        *
+#* distributed under the License is distributed on an "AS IS" BASIS,          *
+#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.   *
+#* See the License for the specific language governing permissions and        *
+#* limitations under the License.                                             *
+#*                                                                            *
+#* Author:  BioFoundation Contributors                                       *
+#*----------------------------------------------------------------------------*
+
+"""Structural contracts for the encoder and prediction-head split.
+
+BioFoundation carries two model shapes. The five original families bundle their
+output layer into the model itself, selecting it from ``num_classes`` at
+construction time. Newer families separate the two: an encoder that produces token
+embeddings, and a prediction head that consumes them, each instantiated from its own
+Hydra group.
+
+The protocols below describe the second shape. They are :class:`typing.Protocol`
+definitions, so conformance is structural: a module satisfies one by having a
+matching ``forward``, with no base class to inherit and no registration step. That
+keeps the bundled families valid exactly as they are while giving the split families
+a contract a type checker can verify.
+
+Tensors are annotated only under :data:`typing.TYPE_CHECKING`. This module is
+imported by the fast contract test suite, which runs without PyTorch installed, so
+nothing here may import ``torch`` at runtime.
+"""
+
+from typing import TYPE_CHECKING, Any, Optional, Protocol, runtime_checkable
+
+if TYPE_CHECKING:  # pragma: no cover - typing only
+    from torch import Tensor
+else:  # pragma: no cover - runtime fallback keeps this module torch-free
+    Tensor = Any
+
+
+@runtime_checkable
+class SignalEncoder(Protocol):
+    """A model that turns a patched biosignal into contextualised token embeddings.
+
+    The encoder owns tokenisation, positional and channel embeddings, and the
+    backbone. It owns no task-specific output layer, so one pre-trained encoder can
+    be paired with any :class:`PredictionHead` without being rebuilt.
+
+    Implementations accept ``(batch, channels, patches, patch_size)`` and return
+    ``(batch, channels * patches, embed_dim)``. Keyword arguments beyond the input
+    carry whatever batch metadata the family requires, declared in its
+    :class:`~biofoundation.core.batch.BatchRequirements`.
+    """
+
+    def forward(self, x: "Tensor", *args: Any, **kwargs: Any) -> "Tensor":
+        """Encode a patched biosignal into token embeddings."""
+        ...
+
+
+@runtime_checkable
+class PredictionHead(Protocol):
+    """A module that turns encoder token embeddings into a task prediction.
+
+    Heads accept ``(batch, num_tokens, embed_dim)``. The output shape is the head's
+    own concern: class logits, a scalar per window, or one reconstructed patch per
+    token. A head never reads the raw waveform and never receives batch metadata, so
+    swapping the task means swapping the head alone.
+    """
+
+    def forward(self, x: "Tensor", *args: Any, **kwargs: Any) -> "Tensor":
+        """Map token embeddings to a prediction."""
+        ...
+
+
+__all__ = ["PredictionHead", "SignalEncoder"]