diff --git a/docs/api/models.rst b/docs/api/models.rst index e881e7894..59bb22d26 100644 --- a/docs/api/models.rst +++ b/docs/api/models.rst @@ -254,3 +254,5 @@ API Reference models/pyhealth.models.califorest models/pyhealth.models.XGBoostModel models/pyhealth.models.survival_mdn + models/pyhealth.models.ProtoECGNet + diff --git a/docs/api/models/pyhealth.models.ProtoECGNet.rst b/docs/api/models/pyhealth.models.ProtoECGNet.rst new file mode 100644 index 000000000..f3ba1e707 --- /dev/null +++ b/docs/api/models/pyhealth.models.ProtoECGNet.rst @@ -0,0 +1,26 @@ +pyhealth.models.ProtoECGNet +=================================== + +ProtoECGNet: case-based interpretable multi-label ECG classification with a +prototype loss that includes a label co-occurrence contrastive term +(`Sethi et al., MLHC 2025 `_). + +.. autoclass:: pyhealth.models.ProtoECGNet + :members: + :undoc-members: + :show-inheritance: + +.. autoclass:: pyhealth.models.PrototypeBranch + :members: + :undoc-members: + :show-inheritance: + +.. autoclass:: pyhealth.models.ResNet1dBackbone + :members: + :undoc-members: + :show-inheritance: + +.. autoclass:: pyhealth.models.ResNet2dBackbone + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/tasks.rst b/docs/api/tasks.rst index e286e1b80..03f14872f 100644 --- a/docs/api/tasks.rst +++ b/docs/api/tasks.rst @@ -215,6 +215,7 @@ Available Tasks Cardiology Detection COVID-19 CXR Classification COVID-19 CXR Image Generation + PTB-XL Multilabel Classification DKA Prediction (MIMIC-IV) Drug Recommendation EHR Generation diff --git a/docs/api/tasks/pyhealth.tasks.PTBXLMultilabelClassification.rst b/docs/api/tasks/pyhealth.tasks.PTBXLMultilabelClassification.rst new file mode 100644 index 000000000..11702b130 --- /dev/null +++ b/docs/api/tasks/pyhealth.tasks.PTBXLMultilabelClassification.rst @@ -0,0 +1,10 @@ +pyhealth.tasks.PTBXLMultilabelClassification +============================================ + +Multi-label classification of PTB-XL 12-lead ECGs (all 71 SCP statements, or the +diagnostic, sub/super-diagnostic, form or rhythm label sets). + +.. autoclass:: pyhealth.tasks.PTBXLMultilabelClassification + :members: + :undoc-members: + :show-inheritance: diff --git a/examples/ecg/ptbxl/ptbxl_multilabel_classification_protoecgnet.py b/examples/ecg/ptbxl/ptbxl_multilabel_classification_protoecgnet.py new file mode 100644 index 000000000..014444342 --- /dev/null +++ b/examples/ecg/ptbxl/ptbxl_multilabel_classification_protoecgnet.py @@ -0,0 +1,380 @@ +"""ProtoECGNet on PTB-XL: full pipeline and ablation study. + +Pipeline (PTBXLDataset -> PTBXLMultilabelClassification -> ProtoECGNet): + +1. Load PTB-XL and build the multi-label task (12 leads, 100 Hz). +2. Split by the recommended ``strat_fold`` (1-8 train, 9 validation, 10 test). +3. Train in the reference stages: ``feature_extractor`` -> (``joint`` -> + prototype projection) x ``--rounds`` -> ``classifier`` (-> ``fusion``). +4. Report test macro AUROC (the paper's primary metric). + +``--paper-groups`` uses the paper's three branches on all 71 labels: a 1D +ResNet with global prototypes for the 16 rhythm labels, and 2D ResNet-18 +branches for the 52 morphology labels (partial prototypes) and the 3 global +labels, joined by a learned fusion layer. + +Ablation (``--ablation``) varies settings that the paper does not sweep: +(a) the contrastive prototype loss on/off, (b) cosine vs L2 similarity, +(c) the number of prototypes per label. Each variant uses identical data and +seeds; the table is saved as markdown. + +Usage: + # Quick check on synthetic ECGs (no download needed, runs on CPU): + python examples/ecg/ptbxl/ptbxl_multilabel_classification_protoecgnet.py --synthetic + + # Real data (open access; only the 100 Hz records, ptbxl_database.csv and + # scp_statements.csv are needed): + # wget -r -N -c -np https://physionet.org/files/ptb-xl/1.0.3/ + python examples/ecg/ptbxl/ptbxl_multilabel_classification_protoecgnet.py \ + --root physionet.org/files/ptb-xl/1.0.3 --label-set superdiagnostic \ + --epochs 20 --width 64 --ablation + + # Paper setting (all 71 labels, three branches, learned fusion): + python examples/ecg/ptbxl/ptbxl_multilabel_classification_protoecgnet.py \ + --root physionet.org/files/ptb-xl/1.0.3 --label-set all --paper-groups + +Note: + Results on ``--synthetic`` data only demonstrate that the pipeline runs; + they say nothing about performance on real ECGs. For small compute budgets + (e.g. a free Colab T4) use ``--dev`` and fewer epochs, and state the + compute limit when comparing with the paper. +""" + +import argparse +import itertools +import os +import random +import shutil +import tempfile +import warnings + +import numpy as np +import pandas as pd +import torch +from sklearn.metrics import roc_auc_score + +from pyhealth.datasets import PTBXLDataset, get_dataloader +from pyhealth.datasets.splitter import split_by_strat_fold +from pyhealth.models import PTBXL_BRANCH_KWARGS, PTBXL_BRANCH_LABELS, ProtoECGNet +from pyhealth.tasks import PTBXLMultilabelClassification +from pyhealth.trainer import Trainer + +LEADS = ["I", "II", "III", "aVR", "aVL", "aVF", "V1", "V2", "V3", "V4", "V5", "V6"] +SYNTHETIC_CODES = ["NORM", "MI", "STTC", "CD", "HYP"] + + +def safe_roc_auc_macro(y_true: np.ndarray, y_prob: np.ndarray) -> float: + """Macro AUROC over labels that have at least one positive in y_true. + + sklearn raises ValueError when a label has no positives (e.g. rare SCP + codes absent from the validation or test fold), so we skip those labels + and warn instead. + """ + has_pos = y_true.sum(axis=0) > 0 + n_skip = int((~has_pos).sum()) + if n_skip: + warnings.warn( + f"{n_skip} label(s) with no positives in the eval fold are " + "excluded from macro AUROC." + ) + if not has_pos.any(): + return float("nan") + return float(roc_auc_score(y_true[:, has_pos], y_prob[:, has_pos], average="macro")) + + +def build_synthetic_ptbxl(root: str, num_patients: int = 60, seed: int = 0) -> None: + """Writes a small PTB-XL-like dataset whose labels are learnable. + + Label ``k`` adds a sine wave of frequency ``k + 1`` Hz to lead ``k``, so a + classifier can recover the labels from the signal. The layout matches what + :class:`~pyhealth.datasets.PTBXLDataset` expects (v1.0.3 column set). + + Args: + root: Output directory. + num_patients: Number of synthetic patients (1-2 ECGs each). + seed: Random seed. + """ + rng = np.random.default_rng(seed) + t = np.arange(1000) / 100.0 + rows, ecg_id = [], 0 + for patient in range(num_patients): + for _ in range(1 + patient % 2): + ecg_id += 1 + labels = list( + rng.choice(len(SYNTHETIC_CODES), size=1 + rng.integers(0, 2), + replace=False) + ) + signal = 0.3 * rng.normal(size=(1000, 12)) + for k in labels: + signal[:, k] += np.sin(2 * np.pi * (k + 1) * t) + rel = f"records100/00000/{ecg_id:05d}_lr" + os.makedirs(os.path.join(root, os.path.dirname(rel)), exist_ok=True) + np.round(signal * 1000).astype(" torch.Tensor: + """Collects the full multi-hot label matrix in one pass over the dataloader. + + ``compute_label_statistics`` requires a single tensor rather than a + dataloader, so we accumulate all batches here before calling it. + """ + loader = get_dataloader(dataset, batch_size=256, shuffle=False) + return torch.cat([batch["labels"] for batch in loader]) + + +def to_markdown_table(rows: list[dict[str, object]]) -> str: + """Formats a list of dicts as a markdown table (no extra dependencies).""" + columns = list(rows[0]) + lines = [ + "| " + " | ".join(columns) + " |", + "| " + " | ".join("---" for _ in columns) + " |", + ] + for row in rows: + lines.append("| " + " | ".join(str(row[c]) for c in columns) + " |") + return "\n".join(lines) + + +def run_pipeline( + samples, + splits, + args: argparse.Namespace, + use_contrastive: bool = True, + similarity: str = "cosine", + prototypes_per_label: int = 0, +) -> dict[str, float]: + """Trains ProtoECGNet in stages and evaluates it on the test fold. + + Args: + samples: The full SampleDataset (defines the label vocabulary). + splits: ``(train, val, test)`` datasets. + args: Parsed command-line arguments. + use_contrastive: Whether to use the contrastive prototype loss. + similarity: ``"cosine"`` or ``"l2"``. + prototypes_per_label: Prototypes per label (0 uses the CLI value). + + Returns: + A dict with the test ``roc_auc_macro`` and ``loss``. + """ + # Seed every RNG that affects training/ablation reproducibility. + random.seed(args.seed) + np.random.seed(args.seed) + torch.manual_seed(args.seed) + train, val, test = splits + loaders = [ + get_dataloader(d, batch_size=args.batch_size, shuffle=(i == 0)) + for i, d in enumerate((train, val, test)) + ] + projection_loader = get_dataloader(train, batch_size=args.batch_size, shuffle=False) + weights, cooc = ProtoECGNet.compute_label_statistics(collect_labels(train)) + branch_kwargs = None + if args.paper_groups: + # The paper uses 5 / 18 / 7 prototypes per label for the three + # branches (heavy for a free GPU). Keep them only with + # --paper-prototypes and when no ablation overrides the count. + keep = args.paper_prototypes and not prototypes_per_label + branch_kwargs = { + name: {k: v for k, v in kwargs.items() + if keep or k != "prototypes_per_label"} + for name, kwargs in PTBXL_BRANCH_KWARGS.items() + } + model = ProtoECGNet( + samples, + label_groups=PTBXL_BRANCH_LABELS if args.paper_groups else None, + branch_kwargs=branch_kwargs, + fusion="learned" if args.paper_groups else "average", + proto_dim=args.proto_dim, + prototypes_per_label=prototypes_per_label or args.prototypes_per_label, + similarity=similarity, + topk=args.topk, + base_width=args.width, + lam_cnrst=300.0 if use_contrastive else 0.0, + lam_l1=args.lam_l1, + class_weights=weights, + label_cooccurrence=cooc, + ) + trainer = Trainer( + model=model, + metrics=[], + device=args.device, + enable_logging=False, + ) + + def stage(name: str, epochs: int) -> None: + model.set_stage(name) + trainer.train( + train_dataloader=loaders[0], + val_dataloader=loaders[1], + epochs=epochs, + monitor="loss", + monitor_criterion="min", + optimizer_params={"lr": args.lr}, + ) + + stage("feature_extractor", args.epochs) + for _ in range(args.rounds): # the paper repeats joint training + projection + stage("joint", args.epochs) + model.project_prototypes(projection_loader) + stage("classifier", max(1, args.epochs // 2)) + if args.paper_groups: + stage("fusion", max(1, args.epochs // 2)) + y_true, y_prob, loss = trainer.inference(loaders[2])[:3] + return {"roc_auc_macro": safe_roc_auc_macro(y_true, y_prob), "loss": loss} + + +def main() -> None: + """Parses arguments, runs the pipeline and optionally the ablation.""" + parser = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + parser.add_argument("--root", type=str, help="PTB-XL 1.0.3 directory") + parser.add_argument("--synthetic", action="store_true", help="use fake ECGs") + parser.add_argument("--dev", action="store_true", help="load a small subset") + parser.add_argument("--label-set", default="superdiagnostic") + parser.add_argument("--epochs", type=int, default=None) + parser.add_argument("--width", type=int, default=None) + parser.add_argument("--proto-dim", type=int, default=None) + parser.add_argument("--prototypes-per-label", type=int, default=3) + parser.add_argument("--topk", type=int, default=1, help="top-k window pooling") + parser.add_argument("--rounds", type=int, default=1, help="joint+projection rounds") + parser.add_argument("--lam-l1", type=float, default=0.0, + help="L1 sparsity weight for classifier/fusion layers " + "(off by default; try 1e-4 to encourage each " + "prototype to predict only its assigned labels)") + parser.add_argument("--paper-prototypes", action="store_true", + help="with --paper-groups, use the paper's 5/18/7 " + "prototypes per label instead of " + "--prototypes-per-label") + parser.add_argument("--paper-groups", action="store_true", + help="3 branches / 71 labels as in the paper (real data only)") + parser.add_argument("--no-contrastive", action="store_true", + help="train the main run without the contrastive " + "prototype loss (for H2-style comparisons)") + parser.add_argument("--remove-baseline", action="store_true", + help="0.5 Hz high-pass filter (as in the paper)") + parser.add_argument("--batch-size", type=int, default=32) + parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--cache-dir", type=str, default=None, + help="persistent directory for the processed dataset " + "cache (default: a temp dir cleaned up after the " + "run; set this to avoid re-preprocessing several " + "GB of PTB-XL on every invocation)") + parser.add_argument("--ablation", action="store_true") + parser.add_argument("--output", default="protoecgnet_ablation.md") + parser.add_argument("--device", choices=["auto", "cpu", "cuda", "mps"], + default="auto", + help="training device (default: auto picks " + "CUDA > MPS > CPU; pass 'cpu' to force CPU)") + args = parser.parse_args() + if args.device == "auto": + if torch.cuda.is_available(): + args.device = "cuda" + elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available(): + args.device = "mps" + else: + args.device = "cpu" + + # Validate and set defaults before creating the scratch directory so that + # parser.error() exits cleanly without leaving a temp dir behind. + if args.synthetic: + args.epochs = args.epochs or 3 + args.width = args.width or 8 + args.proto_dim = args.proto_dim or 16 + args.label_set = "superdiagnostic" + elif not args.root: + parser.error("pass --root /path/to/ptb-xl/1.0.3 or use --synthetic") + if args.paper_groups and (args.synthetic or args.label_set != "all"): + parser.error("--paper-groups needs real PTB-XL data and --label-set all") + args.epochs = args.epochs or 20 + args.width = args.width or 64 + args.proto_dim = args.proto_dim or 128 + + scratch = tempfile.mkdtemp() + try: + if args.synthetic: + args.root = os.path.join(scratch, "ptbxl") + os.makedirs(args.root) + build_synthetic_ptbxl(args.root) + + cache_dir = args.cache_dir or os.path.join(scratch, "cache") + dataset = PTBXLDataset( + root=args.root, cache_dir=cache_dir, dev=args.dev + ) + task = PTBXLMultilabelClassification( + scp_statements_path=dataset.scp_statements_path, + label_set=args.label_set, + remove_baseline=args.remove_baseline, + ) + samples = dataset.set_task(task) + splits = split_by_strat_fold(samples) + print(f"samples: train={len(splits[0])} val={len(splits[1])} test={len(splits[2])}") + + baseline = run_pipeline( + samples, splits, args, use_contrastive=not args.no_contrastive + ) + name = "no-contrastive" if args.no_contrastive else "baseline" + print(f"{name} test macro AUROC: {baseline['roc_auc_macro']:.4f}") + + if args.ablation: + rows = [] + for contrastive, similarity, per_label in itertools.product( + [True, False], ["cosine", "l2"], [1, 3] + ): + result = run_pipeline( + samples, splits, args, contrastive, similarity, per_label + ) + rows.append( + { + "contrastive": contrastive, + "similarity": similarity, + "prototypes_per_label": per_label, + "test_roc_auc_macro": round(result["roc_auc_macro"], 4), + } + ) + print(rows[-1]) + table = to_markdown_table(rows) + with open(args.output, "w", encoding="utf-8") as handle: + handle.write(table + "\n") + print(table) + finally: + shutil.rmtree(scratch, ignore_errors=True) + + +if __name__ == "__main__": + main() diff --git a/pyhealth/datasets/__init__.py b/pyhealth/datasets/__init__.py index 5cc4cfc4f..3a33ed9d5 100644 --- a/pyhealth/datasets/__init__.py +++ b/pyhealth/datasets/__init__.py @@ -52,6 +52,7 @@ def __init__(self, *args, **kwargs): from .clinvar import ClinVarDataset from .cosmic import COSMICDataset from .covid19_cxr import COVID19CXRDataset +from .ptbxl import PTBXLDataset as PTBXLDataset from .dreamt import DREAMTDataset from .ehrshot import EHRShotDataset from .eicu import eICUDataset @@ -88,7 +89,6 @@ def __init__(self, *args, **kwargs): split_by_visit_conformal, ) from .eegbci import EEGBCIDataset as EEGBCIDataset # noqa: E402 -from .ptbxl import PTBXLDataset as PTBXLDataset from .tuab import TUABDataset from .tuev import TUEVDataset from .utils import ( diff --git a/pyhealth/models/__init__.py b/pyhealth/models/__init__.py index a9e6f90b2..fcfd851d5 100644 --- a/pyhealth/models/__init__.py +++ b/pyhealth/models/__init__.py @@ -50,6 +50,14 @@ from .califorest import CaliForest from .gradient_boosted_trees import GradientBoostedTreeModel as GradientBoostedTreeModel from .xgboost_model import XGBoostModel as XGBoostModel +from .protoecgnet import ( + PTBXL_BRANCH_KWARGS as PTBXL_BRANCH_KWARGS, + PTBXL_BRANCH_LABELS as PTBXL_BRANCH_LABELS, + ProtoECGNet as ProtoECGNet, + PrototypeBranch as PrototypeBranch, + ResNet1dBackbone as ResNet1dBackbone, + ResNet2dBackbone as ResNet2dBackbone, +) from .generators.halo import HALO from .generators.gpt2 import GPT2 from .generators.promptehr import PromptEHR diff --git a/pyhealth/models/protoecgnet.py b/pyhealth/models/protoecgnet.py new file mode 100644 index 000000000..9971d223e --- /dev/null +++ b/pyhealth/models/protoecgnet.py @@ -0,0 +1,1075 @@ +"""ProtoECGNet: case-based interpretable multi-label ECG classification. + +Re-implementation of the model in "ProtoECGNet: Case-Based Interpretable Deep +Learning for Multi-Label ECG Classification with Contrastive Learning" +(Sethi et al., MLHC 2025, https://arxiv.org/abs/2504.08713). The reference +code is at https://github.com/bbj-lab/protoecgnet . + +A 1D ResNet turns a 12-lead ECG into a feature map. Each *prototype branch* +compares (windows of) that map with learned prototype vectors that are assigned +to labels, and a last linear layer turns prototype similarities into logits. +Calling :meth:`project_prototypes` replaces every prototype with the +feature-space representation of a real training ECG segment, and returns the +source-sample and window indices to the caller. In global mode +(``proto_time_len=None``) the stored vector is a time-averaged feature rather +than a raw segment, so case-based explanations should display the matching +window from the source sample rather than the prototype vector itself. Store +the returned mapping alongside the checkpoint to recover those windows after +reloading. +""" + +from collections.abc import Sequence +from itertools import combinations + +import torch +import torch.nn.functional as F +from torch import nn +from torch.utils.data import RandomSampler + +from ..datasets import SampleDataset +from .base_model import BaseModel + +STAGES = ( + "all", + "feature_extractor", + "prototypes", + "joint", + "classifier", + "fusion", +) + +#: Physician-defined grouping of the 71 PTB-XL labels into the three prototype +#: branches of the paper (Appendix G): 16 rhythm, 52 morphology, 3 global. +PTBXL_BRANCH_LABELS: dict[str, list[str]] = { + "rhythm": [ + "1AVB", "2AVB", "3AVB", "AFIB", "AFLT", "BIGU", "IVCD", "PACE", "PSVT", + "SARRH", "SBRAD", "SR", "STACH", "SVARR", "SVTAC", "TRIGU", + ], + "morphology": [ + "ABQRS", "ALMI", "AMI", "ANEUR", "ASMI", "CLBBB", "CRBBB", "HVOLT", + "ILBBB", "ILMI", "IMI", "INJAL", "INJAS", "INJIL", "INJIN", "INJLA", + "INVT", "IPLMI", "IPMI", "IRBBB", "ISCAL", "ISCAN", "ISCAS", "ISCIL", + "ISCIN", "ISCLA", "ISC_", "LAFB", "LAO/LAE", "LMI", "LNGQT", "LOWT", + "LPFB", "LPR", "LVH", "LVOLT", "NDT", "NST_", "NT_", "PAC", "PMI", + "PRC(S)", "PVC", "QWAVE", "RAO/RAE", "RVH", "SEHYP", "STD_", "STE_", + "TAB_", "VCLVH", "WPW", + ], + "global": ["DIG", "EL", "NORM"], +} + +#: Per-branch settings matching the paper (Sec. 3.3 and Appendix F): a 1D +#: ResNet-18 with global prototypes (5 per label) for rhythm; a 2D ResNet-18 +#: with partial prototypes of 3 latent steps, top-5 pooling and 18 prototypes +#: per label for morphology; and a 2D ResNet-18 whose global prototypes span +#: the full latent time axis (32 steps for a 12 x 1000 ECG, i.e. a +#: 512 x 1 x 32 prototype) with 7 prototypes per label for the global labels. +PTBXL_BRANCH_KWARGS: dict[str, dict] = { + "rhythm": { + "backbone": "resnet1d", + "proto_time_len": None, + "prototypes_per_label": 5, + }, + "morphology": { + "backbone": "resnet2d", + "proto_time_len": 3, + "topk": 5, + "prototypes_per_label": 18, + }, + "global": { + "backbone": "resnet2d", + "proto_time_len": 32, + "prototypes_per_label": 7, + }, +} + + +class BasicBlock1d(nn.Module): + """Residual block with two 1D convolutions. + + Args: + in_channels: Number of input channels. + out_channels: Number of output channels. + stride: Stride of the first convolution. + kernel_size: Convolution kernel size. + + Examples: + >>> import torch + >>> block = BasicBlock1d(64, 128, stride=2) + >>> block(torch.randn(2, 64, 100)).shape + torch.Size([2, 128, 50]) + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + stride: int = 1, + kernel_size: int = 3, + ) -> None: + super().__init__() + pad = (kernel_size - 1) // 2 + self.conv1 = nn.Conv1d( + in_channels, out_channels, kernel_size, stride, pad, bias=False + ) + self.bn1 = nn.BatchNorm1d(out_channels) + self.relu = nn.ReLU(inplace=True) + self.conv2 = nn.Conv1d( + out_channels, out_channels, kernel_size, 1, pad, bias=False + ) + self.bn2 = nn.BatchNorm1d(out_channels) + self.downsample: nn.Module | None = None + if stride != 1 or in_channels != out_channels: + self.downsample = nn.Sequential( + nn.Conv1d(in_channels, out_channels, 1, stride, bias=False), + nn.BatchNorm1d(out_channels), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Adds a skip connection that is downsampled when stride > 1 or channels change.""" + residual = x if self.downsample is None else self.downsample(x) + out = self.relu(self.bn1(self.conv1(x))) + out = self.bn2(self.conv2(out)) + return self.relu(out + residual) + + +class ResNet1dBackbone(nn.Module): + """1D ResNet feature extractor that keeps the temporal axis. + + With the defaults (ResNet-18 layout) a 12 x 1000 input gives a + ``(batch, 512, 32)`` feature map. + + Args: + in_channels: Number of ECG leads. + base_width: Channels of the first stage; later stages use 2x, 4x, 8x. + layers: Number of residual blocks in each of the four stages. + kernel_size: Convolution kernel size inside the blocks. + + Attributes: + out_channels: Number of channels of the output feature map. + + Examples: + >>> import torch + >>> backbone = ResNet1dBackbone() + >>> backbone(torch.randn(2, 12, 1000)).shape + torch.Size([2, 512, 32]) + """ + + def __init__( + self, + in_channels: int = 12, + base_width: int = 64, + layers: Sequence[int] = (2, 2, 2, 2), + kernel_size: int = 3, + ) -> None: + super().__init__() + self.stem = nn.Sequential( + nn.Conv1d(in_channels, base_width, 7, 2, 3, bias=False), + nn.BatchNorm1d(base_width), + nn.ReLU(inplace=True), + nn.MaxPool1d(3, 2, 1), + ) + blocks: list[nn.Module] = [] + channels = base_width + for stage, count in enumerate(layers): + width = base_width * 2**stage + for index in range(count): + stride = 2 if (stage > 0 and index == 0) else 1 + blocks.append(BasicBlock1d(channels, width, stride, kernel_size)) + channels = width + self.blocks = nn.Sequential(*blocks) + self.out_channels = channels + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Passes the signal through the stem (conv-7, BN, ReLU, maxpool) and then all residual stages, preserving the temporal axis.""" + return self.blocks(self.stem(x)) + + +class ResNet2dBackbone(nn.Module): + """ResNet-18 that reads a 12-lead ECG as a one-channel image. + + The input ``(batch, leads, time)`` is treated as a ``1 x leads x time`` + image. The first convolution has a ``leads x 7`` kernel (so all leads are + mixed immediately) and the global pooling is dropped to keep the temporal + axis, giving a ``(batch, 512, time')`` map (``time' = 32`` for 1000 samples). + + Args: + in_leads: Number of ECG leads. + pretrained: If True, initializes the residual stages from ImageNet + weights (requires a download); the first convolution is always + trained from scratch because its shape differs. + + Attributes: + out_channels: Channels of the output feature map (512). + + Examples: + >>> import torch + >>> backbone = ResNet2dBackbone() + >>> backbone(torch.randn(2, 12, 1000)).shape + torch.Size([2, 512, 32]) + """ + + def __init__(self, in_leads: int = 12, pretrained: bool = False) -> None: + super().__init__() + from torchvision.models import ResNet18_Weights, resnet18 + + net = resnet18(weights=ResNet18_Weights.DEFAULT if pretrained else None) + net.conv1 = nn.Conv2d(1, 64, (in_leads, 7), (1, 2), (0, 3), bias=False) + # avgpool and fc are not used; replace with no-ops to drop their parameters. + net.avgpool = nn.Identity() + net.fc = nn.Identity() + self.net = net + self.out_channels = 512 + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Unsqueezes to a 1-channel image, runs the ResNet-18 stem and four residual stages, then squeezes the lead axis to restore a 1D temporal map.""" + net = self.net + x = net.maxpool(net.relu(net.bn1(net.conv1(x.unsqueeze(1))))) + x = net.layer4(net.layer3(net.layer2(net.layer1(x)))) + return x.squeeze(2) + + +class PrototypeBranch(nn.Module): + """One prototype branch predicting a subset of the labels. + + Prototypes are assigned to a single label (``prototypes_per_label`` each) + or to a pair of labels (``joint_prototypes_per_pair`` each). In ``global`` + mode (``proto_time_len=None``) the feature map is average-pooled over + time; otherwise a prototype covers ``proto_time_len`` consecutive time + steps and is matched against every window of that length (``partial`` + prototypes). The similarity of a prototype to an ECG is the mean of its + ``topk`` best window similarities. + + Args: + in_channels: Number of ECG leads. + num_labels: Number of labels handled by this branch. + proto_dim: Channels of the prototype space. + prototypes_per_label: Prototypes assigned to each single label. + joint_prototypes_per_pair: Prototypes assigned to each label pair. + proto_time_len: Temporal length of a prototype, or None for global. + backbone: ``"resnet1d"`` (ResNet-18 layout over time) or + ``"resnet2d"`` (torchvision ResNet-18 over a leads x time image). + pretrained: ImageNet initialization of the 2D backbone. + base_width: Width of the 1D backbone's first stage. + layers: Residual blocks per stage of the 1D backbone. + similarity: ``"cosine"`` (scaled cosine, as in the paper) or ``"l2"`` + (ProtoPNet-style ``log((d + 1) / (d + eps))`` of the L2 distance). + topk: Number of best windows averaged into the similarity. + similarity_scale: Scale ``a`` of the cosine similarity. Defaults to + ``sqrt(proto_dim * window)``. + activation: ``"log"`` or ``"linear"`` map used when ``similarity="l2"``. + last_layer_connection_weight: Initial weight between a prototype and + its own label(s); other connections start at ``-0.5``. + epsilon: Small constant of the log activation. + + Examples: + >>> import torch + >>> branch = PrototypeBranch(in_channels=12, num_labels=5) + >>> out = branch(torch.randn(2, 12, 1000)) + >>> out["logits"].shape + torch.Size([2, 5]) + """ + + def __init__( + self, + in_channels: int, + num_labels: int, + proto_dim: int = 512, + prototypes_per_label: int = 5, + joint_prototypes_per_pair: int = 0, + proto_time_len: int | None = None, + backbone: str = "resnet1d", + pretrained: bool = False, + base_width: int = 64, + layers: Sequence[int] = (2, 2, 2, 2), + similarity: str = "cosine", + topk: int = 1, + similarity_scale: float | None = None, + activation: str = "log", + last_layer_connection_weight: float = 1.0, + epsilon: float = 1e-4, + ) -> None: + super().__init__() + if similarity not in ("cosine", "l2"): + raise ValueError("similarity must be 'cosine' or 'l2'") + if activation not in ("log", "linear"): + raise ValueError("activation must be 'log' or 'linear'") + if backbone not in ("resnet1d", "resnet2d"): + raise ValueError("backbone must be 'resnet1d' or 'resnet2d'") + if proto_time_len is not None and proto_time_len < 1: + raise ValueError("proto_time_len must be a positive integer or None") + if topk < 1: + raise ValueError("topk must be at least 1") + self.num_labels = num_labels + self.proto_dim = proto_dim + self.proto_time_len = proto_time_len + self.similarity = similarity + self.topk = topk + self.activation = activation + self.epsilon = epsilon + self.window = 1 if proto_time_len is None else proto_time_len + self.scale = ( + float(similarity_scale) + if similarity_scale is not None + else float((proto_dim * self.window) ** 0.5) + ) + + if backbone == "resnet1d": + self.backbone: nn.Module = ResNet1dBackbone( + in_channels, base_width, layers + ) + else: + self.backbone = ResNet2dBackbone(in_channels, pretrained) + feat = self.backbone.out_channels + self.add_on = nn.Sequential( + nn.Conv1d(feat, proto_dim, 1), + nn.ReLU(), + nn.Conv1d(proto_dim, proto_dim, 1), + nn.Sigmoid(), + ) + self.pretrain_head = nn.Linear(feat, num_labels) + + identity = self._make_identity( + num_labels, prototypes_per_label, joint_prototypes_per_pair + ) + self.register_buffer("prototype_class_identity", identity) + self.num_prototypes = identity.shape[0] + self.prototype_vectors = nn.Parameter( + torch.randn(self.num_prototypes, proto_dim, self.window) + ) + self.classifier = nn.Linear(self.num_prototypes, num_labels, bias=False) + self.set_last_layer_connections(last_layer_connection_weight) + + @staticmethod + def _make_identity( + num_labels: int, per_label: int, per_pair: int + ) -> torch.Tensor: + """Builds the binary prototype-to-label assignment matrix, with per-label rows first and then a shared row for every label pair.""" + pairs = list(combinations(range(num_labels), 2)) + total = num_labels * per_label + len(pairs) * per_pair + identity = torch.zeros(total, num_labels) + for label in range(num_labels): + identity[label * per_label : (label + 1) * per_label, label] = 1.0 + offset = num_labels * per_label + for index, (first, second) in enumerate(pairs): + start = offset + index * per_pair + identity[start : start + per_pair, first] = 1.0 + identity[start : start + per_pair, second] = 1.0 + return identity + + def set_last_layer_connections( + self, correct: float = 1.0, incorrect: float = -0.5 + ) -> None: + """Initializes the classifier from the prototype-label assignment. + + Args: + correct: Weight between a prototype and its own label(s). + incorrect: Weight between a prototype and all other labels. + """ + own = self.prototype_class_identity.t() # (labels, prototypes) + with torch.no_grad(): + self.classifier.weight.copy_(correct * own + incorrect * (1 - own)) + + def feature_map(self, x: torch.Tensor) -> torch.Tensor: + """Projects backbone features to prototype space; collapses the time axis to a single step when using global prototypes (``proto_time_len=None``).""" + fmap = self.add_on(self.backbone(x)) + if self.proto_time_len is None: + fmap = fmap.mean(dim=2, keepdim=True) + return fmap + + def windows(self, fmap: torch.Tensor) -> torch.Tensor: + """Slides a window of ``proto_time_len`` steps over the time axis and flattens channel and time into a single vector per window, matching the prototype shape for dot-product similarity.""" + if fmap.shape[2] < self.window: + raise ValueError( + f"Feature map has {fmap.shape[2]} time steps but proto_time_len " + f"is {self.window}; use a smaller proto_time_len or longer ECGs." + ) + cut = fmap.unfold(2, self.window, 1) # (B, D, W, L) + return cut.permute(0, 2, 1, 3).reshape(fmap.shape[0], cut.shape[2], -1) + + def distances(self, windows: torch.Tensor) -> torch.Tensor: + """Computes pairwise L2 distances using ||a-b||²=||a||²-2aᵀb+||b||², avoiding an explicit loop over the P prototypes.""" + protos = self.prototype_vectors.flatten(1) + d2 = ( + windows.pow(2).sum(-1, keepdim=True) + - 2 * windows @ protos.t() + + protos.pow(2).sum(-1) + ) + return torch.sqrt(d2.clamp_min(0.0) + 1e-8) + + def similarity_windows(self, windows: torch.Tensor) -> torch.Tensor: + """Similarity of every window to every prototype ``(batch, W, P)``. + + Larger values mean a closer match for both similarity types. + """ + if self.similarity == "cosine": + protos = F.normalize(self.prototype_vectors.flatten(1), dim=1) + return self.scale * (F.normalize(windows, dim=-1) @ protos.t()) + distance = self.distances(windows) + if self.activation == "log": + return torch.log((distance + 1) / (distance + self.epsilon)) + return -distance + + def forward( + self, x: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Runs the branch. + + Args: + x: ECG batch of shape ``(batch, leads, time)``. + + Returns: + ``(logits, similarity, best_window)``: logits ``(batch, + num_labels)``, the pooled prototype similarities ``(batch, + num_prototypes)`` and, for explanations, the index of the + best-matching window per prototype ``(batch, num_prototypes)``. + """ + sims = self.similarity_windows(self.windows(self.feature_map(x))) + k = min(self.topk, sims.shape[1]) + pooled = sims.topk(k, dim=1).values.mean(dim=1) + best_window = sims.argmax(dim=1) + return self.classifier(pooled), pooled, best_window + + def pretrain_logits(self, x: torch.Tensor) -> torch.Tensor: + """Logits of the plain ResNet head used by the feature-extractor stage.""" + return self.pretrain_head(self.backbone(x).mean(dim=2)) + + def prototype_terms( + self, + similarity: torch.Tensor, + y: torch.Tensor, + cooccurrence: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Prototype regularizers (Eqs. 4-7 of the paper). + + Args: + similarity: Prototype similarities ``(batch, num_prototypes)``. + y: Multi-hot labels of this branch ``(batch, num_labels)``. + cooccurrence: Optional label co-occurrence matrix of this branch + ``(num_labels, num_labels)``; enables the contrastive term. + + Returns: + A dict with the ``clst`` (clustering), ``sep`` (separation), + ``div`` (diversity) and ``cnrst`` (contrastive) terms. + """ + identity = self.prototype_class_identity + # A prototype is "positive" for a sample if it is assigned to at least + # one label present in the sample; otherwise all its labels are absent + # and it is "negative" (Eqs. 4-5 of the paper). + positive = (y.float() @ identity.t()) > 0 # (B, P) + negative = ~positive + neg_inf = torch.finfo(similarity.dtype).min + + def masked_max(mask: torch.Tensor) -> torch.Tensor: + """Per-sample max similarity over ``mask``; 0 where mask is empty.""" + best = similarity.masked_fill(~mask, neg_inf).max(dim=1).values + return torch.where(mask.any(dim=1), best, torch.zeros_like(best)) + + terms = { + "clst": -masked_max(positive).mean(), + "sep": masked_max(negative).mean(), + } + protos = F.normalize(self.prototype_vectors.flatten(1), dim=1) + gram = protos @ protos.t() + count = self.num_prototypes + eye = torch.eye(count, device=gram.device) + terms["div"] = (gram - eye).pow(2).sum() / count**2 + if cooccurrence is None: + terms["cnrst"] = gram.new_zeros(()) + else: + cos = gram.clamp(-1 + 1e-6, 1 - 1e-6) + pos = (identity @ cooccurrence @ identity.t()).clamp(max=1.0) + neg = 1.0 - pos + pos = pos / (pos.sum() + 1e-6) + neg = neg / (neg.sum() + 1e-6) + terms["cnrst"] = -((cos * pos).sum() - (cos * neg).sum()) / count**0.5 + return terms + + +class ProtoECGNet(BaseModel): + """Case-based interpretable multi-label ECG classifier. + + The model contains one or several :class:`PrototypeBranch` modules. Each + branch is responsible for a disjoint group of labels (for example rhythm + labels with global prototypes and morphology labels with partial + prototypes). Branch outputs are combined by scattering their logits into + the full label space (``fusion="average"``) or by a learned linear layer + over all prototype similarities (``fusion="learned"``, trained in the + ``"fusion"`` stage). + + Training follows the reference procedure and can be driven by PyHealth's + ``Trainer`` by switching stages with :meth:`set_stage`: + + 1. ``"feature_extractor"``: train the backbone with a plain linear head. + 2. ``"joint"``: train backbone, add-on layers and prototypes. + 3. :meth:`project_prototypes`: replace every prototype by the most similar + real training segment of its label(s). Steps 2-3 can be repeated. + 4. ``"classifier"``: train only the last layer on the projected + prototypes. For ``fusion="learned"`` finish with ``"fusion"``. + + ``"all"`` trains everything with one loss and is a quick baseline. In the + ``"classifier"`` and ``"fusion"`` stages an L1 penalty (``lam_l1``) on the + weights between a prototype and labels it is *not* assigned to keeps the + final layers sparse, as in the paper. + + Args: + dataset: A fitted :class:`~pyhealth.datasets.SampleDataset` with a + ``signal`` tensor input and a multilabel output. + label_groups: Optional mapping from branch name to the label names the + branch handles (see :data:`PTBXL_BRANCH_LABELS`). Names absent + from the dataset vocabulary are dropped, and the groups must then + cover every label exactly once. Defaults to one branch with all + labels. + branch_kwargs: Optional per-branch overrides of + :class:`PrototypeBranch` arguments, keyed by branch name (see + :data:`PTBXL_BRANCH_KWARGS`). + in_channels: Number of ECG leads. Defaults to 12. + proto_dim: Channels of the prototype space. Defaults to 512. + prototypes_per_label: Prototypes per label. Defaults to 5. + joint_prototypes_per_pair: Prototypes per label pair. Defaults to 0. + proto_time_len: Prototype length (None means global). Defaults to None. + similarity: ``"cosine"`` (default) or ``"l2"``. + topk: Windows averaged into a similarity. Defaults to 1. + base_width: Width of the 1D ResNet's first stage. Defaults to 64. + layers: Residual blocks per 1D stage. Defaults to ``(2, 2, 2, 2)``. + fusion: ``"average"`` (scatter each branch's logits into the full + label space with no cross-branch interaction) or ``"learned"`` + (route logits through a linear layer over all prototype + similarities; see :meth:`_uses_fusion_head` — active in the + ``"all"`` and ``"fusion"`` stages, and dedicated-trained in the + ``"fusion"`` stage). Defaults to ``"average"``. + lam_clst: Weight of the clustering loss. Defaults to 0.004. + lam_sep: Weight of the separation loss. Defaults to 0.0004. + lam_div: Weight of the diversity loss. Defaults to 250. + lam_cnrst: Weight of the contrastive loss. Defaults to 300. It is + only active once a co-occurrence matrix has been set. + lam_l1: Weight of the L1 penalty of the classifier/fusion stages. + Defaults to 0 (off). + class_weights: Optional positive-class weights ``(num_labels,)``. + label_cooccurrence: Optional label co-occurrence matrix. + + Note: + ``ProtoECGNet`` applies loss weighting via its own ``class_weights`` + buffer. :meth:`set_pos_weight` is overridden to route the inherited + :class:`~pyhealth.models.BaseModel` argument into ``class_weights`` + (including the ``"balanced"`` option), so callers like autoimprove + work unchanged. + + Examples: + >>> from pyhealth.datasets import create_sample_dataset, get_dataloader + >>> samples = [ + ... {"patient_id": "a", "signal": [[0.0] * 256] * 12, "labels": ["X"]}, + ... {"patient_id": "b", "signal": [[1.0] * 256] * 12, "labels": ["Y"]}, + ... ] + >>> dataset = create_sample_dataset( + ... samples=samples, + ... input_schema={"signal": "tensor"}, + ... output_schema={"labels": "multilabel"}, + ... ) + >>> model = ProtoECGNet(dataset, base_width=4, layers=(1, 1, 1, 1)) + >>> batch = next(iter(get_dataloader(dataset, batch_size=2))) + >>> out = model(**batch) + >>> out["logit"].shape + torch.Size([2, 2]) + """ + + def __init__( + self, + dataset: SampleDataset, + label_groups: dict[str, list[str]] | None = None, + branch_kwargs: dict[str, dict] | None = None, + in_channels: int = 12, + proto_dim: int = 512, + prototypes_per_label: int = 5, + joint_prototypes_per_pair: int = 0, + proto_time_len: int | None = None, + similarity: str = "cosine", + topk: int = 1, + base_width: int = 64, + layers: Sequence[int] = (2, 2, 2, 2), + fusion: str = "average", + lam_clst: float = 0.004, + lam_sep: float = 0.0004, + lam_div: float = 250.0, + lam_cnrst: float = 300.0, + lam_l1: float = 0.0, + class_weights: torch.Tensor | None = None, + label_cooccurrence: torch.Tensor | None = None, + ) -> None: + super().__init__(dataset) + if len(self.feature_keys) != 1 or len(self.label_keys) != 1: + raise ValueError("ProtoECGNet expects one signal input and one label.") + if self.mode != "multilabel": + raise ValueError("ProtoECGNet requires a multilabel output schema.") + if fusion not in ("average", "learned"): + raise ValueError("fusion must be 'average' or 'learned'") + self.signal_key = self.feature_keys[0] + self.label_key = self.label_keys[0] + self.fusion = fusion + self.stage = "all" + self.lam_clst, self.lam_sep = lam_clst, lam_sep + self.lam_div, self.lam_cnrst, self.lam_l1 = lam_div, lam_cnrst, lam_l1 + + vocab = dataset.output_processors[self.label_key].label_vocab + self.label_names = [ + name for name, _ in sorted(vocab.items(), key=lambda i: i[1]) + ] + self.num_labels = len(self.label_names) + groups = self._prepare_groups(label_groups) + if any("." in name for name in groups): + raise ValueError("branch names must not contain '.'") + index_of = {name: i for i, name in enumerate(self.label_names)} + + self.branch_names = list(groups) + self.branches = nn.ModuleDict() + for name, members in groups.items(): + self.register_buffer( + f"index_{name}", torch.tensor([index_of[m] for m in members]) + ) + options = { + "proto_dim": proto_dim, + "prototypes_per_label": prototypes_per_label, + "joint_prototypes_per_pair": joint_prototypes_per_pair, + "proto_time_len": proto_time_len, + "similarity": similarity, + "topk": topk, + "base_width": base_width, + "layers": layers, + } + options.update((branch_kwargs or {}).get(name, {})) + self.branches[name] = PrototypeBranch( + in_channels, len(members), **options + ) + + total = sum(b.num_prototypes for b in self.branches.values()) + self.fusion_classifier: nn.Linear | None = ( + nn.Linear(total, self.num_labels) if fusion == "learned" else None + ) + self.register_buffer("fusion_mask", self._build_fusion_mask(total)) + # Real-tensor defaults so state_dict keys are stable regardless of whether + # label statistics were provided. ones(class_weights) is BCE's identity; + # zeros(label_cooccurrence) + the .any() gate in forward keeps the + # contrastive term off until a real matrix is supplied. + self.register_buffer("class_weights", torch.ones(self.num_labels)) + self.register_buffer( + "label_cooccurrence", torch.zeros(self.num_labels, self.num_labels) + ) + if class_weights is not None: + self.class_weights = class_weights.float() + if label_cooccurrence is not None: + self.label_cooccurrence = label_cooccurrence.float() + # Modules whose BN running stats must stay fixed after projection. + # Populated by set_stage; consumed by train(). + self._frozen_modules: set[nn.Module] = set() + # Set by project_prototypes; not in state_dict — persist alongside ckpt. + self.projection_info: dict[str, list[dict | None]] | None = None + + def _prepare_groups( + self, groups: dict[str, list[str]] | None + ) -> dict[str, list[str]]: + """Drops unknown labels and validates that groups partition the labels.""" + if groups is None: + return {"all": list(self.label_names)} + known = set(self.label_names) + cleaned = { + name: [m for m in members if m in known] + for name, members in groups.items() + } + cleaned = {name: members for name, members in cleaned.items() if members} + flat = [m for members in cleaned.values() for m in members] + if sorted(flat) != sorted(self.label_names): + raise ValueError( + "label_groups must contain every label exactly once; got " + f"{len(flat)} known entries for {len(self.label_names)} labels." + ) + return cleaned + + def _build_fusion_mask(self, total: int) -> torch.Tensor: + """Assembles the cross-branch prototype-to-label assignment mask, used to identify non-owned connections for the L1 sparsity penalty in the fusion stage.""" + mask = torch.zeros(self.num_labels, total) + offset = 0 + for name in self.branch_names: + branch = self.branches[name] + index = self.branch_index(name) + mask[index, offset : offset + branch.num_prototypes] = ( + branch.prototype_class_identity.t() + ) + offset += branch.num_prototypes + return mask + + def branch_index(self, name: str) -> torch.Tensor: + """Returns the global label indices owned by branch ``name``; used for scattering branch logits into the full label space and slicing the co-occurrence matrix.""" + return getattr(self, f"index_{name}") + + def set_pos_weight( + self, + pos_weight, + dataset=None, + ) -> None: + """Honor :meth:`BaseModel.set_pos_weight` by routing it to ``class_weights``. + + ``ProtoECGNet`` applies positive-class weighting through its own + ``class_weights`` buffer, so the inherited method would otherwise be a + no-op (silently dropping e.g. ``pos_weight="balanced"`` from + ``autoimprove``). This override delegates to the base implementation + to parse the argument (scalar, sequence, tensor, ``"balanced"``, or + ``None``) and then mirrors ``self._pos_weight`` into + ``self.class_weights`` so BCE picks it up. A scalar is broadcast to a + ``(num_labels,)`` tensor so the buffer's shape is stable across calls + (keeps strict ``load_state_dict`` working after a scalar + ``set_pos_weight``). ``None`` resets ``class_weights`` to ones. + + Args: + pos_weight: Same semantics as :meth:`BaseModel.set_pos_weight`. + dataset: Required when ``pos_weight="balanced"``. + + Raises: + ValueError: If the parsed weight is neither length-1 (broadcastable) + nor length-``num_labels``. + """ + super().set_pos_weight(pos_weight, dataset) + if self._pos_weight is None: + self.class_weights = torch.ones(self.num_labels, device=self.device) + return + weight = self._pos_weight.to(self.device).float() + if weight.numel() == 1: + weight = weight.expand(self.num_labels).clone() + elif weight.numel() != self.num_labels: + raise ValueError( + "pos_weight must be a scalar / length-1 tensor or have one " + f"entry per label ({self.num_labels}); got {weight.numel()}." + ) + self.class_weights = weight + + def set_label_statistics( + self, + class_weights: torch.Tensor | None = None, + label_cooccurrence: torch.Tensor | None = None, + ) -> None: + """Sets class weights and the label co-occurrence matrix. + + Both are backed by persistent buffers initialised to ``ones`` (identity + for BCE's ``pos_weight``) and ``zeros`` (gates the contrastive term off + via an ``.any()`` check in :meth:`forward`); supplying a real + ``label_cooccurrence`` enables the contrastive term, and supplying an + all-zeros matrix explicitly also leaves it off. + + Args: + class_weights: Positive-class weights ``(num_labels,)``. + label_cooccurrence: Co-occurrence matrix ``(num_labels, num_labels)``. + """ + if class_weights is not None: + self.class_weights = class_weights.float().to(self.device) + if label_cooccurrence is not None: + self.label_cooccurrence = label_cooccurrence.float().to(self.device) + + @staticmethod + def compute_label_statistics( + labels: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Computes class weights and Jaccard label co-occurrence. + + Args: + labels: Multi-hot training labels ``(num_samples, num_labels)``. + + Returns: + ``(class_weights, cooccurrence)``. Class weights follow the paper's + reference code, ``(total - count) / count`` where ``total`` is the + sum of all per-label positive counts (``counts.sum()``). + + Note: + The ``total`` here differs from + :meth:`~pyhealth.models.BaseModel.set_pos_weight`'s ``"balanced"`` + option, which uses the **sample count** instead of the label-count + total. Both paths end up in the same ``class_weights`` buffer, so + mix them deliberately (e.g. the model used autoimprove-style + ``"balanced"`` or the paper's reference weights, not both). + """ + y = labels.float() + counts = y.sum(dim=0).clamp_min(1.0) + weights = (counts.sum() - counts) / counts + inter = y.t() @ y + union = (counts.unsqueeze(0) + counts.unsqueeze(1) - inter).clamp_min(1e-6) + return weights, inter / union + + def get_extra_state(self) -> dict: + """Persists non-tensor state (currently ``self.stage``) in ``state_dict``. + + ``stage`` selects the forward code path (branch-scatter logits vs. + ``fusion_classifier`` output) and the loss terms, so without this a + checkpoint saved after :meth:`set_stage` would silently reload in + ``"all"`` and return different logits — in particular, a learned-fusion + model would skip ``fusion_classifier`` entirely. + """ + return {"stage": self.stage} + + def set_extra_state(self, state) -> None: + """Restores non-tensor state saved by :meth:`get_extra_state`.""" + if isinstance(state, dict) and "stage" in state: + self.set_stage(state["stage"]) + + def set_stage(self, stage: str) -> None: + """Selects which parameters are trained. + + Args: + stage: One of ``"all"``, ``"feature_extractor"``, ``"prototypes"``, + ``"joint"``, ``"classifier"`` or ``"fusion"``. + """ + if stage not in STAGES: + raise ValueError(f"stage must be one of {STAGES}") + if stage == "fusion" and self.fusion_classifier is None: + raise ValueError("the 'fusion' stage requires fusion='learned'") + trainable = { + "feature_extractor": ("backbone", "pretrain_head"), + "prototypes": ("prototype_vectors",), + "joint": ("backbone", "add_on", "prototype_vectors"), + "classifier": ("classifier",), + } + for name, param in self.named_parameters(): + if stage == "all": + param.requires_grad = True + elif stage == "fusion": + param.requires_grad = name.startswith("fusion_classifier") + else: + parts = name.split(".") + param.requires_grad = ( + parts[0] == "branches" and parts[2] in trainable[stage] + ) + self.stage = stage + # Recompute once here so train() is O(|frozen|) with no allocation. + # Use recurse=True so a PrototypeBranch whose backbone has trainable + # children is never added even though prototype_vectors itself is frozen. + self._frozen_modules = { + m for m in self.modules() + if any(True for _ in m.parameters(recurse=False)) + and not any(p.requires_grad for p in m.parameters(recurse=True)) + } + + def train(self, mode: bool = True) -> "ProtoECGNet": + super().train(mode) + if mode: + for m in self._frozen_modules: + m.eval() + return self + + def _uses_fusion_head(self) -> bool: + """Whether ``forward`` should route logits through ``fusion_classifier``. + + True for ``fusion="learned"`` in the ``"all"`` and ``"fusion"`` stages + (the first because the stage trains every parameter, the second because + it is the paper's dedicated fine-tune on top of frozen branches). The + paper's staged protocol goes through the per-branch classifiers in + ``"joint"`` / ``"classifier"`` / ``"prototypes"``. + """ + return ( + self.fusion_classifier is not None + and self.stage in ("all", "fusion") + ) + + def _scatter(self, per_branch: dict[str, torch.Tensor], batch: int) -> torch.Tensor: + """Writes branch logits into the full label space.""" + full = torch.zeros(batch, self.num_labels, device=self.device) + for name, logits in per_branch.items(): + full = full.index_copy(1, self.branch_index(name), logits) + return full + + def _l1_penalty(self) -> torch.Tensor: + """L1 norm of the weights between prototypes and non-assigned labels.""" + if self.stage == "fusion": + weight = self.fusion_classifier.weight + return (weight.abs() * (1.0 - self.fusion_mask)).sum() + total = self.fusion_mask.new_zeros(()) + for name in self.branch_names: + branch = self.branches[name] + own = branch.prototype_class_identity.t() + total = total + (branch.classifier.weight.abs() * (1.0 - own)).sum() + return total + + def forward(self, **kwargs) -> dict[str, torch.Tensor]: + """Forward pass. + + Args: + **kwargs: Must contain the ECG tensor under the signal key + ``(batch, leads, time)``; may contain the multi-hot labels. + + Returns: + A dict with ``logit`` and ``y_prob`` ``(batch, num_labels)``, + ``similarity`` (all prototype similarities, concatenated over + branches) and, when labels are given, ``loss`` and ``y_true``. + """ + x = kwargs[self.signal_key].to(self.device).float() + y = kwargs.get(self.label_key) + y = None if y is None else y.to(self.device).float() + batch = x.shape[0] + extra: list[torch.Tensor] = [] + + if self.stage == "feature_extractor": + logits = self._scatter( + {n: self.branches[n].pretrain_logits(x) for n in self.branch_names}, + batch, + ) + sims = x.new_zeros(batch, 0) + else: + outputs = {n: self.branches[n](x) for n in self.branch_names} + sims = torch.cat([o[1] for o in outputs.values()], dim=1) + # Use the learned fusion head in "all" (so it actually gets gradients + # when the stage claims to train everything) and in the "fusion" + # stage (the paper's dedicated fine-tune). Other stages — "joint", + # "classifier", "prototypes" — follow the paper's staged protocol + # and go through the per-branch classifiers. + if self._uses_fusion_head(): + logits = self.fusion_classifier(sims) + else: + logits = self._scatter({n: o[0] for n, o in outputs.items()}, batch) + if y is not None and self.stage in ("all", "joint", "prototypes"): + # Hoist the one-time .any() check out of the per-branch loop. + cnrst_on = self.lam_cnrst > 0 and bool(self.label_cooccurrence.any()) + for name, (_, sim, _) in outputs.items(): + index = self.branch_index(name) + cooc = ( + self.label_cooccurrence[index][:, index] if cnrst_on else None + ) + t = self.branches[name].prototype_terms(sim, y[:, index], cooc) + extra.append( + self.lam_clst * t["clst"] + + self.lam_sep * t["sep"] + + self.lam_div * t["div"] + + self.lam_cnrst * t["cnrst"] + ) + if y is not None and self.lam_l1 > 0 and self.stage in ( + "classifier", + "fusion", + ): + extra.append(self.lam_l1 * self._l1_penalty()) + + result = {"logit": logits, "y_prob": torch.sigmoid(logits), "similarity": sims} + if y is not None: + loss = F.binary_cross_entropy_with_logits( + logits, y, pos_weight=self.class_weights + ) + for term in extra: + loss = loss + term + result["loss"] = loss + result["y_true"] = y + return result + + @torch.no_grad() + def project_prototypes(self, dataloader) -> dict[str, list[dict | None]]: + """Replaces each prototype with the most similar real training segment. + + For every prototype only samples that carry all of its label(s) are + considered. The prototype is set to the feature window with the + highest similarity, so that it corresponds to an actual ECG segment. + In global mode (``proto_time_len=None``), the stored vector is a + time-averaged feature rather than a raw ECG segment. + + The result is also stored on ``self.projection_info`` for the current + session. It is **not** part of ``state_dict``; save it separately + alongside the model checkpoint to support case-based explanations after + reloading (e.g. ``torch.save(model.projection_info, "proj.pt")``). + + Args: + dataloader: Training dataloader (``shuffle=False``) yielding + batches with the signal and the multi-hot labels. + + Returns: + For each branch a list (one entry per prototype) with the + ``sample`` index within the loader order, the ``window`` index and + the ``similarity`` of the chosen segment, or None if no sample + carried the prototype's labels (the prototype is then unchanged). + + Raises: + ValueError: If ``dataloader`` uses a :class:`RandomSampler`. + Sample index tracking requires a deterministic order. + """ + if getattr(dataloader, "shuffle_samples", False) or isinstance( + getattr(dataloader, "sampler", None), RandomSampler + ): + raise ValueError( + "project_prototypes requires a deterministic loader " + "(shuffle=False). Sample index mapping is wrong with a " + "shuffled loader." + ) + was_training = self.training + self.eval() + best = { + name: [None] * self.branches[name].num_prototypes + for name in self.branch_names + } + offset = 0 + for batch in dataloader: + x = batch[self.signal_key].to(self.device).float() + y = batch[self.label_key].to(self.device).float() + for name in self.branch_names: + branch = self.branches[name] + index = self.branch_index(name) + windows = branch.windows(branch.feature_map(x)) # (B, W, N) + sim = branch.similarity_windows(windows) # (B, W, P) + own = branch.prototype_class_identity # (P, L) + needed = own.sum(dim=1) + have = (y[:, index] @ own.t()) >= needed - 1e-6 # (B, P) + for p in range(branch.num_prototypes): + valid = have[:, p].nonzero(as_tuple=True)[0] + if len(valid) == 0: + continue + scores = sim[valid, :, p] # (V, W) + value, flat = scores.flatten().max(dim=0) + row, win = flat // scores.shape[1], flat % scores.shape[1] + current = best[name][p] + if current is None or value.item() > current["similarity"]: + best[name][p] = { + "sample": offset + valid[row].item(), + "window": win.item(), + "similarity": value.item(), + "vector": windows[valid[row], win].clone(), + } + offset += x.shape[0] + for name in self.branch_names: + branch = self.branches[name] + for p, entry in enumerate(best[name]): + if entry is not None: + vector = entry.pop("vector").view( + branch.proto_dim, branch.window + ) + branch.prototype_vectors.data[p] = vector + self.train(was_training) + self.projection_info = best + return best + + @torch.no_grad() + def prototype_contributions( + self, signal: torch.Tensor + ) -> dict[str, torch.Tensor]: + """Per-prototype contributions to each logit, for explanations. + + Args: + signal: ECG batch ``(batch, leads, time)``. + + Returns: + For every branch a tensor equal to similarity times classifier + weight. Shape is ``(batch, num_labels, num_branch_prototypes)`` + when the forward path goes through ``fusion_classifier`` + (``stage in {"all", "fusion"}`` with ``fusion="learned"``, mirroring + :meth:`_uses_fusion_head`; summing over all branches and prototypes + gives ``logit - fusion_classifier.bias`` because the branch + classifiers have no bias but the fusion classifier does), or + ``(batch, num_branch_labels, num_branch_prototypes)`` otherwise + (branch classifier weights, which have no bias so contributions + sum exactly to the branch logits). + """ + was_training = self.training + self.eval() + x = signal.to(self.device).float() + use_fusion = self._uses_fusion_head() + out = {} + offset = 0 + for name in self.branch_names: + branch = self.branches[name] + _, sim, _ = branch(x) + if use_fusion: + w = self.fusion_classifier.weight[:, offset : offset + branch.num_prototypes] + out[name] = sim.unsqueeze(1) * w.unsqueeze(0) + offset += branch.num_prototypes + else: + out[name] = sim.unsqueeze(1) * branch.classifier.weight.unsqueeze(0) + self.train(was_training) + return out diff --git a/pyhealth/tasks/__init__.py b/pyhealth/tasks/__init__.py index a8546d1a6..6d4e33a79 100644 --- a/pyhealth/tasks/__init__.py +++ b/pyhealth/tasks/__init__.py @@ -15,6 +15,7 @@ from .covid19_cxr_generation import ( COVID19CXRImageGeneration as COVID19CXRImageGeneration, ) +from .ptbxl_multilabel_classification import PTBXLMultilabelClassification as PTBXLMultilabelClassification from .deid_ner import DeIDNERTask from .dka import DKAPredictionMIMIC4, T1DDKAPredictionMIMIC4 # New exports must use the redundant `X as X` form: this module has no diff --git a/pyhealth/tasks/ptbxl_multilabel_classification.py b/pyhealth/tasks/ptbxl_multilabel_classification.py new file mode 100644 index 000000000..caa13f4e7 --- /dev/null +++ b/pyhealth/tasks/ptbxl_multilabel_classification.py @@ -0,0 +1,256 @@ +"""Multi-label ECG classification task for the PTB-XL dataset.""" + +import functools +import os +from pathlib import Path +from typing import Any, ClassVar + +import numpy as np +import pandas as pd + +from .base_task import BaseTask + +_LABEL_SETS = ( + "all", + "diagnostic", + "subdiagnostic", + "superdiagnostic", + "form", + "rhythm", +) + + +@functools.cache +def _load_scp_statements(path: str) -> pd.DataFrame: + """Load and cache SCP statements CSV; ``path`` must be the resolved str.""" + return pd.read_csv(path, index_col=0) + + +def _aggregate_labels( + scp_codes: dict[str, float], + scp_statements: pd.DataFrame, + label_set: str, + min_likelihood: float, +) -> list[str]: + """Aggregate SCP codes to the chosen label set with likelihood filtering. + + Generalises :func:`~pyhealth.tasks.ptbxl.aggregate_diagnostic_superclasses` + (which only maps diagnostic codes to the 5 ``diagnostic_class`` superclasses + and ignores the per-code likelihood field) to all six SCP-ECG label sets + (``all``, ``diagnostic``, ``subdiagnostic``, ``superdiagnostic``, ``form``, + ``rhythm``) and applies per-code ``min_likelihood`` filtering. A direct + call-through isn't viable because the public helper doesn't take a label + set or a likelihood threshold, so this reimplementation lives next to the + task. + + For ``"subdiagnostic"`` and ``"superdiagnostic"`` the max likelihood of + all contributing codes is used per parent class (matching the original + PTB-XL paper protocol). + + Args: + scp_codes: ``{code: likelihood}`` from a single ECG record. + scp_statements: Loaded ``scp_statements.csv``, indexed by SCP code. + label_set: One of the six sets in :data:`_LABEL_SETS`. + min_likelihood: Labels whose effective likelihood is below this are + dropped. + + Returns: + Sorted list of unique label strings. + """ + per_class: dict[str, float] = {} + for code, likelihood in scp_codes.items(): + if code not in scp_statements.index: + continue + likelihood = float(likelihood) + row = scp_statements.loc[code] + if label_set == "all" or (label_set == "diagnostic" and row.get("diagnostic") == 1): + per_class[code] = max(per_class.get(code, float("-inf")), likelihood) + elif label_set == "subdiagnostic" and row.get("diagnostic") == 1: + sub = row.get("diagnostic_subclass") + if isinstance(sub, str) and sub: + per_class[sub] = max(per_class.get(sub, float("-inf")), likelihood) + elif label_set == "superdiagnostic" and row.get("diagnostic") == 1: + sup = row.get("diagnostic_class") + if isinstance(sup, str) and sup: + per_class[sup] = max(per_class.get(sup, float("-inf")), likelihood) + elif (label_set == "form" and row.get("form") == 1) or ( + label_set == "rhythm" and row.get("rhythm") == 1 + ): + per_class[code] = max(per_class.get(code, float("-inf")), likelihood) + return sorted(k for k, v in per_class.items() if v >= min_likelihood) + + +class PTBXLMultilabelClassification(BaseTask): + """Multi-label classification of 12-lead ECGs from the PTB-XL dataset. + + Compatible with :class:`~pyhealth.datasets.PTBXLDataset`. Label + aggregation is computed on-the-fly from the raw ``scp_codes`` stored in + each event; ``scp_statements.csv`` is needed to map codes to the chosen + label set. Use ``dataset.scp_statements_path`` to pass it. + + Compared to :class:`~pyhealth.tasks.PTBXLSuperclassClassification`, this + task exposes all six SCP-ECG label sets (diagnostic, subdiagnostic, + superdiagnostic, form, rhythm, all 71 statements) and the full signal + processing pipeline used by ProtoECGNet (optional 0.5 Hz baseline removal, + per-lead z-scoring). + + Args: + scp_statements_path: Path to ``scp_statements.csv`` from the PTB-XL + download. Use ``dataset.scp_statements_path``. + label_set: Which labels to predict. One of ``"all"`` (71 SCP + statements), ``"diagnostic"`` (44 statements), + ``"subdiagnostic"`` (23 sub-classes), ``"superdiagnostic"`` + (5 super-classes, same as :class:`PTBXLSuperclassClassification`), + ``"form"`` or ``"rhythm"``. Defaults to ``"diagnostic"``. + min_likelihood: Keep statements whose annotated likelihood is at + least this value (0 keeps everything, including + unknown-confidence form/rhythm statements). Defaults to 0.0. + standardize: ``"record"`` applies per-lead z-scoring within each + record; ``None`` keeps raw millivolts. Defaults to ``"record"``. + remove_baseline: If True, applies a 0.5 Hz high-pass Butterworth + filter to remove baseline wander (used in the ProtoECGNet paper). + Defaults to False. + + Attributes: + task_name (str): ``"PTBXLMultilabelClassification"``. + input_schema (dict): ``{"signal": "tensor"}``. + output_schema (dict): ``{"labels": "multilabel"}``. + + Note: + The ``sampling_rate`` is read from ``event.sampling_rate`` (set by + :class:`~pyhealth.datasets.PTBXLDataset`), so it follows whatever rate + the dataset was constructed with (100 or 500 Hz). There is no + ``sampling_rate`` constructor argument. + + Examples: + >>> from pyhealth.datasets import PTBXLDataset + >>> from pyhealth.tasks import PTBXLMultilabelClassification + >>> dataset = PTBXLDataset(root="/path/to/ptb-xl/1.0.3") # doctest: +SKIP + >>> task = PTBXLMultilabelClassification( # doctest: +SKIP + ... scp_statements_path=dataset.scp_statements_path, + ... label_set="superdiagnostic", + ... ) + >>> samples = dataset.set_task(task) # doctest: +SKIP + >>> print(samples[0]["signal"].shape) # doctest: +SKIP + torch.Size([12, 1000]) + """ + + task_name: str = "PTBXLMultilabelClassification" + input_schema: ClassVar[dict[str, str]] = {"signal": "tensor"} + output_schema: ClassVar[dict[str, str]] = {"labels": "multilabel"} + + def __init__( + self, + scp_statements_path: str | os.PathLike, + label_set: str = "diagnostic", + min_likelihood: float = 0.0, + standardize: str | None = "record", + remove_baseline: bool = False, + ) -> None: + if label_set not in _LABEL_SETS: + raise ValueError(f"label_set must be one of {_LABEL_SETS}") + if standardize not in ("record", None): + raise ValueError("standardize must be 'record' or None") + # Resolve to a canonical string so the cache key is stable. + self.scp_statements_path = str(Path(scp_statements_path).resolve()) + self.label_set = label_set + self.min_likelihood = min_likelihood + self.standardize = standardize + self.remove_baseline = remove_baseline + super().__init__() + + def _prepare_signal(self, signal_file: str, sampling_rate: int) -> np.ndarray: + """Loads, crops/pads and normalizes one ECG. + + Signal loading goes through + :func:`~pyhealth.datasets.ptbxl.load_ptbxl_record` — the canonical + PTB-XL reader also used by + :class:`~pyhealth.tasks.PTBXLSuperclassClassification` — so there is a + single decoding path. ``wfdb`` is a hard requirement here; install + the ``pyhealth[ptbxl]`` extra. + + Args: + signal_file: Record path without extension (as stored in the event). + sampling_rate: Waveform sampling rate in Hz (100 or 500). + + Returns: + Float32 array of shape ``(12, 10 * sampling_rate)``. + """ + from pyhealth.datasets.ptbxl import load_ptbxl_record + + signal = load_ptbxl_record(signal_file) # (leads, samples) + length = 10 * sampling_rate + if signal.shape[1] < length: + signal = np.pad(signal, ((0, 0), (0, length - signal.shape[1]))) + signal = signal[:, :length] + if self.remove_baseline: + from scipy.signal import butter, filtfilt + + b, a = butter(1, 0.5 / (sampling_rate / 2), btype="high") + signal = filtfilt(b, a, signal, axis=1) + if self.standardize == "record": + mean = signal.mean(axis=1, keepdims=True) + std = signal.std(axis=1, keepdims=True) + signal = (signal - mean) / (std + 1e-6) + return np.ascontiguousarray(signal, dtype=np.float32) + + def __call__(self, patient: Any) -> list[dict[str, Any]]: + """Builds one sample per labelled ECG of a patient. + + For the ``superdiagnostic`` label set with ``min_likelihood == 0`` the + aggregation delegates to + :func:`~pyhealth.tasks.ptbxl.aggregate_diagnostic_superclasses` (used + by :class:`~pyhealth.tasks.PTBXLSuperclassClassification`), so the + labels produced by both tasks are identical for that configuration. + Other label sets or non-zero likelihood thresholds use + :func:`_aggregate_labels`. + + Args: + patient: A patient object with ``records`` events from + :class:`~pyhealth.datasets.PTBXLDataset`. + + Returns: + A list of dicts with ``patient_id``, ``strat_fold``, + ``signal`` (float32 array ``(12, T)``) and ``labels`` (list of + label codes). ECGs with no label in the chosen set are skipped. + """ + from pyhealth.datasets.ptbxl import parse_scp_codes + + if self.label_set == "superdiagnostic" and self.min_likelihood == 0.0: + from pyhealth.tasks.ptbxl import ( + aggregate_diagnostic_superclasses, + load_diagnostic_class_map, + ) + + class_map = load_diagnostic_class_map(self.scp_statements_path) + + def aggregate(scp_codes_raw: Any) -> list[str]: + return aggregate_diagnostic_superclasses(scp_codes_raw, class_map) + else: + statements = _load_scp_statements(self.scp_statements_path) + + def aggregate(scp_codes_raw: Any) -> list[str]: + return _aggregate_labels( + parse_scp_codes(scp_codes_raw), + statements, + self.label_set, + self.min_likelihood, + ) + + samples = [] + for event in patient.get_events(event_type="records"): + labels = aggregate(event.scp_codes) + if not labels: + continue + sampling_rate = int(event.sampling_rate) + samples.append( + { + "patient_id": patient.patient_id, + # int(float(...)) matches PTBXLSuperclassClassification and + # tolerates "9.0" strings from the CSV. + "strat_fold": int(float(event.strat_fold)), + "signal": self._prepare_signal(event.signal_file, sampling_rate), + "labels": labels, + } + ) + return samples diff --git a/tests/core/test_protoecgnet.py b/tests/core/test_protoecgnet.py new file mode 100644 index 000000000..f6eaf2688 --- /dev/null +++ b/tests/core/test_protoecgnet.py @@ -0,0 +1,735 @@ +"""Tests for ProtoECGNet using tiny synthetic ECG tensors (no real data).""" + +import unittest + +import torch + +from pyhealth.datasets import create_sample_dataset, get_dataloader +from pyhealth.models import ( + PTBXL_BRANCH_KWARGS, + PTBXL_BRANCH_LABELS, + ProtoECGNet, + PrototypeBranch, + ResNet1dBackbone, + ResNet2dBackbone, +) + +LABELS = ["A", "B", "C", "D", "E", "F"] +TINY = {"proto_dim": 8, "base_width": 4, "layers": (1, 1, 1, 1)} + + +def make_dataset(num_samples: int = 12, length: int = 192): + """Creates a small multilabel ECG dataset with random 12-lead signals.""" + torch.manual_seed(0) + samples = [] + for i in range(num_samples): + labels = [LABELS[i % 6]] + if i % 3 == 0: + labels.append(LABELS[(i + 2) % 6]) + samples.append( + { + "patient_id": f"p{i}", + "signal": torch.randn(12, length).tolist(), + "labels": labels, + } + ) + return create_sample_dataset( + samples=samples, + input_schema={"signal": "tensor"}, + output_schema={"labels": "multilabel"}, + ) + + +class TestPaperConstants(unittest.TestCase): + """The physician-defined PTB-XL label groups of the paper.""" + + def test_group_sizes_and_disjointness(self): + sizes = {k: len(v) for k, v in PTBXL_BRANCH_LABELS.items()} + self.assertEqual(sizes, {"rhythm": 16, "morphology": 52, "global": 3}) + flat = [c for codes in PTBXL_BRANCH_LABELS.values() for c in codes] + self.assertEqual(len(flat), 71) + self.assertEqual(len(set(flat)), 71) + + def test_branch_kwargs_cover_groups(self): + self.assertEqual(set(PTBXL_BRANCH_KWARGS), set(PTBXL_BRANCH_LABELS)) + self.assertEqual(PTBXL_BRANCH_KWARGS["morphology"]["proto_time_len"], 3) + + def test_branch_kwargs_match_paper(self): + per_label = { + k: v["prototypes_per_label"] for k, v in PTBXL_BRANCH_KWARGS.items() + } + self.assertEqual(per_label, {"rhythm": 5, "morphology": 18, "global": 7}) + self.assertEqual(PTBXL_BRANCH_KWARGS["morphology"]["topk"], 5) + self.assertEqual(PTBXL_BRANCH_KWARGS["global"]["proto_time_len"], 32) + + +class TestPaperShapes(unittest.TestCase): + """Prototype shapes of the paper's branches on a 12 x 1000 ECG.""" + + def test_global_2d_prototype_spans_full_time_axis(self): + branch = PrototypeBranch( + 12, 2, proto_dim=8, prototypes_per_label=1, proto_time_len=32, + backbone="resnet2d", + ) + self.assertEqual(tuple(branch.prototype_vectors.shape[1:]), (8, 32)) + logits, sim, best = branch(torch.randn(1, 12, 1000)) + self.assertEqual(logits.shape, (1, 2)) + self.assertEqual(sim.shape, (1, 2)) + self.assertTrue(torch.all(best == 0)) # proto_time_len=32 spans the full 32-step map for a 1000-sample input, so unfold yields exactly one window position + + +class TestPrototypeTerms(unittest.TestCase): + """Clustering and separation follow Eqs. 4-5 of the paper.""" + + def make_branch(self): + return PrototypeBranch( + 12, 3, proto_dim=4, prototypes_per_label=1, base_width=4, + layers=(1, 1, 1, 1), + ) + + def test_clustering_and_separation_use_per_sample_prototype_sets(self): + branch = self.make_branch() # 3 labels, 1 prototype each → prototype index equals label index + sim = torch.tensor([[5.0, 3.0, 1.0]]) + y = torch.tensor([[1.0, 1.0, 0.0]]) # labels 0 and 1 present + terms = branch.prototype_terms(sim, y) + self.assertAlmostEqual(float(terms["clst"]), -5.0) # max over {0, 1} + # Separation only sees prototype 2; prototypes of the other present + # label must not be penalized. + self.assertAlmostEqual(float(terms["sep"]), 1.0) + + def test_empty_sets_contribute_zero(self): + branch = self.make_branch() + sim = torch.tensor([[2.0, 2.0, 2.0], [2.0, 2.0, 2.0]]) + y = torch.tensor([[1.0, 1.0, 1.0], [0.0, 0.0, 0.0]]) + terms = branch.prototype_terms(sim, y) + self.assertAlmostEqual(float(terms["clst"]), -1.0) # (-2 + 0) / 2 + self.assertAlmostEqual(float(terms["sep"]), 1.0) # (0 + 2) / 2 + + +class TestBackbones(unittest.TestCase): + """Feature-map shapes of the two backbones.""" + + def test_backbone_1d_output_shape(self): + backbone = ResNet1dBackbone(base_width=4, layers=(1, 1, 1, 1)) + out = backbone(torch.randn(2, 12, 256)) + self.assertEqual(out.shape[:2], (2, backbone.out_channels)) + self.assertEqual(out.shape[2], 8) # 256 / (2 * 2 * 2 * 2 * 2) + + def test_backbone_2d_keeps_time_axis(self): + backbone = ResNet2dBackbone() + out = backbone(torch.randn(2, 12, 256)) + self.assertEqual(out.shape, (2, 512, 8)) + + +class TestPrototypeBranch(unittest.TestCase): + """Prototype assignment, similarity types and pooling.""" + + def kwargs(self, **extra): + base = {"proto_dim": 8, "base_width": 4, "layers": (1, 1, 1, 1)} + base.update(extra) + return base + + def test_prototype_assignment_counts(self): + branch = PrototypeBranch( + 12, 4, prototypes_per_label=3, joint_prototypes_per_pair=2, + **self.kwargs(), + ) + self.assertEqual(branch.num_prototypes, 4 * 3 + 6 * 2) + identity = branch.prototype_class_identity + self.assertEqual(identity.shape, (24, 4)) + self.assertTrue(torch.all(identity[:12].sum(dim=1) == 1)) + self.assertTrue(torch.all(identity[12:].sum(dim=1) == 2)) + + def test_last_layer_initialization(self): + branch = PrototypeBranch(12, 3, prototypes_per_label=2, **self.kwargs()) + own = branch.prototype_class_identity.t() + weight = branch.classifier.weight + self.assertTrue(torch.all(weight[own == 1] == 1.0)) + self.assertTrue(torch.all(weight[own == 0] == -0.5)) + + def test_cosine_similarity_is_bounded_by_scale(self): + branch = PrototypeBranch( + 12, 3, prototypes_per_label=2, proto_time_len=2, **self.kwargs() + ) + x = torch.randn(3, 12, 128) + sims = branch.similarity_windows(branch.windows(branch.feature_map(x))) + self.assertLessEqual(float(sims.max()), branch.scale + 1e-4) + self.assertGreaterEqual(float(sims.min()), -branch.scale - 1e-4) + self.assertAlmostEqual(branch.scale, (8 * 2) ** 0.5) + + def test_l2_similarity_is_finite(self): + branch = PrototypeBranch( + 12, 3, prototypes_per_label=1, similarity="l2", **self.kwargs() + ) + logits, sim, window = branch(torch.randn(2, 12, 128)) + self.assertTrue(torch.isfinite(logits).all()) + self.assertTrue(torch.isfinite(sim).all()) + self.assertEqual(window.shape, (2, 3)) + + def test_topk_pooling_is_mean_of_best_windows(self): + torch.manual_seed(1) + branch = PrototypeBranch( + 12, 2, prototypes_per_label=1, proto_time_len=2, topk=2, **self.kwargs() + ).eval() + x = torch.randn(2, 12, 192) + windows = branch.windows(branch.feature_map(x)) + sims = branch.similarity_windows(windows) + expected = sims.topk(2, dim=1).values.mean(dim=1) + _, pooled, _ = branch(x) + torch.testing.assert_close(pooled, expected) + single = PrototypeBranch( + 12, 2, prototypes_per_label=1, proto_time_len=2, topk=1, **self.kwargs() + ) + self.assertEqual(single.topk, 1) + + def test_invalid_arguments(self): + with self.assertRaises(ValueError): + PrototypeBranch(12, 2, similarity="dot", **self.kwargs()) + with self.assertRaises(ValueError): + PrototypeBranch(12, 2, backbone="vit", **self.kwargs()) + with self.assertRaises(ValueError): + PrototypeBranch(12, 2, topk=0, **self.kwargs()) + with self.assertRaises(ValueError): + PrototypeBranch(12, 2, proto_time_len=0, **self.kwargs()) + + def test_2d_branch_forward(self): + branch = PrototypeBranch( + 12, 3, proto_dim=8, prototypes_per_label=2, proto_time_len=3, + backbone="resnet2d", topk=2, + ) + logits, sim, _ = branch(torch.randn(2, 12, 256)) + self.assertEqual(logits.shape, (2, 3)) + self.assertEqual(sim.shape, (2, 6)) + + +class TestProtoECGNet(unittest.TestCase): + """Instantiation, forward, loss, stages, projection and fusion.""" + + @classmethod + def setUpClass(cls): + cls.dataset = make_dataset() + cls.loader = get_dataloader(cls.dataset, batch_size=4, shuffle=False) + cls.batch = next(iter(cls.loader)) + + def test_forward_shapes_and_loss(self): + model = ProtoECGNet( + self.dataset, prototypes_per_label=2, proto_time_len=2, **TINY + ) + out = model(**self.batch) + self.assertEqual(out["logit"].shape, (4, 6)) + self.assertEqual(out["y_prob"].shape, (4, 6)) + self.assertEqual(out["y_true"].shape, (4, 6)) + self.assertEqual(out["similarity"].shape, (4, 12)) + self.assertTrue(torch.isfinite(out["loss"])) + self.assertTrue(((out["y_prob"] >= 0) & (out["y_prob"] <= 1)).all()) + + def test_l2_model_variant_runs(self): + model = ProtoECGNet( + self.dataset, prototypes_per_label=1, similarity="l2", **TINY + ) + self.assertTrue(torch.isfinite(model(**self.batch)["loss"])) + + def test_forward_without_labels_has_no_loss(self): + model = ProtoECGNet(self.dataset, prototypes_per_label=1, **TINY) + out = model(signal=self.batch["signal"]) + self.assertNotIn("loss", out) + + def test_gradients_reach_backbone_and_prototypes(self): + model = ProtoECGNet(self.dataset, prototypes_per_label=2, **TINY) + model(**self.batch)["loss"].backward() + branch = model.branches["all"] + self.assertGreater(float(branch.prototype_vectors.grad.abs().sum()), 0.0) + self.assertTrue(torch.isfinite(branch.backbone.stem[0].weight.grad).all()) + + def test_contrastive_term_activates_with_cooccurrence(self): + model = ProtoECGNet( + self.dataset, prototypes_per_label=2, lam_cnrst=300.0, **TINY + ).eval() + base = model(**self.batch)["loss"].item() + labels = torch.stack([self.dataset[i]["labels"] for i in range(12)]) + _, cooc = ProtoECGNet.compute_label_statistics(labels) + model.set_label_statistics(label_cooccurrence=cooc) + with_cnrst = model(**self.batch)["loss"].item() + self.assertNotEqual(base, with_cnrst) + self.assertTrue(torch.isfinite(torch.tensor(with_cnrst))) + + def test_label_statistics_values(self): + labels = torch.tensor([[1, 1, 0], [1, 0, 0], [0, 1, 1], [1, 1, 0.0]]) + weights, cooc = ProtoECGNet.compute_label_statistics(labels) + counts = torch.tensor([3.0, 3.0, 1.0]) + torch.testing.assert_close(weights, (counts.sum() - counts) / counts) + self.assertAlmostEqual(float(cooc[0, 1]), 2 / 4) # Jaccard(A, B) + self.assertAlmostEqual(float(cooc[0, 0]), 1.0) + + def test_stage_freezing(self): + model = ProtoECGNet(self.dataset, prototypes_per_label=2, **TINY) + branch = model.branches["all"] + model.set_stage("joint") + self.assertTrue(branch.prototype_vectors.requires_grad) + self.assertTrue(branch.backbone.stem[0].weight.requires_grad) + self.assertFalse(branch.classifier.weight.requires_grad) + model.set_stage("classifier") + self.assertTrue(branch.classifier.weight.requires_grad) + self.assertFalse(branch.prototype_vectors.requires_grad) + self.assertFalse(branch.backbone.stem[0].weight.requires_grad) + model.set_stage("prototypes") + self.assertTrue(branch.prototype_vectors.requires_grad) + self.assertFalse(branch.add_on[0].weight.requires_grad) + model.set_stage("all") + self.assertTrue(branch.classifier.weight.requires_grad) + + def test_feature_extractor_stage_trains_head_only_path(self): + model = ProtoECGNet(self.dataset, prototypes_per_label=1, **TINY) + model.set_stage("feature_extractor") + model(**self.batch)["loss"].backward() + branch = model.branches["all"] + self.assertIsNotNone(branch.pretrain_head.weight.grad) + self.assertIsNone(branch.prototype_vectors.grad) + + def test_backbone_bn_stays_in_train_mode_in_feature_extractor_stage(self): + """Backbone BatchNorm must stay in train mode so running stats update. + + Regression: _frozen_modules previously used recurse=False, which + mistakenly included the whole PrototypeBranch (because + prototype_vectors has requires_grad=False), causing train() to call + .eval() on the branch and freeze BN statistics. + """ + import torch.nn as nn + + model = ProtoECGNet(self.dataset, prototypes_per_label=1, **TINY) + model.set_stage("feature_extractor") + model.train() + + branch = model.branches["all"] + bn = next( + ( + m + for m in branch.backbone.modules() + if isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d)) + ), + None, + ) + self.assertIsNotNone(bn, "backbone must contain a BatchNorm layer") + self.assertTrue(bn.training, "BatchNorm must be in train mode") + + mean_before = bn.running_mean.clone() + model(**self.batch) + self.assertFalse( + torch.equal(bn.running_mean, mean_before), + "BatchNorm running_mean must update when backbone is in train mode", + ) + + def test_invalid_stage_and_fusion(self): + model = ProtoECGNet(self.dataset, prototypes_per_label=1, **TINY) + with self.assertRaises(ValueError): + model.set_stage("warmup") + with self.assertRaises(ValueError): + model.set_stage("fusion") # fusion='average' has no fusion head + with self.assertRaises(ValueError): + ProtoECGNet(self.dataset, fusion="max", **TINY) + + def test_label_groups_must_partition(self): + with self.assertRaises(ValueError): + ProtoECGNet(self.dataset, label_groups={"x": ["A", "B"]}, **TINY) + with self.assertRaises(ValueError): + ProtoECGNet( + self.dataset, label_groups={"x": LABELS, "y": ["A"]}, **TINY + ) + + def test_unknown_labels_in_groups_are_dropped(self): + groups = {"r": ["A", "B", "C", "ZZZ"], "m": ["D", "E", "F"], "empty": ["QQ"]} + model = ProtoECGNet(self.dataset, label_groups=groups, **TINY) + self.assertEqual(model.branch_names, ["r", "m"]) + self.assertEqual(model.branches["r"].num_labels, 3) + + def test_two_branches_with_learned_fusion(self): + groups = {"rhythm": ["A", "B", "C"], "morph": ["D", "E", "F"]} + model = ProtoECGNet( + self.dataset, + label_groups=groups, + branch_kwargs={"morph": {"proto_time_len": 3}}, + prototypes_per_label=2, + fusion="learned", + **TINY, + ) + self.assertEqual(model.branches["rhythm"].prototype_vectors.shape[2], 1) + self.assertEqual(model.branches["morph"].prototype_vectors.shape[2], 3) + out = model(**self.batch) + self.assertEqual(out["logit"].shape, (4, 6)) + self.assertEqual(out["similarity"].shape, (4, 12)) + model.set_stage("fusion") + model(**self.batch)["loss"].backward() + self.assertIsNotNone(model.fusion_classifier.weight.grad) + self.assertIsNone(model.branches["morph"].prototype_vectors.grad) + + def test_fusion_head_affects_logits_in_all_and_fusion_stages(self): + """fusion_classifier is used in "all" (gets gradients, train everything) + and "fusion" stages; the branch classifiers are used in the staged + protocol's "joint" / "classifier" / "prototypes". Perturbing the fusion + head's weights must change the logits iff the stage uses it. + """ + groups = {"rhythm": ["A", "B", "C"], "morph": ["D", "E", "F"]} + model = ProtoECGNet( + self.dataset, label_groups=groups, prototypes_per_label=2, + fusion="learned", **TINY, + ).eval() + + def logits(stage: str) -> torch.Tensor: + model.set_stage(stage) + with torch.no_grad(): + return model(**self.batch)["logit"].clone() + + for stage in ("all", "fusion", "joint", "classifier"): + before = logits(stage) + with torch.no_grad(): + model.fusion_classifier.weight.add_(1.0) + after = logits(stage) + with torch.no_grad(): + model.fusion_classifier.weight.sub_(1.0) # restore + changed = not torch.allclose(before, after) + if stage in ("all", "fusion"): + self.assertTrue( + changed, f"fusion head must affect logits in stage {stage}" + ) + else: + self.assertFalse( + changed, + f"fusion head must NOT affect logits in stage {stage}", + ) + + def test_all_stage_trains_fusion_head_with_learned_fusion(self): + """Backprop in "all" with fusion="learned" must give the fusion head + a gradient — otherwise its claim to "train everything" is false. + """ + groups = {"rhythm": ["A", "B", "C"], "morph": ["D", "E", "F"]} + model = ProtoECGNet( + self.dataset, label_groups=groups, prototypes_per_label=2, + fusion="learned", **TINY, + ) + # "all" is the default stage; be explicit for the test's intent. + model.set_stage("all") + model(**self.batch)["loss"].backward() + self.assertIsNotNone(model.fusion_classifier.weight.grad) + self.assertFalse(torch.all(model.fusion_classifier.weight.grad == 0)) + + def test_mixed_1d_and_2d_branches(self): + groups = {"rhythm": ["A", "B", "C"], "morph": ["D", "E", "F"]} + model = ProtoECGNet( + self.dataset, + label_groups=groups, + branch_kwargs={ + "morph": {"backbone": "resnet2d", "proto_time_len": 3, "topk": 2} + }, + prototypes_per_label=1, + proto_dim=8, + base_width=4, + layers=(1, 1, 1, 1), + ) + out = model(**self.batch) + self.assertEqual(out["logit"].shape, (4, 6)) + self.assertTrue(torch.isfinite(out["loss"])) + + def test_fusion_mask_matches_assignment(self): + groups = {"rhythm": ["A", "B", "C"], "morph": ["D", "E", "F"]} + model = ProtoECGNet( + self.dataset, label_groups=groups, prototypes_per_label=2, + fusion="learned", **TINY, + ) + mask = model.fusion_mask + self.assertEqual(mask.shape, (6, 12)) + self.assertTrue(torch.all(mask.sum(dim=0) == 1)) # one label per prototype + self.assertTrue(torch.all(mask.sum(dim=1) == 2)) # 2 prototypes per label + + def test_l1_penalty_only_in_final_stages(self): + torch.manual_seed(0) + plain = ProtoECGNet(self.dataset, prototypes_per_label=2, **TINY).eval() + sparse = ProtoECGNet( + self.dataset, prototypes_per_label=2, lam_l1=1.0, **TINY + ).eval() + sparse.load_state_dict(plain.state_dict()) + for stage in ("joint", "classifier"): + plain.set_stage(stage) + sparse.set_stage(stage) + diff = float(sparse(**self.batch)["loss"] - plain(**self.batch)["loss"]) + if stage == "classifier": + self.assertGreater(diff, 0.0) + else: + self.assertAlmostEqual(diff, 0.0, places=5) + + def test_too_long_prototype_raises(self): + model = ProtoECGNet( + self.dataset, prototypes_per_label=1, proto_time_len=99, **TINY + ) + with self.assertRaises(ValueError): + model(**self.batch) + + def test_projection_snaps_prototypes_to_training_segments(self): + model = ProtoECGNet( + self.dataset, prototypes_per_label=1, proto_time_len=2, **TINY + ) + info = model.project_prototypes(self.loader)["all"] + self.assertTrue(all(entry is not None for entry in info)) + branch = model.branches["all"] + model.eval() + best = [] + with torch.no_grad(): + for batch in self.loader: + windows = branch.windows(branch.feature_map(batch["signal"].float())) + best.append(branch.similarity_windows(windows).amax(dim=(0, 1))) + # After projection every prototype exactly matches a training window, + # i.e. attains the maximal cosine similarity (equal to the scale). + self.assertTrue(torch.all(torch.stack(best).amax(dim=0) > branch.scale - 1e-3)) + + def test_projection_respects_labels(self): + model = ProtoECGNet(self.dataset, prototypes_per_label=1, **TINY) + info = model.project_prototypes(self.loader)["all"] + for label_index, entry in enumerate(info): + sample = self.dataset[entry["sample"]] + self.assertEqual(float(sample["labels"][label_index]), 1.0) + + def test_prototype_contributions_shape(self): + model = ProtoECGNet(self.dataset, prototypes_per_label=2, **TINY) + out = model.prototype_contributions(self.batch["signal"]) + self.assertEqual(out["all"].shape, (4, 6, 12)) + + def test_frozen_bn_stats_unchanged_in_classifier_stage(self): + model = ProtoECGNet(self.dataset, prototypes_per_label=2, **TINY) + model.set_stage("classifier") + branch = model.branches["all"] + bn = next( + m for m in branch.backbone.modules() + if isinstance(m, torch.nn.BatchNorm1d) + ) + mean_before = bn.running_mean.clone() + var_before = bn.running_var.clone() + model.train() + model(**self.batch) + torch.testing.assert_close(bn.running_mean, mean_before) + torch.testing.assert_close(bn.running_var, var_before) + + def test_frozen_bn_stats_unchanged_in_fusion_stage(self): + groups = {"rhythm": ["A", "B", "C"], "morph": ["D", "E", "F"]} + model = ProtoECGNet( + self.dataset, label_groups=groups, prototypes_per_label=2, + fusion="learned", **TINY, + ) + model.set_stage("fusion") + bn = next( + m for m in model.branches["rhythm"].backbone.modules() + if isinstance(m, torch.nn.BatchNorm1d) + ) + mean_before = bn.running_mean.clone() + model.train() + model(**self.batch) + torch.testing.assert_close(bn.running_mean, mean_before) + + def test_prototype_contributions_sum_to_logit_fusion_stage(self): + groups = {"rhythm": ["A", "B", "C"], "morph": ["D", "E", "F"]} + model = ProtoECGNet( + self.dataset, label_groups=groups, prototypes_per_label=2, + fusion="learned", **TINY, + ) + model.set_stage("fusion") + signal = self.batch["signal"] + contribs = model.prototype_contributions(signal) + # Contributions (without bias) must sum to logit - fusion bias. + total = sum(contribs[n].sum(dim=-1) for n in model.branch_names) + with torch.no_grad(): + model.eval() + logit = model(**self.batch)["logit"] + expected = logit - model.fusion_classifier.bias + torch.testing.assert_close(total, expected, atol=1e-4, rtol=1e-4) + + + def test_set_pos_weight_routes_to_class_weights(self): + """ProtoECGNet.set_pos_weight must mirror the parsed weight into + ``class_weights`` with a stable ``(num_labels,)`` shape (regression: + previously a silent no-op, then a scalar left a ``(1,)`` buffer that + broke strict ``load_state_dict``).""" + model = ProtoECGNet(self.dataset, prototypes_per_label=1, **TINY) + n = model.num_labels + + # Scalar: broadcast to (num_labels,) filled with the scalar. + model.set_pos_weight(4.0) + self.assertEqual(model.class_weights.shape, (n,)) + torch.testing.assert_close( + model.class_weights.detach().cpu(), torch.full((n,), 4.0) + ) + + # Per-label sequence: length-num_labels vector. + weights = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + model.set_pos_weight(weights) + torch.testing.assert_close( + model.class_weights.detach().cpu(), torch.tensor(weights) + ) + self.assertEqual(model.class_weights.shape, (n,)) + + # Wrong length (not 1, not num_labels): clear error. + with self.assertRaises(ValueError): + model.set_pos_weight([1.0, 2.0]) + + # None resets to the ones default (identity for BCE pos_weight). + model.set_pos_weight(None) + torch.testing.assert_close( + model.class_weights.detach().cpu(), torch.ones(n) + ) + self.assertEqual(model.class_weights.shape, (n,)) + + # state_dict round-trip after a scalar set_pos_weight must work with + # strict loading — the (num_labels,) shape must survive the save/load. + model.set_pos_weight(2.5) + clone = ProtoECGNet(self.dataset, prototypes_per_label=1, **TINY) + clone.load_state_dict(model.state_dict()) + torch.testing.assert_close(clone.class_weights, model.class_weights) + + def test_state_dict_round_trip_preserves_fusion_stage(self): + """A learned-fusion model saved in the ``fusion`` stage must reload in + that same stage, so forward keeps using ``fusion_classifier`` instead of + falling back to the branch-scatter path. + + Regression: ``self.stage`` was a plain attribute (not in ``state_dict``), + so a reloaded model silently resumed in ``"all"`` and returned the + wrong logits. + """ + torch.manual_seed(0) + groups = {"rhythm": ["A", "B", "C"], "morph": ["D", "E", "F"]} + model = ProtoECGNet( + self.dataset, label_groups=groups, prototypes_per_label=2, + fusion="learned", **TINY, + ) + model.set_stage("fusion") + with torch.no_grad(): + for p in model.parameters(): + if p.requires_grad: + p.add_(0.1 * torch.randn_like(p)) + model.eval() + with torch.no_grad(): + expected = model(**self.batch)["logit"] + expected_contribs = model.prototype_contributions(self.batch["signal"]) + + clone = ProtoECGNet( + self.dataset, label_groups=groups, prototypes_per_label=2, + fusion="learned", **TINY, + ).eval() + # Default stage on a freshly constructed model. + self.assertEqual(clone.stage, "all") + clone.load_state_dict(model.state_dict()) + # Extra state restored the fusion stage. + self.assertEqual(clone.stage, "fusion") + with torch.no_grad(): + got = clone(**self.batch)["logit"] + got_contribs = clone.prototype_contributions(self.batch["signal"]) + torch.testing.assert_close(got, expected, atol=1e-6, rtol=1e-6) + for name in model.branch_names: + torch.testing.assert_close( + got_contribs[name], expected_contribs[name], atol=1e-6, rtol=1e-6 + ) + + def test_set_pos_weight_balanced_uses_dataset(self): + """``set_pos_weight('balanced', dataset)`` fills class_weights from the + dataset's label distribution; checks the string / dataset path end-to-end. + """ + model = ProtoECGNet(self.dataset, prototypes_per_label=1, **TINY) + model.set_pos_weight("balanced", self.dataset) + self.assertEqual(model.class_weights.shape, (model.num_labels,)) + self.assertTrue(torch.isfinite(model.class_weights).all()) + self.assertTrue((model.class_weights > 0).all()) + # 'balanced' needs a dataset; otherwise the base class raises. + with self.assertRaises(ValueError): + model.set_pos_weight("balanced") + + def test_project_prototypes_raises_on_shuffled_loader(self): + model = ProtoECGNet(self.dataset, prototypes_per_label=1, **TINY) + shuffled = get_dataloader(self.dataset, batch_size=4, shuffle=True) + with self.assertRaises(ValueError, msg="shuffled loader must raise"): + model.project_prototypes(shuffled) + + def test_trainer_runs_through_staged_protocol(self): + """One epoch per stage through the PyHealth Trainer, with a projection + between the joint and classifier stages. Catches regressions where the + trainer loop (optimizer / loss / stage switch) breaks even though the + unit tests of individual stages still pass. + """ + from pyhealth.trainer import Trainer + + model = ProtoECGNet(self.dataset, prototypes_per_label=1, **TINY) + loader = get_dataloader(self.dataset, batch_size=4, shuffle=False) + proj_loader = get_dataloader(self.dataset, batch_size=4, shuffle=False) + trainer = Trainer(model=model, metrics=[], enable_logging=False) + + for stage in ("feature_extractor", "joint"): + model.set_stage(stage) + trainer.train( + train_dataloader=loader, + val_dataloader=loader, + epochs=1, + monitor="loss", + monitor_criterion="min", + optimizer_params={"lr": 1e-3}, + ) + + # Projection between joint and classifier mirrors the example's loop. + model.project_prototypes(proj_loader) + + model.set_stage("classifier") + trainer.train( + train_dataloader=loader, + val_dataloader=loader, + epochs=1, + monitor="loss", + monitor_criterion="min", + optimizer_params={"lr": 1e-3}, + ) + + # Final inference must still produce finite logits for all labels. + y_true, y_prob, loss = trainer.inference(loader)[:3] + self.assertEqual(y_prob.shape[1], model.num_labels) + self.assertTrue(torch.isfinite(torch.as_tensor(y_prob)).all()) + self.assertTrue(torch.isfinite(torch.as_tensor(loss))) + + def test_state_dict_round_trip_preserves_forward_output(self): + """state_dict → load_state_dict must reproduce forward outputs exactly. + + Covers prototypes, classifier weights, fusion weights and the buffers + populated by ``set_label_statistics`` (``class_weights``, + ``label_cooccurrence``). ``class_weights`` and ``label_cooccurrence`` + were previously registered as ``None`` buffers and promoted to tensors + only when supplied; strict ``load_state_dict`` then failed because the + saved checkpoint carried keys the fresh model lacked. This test guards + against that regression. + """ + torch.manual_seed(0) + groups = {"rhythm": ["A", "B", "C"], "morph": ["D", "E", "F"]} + model = ProtoECGNet( + self.dataset, label_groups=groups, prototypes_per_label=2, + fusion="learned", **TINY, + ) + # Set non-trivial label statistics so the buffers carry real tensors + # and the contrastive term is active in forward. + weights, cooc = ProtoECGNet.compute_label_statistics( + torch.stack([self.dataset[i]["labels"] for i in range(12)]) + ) + model.set_label_statistics(class_weights=weights, label_cooccurrence=cooc) + # Nudge weights off their initializer so random re-init would be caught. + with torch.no_grad(): + for p in model.parameters(): + p.add_(0.1 * torch.randn_like(p)) + model.eval() + with torch.no_grad(): + expected = model(**self.batch)["logit"] + + clone = ProtoECGNet( + self.dataset, label_groups=groups, prototypes_per_label=2, + fusion="learned", **TINY, + ).eval() + # strict=True by default; fails if keys mismatch. + clone.load_state_dict(model.state_dict()) + with torch.no_grad(): + got = clone(**self.batch)["logit"] + torch.testing.assert_close(got, expected, atol=1e-6, rtol=1e-6) + torch.testing.assert_close(clone.class_weights, model.class_weights) + torch.testing.assert_close(clone.label_cooccurrence, model.label_cooccurrence) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/core/test_ptbxl.py b/tests/core/test_ptbxl.py index 693ac7cfa..dea5c8645 100644 --- a/tests/core/test_ptbxl.py +++ b/tests/core/test_ptbxl.py @@ -717,5 +717,262 @@ def _guard(name, *args, **kwargs): self.assertIn("pyhealth[ptbxl]", str(ctx.exception)) +try: + import wfdb as _wfdb # noqa: F401 + _WFDB_AVAILABLE = True +except ImportError: + _WFDB_AVAILABLE = False + + +@unittest.skipUnless( + _WFDB_AVAILABLE, + "wfdb not installed; install the 'pyhealth[ptbxl]' extra to run these tests", +) +class TestPTBXLMultilabelClassification(unittest.TestCase): + """Tests for PTBXLMultilabelClassification wired to the master dataset.""" + + @classmethod + def setUpClass(cls): + cls.tmp = tempfile.mkdtemp() + cls.root = Path(cls.tmp) / "ptbxl" + cls.root.mkdir() + shutil.copy(FIXTURE_ROOT / "ptbxl_database.csv", cls.root / "ptbxl_database.csv") + shutil.copy(FIXTURE_ROOT / "scp_statements.csv", cls.root / "scp_statements.csv") + for row in pd.read_csv(cls.root / "ptbxl_database.csv").itertuples(): + _write_dummy_wfdb(cls.root / str(row.filename_lr), n_samples=1000) + _write_dummy_wfdb(cls.root / str(row.filename_hr), n_samples=5000, fs=500) + cls.cache = Path(cls.tmp) / "cache" + from pyhealth.datasets.ptbxl import PTBXLDataset + cls.dataset = PTBXLDataset(root=str(cls.root), cache_dir=str(cls.cache)) + cls.scp_path = cls.dataset.scp_statements_path + + @classmethod + def tearDownClass(cls): + shutil.rmtree(cls.tmp, ignore_errors=True) + + def test_superdiagnostic_samples_shape(self): + from pyhealth.tasks import PTBXLMultilabelClassification + task = PTBXLMultilabelClassification( + scp_statements_path=self.scp_path, label_set="superdiagnostic" + ) + samples = self.dataset.set_task(task) + self.assertGreater(len(samples), 0) + item = samples[0] + self.assertEqual(tuple(item["signal"].shape), (12, 1000)) + + def test_rhythm_skips_non_rhythm_ecgs(self): + from pyhealth.tasks import PTBXLMultilabelClassification + task = PTBXLMultilabelClassification( + scp_statements_path=self.scp_path, label_set="rhythm" + ) + samples = self.dataset.set_task(task) + self.assertGreater(len(samples), 0) + for item in samples: + self.assertTrue(len(item["labels"]) > 0) + + def test_invalid_label_set_raises(self): + from pyhealth.tasks import PTBXLMultilabelClassification + with self.assertRaises(ValueError): + PTBXLMultilabelClassification( + scp_statements_path=self.scp_path, label_set="nope" + ) + + def test_invalid_standardize_raises(self): + from pyhealth.tasks import PTBXLMultilabelClassification + with self.assertRaises(ValueError): + PTBXLMultilabelClassification( + scp_statements_path=self.scp_path, standardize="dataset" + ) + + def test_prepare_signal_shape_and_dtype(self): + from pyhealth.tasks import PTBXLMultilabelClassification + task = PTBXLMultilabelClassification( + scp_statements_path=self.scp_path, remove_baseline=True + ) + rec = str(self.root / "records100" / "00000" / "00001_lr") + out = task._prepare_signal(rec, 100) + self.assertEqual(out.shape, (12, 1000)) + self.assertEqual(out.dtype, np.float32) + self.assertTrue(np.isfinite(out).all()) + + def test_prepare_signal_500hz(self): + from pyhealth.tasks import PTBXLMultilabelClassification + task = PTBXLMultilabelClassification(scp_statements_path=self.scp_path) + rec = str(self.root / "records500" / "00000" / "00001_hr") + out = task._prepare_signal(rec, 500) + self.assertEqual(out.shape, (12, 5000)) + + def test_aggregate_labels_likelihood_filter(self): + from pyhealth.tasks.ptbxl_multilabel_classification import ( + _aggregate_labels, + _load_scp_statements, + ) + stmts = _load_scp_statements(str(self.scp_path)) + labels = _aggregate_labels( + {"IMI": 80.0, "NORM": 100.0}, stmts, "superdiagnostic", 0.0 + ) + self.assertIn("MI", labels) + self.assertIn("NORM", labels) + labels_hi = _aggregate_labels( + {"IMI": 80.0, "NORM": 50.0}, stmts, "superdiagnostic", 75.0 + ) + self.assertIn("MI", labels_hi) + self.assertNotIn("NORM", labels_hi) + + def test_all_label_set_includes_rhythm(self): + from pyhealth.tasks.ptbxl_multilabel_classification import ( + _aggregate_labels, + _load_scp_statements, + ) + stmts = _load_scp_statements(str(self.scp_path)) + # IMI is diagnostic; SR is rhythm — both appear under "all" + labels = _aggregate_labels({"IMI": 80.0, "SR": 100.0}, stmts, "all", 0.0) + self.assertIn("IMI", labels) + self.assertIn("SR", labels) + + def test_diagnostic_label_set_excludes_rhythm(self): + from pyhealth.tasks.ptbxl_multilabel_classification import ( + _aggregate_labels, + _load_scp_statements, + ) + stmts = _load_scp_statements(str(self.scp_path)) + labels = _aggregate_labels({"IMI": 80.0, "SR": 100.0}, stmts, "diagnostic", 0.0) + self.assertIn("IMI", labels) + self.assertNotIn("SR", labels) + + def test_subdiagnostic_label_set_maps_to_subclass(self): + from pyhealth.tasks.ptbxl_multilabel_classification import ( + _aggregate_labels, + _load_scp_statements, + ) + stmts = _load_scp_statements(str(self.scp_path)) + # IMI → diagnostic_subclass=IMI, LVH → diagnostic_subclass=LVH + labels = _aggregate_labels( + {"IMI": 100.0, "LVH": 80.0}, stmts, "subdiagnostic", 0.0 + ) + self.assertIn("IMI", labels) + self.assertIn("LVH", labels) + # SR has no diagnostic==1, so should not appear + labels_rhythm = _aggregate_labels({"SR": 100.0}, stmts, "subdiagnostic", 0.0) + self.assertEqual(labels_rhythm, []) + + def test_form_label_set_empty_in_fixture(self): + from pyhealth.tasks.ptbxl_multilabel_classification import ( + _aggregate_labels, + _load_scp_statements, + ) + stmts = _load_scp_statements(str(self.scp_path)) + # Fixture scp_statements has no form==1 codes + labels = _aggregate_labels({"NORM": 100.0, "IMI": 80.0}, stmts, "form", 0.0) + self.assertEqual(labels, []) + + def test_superdiagnostic_matches_superclass_task(self): + """When label_set='superdiagnostic' and min_likelihood=0, labels match + PTBXLSuperclassClassification record-for-record. + + Guards the Blocker-1 refactor that delegates the overlapping path to + aggregate_diagnostic_superclasses. + """ + from pyhealth.tasks import ( + PTBXLMultilabelClassification, + PTBXLSuperclassClassification, + ) + multi = PTBXLMultilabelClassification( + scp_statements_path=self.scp_path, label_set="superdiagnostic" + ) + superclass = PTBXLSuperclassClassification( + scp_statements_path=self.scp_path, drop_empty_labels=True + ) + multi_samples = { + (s["patient_id"], s["strat_fold"]): sorted(s["labels"]) + for s in self.dataset.set_task(multi) + } + # Build a fresh dataset for the superclass task so the two caches don't collide. + from pyhealth.datasets.ptbxl import PTBXLDataset + other_dataset = PTBXLDataset( + root=str(self.root), cache_dir=str(Path(self.tmp) / "cache_sup") + ) + super_samples = { + (s["patient_id"], s["strat_fold"]): sorted(s["labels"]) + for s in other_dataset.set_task(superclass) + } + self.assertEqual(set(multi_samples.keys()), set(super_samples.keys())) + for key, labels in multi_samples.items(): + self.assertEqual(labels, super_samples[key]) + + def test_diagnostic_samples_labels_nonempty(self): + from pyhealth.tasks import PTBXLMultilabelClassification + task = PTBXLMultilabelClassification( + scp_statements_path=self.scp_path, label_set="diagnostic" + ) + samples = self.dataset.set_task(task) + self.assertGreater(len(samples), 0) + for item in samples: + self.assertGreater(len(item["labels"]), 0) + self.assertIn("strat_fold", item) + + def test_standardize_record_versus_none(self): + """'record' forces per-lead mean≈0 and std≈1; None does not.""" + from pyhealth.tasks import PTBXLMultilabelClassification + rec = str(self.root / "records100" / "00000" / "00004_lr") + + task_std = PTBXLMultilabelClassification(scp_statements_path=self.scp_path) + out_std = task_std._prepare_signal(rec, 100) + np.testing.assert_allclose(out_std.mean(axis=1), 0.0, atol=1e-5) + np.testing.assert_allclose(out_std.std(axis=1), 1.0, atol=1e-3) + + task_raw = PTBXLMultilabelClassification( + scp_statements_path=self.scp_path, standardize=None + ) + out_raw = task_raw._prepare_signal(rec, 100) + self.assertEqual(out_raw.shape, (12, 1000)) + # The fixture signals are not centered, so standardize=None keeps a + # per-lead mean clearly away from zero for at least one lead. + self.assertTrue((np.abs(out_raw.mean(axis=1)) > 0.01).any()) + + def test_remove_baseline_changes_signal(self): + """remove_baseline=True must perturb the output vs. remove_baseline=False.""" + from pyhealth.tasks import PTBXLMultilabelClassification + rec = str(self.root / "records100" / "00000" / "00004_lr") + + task_no_bl = PTBXLMultilabelClassification( + scp_statements_path=self.scp_path, + standardize=None, + remove_baseline=False, + ) + task_bl = PTBXLMultilabelClassification( + scp_statements_path=self.scp_path, + standardize=None, + remove_baseline=True, + ) + out_no_bl = task_no_bl._prepare_signal(rec, 100) + out_bl = task_bl._prepare_signal(rec, 100) + self.assertEqual(out_bl.shape, out_no_bl.shape) + self.assertFalse(np.allclose(out_no_bl, out_bl)) + # A high-pass filter removes DC; so the per-lead mean of the filtered + # output is closer to zero than the unfiltered one. + self.assertLess( + float(np.abs(out_bl.mean(axis=1)).max()), + float(np.abs(out_no_bl.mean(axis=1)).max()), + ) + + def test_load_signal_matches_superclass_task_signal(self): + """Both tasks must load the exact same signal tensor from the same record. + + The multi-label task now uses load_ptbxl_record directly (same as the + superclass task), so there is one decoding path and the signals match. + """ + from pyhealth.datasets.ptbxl import load_ptbxl_record + from pyhealth.tasks import PTBXLMultilabelClassification + + task = PTBXLMultilabelClassification( + scp_statements_path=self.scp_path, standardize=None, + ) + rec = str(self.root / "records100" / "00000" / "00001_lr") + prepped = task._prepare_signal(rec, 100) + canonical = load_ptbxl_record(rec)[:, :1000] + np.testing.assert_allclose(prepped, canonical, atol=1e-5) + + if __name__ == "__main__": unittest.main()