Release v0.2.0: custom channel mappings + tolerant EDF label matching

README.md

@@ -65,6 +65,32 @@ tokens, modality_mask, channel_ids = tokenize(tokenizers, meta, signals)
 emb = embed(model, tokens, modality_mask, channel_ids, meta)   # {name: [T, D]}
 ```
 
+### Custom channel labels
+
+Hypnos maps each EDF signal label onto a canonical channel (`C3`, `C4`, `E1`, `E2`, `Chin`,
+`ECG`, `ABD`, `THX`). It already recognizes many common naming conventions out of the box
+(e.g. `EKG`/`ECG L-ECG R` → `ECG`, `C3-M2` → `C3`), plus AASM contralateral re-referencing
+(mastoid equivalents `A1`/`A2` and `TP9`/`TP10` are accepted as `M1`/`M2`) and chin-EMG
+bipolar derivation. Matching is tolerant of case, whitespace and `:`/`/` separators (so
+`c3:m2` resolves like `C3-M2`). A modality whose channel can't be found is simply skipped.
+
+If your recording uses labels Hypnos doesn't recognize, pass `channel_aliases` — a
+`{canonical_name: [extra EDF labels]}` mapping that's merged with the built-ins (your aliases
+take precedence):
+
+```python
+emb = embed_edf(
+    "recording.edf",
+    channel_aliases={
+        "ECG": ["MyDeviceEKG"],     # canonical "ECG" <- EDF label "MyDeviceEKG"
+        "C3": ["EEG_C3_custom"],
+    },
+)
+```
+
+`channel_aliases` is also accepted by `preprocess_edf(...)` in the step-by-step API. The
+built-in alias tables live in `hypnos.data.edf` (`ALT_COLUMNS`).
+
 ### Pooling
 
 Hypnos produces embeddings at 1 Hz for each modality. In our experiments, we found that simple pooling over modalities and timescales works well for downstream tasks. For example, to produce a single embedding per 30-second sleep epoch:

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "hypnos"
-version = "0.1.0"
+version = "0.2.0"
 description = "Minimal inference library for Hypnos: load an EDF, preprocess, and generate sleep embeddings from pretrained checkpoints."
 readme = "README.md"
 requires-python = ">=3.11"

src/hypnos/data/edf.py

@@ -6,6 +6,7 @@
 """
 
 import logging
+from collections.abc import Mapping, Sequence
 from dataclasses import dataclass, field
 from math import gcd
 
@@ -93,8 +94,8 @@
 # Contralateral mastoid referencing (AASM): canonical channel -> required reference electrode.
 CONTRALATERAL_REF: dict[str, str] = {"C3": "M2", "C4": "M1", "E1": "M2", "E2": "M1"}
 
-# A1/A2 are legacy equivalents of M1/M2.
-REFERENCE_ALTS: dict[str, list[str]] = {"M1": ["A1"], "M2": ["A2"]}
+# Mastoid equivalents of M1/M2: A1/A2 (legacy 10-20) and TP9/TP10 (10-10).
+REFERENCE_ALTS: dict[str, list[str]] = {"M1": ["A1", "TP9"], "M2": ["A2", "TP10"]}
 _ALL_REF_NAMES: set[str] = set(REFERENCE_ALTS) | {a for alts in REFERENCE_ALTS.values() for a in alts}
 
 # Pre-computed bipolar labels (already derived in the EDF), checked before components.
@@ -114,13 +115,47 @@
 }
 
 
-def get_column_match(target_col: str, available_cols: list[str]) -> str | None:
-    """Return the EDF label matching ``target_col`` (exact or via ``ALT_COLUMNS``), else None."""
+def _clean_label(label: str) -> str:
+    """Normalize an EDF label for tolerant matching.
+
+    Uppercases, strips whitespace and a few device-specific suffixes, and unifies the
+    separators (``:`` / ``/`` → ``-``) so labels like ``'c3:m2'``, ``'C3-M2_PDS'`` and
+    ``'C3-M2'`` all compare equal.
+    """
+    cleaned = label.strip().upper()
+    for suffix in ("_PDS", "_EG"):
+        if cleaned.endswith(suffix):
+            cleaned = cleaned[: -len(suffix)]
+    return cleaned.replace(":", "-").replace("/", "-").strip()
+
+
+def get_column_match(
+    target_col: str,
+    available_cols: list[str],
+    channel_aliases: Mapping[str, Sequence[str]] | None = None,
+) -> str | None:
+    """Return the EDF label matching ``target_col``, else None.
+
+    An exact label match wins outright. Otherwise the canonical name and its aliases are
+    compared against the available labels under :func:`_clean_label` normalization (case,
+    whitespace, separator and suffix insensitive). Candidate order: the canonical name,
+    then caller-supplied ``channel_aliases`` (precedence over the built-ins), then the
+    built-in ``ALT_COLUMNS`` table. The original (un-normalized) EDF label is returned.
+    """
     if target_col in available_cols:
         return target_col
-    for alt_col in ALT_COLUMNS.get(target_col, ()):
-        if alt_col in available_cols:
-            return alt_col
+
+    candidates: list[str] = [target_col]
+    if channel_aliases is not None:
+        candidates.extend(channel_aliases.get(target_col, ()))
+    candidates.extend(ALT_COLUMNS.get(target_col, ()))
+
+    cleaned_available = [(col, _clean_label(col)) for col in available_cols]
+    for cand in candidates:
+        cand_clean = _clean_label(cand)
+        for col, col_clean in cleaned_available:
+            if col_clean == cand_clean:
+                return col
     return None
 
 
@@ -138,12 +173,17 @@ class ResolvedChannel:
 
 
 def _find_reference_label(ref_name: str, available_labels: list[str]) -> str | None:
-    """Find a reference channel (M1/M2 or A1/A2 equivalent) in available EDF labels."""
+    """Find a reference channel (M1/M2 or A1/A2/TP9/TP10 equivalent) in available EDF labels.
+
+    Matches under :func:`_clean_label` normalization, returning the original EDF label.
+    """
     if ref_name in available_labels:
         return ref_name
-    for alt in REFERENCE_ALTS.get(ref_name, []):
-        if alt in available_labels:
-            return alt
+    for cand in (ref_name, *REFERENCE_ALTS.get(ref_name, [])):
+        cand_clean = _clean_label(cand)
+        for label in available_labels:
+            if _clean_label(label) == cand_clean:
+                return label
     return None
 
 
@@ -171,10 +211,14 @@ def _resample_reference(ref_signal: np.ndarray, ref_fs: int, target_fs: int) ->
 
 
 def _is_pre_referenced(ch_name: str, actual_label: str) -> bool:
-    """Whether the matched EDF label is already referenced (e.g. 'C3-M2', 'E2-M1')."""
+    """Whether the matched EDF label is already referenced (e.g. 'C3-M2', 'E2-M1').
+
+    Case-insensitive; a label that *is* a bare reference (e.g. 'M1') is not pre-referenced.
+    """
     if ch_name not in CONTRALATERAL_REF:
         return False
-    return any(ref in actual_label for ref in _ALL_REF_NAMES)
+    upper = actual_label.upper()
+    return any(ref in upper and upper != ref for ref in _ALL_REF_NAMES)
 
 
 def _read_signal_metadata(f: pyedflib.EdfReader, idx: int) -> tuple[int, str, float, float]:
@@ -191,6 +235,7 @@ def load_psg_channels(
     f: pyedflib.EdfReader,
     channels: list[str],
     drop_unreferenced: bool = False,
+    channel_aliases: Mapping[str, Sequence[str]] | None = None,
 ) -> dict[str, ResolvedChannel]:
     """Load PSG channels from an open EDF with proper referencing and derivations.
 
@@ -206,6 +251,9 @@ def load_psg_channels(
         channels: Canonical channel names (e.g. ['C3', 'E1', 'Chin', 'ECG']).
         drop_unreferenced: If True, skip channels that need contralateral referencing
             but have no reference electrode available (e.g. bare E1 without M2).
+        channel_aliases: Optional ``{canonical_name: [extra EDF labels]}`` mapping, merged
+            with the built-in ``ALT_COLUMNS`` for recordings with non-standard labels.
+            Caller aliases take precedence over the built-ins.
 
     Returns:
         Dict mapping canonical channel name to ResolvedChannel (channels not found are omitted).
@@ -217,7 +265,9 @@ def load_psg_channels(
     ref_signals = _load_reference_signals(f, label_to_idx)
     result: dict[str, ResolvedChannel] = {}
     for ch_name in channels:
-        resolved = _resolve_one_channel(f, ch_name, available, label_to_idx, ref_signals, drop_unreferenced)
+        resolved = _resolve_one_channel(
+            f, ch_name, available, label_to_idx, ref_signals, drop_unreferenced, channel_aliases
+        )
         if resolved is not None:
             result[ch_name] = resolved
     return result
@@ -230,6 +280,7 @@ def _resolve_one_channel(
     label_to_idx: dict[str, int],
     ref_signals: dict[str, tuple[np.ndarray, int]],
     drop_unreferenced: bool,
+    channel_aliases: Mapping[str, Sequence[str]] | None = None,
 ) -> ResolvedChannel | None:
     """Resolve a single canonical channel from an EDF.
 
@@ -283,7 +334,7 @@ def _resolve_one_channel(
         # Fall through to standard resolution (single electrode fallback)
 
     # --- Step 2: Standard name resolution ---
-    actual_name = get_column_match(ch_name, available)
+    actual_name = get_column_match(ch_name, available, channel_aliases)
     if actual_name is None:
         _logger.info(f"Channel {ch_name} not found in EDF")
         return None

src/hypnos/embedding/__init__.py

@@ -19,6 +19,8 @@
 
 from __future__ import annotations
 
+from collections.abc import Mapping, Sequence
+
 import numpy as np
 import torch
 
@@ -81,18 +83,26 @@ def embed_edf(
     causal: bool = True,
     chunk_tokens: int | None = None,
     autocast_dtype: torch.dtype | None = None,
+    channel_aliases: Mapping[str, Sequence[str]] | None = None,
 ) -> dict[str, np.ndarray]:
     """Convenience: load model -> preprocess EDF -> tokenize -> embed, in one call.
 
     ``model_repo_or_path`` defaults to the released model on the Hub. ``notch_freq`` is the
     powerline frequency to filter out — 50 Hz (default, most of the world) or 60 Hz (Americas).
 
+    ``channel_aliases`` lets you map non-standard EDF labels onto the model's canonical
+    channels for recordings the built-in ``ALT_COLUMNS`` table doesn't cover, e.g.
+    ``{"ECG": ["MyECGLabel"], "C3": ["EEG_C3_custom"]}``. Caller aliases take precedence
+    over the built-ins.
+
     Returns a ``{modality_name: [n_seconds, embed_dim]}`` dict of per-modality 1 Hz
     embeddings (only modalities present in the recording). For repeated embedding, call
     :func:`load_model` once and reuse the returned model/tokenizers.
     """
     model, tokenizers, meta = load_model(model_repo_or_path, device=device, dtype=dtype)
-    signals = preprocess_edf(edf_path, meta, notch_freq=notch_freq, causal=causal)
+    signals = preprocess_edf(
+        edf_path, meta, notch_freq=notch_freq, causal=causal, channel_aliases=channel_aliases
+    )
     tokens, modality_mask, channel_ids = tokenize(tokenizers, meta, signals, device=device)
     return embed(
         model,

src/hypnos/embedding/pipeline.py

@@ -15,6 +15,8 @@
 
 import logging
 
+from collections.abc import Mapping, Sequence
+
 import numpy as np
 import pyedflib
 import torch
@@ -35,6 +37,7 @@ def preprocess_edf(
     notch_freq: float = 50.0,
     causal: bool = True,
     tau_seconds: float = 60.0,
+    channel_aliases: Mapping[str, Sequence[str]] | None = None,
 ) -> dict[str, np.ndarray]:
     """Load and preprocess each modality's signal from an EDF.
 
@@ -51,13 +54,19 @@ def preprocess_edf(
         causal: Use the causal preprocessing path (matches the released causal tokenizers).
             Set False only to experiment with the zero-phase ``preprocess_signal``.
         tau_seconds: Rolling-normaliser timescale for the causal path.
+        channel_aliases: Optional ``{canonical_name: [extra EDF labels]}`` mapping for
+            recordings whose channel labels aren't covered by the built-in ``ALT_COLUMNS``.
+            Keys are canonical channel names (e.g. ``"ECG"``, ``"C3"``); caller aliases take
+            precedence over the built-ins.
     """
     # One channel per modality for the released model (in_channels=1 tokenizers); we take
     # the first channel of each modality spec.
     all_channels = sorted({ch for m in metadata.modalities for ch in m.channels})
 
     with pyedflib.EdfReader(edf_path) as f:
-        resolved = load_psg_channels(f, all_channels, drop_unreferenced=True)
+        resolved = load_psg_channels(
+            f, all_channels, drop_unreferenced=True, channel_aliases=channel_aliases
+        )
 
     signals: dict[str, np.ndarray] = {}
     for m in metadata.modalities:

tests/test_pipeline.py

@@ -144,6 +144,54 @@ def test_temporal_only_matches_forward():
     print(f"[A] _temporal_only_forward matches forward (max_diff={max_diff:.2e}) OK")
 
 
+def test_get_column_match_channel_aliases():
+    from hypnos.data.edf import get_column_match
+
+    # Exact match wins.
+    assert get_column_match("ECG", ["ECG", "EKG"]) == "ECG"
+    # Built-in ALT_COLUMNS alias is used when no exact match.
+    assert get_column_match("ECG", ["EKG"]) == "EKG"
+    # Caller-supplied alias resolves labels the built-ins don't cover.
+    assert get_column_match("ECG", ["MyECG"], {"ECG": ["MyECG"]}) == "MyECG"
+    # Caller aliases take precedence over the built-ins.
+    assert get_column_match("ECG", ["EKG", "MyECG"], {"ECG": ["MyECG"]}) == "MyECG"
+    # No match returns None.
+    assert get_column_match("ECG", ["Nope"], {"ECG": ["MyECG"]}) is None
+
+    print("[C] get_column_match honours channel_aliases OK")
+
+
+def test_get_column_match_normalization():
+    from hypnos.data.edf import get_column_match
+
+    # Case / separator / suffix insensitive, returning the original EDF label.
+    assert get_column_match("C3", ["c3-m2"]) == "c3-m2"
+    assert get_column_match("C3", ["C3:M2"]) == "C3:M2"
+    assert get_column_match("C3", ["C3-M2_PDS"]) == "C3-M2_PDS"
+    # Aliases are normalized too.
+    assert get_column_match("ECG", ["ekg"]) == "ekg"
+    # Exact match still wins outright (returned verbatim).
+    assert get_column_match("ECG", ["ECG", "ekg"]) == "ECG"
+
+    print("[D] get_column_match normalization OK")
+
+
+def test_reference_label_resolution():
+    from hypnos.data.edf import _find_reference_label, _is_pre_referenced
+
+    # TP9/TP10 are accepted as M1/M2 equivalents, normalized, original label returned.
+    assert _find_reference_label("M1", ["TP9"]) == "TP9"
+    assert _find_reference_label("M2", ["a2"]) == "a2"
+    assert _find_reference_label("M2", ["Nope"]) is None
+
+    # Pre-reference detection is case-insensitive; a bare reference is not pre-referenced.
+    assert _is_pre_referenced("C3", "c3-m2") is True
+    assert _is_pre_referenced("C3", "C3") is False
+    assert _is_pre_referenced("M2", "M2") is False  # M2 isn't a CONTRALATERAL_REF channel
+
+    print("[E] reference label resolution OK")
+
+
 def make_synthetic_edf(path, duration_sec=120):
     labels = {"C3-M2": 128, "C4-M1": 128, "ECG": 128, "ABD": 32}
     w = pyedflib.EdfWriter(str(path), len(labels))