From 389f435db3d3119aa0f8f786af48227082d3361c Mon Sep 17 00:00:00 2001 From: Jonathan Carter <42900403+joncarter1@users.noreply.github.com> Date: Thu, 25 Jun 2026 09:52:14 +0100 Subject: [PATCH] feat: custom channel mappings and tolerant EDF label matching (v0.2.0) Add an optional `channel_aliases` parameter threaded through embed_edf -> preprocess_edf -> load_psg_channels for recordings whose labels aren't covered by the built-in ALT_COLUMNS; caller aliases take precedence. Make channel and reference matching tolerant of case, whitespace, :/-/ separators and device suffixes (_clean_label), and accept TP9/TP10 as M1/M2 mastoid equivalents. Pre-reference detection is now case-insensitive and no longer flags a bare reference electrode. Bump version to 0.2.0 and document under "Custom channel labels". Co-Authored-By: Claude Opus 4.8 (1M context) --- README.md | 26 ++++++++++ pyproject.toml | 2 +- src/hypnos/data/edf.py | 81 ++++++++++++++++++++++++++------ src/hypnos/embedding/__init__.py | 12 ++++- src/hypnos/embedding/pipeline.py | 11 ++++- tests/test_pipeline.py | 48 +++++++++++++++++++ 6 files changed, 162 insertions(+), 18 deletions(-) diff --git a/README.md b/README.md index 1731618..fd42614 100644 --- a/README.md +++ b/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: diff --git a/pyproject.toml b/pyproject.toml index 73e5af9..f8a60eb 100644 --- a/pyproject.toml +++ b/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" diff --git a/src/hypnos/data/edf.py b/src/hypnos/data/edf.py index 1f91bd6..a880ed7 100644 --- a/src/hypnos/data/edf.py +++ b/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 diff --git a/src/hypnos/embedding/__init__.py b/src/hypnos/embedding/__init__.py index 9f56f35..b9f6666 100644 --- a/src/hypnos/embedding/__init__.py +++ b/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, diff --git a/src/hypnos/embedding/pipeline.py b/src/hypnos/embedding/pipeline.py index 433f479..dbe117b 100644 --- a/src/hypnos/embedding/pipeline.py +++ b/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: diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index d74295c..3f7d55d 100644 --- a/tests/test_pipeline.py +++ b/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))