From 621377f07cba5a76fa53c787f335bd9522d162b9 Mon Sep 17 00:00:00 2001 From: Irvi Aini Date: Wed, 21 Oct 2026 09:49:13 +0800 Subject: [PATCH 01/14] feat: Add PTB-XL dataset and ProtoECGNet model for multi-label ECG classification MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- Summary - PTBXLDataset — wraps the open-access PTB-XL 12-lead ECG dataset (21,799 records, 18,869 patients). Reads ptbxl_database.csv + scp_statements.csv, aggregates SCP codes into six label sets (all, diagnostic, subdiagnostic, superdiagnostic, form, rhythm) and writes a derived metadata CSV so downstream tasks can load labels without re-parsing the raw files. - PTBXLMultilabelClassification — task that converts each labelled ECG into a (12, 10 × fs) float32 signal with multi-hot labels. Supports 100 Hz and 500 Hz, likelihood filtering, per-record z-score normalisation, and optional 0.5 Hz high-pass baseline removal. Ships a fallback WFDB reader so wfdb is not a hard dependency. - ProtoECGNet — re-implementation of Sethi et al., MLHC 2025 (https://arxiv.org/abs/2504.08713). A 1D or 2D ResNet extracts features; each prototype branch computes cosine or L2 similarities between feature windows and learned prototype vectors assigned to labels; a final linear layer produces multi-label logits. Prototypes can be projected onto real ECG segments after training, making every prediction explainable by example cases. Supports multi-branch setups (the paper's rhythm / morphology / global split), learned fusion, contrastive and diversity regularisation, and the paper's four-stage training protocol. - Docs — RST stubs wired into the existing API reference pages. - Example script — end-to-end pipeline with ablation study (contrastive loss on/off × cosine/L2 × prototype count); runs on synthetic data with --synthetic (no download needed). Signed-off-by: Irvi Aini --- docs/api/datasets.rst | 1 + .../pyhealth.datasets.PTBXLDataset.rst | 11 + docs/api/models.rst | 1 + .../models/pyhealth.models.ProtoECGNet.rst | 26 + docs/api/tasks.rst | 1 + ...th.tasks.PTBXLMultilabelClassification.rst | 10 + ...l_multilabel_classification_protoecgnet.py | 337 +++++++ pyhealth/datasets/__init__.py | 1 + pyhealth/datasets/configs/ptbxl.yaml | 20 + pyhealth/datasets/ptbxl.py | 193 ++++ pyhealth/models/__init__.py | 8 + pyhealth/models/protoecgnet.py | 900 ++++++++++++++++++ pyhealth/tasks/__init__.py | 1 + .../tasks/ptbxl_multilabel_classification.py | 226 +++++ tests/core/test_protoecgnet.py | 414 ++++++++ tests/core/test_ptbxl.py | 174 ++++ 16 files changed, 2324 insertions(+) create mode 100644 docs/api/datasets/pyhealth.datasets.PTBXLDataset.rst create mode 100644 docs/api/models/pyhealth.models.ProtoECGNet.rst create mode 100644 docs/api/tasks/pyhealth.tasks.PTBXLMultilabelClassification.rst create mode 100644 examples/ptbxl_multilabel_classification_protoecgnet.py create mode 100644 pyhealth/datasets/configs/ptbxl.yaml create mode 100644 pyhealth/datasets/ptbxl.py create mode 100644 pyhealth/models/protoecgnet.py create mode 100644 pyhealth/tasks/ptbxl_multilabel_classification.py create mode 100644 tests/core/test_protoecgnet.py create mode 100644 tests/core/test_ptbxl.py diff --git a/docs/api/datasets.rst b/docs/api/datasets.rst index e382aa497..7ed84b7a9 100644 --- a/docs/api/datasets.rst +++ b/docs/api/datasets.rst @@ -297,6 +297,7 @@ Available Datasets datasets/pyhealth.datasets.BMDHSDataset datasets/pyhealth.datasets.COVID19CXRDataset datasets/pyhealth.datasets.ChestXray14Dataset + datasets/pyhealth.datasets.PTBXLDataset datasets/pyhealth.datasets.PhysioNetDeIDDataset datasets/pyhealth.datasets.EEGBCIDataset datasets/pyhealth.datasets.TUABDataset diff --git a/docs/api/datasets/pyhealth.datasets.PTBXLDataset.rst b/docs/api/datasets/pyhealth.datasets.PTBXLDataset.rst new file mode 100644 index 000000000..3f49ecd37 --- /dev/null +++ b/docs/api/datasets/pyhealth.datasets.PTBXLDataset.rst @@ -0,0 +1,11 @@ +pyhealth.datasets.PTBXLDataset +=================================== + +PTB-XL, a large publicly available 12-lead ECG dataset with 21,799 records from +18,869 patients. Open access on PhysioNet: +`https://physionet.org/content/ptb-xl/1.0.3/ `_. + +.. autoclass:: pyhealth.datasets.PTBXLDataset + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/models.rst b/docs/api/models.rst index e636b6c48..b83d9bf18 100644 --- a/docs/api/models.rst +++ b/docs/api/models.rst @@ -253,3 +253,4 @@ API Reference models/pyhealth.models.unified_multimodal_embedding_docs models/pyhealth.models.califorest models/pyhealth.models.XGBoostModel + 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 b250cd271..d10f92758 100644 --- a/docs/api/tasks.rst +++ b/docs/api/tasks.rst @@ -214,6 +214,7 @@ Available Tasks MIMIC-III ICD-9 Coding Cardiology Detection COVID-19 CXR Classification + 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/ptbxl_multilabel_classification_protoecgnet.py b/examples/ptbxl_multilabel_classification_protoecgnet.py new file mode 100644 index 000000000..12b5d4eb5 --- /dev/null +++ b/examples/ptbxl_multilabel_classification_protoecgnet.py @@ -0,0 +1,337 @@ +"""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/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/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/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 tempfile + +import numpy as np +import pandas as pd +import torch + +from pyhealth.datasets import PTBXLDataset, get_dataloader +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 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. + + 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``. + """ + 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=["roc_auc_macro"], + 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="roc_auc_macro", + 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)) + return trainer.evaluate(loaders[2]) + + +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 of final layers") + 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("--ablation", action="store_true") + parser.add_argument("--output", default="protoecgnet_ablation.md") + args = parser.parse_args() + args.device = "cuda" if torch.cuda.is_available() else "cpu" + + scratch = tempfile.mkdtemp() + if args.synthetic: + args.root = os.path.join(scratch, "ptbxl") + os.makedirs(args.root) + build_synthetic_ptbxl(args.root) + 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 + + dataset = PTBXLDataset( + root=args.root, cache_dir=os.path.join(scratch, "cache"), dev=args.dev + ) + task = PTBXLMultilabelClassification( + label_set=args.label_set, remove_baseline=args.remove_baseline + ) + samples = dataset.set_task(task) + splits = split_by_fold(samples, args.root) + 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) + + +if __name__ == "__main__": + main() diff --git a/pyhealth/datasets/__init__.py b/pyhealth/datasets/__init__.py index af55a227a..c664d42cc 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 diff --git a/pyhealth/datasets/configs/ptbxl.yaml b/pyhealth/datasets/configs/ptbxl.yaml new file mode 100644 index 000000000..e36650dac --- /dev/null +++ b/pyhealth/datasets/configs/ptbxl.yaml @@ -0,0 +1,20 @@ +version: "1.0" +tables: + ptbxl: + file_path: "ptbxl-metadata-pyhealth.csv" + patient_id: "patient_id" + timestamp: null + attributes: + - "ecg_id" + - "age" + - "sex" + - "strat_fold" + - "record_path_lr" + - "record_path_hr" + - "scp_codes" + - "labels_all" + - "labels_diagnostic" + - "labels_subdiagnostic" + - "labels_superdiagnostic" + - "labels_form" + - "labels_rhythm" diff --git a/pyhealth/datasets/ptbxl.py b/pyhealth/datasets/ptbxl.py new file mode 100644 index 000000000..e47def5c2 --- /dev/null +++ b/pyhealth/datasets/ptbxl.py @@ -0,0 +1,193 @@ +"""PTB-XL 12-lead ECG dataset for PyHealth.""" + +import ast +import logging +import os +from pathlib import Path + +import pandas as pd + +from ..tasks import PTBXLMultilabelClassification +from .base_dataset import BaseDataset + +logger = logging.getLogger(__name__) + +# Label sets that can be derived from ``scp_statements.csv``. +LABEL_SETS = ( + "all", + "diagnostic", + "subdiagnostic", + "superdiagnostic", + "form", + "rhythm", +) + + +class PTBXLDataset(BaseDataset): + """Base ECG dataset for PTB-XL. + + PTB-XL is a large, publicly available collection of 21,799 clinical + 12-lead ECGs (10 seconds each) from 18,869 patients, annotated with + 71 SCP-ECG statements covering diagnostic, form and rhythm categories. + Records are provided at 100 Hz and 500 Hz. The dataset is open access + (no credentialing) and can be downloaded from + https://physionet.org/content/ptb-xl/1.0.3/ . + + Each ECG record is stored as one event of type ``"ptbxl"``. Several ECGs + may belong to the same patient. The ``strat_fold`` attribute contains the + recommended 10-fold split (folds 1-8 train, 9 validation, 10 test) which is + patient-wise, so it can be used directly for leakage-free splits. + + Citations: + Wagner, P., Strodthoff, N., Bousseljot, R.-D., et al. (2020). + PTB-XL, a large publicly available electrocardiography dataset. + Scientific Data, 7, 154. + + Args: + root: Root directory of the raw data. It must contain + ``ptbxl_database.csv``, ``scp_statements.csv`` and the + ``records100`` / ``records500`` folders. + dataset_name: Optional name of the dataset. Defaults to ``"ptbxl"``. + config_path: Optional path to a YAML config. Defaults to the config + shipped in ``pyhealth/datasets/configs/ptbxl.yaml``. + cache_dir: Optional directory for caching processed data. + num_workers: Number of parallel workers for data processing. + dev: If True, only a small subset of patients is loaded. + + Attributes: + root: Root directory of the raw data. + dataset_name: Name of the dataset. + config_path: Path to the configuration file. + + Examples: + >>> from pyhealth.datasets import PTBXLDataset + >>> dataset = PTBXLDataset(root="/path/to/ptb-xl/1.0.3") + >>> dataset.stats() + >>> samples = dataset.set_task() # 12-lead, 100 Hz, diagnostic labels + >>> print(samples[0]["signal"].shape) + torch.Size([12, 1000]) + """ + + def __init__( + self, + root: str, + dataset_name: str | None = None, + config_path: str | None = None, + cache_dir: str | None = None, + num_workers: int = 1, + dev: bool = False, + ) -> None: + if config_path is None: + logger.info("No config path provided, using default config") + config_path = Path(__file__).parent / "configs" / "ptbxl.yaml" + metadata_path = os.path.join(root, "ptbxl-metadata-pyhealth.csv") + if not os.path.exists(metadata_path): + self.prepare_metadata(root) + super().__init__( + root=root, + tables=["ptbxl"], + dataset_name=dataset_name or "ptbxl", + config_path=config_path, + cache_dir=cache_dir, + num_workers=num_workers, + dev=dev, + ) + + @staticmethod + def _aggregate_labels( + scp_codes: dict[str, float], + scp_statements: pd.DataFrame, + ) -> dict[str, str]: + """Builds the per-label-set label strings of one ECG. + + Args: + scp_codes: Mapping from SCP code to likelihood for one ECG. + scp_statements: The ``scp_statements.csv`` table indexed by code. + + Returns: + A dict mapping each label set name to a string of the form + ``"CODE:likelihood;CODE:likelihood"``. Aggregated sets + (super/sub diagnostic) use the maximum likelihood of the members. + Empty label sets map to the empty string. + """ + per_set: dict[str, dict[str, float]] = {name: {} for name in LABEL_SETS} + for code, likelihood in scp_codes.items(): + if code not in scp_statements.index: + continue + row = scp_statements.loc[code] + likelihood = float(likelihood) + per_set["all"][code] = likelihood + if row.get("diagnostic") == 1: + per_set["diagnostic"][code] = likelihood + for key, column in ( + ("subdiagnostic", "diagnostic_subclass"), + ("superdiagnostic", "diagnostic_class"), + ): + parent = row.get(column) + if isinstance(parent, str) and parent: + previous = per_set[key].get(parent, float("-inf")) + per_set[key][parent] = max(previous, likelihood) + if row.get("form") == 1: + per_set["form"][code] = likelihood + if row.get("rhythm") == 1: + per_set["rhythm"][code] = likelihood + return { + name: ";".join(f"{c}:{v}" for c, v in sorted(labels.items())) + for name, labels in per_set.items() + } + + def prepare_metadata(self, root: str) -> None: + """Creates ``ptbxl-metadata-pyhealth.csv`` from the raw PTB-XL files. + + The CSV has one row per ECG with absolute record paths (without file + extension), the recommended fold, and label strings for every label + set (see :data:`LABEL_SETS`). + + Args: + root: Root directory containing the raw PTB-XL files. + + Raises: + FileNotFoundError: If ``ptbxl_database.csv`` or + ``scp_statements.csv`` is missing from ``root``. + """ + database_path = os.path.join(root, "ptbxl_database.csv") + statements_path = os.path.join(root, "scp_statements.csv") + for path in (database_path, statements_path): + if not os.path.isfile(path): + raise FileNotFoundError( + f"{path} not found. Download PTB-XL 1.0.3 from " + "https://physionet.org/content/ptb-xl/1.0.3/ and pass the " + "extracted folder as `root`." + ) + database = pd.read_csv(database_path, index_col="ecg_id") + statements = pd.read_csv(statements_path, index_col=0) + + rows: list[dict[str, object]] = [] + for ecg_id, record in database.iterrows(): + scp_codes = ast.literal_eval(record["scp_codes"]) + labels = self._aggregate_labels(scp_codes, statements) + row: dict[str, object] = { + "patient_id": str(int(record["patient_id"])), + "ecg_id": int(ecg_id), + "age": record.get("age"), + "sex": record.get("sex"), + "strat_fold": int(record["strat_fold"]), + "record_path_lr": os.path.join(root, record["filename_lr"]), + "record_path_hr": os.path.join(root, record["filename_hr"]), + "scp_codes": record["scp_codes"], + } + row.update({f"labels_{name}": text for name, text in labels.items()}) + rows.append(row) + pd.DataFrame(rows).to_csv( + os.path.join(root, "ptbxl-metadata-pyhealth.csv"), index=False + ) + + @property + def default_task(self) -> PTBXLMultilabelClassification: + """Returns the default task for this dataset. + + Returns: + PTBXLMultilabelClassification: Multi-label diagnosis from the + 100 Hz, 12-lead ECG using the ``"diagnostic"`` label set. + """ + return PTBXLMultilabelClassification() diff --git a/pyhealth/models/__init__.py b/pyhealth/models/__init__.py index ddb74249f..332b16ed9 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..3c94bf6c8 --- /dev/null +++ b/pyhealth/models/protoecgnet.py @@ -0,0 +1,900 @@ +"""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. +Because prototypes can be projected onto real training ECG segments, every +prediction can be explained with example cases. +""" + +from collections.abc import Sequence +from itertools import combinations + +import torch +import torch.nn.functional as F +from torch import nn + +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) + 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"`` or ``"learned"``. 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. + + 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)) + self.register_buffer("class_weights", None) + self.register_buffer("label_cooccurrence", None) + if class_weights is not None: + self.class_weights = class_weights.float() + if label_cooccurrence is not None: + self.label_cooccurrence = label_cooccurrence.float() + + 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_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. + + 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 + reference code, ``(total - count) / count`` where ``total`` is the + sum of all label counts. + """ + 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 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 + + 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) + logits = self._scatter({n: o[0] for n, o in outputs.items()}, batch) + if self.stage == "fusion": + logits = self.fusion_classifier(sims) + if y is not None and self.stage in ("all", "joint", "prototypes"): + for name, (_, sim, _) in outputs.items(): + index = self.branch_index(name) + cooc = None + if self.label_cooccurrence is not None and self.lam_cnrst > 0: + cooc = self.label_cooccurrence[index][:, index] + 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. + + 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). + """ + 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) + 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 ``(batch, num_branch_labels, + num_prototypes)`` equal to similarity times classifier weight. + """ + was_training = self.training + self.eval() + x = signal.to(self.device).float() + out = {} + for name in self.branch_names: + branch = self.branches[name] + _, sim, _ = branch(x) + 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 2e7bd6bab..c2e9aa2f5 100644 --- a/pyhealth/tasks/__init__.py +++ b/pyhealth/tasks/__init__.py @@ -12,6 +12,7 @@ from .chestxray14_binary_classification import ChestXray14BinaryClassification from .chestxray14_multilabel_classification import ChestXray14MultilabelClassification from .covid19_cxr_classification import COVID19CXRClassification +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..af5f04280 --- /dev/null +++ b/pyhealth/tasks/ptbxl_multilabel_classification.py @@ -0,0 +1,226 @@ +"""Multi-label ECG classification task for the PTB-XL dataset.""" + +import os +import re +from typing import Any, ClassVar + +import numpy as np + +from .base_task import BaseTask + +_LABEL_SETS = ( + "all", + "diagnostic", + "subdiagnostic", + "superdiagnostic", + "form", + "rhythm", +) +_GAIN_PATTERN = re.compile(r"^([0-9.eE+-]+)(?:\(([-0-9]+)\))?(?:/.*)?$") + + +def read_wfdb_signal(record_path: str) -> np.ndarray: + """Reads a WFDB record and returns its physical signal. + + The ``wfdb`` package is used when it is installed. Otherwise a small + fallback reader that supports the 16-bit format used by PTB-XL is used, so + that no extra dependency is required. + + Args: + record_path: Path of the record without extension, e.g. + ``".../records100/00000/00001_lr"``. + + Returns: + A float32 array of shape ``(num_samples, num_leads)`` in millivolts. + + Raises: + FileNotFoundError: If the header or data file does not exist. + ValueError: If the record uses an unsupported WFDB storage format. + + Examples: + >>> signal = read_wfdb_signal("records100/00000/00001_lr") # doctest: +SKIP + >>> signal.shape # doctest: +SKIP + (1000, 12) + """ + try: + import wfdb # type: ignore + + return wfdb.rdrecord(record_path).p_signal.astype(np.float32) + except ImportError: + pass + + header_path = record_path + ".hea" + if not os.path.isfile(header_path): + raise FileNotFoundError(f"WFDB header not found: {header_path}") + with open(header_path, "r", encoding="utf-8") as handle: + lines = [ln.strip() for ln in handle if ln.strip()] + lines = [ln for ln in lines if not ln.startswith("#")] + first = lines[0].split() + num_leads, num_samples = int(first[1]), int(first[3]) + gains, baselines, files = [], [], set() + for line in lines[1 : 1 + num_leads]: + fields = line.split() + files.add(fields[0]) + if fields[1].split("+")[0] != "16": + raise ValueError( + f"Unsupported WFDB format {fields[1]} in {header_path}; " + "install `wfdb` to read this record." + ) + match = _GAIN_PATTERN.match(fields[2]) + gain = float(match.group(1)) if match else 200.0 + gains.append(gain if gain != 0 else 200.0) + baselines.append( + int(match.group(2)) if match and match.group(2) else int(fields[4]) + ) + if len(files) != 1: + raise ValueError("Only records with a single signal file are supported.") + data_path = os.path.join(os.path.dirname(header_path), files.pop()) + if not os.path.isfile(data_path): + raise FileNotFoundError(f"WFDB data file not found: {data_path}") + digital = np.fromfile(data_path, dtype=">> from pyhealth.datasets import PTBXLDataset + >>> from pyhealth.tasks import PTBXLMultilabelClassification + >>> dataset = PTBXLDataset(root="/path/to/ptb-xl/1.0.3") + >>> task = PTBXLMultilabelClassification(label_set="superdiagnostic") + >>> samples = dataset.set_task(task) + """ + + task_name: str = "PTBXLMultilabelClassification" + input_schema: ClassVar[dict[str, str]] = {"signal": "tensor"} + output_schema: ClassVar[dict[str, str]] = {"labels": "multilabel"} + + def __init__( + self, + label_set: str = "diagnostic", + sampling_rate: int = 100, + 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 sampling_rate not in (100, 500): + raise ValueError("sampling_rate must be 100 or 500") + if standardize not in ("record", None): + raise ValueError("standardize must be 'record' or None") + self.label_set = label_set + self.sampling_rate = sampling_rate + self.min_likelihood = min_likelihood + self.standardize = standardize + self.remove_baseline = remove_baseline + super().__init__() + + @staticmethod + def _parse_labels(text: Any, min_likelihood: float) -> list[str]: + """Parses a ``"CODE:likelihood;CODE:likelihood"`` string. + + Args: + text: The label string, or None / NaN when the set is empty. + min_likelihood: Minimum likelihood to keep a label. + + Returns: + The list of label codes that pass the likelihood threshold. + """ + if not isinstance(text, str) or not text: + return [] + labels = [] + for item in text.split(";"): + code, _, likelihood = item.partition(":") + if code and float(likelihood or 0.0) >= min_likelihood: + labels.append(code) + return labels + + def _prepare_signal(self, record_path: str) -> np.ndarray: + """Loads, crops/pads and normalizes one ECG. + + Args: + record_path: Record path without extension. + + Returns: + A float32 array of shape ``(12, 10 * sampling_rate)``. + """ + signal = read_wfdb_signal(record_path).T # (leads, samples) + length = 10 * self.sampling_rate + if signal.shape[1] < length: + pad = length - signal.shape[1] + signal = np.pad(signal, ((0, 0), (0, pad))) + signal = signal[:, :length] + if self.remove_baseline: + from scipy.signal import butter, filtfilt + + b, a = butter(1, 0.5 / (self.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. + + Args: + patient: A patient object with ``ptbxl`` events. + + Returns: + A list of dicts with ``patient_id``, ``ecg_id``, ``strat_fold``, + ``signal`` (12 x T float32 array) and ``labels`` (list of codes). + ECGs without any label in the chosen set are skipped. + """ + samples = [] + for event in patient.get_events(event_type="ptbxl"): + labels = self._parse_labels( + getattr(event, f"labels_{self.label_set}", None), + self.min_likelihood, + ) + if not labels: + continue + path = ( + event.record_path_lr + if self.sampling_rate == 100 + else event.record_path_hr + ) + samples.append( + { + "patient_id": patient.patient_id, + "ecg_id": int(event.ecg_id), + "strat_fold": int(event.strat_fold), + "signal": self._prepare_signal(path), + "labels": labels, + } + ) + return samples diff --git a/tests/core/test_protoecgnet.py b/tests/core/test_protoecgnet.py new file mode 100644 index 000000000..3b596c5fe --- /dev/null +++ b/tests/core/test_protoecgnet.py @@ -0,0 +1,414 @@ +"""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_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_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)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/core/test_ptbxl.py b/tests/core/test_ptbxl.py new file mode 100644 index 000000000..8e34bc074 --- /dev/null +++ b/tests/core/test_ptbxl.py @@ -0,0 +1,174 @@ +"""Tests for the PTB-XL dataset and its multi-label classification task. + +All tests use a tiny synthetic PTB-XL-like folder (5 ECGs from 4 patients) +written to a temporary directory; no real data is read. +""" + +import os +import shutil +import tempfile +import unittest + +import numpy as np +import pandas as pd + +from pyhealth.datasets import PTBXLDataset +from pyhealth.tasks import PTBXLMultilabelClassification +from pyhealth.tasks.ptbxl_multilabel_classification import read_wfdb_signal + +LEADS = ["I", "II", "III", "aVR", "aVL", "aVF", "V1", "V2", "V3", "V4", "V5", "V6"] + +# Chosen to cover: likelihood filtering, super/sub-diagnostic aggregation, and +# the edge case where a record has rhythm codes only (ECG 5, no diagnostic label). +ECGS = [ + (1, 10, {"NORM": 100.0, "SR": 0.0}, 1), + (2, 10, {"IMI": 80.0, "ABQRS": 0.0}, 1), + (3, 11, {"LVH": 50.0, "SR": 0.0}, 9), + (4, 12, {"NORM": 100.0}, 10), + (5, 13, {"SR": 0.0}, 2), # rhythm only: no diagnostic labels +] + +STATEMENTS = pd.DataFrame( + { + "diagnostic": [1, 1, 1, 0, 0], + "form": [0, 0, 0, 0, 1], + "rhythm": [0, 0, 0, 1, 0], + "diagnostic_class": ["NORM", "MI", "HYP", None, None], + "diagnostic_subclass": ["NORM", "IMI", "LVH", None, None], + }, + index=pd.Index(["NORM", "IMI", "LVH", "SR", "ABQRS"], name="scp_code"), +) + + +def write_record(path_without_ext: str, signal: np.ndarray) -> None: + """Writes a 16-bit WFDB record (PTB-XL layout, gain 1000/mV).""" + os.makedirs(os.path.dirname(path_without_ext), exist_ok=True) + num_samples = signal.shape[0] + digital = np.round(signal * 1000).astype(" None: + """Creates a minimal PTB-XL folder under ``root``.""" + rng = np.random.default_rng(0) + rows = [] + for ecg_id, patient_id, codes, fold in ECGS: + rel = f"records100/00000/{ecg_id:05d}_lr" + write_record(os.path.join(root, rel), rng.normal(size=(1000, 12))) + rows.append( + { + "ecg_id": ecg_id, + "patient_id": float(patient_id), + "age": 60.0, + "sex": 1, + "scp_codes": str(codes), + "strat_fold": fold, + "filename_lr": rel, + "filename_hr": rel.replace("records100", "records500"), + } + ) + pd.DataFrame(rows).to_csv(os.path.join(root, "ptbxl_database.csv"), index=False) + STATEMENTS.to_csv(os.path.join(root, "scp_statements.csv")) + + +class TestPTBXL(unittest.TestCase): + """Dataset, task and reader tests on synthetic PTB-XL files.""" + + @classmethod + def setUpClass(cls): + cls.root = tempfile.mkdtemp() + cls.cache = tempfile.mkdtemp() + build_synthetic_ptbxl(cls.root) + cls.dataset = PTBXLDataset(root=cls.root, cache_dir=cls.cache) + + @classmethod + def tearDownClass(cls): + shutil.rmtree(cls.root, ignore_errors=True) + shutil.rmtree(cls.cache, ignore_errors=True) + + def test_metadata_file_and_label_strings(self): + meta = pd.read_csv(os.path.join(self.root, "ptbxl-metadata-pyhealth.csv")) + self.assertEqual(len(meta), 5) + self.assertEqual(meta.loc[0, "patient_id"], 10) + # IMI aggregates to the MI super-class and the IMI sub-class. + second = meta.set_index("ecg_id").loc[2] + self.assertEqual(second["labels_superdiagnostic"], "MI:80.0") + self.assertEqual(second["labels_subdiagnostic"], "IMI:80.0") + self.assertEqual(second["labels_form"], "ABQRS:0.0") + self.assertEqual(second["labels_all"], "ABQRS:0.0;IMI:80.0") + + def test_missing_files_raise(self): + with tempfile.TemporaryDirectory() as empty, self.assertRaises(FileNotFoundError): + PTBXLDataset(root=empty) + + def test_patients_and_events(self): + self.assertEqual(len(self.dataset.unique_patient_ids), 4) + patient = self.dataset.get_patient("10") + self.assertEqual(len(patient.get_events(event_type="ptbxl")), 2) + + def test_default_task(self): + self.assertIsInstance( + self.dataset.default_task, PTBXLMultilabelClassification + ) + + def test_task_samples_superdiagnostic(self): + task = PTBXLMultilabelClassification(label_set="superdiagnostic") + samples = self.dataset.set_task(task) + self.assertEqual(len(samples), 4) # ECG 5 has no diagnostic label + item = samples[0] + self.assertEqual(tuple(item["signal"].shape), (12, 1000)) + self.assertEqual(item["labels"].shape[0], 3) # NORM, MI, HYP vocab + self.assertAlmostEqual(float(item["signal"].mean()), 0.0, places=4) + + def test_task_rhythm_skips_unlabelled(self): + task = PTBXLMultilabelClassification(label_set="rhythm") + samples = self.dataset.set_task(task) + self.assertEqual(len(samples), 3) # ECGs 1, 3 and 5 have SR + + def test_parse_labels_likelihood_filter(self): + parse = PTBXLMultilabelClassification._parse_labels + self.assertEqual(parse("A:100.0;B:0.0", 0.0), ["A", "B"]) + self.assertEqual(parse("A:100.0;B:0.0", 50.0), ["A"]) + self.assertEqual(parse(None, 0.0), []) + self.assertEqual(parse(float("nan"), 0.0), []) + + def test_invalid_arguments(self): + with self.assertRaises(ValueError): + PTBXLMultilabelClassification(label_set="nope") + with self.assertRaises(ValueError): + PTBXLMultilabelClassification(sampling_rate=250) + with self.assertRaises(ValueError): + PTBXLMultilabelClassification(standardize="dataset") + + def test_read_wfdb_signal_matches_written_values(self): + with tempfile.TemporaryDirectory() as tmp: + signal = np.linspace(-1.0, 1.0, 24).reshape(2, 12) + path = os.path.join(tmp, "rec") + write_record(path, signal) + read = read_wfdb_signal(path) + self.assertEqual(read.shape, (2, 12)) + np.testing.assert_allclose(read, signal, atol=1e-3) + + def test_read_wfdb_signal_missing_file(self): + with self.assertRaises(FileNotFoundError): + read_wfdb_signal("/nonexistent/record") + + def test_prepare_signal_pads_and_standardizes(self): + task = PTBXLMultilabelClassification(remove_baseline=True) + with tempfile.TemporaryDirectory() as tmp: + path = os.path.join(tmp, "short") + write_record(path, np.random.default_rng(1).normal(size=(500, 12))) + out = task._prepare_signal(path) + self.assertEqual(out.shape, (12, 1000)) + self.assertEqual(out.dtype, np.float32) + self.assertTrue(np.isfinite(out).all()) + + +if __name__ == "__main__": + unittest.main() From acded1316bab1d6e7eaff3e789f9f7a2245093f1 Mon Sep 17 00:00:00 2001 From: Irvi Aini Date: Sat, 24 Oct 2026 12:25:29 +0800 Subject: [PATCH 02/14] fix: address AI review feedback on ProtoECGNet and PTB-XL Blocking fixes: - Override train() to keep frozen backbone/add-on BatchNorm layers in eval mode during classifier and fusion stages; frozen module set is computed once in set_stage so train() has no per-call allocation - Fix prototype_contributions() to use fusion_classifier weights when stage="fusion" so contributions sum to the logit (branch classifier weights were used unconditionally before) Optional: - Add refresh_metadata parameter to PTBXLDataset for forced CSV rebuild - Add MPS device support in the example script for Apple Silicon - Clarify project_prototypes docstring: global mode stores a time-averaged feature vector, not a raw ECG segment Tests: - Verify BN running stats are unchanged during classifier and fusion stages after a train()+forward pass - Verify prototype_contributions sums to logit minus fusion bias - Add 500 Hz path test and remove_baseline DC-suppression value test - Add refresh_metadata round-trip test Signed-off-by: Irvi Aini --- ...l_multilabel_classification_protoecgnet.py | 7 ++- pyhealth/datasets/ptbxl.py | 6 ++- pyhealth/models/protoecgnet.py | 34 +++++++++++-- tests/core/test_protoecgnet.py | 48 +++++++++++++++++++ tests/core/test_ptbxl.py | 30 ++++++++++++ 5 files changed, 120 insertions(+), 5 deletions(-) diff --git a/examples/ptbxl_multilabel_classification_protoecgnet.py b/examples/ptbxl_multilabel_classification_protoecgnet.py index 12b5d4eb5..a3fb342a5 100644 --- a/examples/ptbxl_multilabel_classification_protoecgnet.py +++ b/examples/ptbxl_multilabel_classification_protoecgnet.py @@ -275,7 +275,12 @@ def main() -> None: parser.add_argument("--ablation", action="store_true") parser.add_argument("--output", default="protoecgnet_ablation.md") args = parser.parse_args() - args.device = "cuda" if torch.cuda.is_available() else "cpu" + 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" scratch = tempfile.mkdtemp() if args.synthetic: diff --git a/pyhealth/datasets/ptbxl.py b/pyhealth/datasets/ptbxl.py index e47def5c2..e3fabeb54 100644 --- a/pyhealth/datasets/ptbxl.py +++ b/pyhealth/datasets/ptbxl.py @@ -53,6 +53,9 @@ class PTBXLDataset(BaseDataset): cache_dir: Optional directory for caching processed data. num_workers: Number of parallel workers for data processing. dev: If True, only a small subset of patients is loaded. + refresh_metadata: If True, recreates ``ptbxl-metadata-pyhealth.csv`` + even when it already exists. Useful after upgrading PyHealth if the + label aggregation logic has changed. Attributes: root: Root directory of the raw data. @@ -76,12 +79,13 @@ def __init__( cache_dir: str | None = None, num_workers: int = 1, dev: bool = False, + refresh_metadata: bool = False, ) -> None: if config_path is None: logger.info("No config path provided, using default config") config_path = Path(__file__).parent / "configs" / "ptbxl.yaml" metadata_path = os.path.join(root, "ptbxl-metadata-pyhealth.csv") - if not os.path.exists(metadata_path): + if refresh_metadata or not os.path.exists(metadata_path): self.prepare_metadata(root) super().__init__( root=root, diff --git a/pyhealth/models/protoecgnet.py b/pyhealth/models/protoecgnet.py index 3c94bf6c8..92f245a36 100644 --- a/pyhealth/models/protoecgnet.py +++ b/pyhealth/models/protoecgnet.py @@ -629,6 +629,9 @@ def __init__( 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() def _prepare_groups( self, groups: dict[str, list[str]] | None @@ -732,6 +735,19 @@ def set_stage(self, stage: str) -> None: parts[0] == "branches" and parts[2] in trainable[stage] ) self.stage = stage + # Recompute once here so train() is O(|frozen|) with no allocation. + self._frozen_modules = { + m for m in self.modules() + if any(True for _ in m.parameters(recurse=False)) + and all(not p.requires_grad for p in m.parameters(recurse=False)) + } + + def train(self, mode: bool = True) -> "ProtoECGNet": + super().train(mode) + if mode: + for m in self._frozen_modules: + m.eval() + return self def _scatter(self, per_branch: dict[str, torch.Tensor], batch: int) -> torch.Tensor: """Writes branch logits into the full label space.""" @@ -819,6 +835,8 @@ def project_prototypes(self, dataloader) -> dict[str, list[dict | None]]: 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. Args: dataloader: Training dataloader (``shuffle=False``) yielding @@ -885,16 +903,26 @@ def prototype_contributions( signal: ECG batch ``(batch, leads, time)``. Returns: - For every branch a tensor ``(batch, num_branch_labels, - num_prototypes)`` equal to similarity times classifier weight. + For every branch a tensor equal to similarity times classifier + weight. Shape is ``(batch, num_labels, num_branch_prototypes)`` + when ``stage="fusion"`` (fusion classifier weights, contributions + sum to the logit), or ``(batch, num_branch_labels, + num_branch_prototypes)`` otherwise (branch classifier weights). """ was_training = self.training self.eval() x = signal.to(self.device).float() + use_fusion = self.stage == "fusion" and self.fusion_classifier is not None out = {} + offset = 0 for name in self.branch_names: branch = self.branches[name] _, sim, _ = branch(x) - out[name] = sim.unsqueeze(1) * branch.classifier.weight.unsqueeze(0) + 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/tests/core/test_protoecgnet.py b/tests/core/test_protoecgnet.py index 3b596c5fe..449679b3b 100644 --- a/tests/core/test_protoecgnet.py +++ b/tests/core/test_protoecgnet.py @@ -409,6 +409,54 @@ def test_prototype_contributions_shape(self): 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) + if __name__ == "__main__": unittest.main() diff --git a/tests/core/test_ptbxl.py b/tests/core/test_ptbxl.py index 8e34bc074..1893e9ea5 100644 --- a/tests/core/test_ptbxl.py +++ b/tests/core/test_ptbxl.py @@ -169,6 +169,36 @@ def test_prepare_signal_pads_and_standardizes(self): self.assertEqual(out.dtype, np.float32) self.assertTrue(np.isfinite(out).all()) + def test_500hz_path_yields_5000_samples(self): + task = PTBXLMultilabelClassification(sampling_rate=500) + with tempfile.TemporaryDirectory() as tmp: + path = os.path.join(tmp, "rec500") + write_record(path, np.random.default_rng(2).normal(size=(5000, 12))) + out = task._prepare_signal(path) + self.assertEqual(out.shape, (12, 5000)) + self.assertEqual(out.dtype, np.float32) + + def test_remove_baseline_suppresses_dc(self): + # A constant (DC-only) signal must be driven near zero by the 0.5 Hz + # high-pass filter; standardize=None isolates the filter's effect. + task = PTBXLMultilabelClassification(remove_baseline=True, standardize=None) + with tempfile.TemporaryDirectory() as tmp: + path = os.path.join(tmp, "dc") + write_record(path, np.ones((1000, 12), dtype=np.float32)) + out = task._prepare_signal(path) + self.assertEqual(out.shape, (12, 1000)) + np.testing.assert_allclose(out.mean(axis=1), 0.0, atol=0.05) + + def test_refresh_metadata_recreates_file(self): + # Patch the metadata with stale content; refresh_metadata must overwrite it. + meta_path = os.path.join(self.root, "ptbxl-metadata-pyhealth.csv") + with open(meta_path, "w") as handle: + handle.write("stale\n") + PTBXLDataset(root=self.root, cache_dir=self.cache, refresh_metadata=True) + updated = pd.read_csv(meta_path) + self.assertIn("patient_id", updated.columns) + self.assertEqual(len(updated), 5) + if __name__ == "__main__": unittest.main() From 9a4bc14c3bbb5c09ff3394c0d8ae93d0e200ad9e Mon Sep 17 00:00:00 2001 From: Irvi Aini Date: Sat, 24 Oct 2026 12:28:58 +0800 Subject: [PATCH 03/14] docs: include survival_mdn and ProtoECGNet in models API reference --- docs/api/models.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/api/models.rst b/docs/api/models.rst index b83d9bf18..e8d211056 100644 --- a/docs/api/models.rst +++ b/docs/api/models.rst @@ -253,4 +253,5 @@ API Reference models/pyhealth.models.unified_multimodal_embedding_docs models/pyhealth.models.califorest models/pyhealth.models.XGBoostModel + models/pyhealth.models.survival_mdn models/pyhealth.models.ProtoECGNet From b9a9d3a18b0914aad582f4ed631d99308cb7e672 Mon Sep 17 00:00:00 2001 From: Irvi Aini Date: Mon, 26 Oct 2026 11:07:37 +0800 Subject: [PATCH 04/14] fix: address optional review feedback on ProtoECGNet and PTB-XL - prototype_contributions docstring: clarify that in fusion mode contributions sum to logit minus fusion_classifier.bias (fusion classifier has a bias; branch classifiers do not) - ResNet2dBackbone: replace unused avgpool and fc with nn.Identity() to drop ~512K dead parameters - Example: wrap dataset/training in try/finally so the scratch directory (potentially multi-GB litdata cache) is always cleaned up; move parser.error() calls before mkdtemp so they exit cleanly - Example: expand --lam-l1 help text to explain when to enable it - ptbxl.py: use os.path.abspath(root) in prepare_metadata so record paths in the metadata CSV survive os.chdir Signed-off-by: Irvi Aini --- ...l_multilabel_classification_protoecgnet.py | 122 +++++++++++------- pyhealth/datasets/ptbxl.py | 1 + pyhealth/models/protoecgnet.py | 32 ++++- 3 files changed, 105 insertions(+), 50 deletions(-) diff --git a/examples/ptbxl_multilabel_classification_protoecgnet.py b/examples/ptbxl_multilabel_classification_protoecgnet.py index a3fb342a5..f6583795a 100644 --- a/examples/ptbxl_multilabel_classification_protoecgnet.py +++ b/examples/ptbxl_multilabel_classification_protoecgnet.py @@ -43,11 +43,14 @@ import argparse import itertools import os +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.models import PTBXL_BRANCH_KWARGS, PTBXL_BRANCH_LABELS, ProtoECGNet @@ -58,6 +61,25 @@ 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. @@ -219,7 +241,7 @@ def run_pipeline( ) trainer = Trainer( model=model, - metrics=["roc_auc_macro"], + metrics=[], device=args.device, enable_logging=False, ) @@ -230,7 +252,8 @@ def stage(name: str, epochs: int) -> None: train_dataloader=loaders[0], val_dataloader=loaders[1], epochs=epochs, - monitor="roc_auc_macro", + monitor="loss", + monitor_criterion="min", optimizer_params={"lr": args.lr}, ) @@ -241,7 +264,8 @@ def stage(name: str, epochs: int) -> None: stage("classifier", max(1, args.epochs // 2)) if args.paper_groups: stage("fusion", max(1, args.epochs // 2)) - return trainer.evaluate(loaders[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: @@ -257,7 +281,10 @@ def main() -> 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 of final layers") + 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 " @@ -282,11 +309,9 @@ def main() -> None: else: args.device = "cpu" - scratch = tempfile.mkdtemp() + # 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.root = os.path.join(scratch, "ptbxl") - os.makedirs(args.root) - build_synthetic_ptbxl(args.root) args.epochs = args.epochs or 3 args.width = args.width or 8 args.proto_dim = args.proto_dim or 16 @@ -299,43 +324,52 @@ def main() -> None: args.width = args.width or 64 args.proto_dim = args.proto_dim or 128 - dataset = PTBXLDataset( - root=args.root, cache_dir=os.path.join(scratch, "cache"), dev=args.dev - ) - task = PTBXLMultilabelClassification( - label_set=args.label_set, remove_baseline=args.remove_baseline - ) - samples = dataset.set_task(task) - splits = split_by_fold(samples, args.root) - print(f"samples: train={len(splits[0])} val={len(splits[1])} test={len(splits[2])}") + scratch = tempfile.mkdtemp() + try: + if args.synthetic: + args.root = os.path.join(scratch, "ptbxl") + os.makedirs(args.root) + build_synthetic_ptbxl(args.root) + + dataset = PTBXLDataset( + root=args.root, cache_dir=os.path.join(scratch, "cache"), dev=args.dev + ) + task = PTBXLMultilabelClassification( + label_set=args.label_set, remove_baseline=args.remove_baseline + ) + samples = dataset.set_task(task) + splits = split_by_fold(samples, args.root) + 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) + 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__": diff --git a/pyhealth/datasets/ptbxl.py b/pyhealth/datasets/ptbxl.py index e3fabeb54..aad900f2d 100644 --- a/pyhealth/datasets/ptbxl.py +++ b/pyhealth/datasets/ptbxl.py @@ -163,6 +163,7 @@ def prepare_metadata(self, root: str) -> None: "https://physionet.org/content/ptb-xl/1.0.3/ and pass the " "extracted folder as `root`." ) + root = os.path.abspath(root) database = pd.read_csv(database_path, index_col="ecg_id") statements = pd.read_csv(statements_path, index_col=0) diff --git a/pyhealth/models/protoecgnet.py b/pyhealth/models/protoecgnet.py index 92f245a36..b9fecb2ef 100644 --- a/pyhealth/models/protoecgnet.py +++ b/pyhealth/models/protoecgnet.py @@ -8,8 +8,10 @@ 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. -Because prototypes can be projected onto real training ECG segments, every -prediction can be explained with example cases. +Calling :meth:`project_prototypes` replaces every prototype with the most +similar real training segment and returns the source-sample mapping to the +caller; store that alongside the checkpoint to support case-based explanations +after reloading. """ from collections.abc import Sequence @@ -207,6 +209,9 @@ def __init__(self, in_leads: int = 12, pretrained: bool = False) -> None: 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 @@ -521,7 +526,10 @@ class ProtoECGNet(BaseModel): 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"`` or ``"learned"``. Defaults to ``"average"``. + fusion: ``"average"`` (scatter each branch's logits into the full + label space with no cross-branch interaction) or ``"learned"`` + (train a linear layer over all prototype similarities 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. @@ -632,6 +640,8 @@ def __init__( # 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 @@ -838,6 +848,11 @@ def project_prototypes(self, dataloader) -> dict[str, list[dict | None]]: 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. @@ -891,6 +906,7 @@ def project_prototypes(self, dataloader) -> dict[str, list[dict | None]]: ) branch.prototype_vectors.data[p] = vector self.train(was_training) + self.projection_info = best return best @torch.no_grad() @@ -905,9 +921,13 @@ def prototype_contributions( Returns: For every branch a tensor equal to similarity times classifier weight. Shape is ``(batch, num_labels, num_branch_prototypes)`` - when ``stage="fusion"`` (fusion classifier weights, contributions - sum to the logit), or ``(batch, num_branch_labels, - num_branch_prototypes)`` otherwise (branch classifier weights). + when ``stage="fusion"`` (fusion classifier weights; 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() From 7bdd110b65adb963aae57a7b6ca456d5143768cb Mon Sep 17 00:00:00 2001 From: Irvi Aini Date: Mon, 26 Oct 2026 11:30:11 +0800 Subject: [PATCH 05/14] fix: address optional review feedback on ProtoECGNet and PTB-XL - prototype_contributions docstring: clarify that in fusion mode contributions sum to logit minus fusion_classifier.bias (fusion classifier has a bias; branch classifiers do not) - ResNet2dBackbone: replace unused avgpool and fc with nn.Identity() to drop ~512K dead parameters - Example: wrap dataset/training in try/finally so the scratch directory (potentially multi-GB litdata cache) is always cleaned up; move parser.error() calls before mkdtemp so they exit cleanly - Example: add --cache-dir option so repeated real-data runs reuse the processed cache instead of re-preprocessing several GB each time - Example: expand --lam-l1 help text to explain when to enable it - ptbxl.py: use os.path.abspath(root) in prepare_metadata so record paths in the metadata CSV survive os.chdir Signed-off-by: Irvi Aini --- examples/ptbxl_multilabel_classification_protoecgnet.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/examples/ptbxl_multilabel_classification_protoecgnet.py b/examples/ptbxl_multilabel_classification_protoecgnet.py index f6583795a..c28fa97a5 100644 --- a/examples/ptbxl_multilabel_classification_protoecgnet.py +++ b/examples/ptbxl_multilabel_classification_protoecgnet.py @@ -299,6 +299,11 @@ def main() -> None: 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") args = parser.parse_args() @@ -331,8 +336,9 @@ def main() -> None: 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=os.path.join(scratch, "cache"), dev=args.dev + root=args.root, cache_dir=cache_dir, dev=args.dev ) task = PTBXLMultilabelClassification( label_set=args.label_set, remove_baseline=args.remove_baseline From 8916bba903589f873f7ee3202b0bc64a680b4829 Mon Sep 17 00:00:00 2001 From: Irvi Aini Date: Mon, 26 Oct 2026 12:31:24 +0800 Subject: [PATCH 06/14] fix: write PTB-XL metadata CSV to cache_dir to support read-only roots When cache_dir is provided, ptbxl-metadata-pyhealth.csv is now written there rather than into root, so PTBXLDataset works on Kaggle /kaggle/input mounts and other read-only data directories. A patched YAML config with a relative path is generated in cache_dir so BaseDataset can still locate the file. self.metadata_path exposes the resolved path to callers. Also adds: shuffled-loader guard in project_prototypes, CardiologyDataset overlap note in PTBXLDataset docstring, set_pos_weight note in ProtoECGNet docstring, and one new test for each guard. Signed-off-by: Irvi Aini --- ...l_multilabel_classification_protoecgnet.py | 9 +-- pyhealth/datasets/ptbxl.py | 62 ++++++++++++++++--- pyhealth/models/protoecgnet.py | 19 ++++++ tests/core/test_protoecgnet.py | 7 +++ tests/core/test_ptbxl.py | 13 +++- 5 files changed, 97 insertions(+), 13 deletions(-) diff --git a/examples/ptbxl_multilabel_classification_protoecgnet.py b/examples/ptbxl_multilabel_classification_protoecgnet.py index c28fa97a5..a1f357056 100644 --- a/examples/ptbxl_multilabel_classification_protoecgnet.py +++ b/examples/ptbxl_multilabel_classification_protoecgnet.py @@ -142,17 +142,18 @@ def build_synthetic_ptbxl(root: str, num_patients: int = 60, seed: int = 0) -> N statements.to_csv(os.path.join(root, "scp_statements.csv")) -def split_by_fold(sample_dataset, root: str): +def split_by_fold(sample_dataset, metadata_path: str): """Splits a PTB-XL SampleDataset by the recommended ``strat_fold``. Args: sample_dataset: The dataset returned by ``set_task``. - root: PTB-XL root containing ``ptbxl-metadata-pyhealth.csv``. + metadata_path: Path to ``ptbxl-metadata-pyhealth.csv`` + (use ``dataset.metadata_path``). Returns: ``(train, val, test)`` datasets (folds 1-8, 9 and 10). """ - meta = pd.read_csv(os.path.join(root, "ptbxl-metadata-pyhealth.csv")) + meta = pd.read_csv(metadata_path) fold_of = dict(zip(meta["patient_id"].astype(str), meta["strat_fold"])) buckets: dict[str, list[int]] = {"train": [], "val": [], "test": []} for patient, indices in sample_dataset.patient_to_index.items(): @@ -344,7 +345,7 @@ def main() -> None: label_set=args.label_set, remove_baseline=args.remove_baseline ) samples = dataset.set_task(task) - splits = split_by_fold(samples, args.root) + splits = split_by_fold(samples, dataset.metadata_path) print(f"samples: train={len(splits[0])} val={len(splits[1])} test={len(splits[2])}") baseline = run_pipeline( diff --git a/pyhealth/datasets/ptbxl.py b/pyhealth/datasets/ptbxl.py index aad900f2d..a72b4e2d8 100644 --- a/pyhealth/datasets/ptbxl.py +++ b/pyhealth/datasets/ptbxl.py @@ -6,12 +6,33 @@ from pathlib import Path import pandas as pd +import yaml from ..tasks import PTBXLMultilabelClassification from .base_dataset import BaseDataset logger = logging.getLogger(__name__) + +def _patch_ptbxl_config(source: str, rel_metadata: str, cache_dir: str) -> str: + """Writes a patched copy of the YAML config to cache_dir. + + Replaces ``file_path`` in the ``ptbxl`` table entry with ``rel_metadata`` + so :class:`~pyhealth.datasets.BaseDataset` can find the CSV when it lives + outside ``root``. + + Returns: + Absolute path to the patched YAML in ``cache_dir``. + """ + with open(source) as fh: + cfg = yaml.safe_load(fh) + cfg["tables"]["ptbxl"]["file_path"] = rel_metadata + patched = os.path.join(cache_dir, "ptbxl.yaml") + with open(patched, "w") as fh: + yaml.dump(cfg, fh, default_flow_style=False) + return patched + + # Label sets that can be derived from ``scp_statements.csv``. LABEL_SETS = ( "all", @@ -61,6 +82,16 @@ class PTBXLDataset(BaseDataset): root: Root directory of the raw data. dataset_name: Name of the dataset. config_path: Path to the configuration file. + metadata_path: Absolute path to the derived + ``ptbxl-metadata-pyhealth.csv``. Written to ``cache_dir`` when + provided (so it works on read-only roots such as + ``/kaggle/input``), otherwise written to ``root``. + + Note: + PTB-XL is also available through + :class:`~pyhealth.datasets.CardiologyDataset` via the PhysioNet + Challenge 2020 subset. That version uses SNOMED-CT codes and a + different fold structure; the two datasets are not directly comparable. Examples: >>> from pyhealth.datasets import PTBXLDataset @@ -84,9 +115,22 @@ def __init__( if config_path is None: logger.info("No config path provided, using default config") config_path = Path(__file__).parent / "configs" / "ptbxl.yaml" - metadata_path = os.path.join(root, "ptbxl-metadata-pyhealth.csv") + # Prefer cache_dir for the derived CSV so PTBXLDataset works on + # read-only roots (e.g. /kaggle/input mounts). When cache_dir is given, + # a patched YAML config is written there so BaseDataset can find the CSV + # via a relative path from root. + if cache_dir is not None: + os.makedirs(cache_dir, exist_ok=True) + metadata_path = os.path.join(cache_dir, "ptbxl-metadata-pyhealth.csv") + rel_to_root = os.path.relpath( + os.path.abspath(metadata_path), os.path.abspath(root) + ) + config_path = _patch_ptbxl_config(str(config_path), rel_to_root, cache_dir) + else: + metadata_path = os.path.join(root, "ptbxl-metadata-pyhealth.csv") + self.metadata_path = metadata_path if refresh_metadata or not os.path.exists(metadata_path): - self.prepare_metadata(root) + self.prepare_metadata(root, metadata_path) super().__init__( root=root, tables=["ptbxl"], @@ -140,8 +184,10 @@ def _aggregate_labels( for name, labels in per_set.items() } - def prepare_metadata(self, root: str) -> None: - """Creates ``ptbxl-metadata-pyhealth.csv`` from the raw PTB-XL files. + def prepare_metadata( + self, root: str, metadata_path: str | None = None + ) -> None: + """Creates the metadata CSV from the raw PTB-XL files. The CSV has one row per ECG with absolute record paths (without file extension), the recommended fold, and label strings for every label @@ -149,6 +195,8 @@ def prepare_metadata(self, root: str) -> None: Args: root: Root directory containing the raw PTB-XL files. + metadata_path: Where to write the CSV. Defaults to + ``ptbxl-metadata-pyhealth.csv`` inside ``root``. Raises: FileNotFoundError: If ``ptbxl_database.csv`` or @@ -164,6 +212,8 @@ def prepare_metadata(self, root: str) -> None: "extracted folder as `root`." ) root = os.path.abspath(root) + if metadata_path is None: + metadata_path = os.path.join(root, "ptbxl-metadata-pyhealth.csv") database = pd.read_csv(database_path, index_col="ecg_id") statements = pd.read_csv(statements_path, index_col=0) @@ -183,9 +233,7 @@ def prepare_metadata(self, root: str) -> None: } row.update({f"labels_{name}": text for name, text in labels.items()}) rows.append(row) - pd.DataFrame(rows).to_csv( - os.path.join(root, "ptbxl-metadata-pyhealth.csv"), index=False - ) + pd.DataFrame(rows).to_csv(metadata_path, index=False) @property def default_task(self) -> PTBXLMultilabelClassification: diff --git a/pyhealth/models/protoecgnet.py b/pyhealth/models/protoecgnet.py index b9fecb2ef..434cea7c7 100644 --- a/pyhealth/models/protoecgnet.py +++ b/pyhealth/models/protoecgnet.py @@ -20,6 +20,7 @@ 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 @@ -540,6 +541,12 @@ class ProtoECGNet(BaseModel): 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`` + parameter. The inherited + :meth:`~pyhealth.models.BaseModel.set_pos_weight` has no effect on + this model. + Examples: >>> from pyhealth.datasets import create_sample_dataset, get_dataloader >>> samples = [ @@ -862,7 +869,19 @@ def project_prototypes(self, dataloader) -> dict[str, list[dict | None]]: ``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 = { diff --git a/tests/core/test_protoecgnet.py b/tests/core/test_protoecgnet.py index 449679b3b..e27157c57 100644 --- a/tests/core/test_protoecgnet.py +++ b/tests/core/test_protoecgnet.py @@ -458,5 +458,12 @@ def test_prototype_contributions_sum_to_logit_fusion_stage(self): torch.testing.assert_close(total, expected, atol=1e-4, rtol=1e-4) + 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) + + if __name__ == "__main__": unittest.main() diff --git a/tests/core/test_ptbxl.py b/tests/core/test_ptbxl.py index 1893e9ea5..8ff946718 100644 --- a/tests/core/test_ptbxl.py +++ b/tests/core/test_ptbxl.py @@ -93,7 +93,7 @@ def tearDownClass(cls): shutil.rmtree(cls.cache, ignore_errors=True) def test_metadata_file_and_label_strings(self): - meta = pd.read_csv(os.path.join(self.root, "ptbxl-metadata-pyhealth.csv")) + meta = pd.read_csv(self.dataset.metadata_path) self.assertEqual(len(meta), 5) self.assertEqual(meta.loc[0, "patient_id"], 10) # IMI aggregates to the MI super-class and the IMI sub-class. @@ -189,9 +189,18 @@ def test_remove_baseline_suppresses_dc(self): self.assertEqual(out.shape, (12, 1000)) np.testing.assert_allclose(out.mean(axis=1), 0.0, atol=0.05) + def test_metadata_written_to_cache_dir(self): + """When cache_dir is provided, the metadata CSV is written there, not root.""" + expected = os.path.join(self.cache, "ptbxl-metadata-pyhealth.csv") + self.assertEqual(self.dataset.metadata_path, expected) + self.assertTrue(os.path.exists(expected)) + self.assertFalse( + os.path.exists(os.path.join(self.root, "ptbxl-metadata-pyhealth.csv")) + ) + def test_refresh_metadata_recreates_file(self): # Patch the metadata with stale content; refresh_metadata must overwrite it. - meta_path = os.path.join(self.root, "ptbxl-metadata-pyhealth.csv") + meta_path = os.path.join(self.cache, "ptbxl-metadata-pyhealth.csv") with open(meta_path, "w") as handle: handle.write("stale\n") PTBXLDataset(root=self.root, cache_dir=self.cache, refresh_metadata=True) From f1610fdb4ccc32e992023bee31d27c1bbc55330b Mon Sep 17 00:00:00 2001 From: Irvi Aini Date: Tue, 27 Oct 2026 09:05:48 +0800 Subject: [PATCH 07/14] Four blocking issues fixed: MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 1. examples: split_by_strat_fold(samples) — removed incorrect folds=[range(1,9),[9],[10]] argument that crashed at runtime; the default 1-8/9/10 split is correct. 2. models/protoecgnet: _frozen_modules now checks recurse=True so a PrototypeBranch whose backbone has trainable children is never put in eval() during the feature_extractor stage (backbone BatchNorm was not updating its running stats). Added test_backbone_bn_stays_in_train_mode_in_feature_extractor_stage to confirm. 3. tests: added 8 tests to TestPTBXLMultilabelClassification covering the all, diagnostic, subdiagnostic, form label sets, the missing-header error path, standardize=None, and strat_fold field presence. Task coverage now >80%. 4. tasks/ptbxl_multilabel_classification: read_wfdb_signal delegates to load_ptbxl_record (master's canonical reader) when wfdb is installed, falling back to the built-in 16-bit reader otherwise — eliminates duplicate logic for the wfdb-available path. _aggregate_labels docstring notes it extends aggregate_diagnostic_superclasses to all six label sets. Two optional fixes: removed duplicate PTBXLDataset import in datasets/__init__.py and duplicate toctree entry in docs/api/datasets.rst. Signed-off-by: Irvi Aini --- docs/api/datasets.rst | 1 - ...l_multilabel_classification_protoecgnet.py | 2 +- pyhealth/datasets/__init__.py | 1 - pyhealth/models/protoecgnet.py | 4 +- .../tasks/ptbxl_multilabel_classification.py | 18 +++-- tests/core/test_protoecgnet.py | 33 +++++++++ tests/core/test_ptbxl.py | 74 +++++++++++++++++++ 7 files changed, 124 insertions(+), 9 deletions(-) diff --git a/docs/api/datasets.rst b/docs/api/datasets.rst index 3a82b7b0d..41f795d22 100644 --- a/docs/api/datasets.rst +++ b/docs/api/datasets.rst @@ -308,7 +308,6 @@ Available Datasets datasets/pyhealth.datasets.PTBXLDataset datasets/pyhealth.datasets.PhysioNetDeIDDataset datasets/pyhealth.datasets.EEGBCIDataset - PTB-XL Dataset datasets/pyhealth.datasets.TUABDataset datasets/pyhealth.datasets.TUEVDataset datasets/pyhealth.datasets.ClinVarDataset diff --git a/examples/ptbxl_multilabel_classification_protoecgnet.py b/examples/ptbxl_multilabel_classification_protoecgnet.py index b91316cf6..6ff1045c7 100644 --- a/examples/ptbxl_multilabel_classification_protoecgnet.py +++ b/examples/ptbxl_multilabel_classification_protoecgnet.py @@ -333,7 +333,7 @@ def main() -> None: remove_baseline=args.remove_baseline, ) samples = dataset.set_task(task) - splits = split_by_strat_fold(samples, folds=[range(1, 9), [9], [10]]) + splits = split_by_strat_fold(samples) print(f"samples: train={len(splits[0])} val={len(splits[1])} test={len(splits[2])}") baseline = run_pipeline( diff --git a/pyhealth/datasets/__init__.py b/pyhealth/datasets/__init__.py index 3fab22998..3a33ed9d5 100644 --- a/pyhealth/datasets/__init__.py +++ b/pyhealth/datasets/__init__.py @@ -89,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/protoecgnet.py b/pyhealth/models/protoecgnet.py index 434cea7c7..e63ff13e6 100644 --- a/pyhealth/models/protoecgnet.py +++ b/pyhealth/models/protoecgnet.py @@ -753,10 +753,12 @@ def set_stage(self, stage: str) -> None: ) 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 all(not p.requires_grad for p 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": diff --git a/pyhealth/tasks/ptbxl_multilabel_classification.py b/pyhealth/tasks/ptbxl_multilabel_classification.py index 3ef871368..b6911e935 100644 --- a/pyhealth/tasks/ptbxl_multilabel_classification.py +++ b/pyhealth/tasks/ptbxl_multilabel_classification.py @@ -25,9 +25,10 @@ def read_wfdb_signal(record_path: str) -> np.ndarray: """Reads a WFDB record and returns its physical signal. - The ``wfdb`` package is used when it is installed. Otherwise a small - fallback reader that supports the 16-bit format used by PTB-XL is used, so - that no extra dependency is required. + Delegates to :func:`~pyhealth.datasets.ptbxl.load_ptbxl_record` when + the optional ``wfdb`` extra is installed. Otherwise a small built-in + reader that supports the 16-bit format used by PTB-XL is used, so no + extra dependency is required for the task itself. Args: record_path: Path of the record without extension, e.g. @@ -45,13 +46,16 @@ def read_wfdb_signal(record_path: str) -> np.ndarray: >>> signal.shape # doctest: +SKIP (1000, 12) """ + # Prefer load_ptbxl_record (master's canonical reader) when wfdb is available. try: - import wfdb # type: ignore + from pyhealth.datasets.ptbxl import load_ptbxl_record - return wfdb.rdrecord(record_path).p_signal.astype(np.float32) + # load_ptbxl_record returns (n_leads, n_samples); transpose to compat shape. + return load_ptbxl_record(record_path).T except ImportError: pass + # Fallback: read 16-bit PTB-XL WFDB records without the wfdb package. header_path = record_path + ".hea" if not os.path.isfile(header_path): raise FileNotFoundError(f"WFDB header not found: {header_path}") @@ -101,6 +105,10 @@ def _aggregate_labels( ) -> list[str]: """Aggregate SCP codes to the chosen label set with likelihood filtering. + Extends :func:`~pyhealth.tasks.ptbxl.aggregate_diagnostic_superclasses` to + all six SCP-ECG label sets (diagnostic, subdiagnostic, superdiagnostic, + form, rhythm, all). + For ``"subdiagnostic"`` and ``"superdiagnostic"`` the max likelihood of all contributing codes is used per parent class (matching the original PTB-XL paper protocol). diff --git a/tests/core/test_protoecgnet.py b/tests/core/test_protoecgnet.py index e27157c57..54c962248 100644 --- a/tests/core/test_protoecgnet.py +++ b/tests/core/test_protoecgnet.py @@ -286,6 +286,39 @@ def test_feature_extractor_stage_trains_head_only_path(self): 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): diff --git a/tests/core/test_ptbxl.py b/tests/core/test_ptbxl.py index 51e7a549e..af01117f4 100644 --- a/tests/core/test_ptbxl.py +++ b/tests/core/test_ptbxl.py @@ -815,6 +815,80 @@ def test_aggregate_labels_likelihood_filter(self): 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_read_wfdb_signal_missing_header_raises(self): + from pyhealth.tasks.ptbxl_multilabel_classification import read_wfdb_signal + with self.assertRaises(FileNotFoundError): + read_wfdb_signal(str(self.root / "records100" / "00000" / "nonexistent")) + + 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_none_keeps_raw_scale(self): + from pyhealth.tasks import PTBXLMultilabelClassification + task = PTBXLMultilabelClassification( + scp_statements_path=self.scp_path, standardize=None + ) + rec = str(self.root / "records100" / "00000" / "00004_lr") + out = task._prepare_signal(rec, 100) + self.assertEqual(out.shape, (12, 1000)) + # Without standardization the mean is not forced to zero + self.assertIsNotNone(out) + if __name__ == "__main__": unittest.main() From ca0f9a3eaeeb7a227c676856034eb8fd5ee1efda Mon Sep 17 00:00:00 2001 From: Irvi Aini Date: Tue, 27 Oct 2026 09:30:49 +0800 Subject: [PATCH 08/14] chore: tighten PR per self-review before next bot cycle MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - examples: add --device {auto,cpu,cuda,mps} flag (default auto picks CUDA > MPS > CPU) so MacBook CPU-only results are reproducible with --device cpu instead of being accidentally overridden by MPS. - examples: seed random and numpy alongside torch at the start of run_pipeline so the 8-cell ablation grid is reproducible, not just the first run. - tests: add test_state_dict_round_trip_preserves_forward_output — nudges parameters off initializer, saves state_dict into a fresh model, and asserts forward outputs match exactly. Covers prototypes, classifier/fusion weights, and buffers. Not fixed in this commit (documented instead): projection_info is still not part of state_dict. The "every prediction explained by example" claim is softened in the PR description and the project_prototypes docstring (already) explains the save-separately workflow. The paper --paper-groups path remains unbenchmarked on real PTB-XL; the PR description now flags that explicitly. Signed-off-by: Irvi Aini --- ...l_multilabel_classification_protoecgnet.py | 21 ++++++++++---- tests/core/test_protoecgnet.py | 29 +++++++++++++++++++ 2 files changed, 44 insertions(+), 6 deletions(-) diff --git a/examples/ptbxl_multilabel_classification_protoecgnet.py b/examples/ptbxl_multilabel_classification_protoecgnet.py index 6ff1045c7..21c4d5e04 100644 --- a/examples/ptbxl_multilabel_classification_protoecgnet.py +++ b/examples/ptbxl_multilabel_classification_protoecgnet.py @@ -43,6 +43,7 @@ import argparse import itertools import os +import random import shutil import tempfile import warnings @@ -192,6 +193,9 @@ def run_pipeline( 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 = [ @@ -293,13 +297,18 @@ def main() -> None: "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 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" + 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. diff --git a/tests/core/test_protoecgnet.py b/tests/core/test_protoecgnet.py index 54c962248..db6be5afb 100644 --- a/tests/core/test_protoecgnet.py +++ b/tests/core/test_protoecgnet.py @@ -497,6 +497,35 @@ def test_project_prototypes_raises_on_shuffled_loader(self): with self.assertRaises(ValueError, msg="shuffled loader must raise"): model.project_prototypes(shuffled) + 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 buffers + (class_weights, label_cooccurrence, fusion_mask). + """ + 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, + ) + # 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() + 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) + if __name__ == "__main__": unittest.main() From be90e67dcc4f4e8b62495db909da840b7286ed36 Mon Sep 17 00:00:00 2001 From: Irvi Aini Date: Tue, 27 Oct 2026 10:13:11 +0800 Subject: [PATCH 09/14] fix: address bot review on class_weights buffer, interpretability claim, test coverage and dedup justifications MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - protoecgnet: register class_weights and label_cooccurrence as real-tensor buffers (ones/zeros with known shape) instead of None. Previously a checkpoint saved from a model with label statistics carried keys that a freshly constructed model lacked, so strict load_state_dict failed. Forward uses label_cooccurrence.any() to gate the contrastive term so defaulting to zeros is a no-op. - test_state_dict_round_trip_preserves_forward_output now actually calls set_label_statistics before saving, so the round-trip covers the two buffers its docstring claims to cover. - protoecgnet module docstring: correct the interpretability claim. Projected prototypes are feature-space representations; in global mode (proto_time_len=None) specifically, the stored vector is a time-averaged feature rather than a raw segment. Case-based explanations should display the matching window from the source sample, retrieved via the returned projection mapping. - ptbxl_multilabel_classification: explain in the docstrings why read_wfdb_signal keeps a fallback (so tests + --synthetic run without the optional wfdb extra) and why _aggregate_labels doesn't call through to aggregate_diagnostic_superclasses (the public helper doesn't take a label set or likelihood threshold). - tests: replace the weak test_standardize_none_keeps_raw_scale (asserted assertIsNotNone) with test_standardize_record_versus_none that checks per-lead mean≈0 / std≈1 for 'record' and non-zero mean for None. Add test_remove_baseline_changes_signal and three wfdb-import-blocked tests (unsupported format, missing .dat, successful roundtrip) that exercise the fallback reader branches CI's wfdb install would otherwise skip. Task coverage expected to cross 80% once CI re-runs. Signed-off-by: Irvi Aini --- pyhealth/models/protoecgnet.py | 24 ++-- .../tasks/ptbxl_multilabel_classification.py | 23 ++-- tests/core/test_protoecgnet.py | 18 ++- tests/core/test_ptbxl.py | 116 +++++++++++++++++- 4 files changed, 159 insertions(+), 22 deletions(-) diff --git a/pyhealth/models/protoecgnet.py b/pyhealth/models/protoecgnet.py index e63ff13e6..aea75ca31 100644 --- a/pyhealth/models/protoecgnet.py +++ b/pyhealth/models/protoecgnet.py @@ -8,10 +8,14 @@ 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 most -similar real training segment and returns the source-sample mapping to the -caller; store that alongside the checkpoint to support case-based explanations -after reloading. +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 @@ -638,8 +642,14 @@ def __init__( nn.Linear(total, self.num_labels) if fusion == "learned" else None ) self.register_buffer("fusion_mask", self._build_fusion_mask(total)) - self.register_buffer("class_weights", None) - self.register_buffer("label_cooccurrence", None) + # 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: @@ -821,7 +831,7 @@ def forward(self, **kwargs) -> dict[str, torch.Tensor]: for name, (_, sim, _) in outputs.items(): index = self.branch_index(name) cooc = None - if self.label_cooccurrence is not None and self.lam_cnrst > 0: + if self.lam_cnrst > 0 and self.label_cooccurrence.any(): cooc = self.label_cooccurrence[index][:, index] t = self.branches[name].prototype_terms(sim, y[:, index], cooc) extra.append( diff --git a/pyhealth/tasks/ptbxl_multilabel_classification.py b/pyhealth/tasks/ptbxl_multilabel_classification.py index b6911e935..8bab54711 100644 --- a/pyhealth/tasks/ptbxl_multilabel_classification.py +++ b/pyhealth/tasks/ptbxl_multilabel_classification.py @@ -25,10 +25,14 @@ def read_wfdb_signal(record_path: str) -> np.ndarray: """Reads a WFDB record and returns its physical signal. - Delegates to :func:`~pyhealth.datasets.ptbxl.load_ptbxl_record` when - the optional ``wfdb`` extra is installed. Otherwise a small built-in - reader that supports the 16-bit format used by PTB-XL is used, so no - extra dependency is required for the task itself. + Delegates to :func:`~pyhealth.datasets.ptbxl.load_ptbxl_record` (the + canonical PTB-XL reader in ``pyhealth.datasets.ptbxl``) when the optional + ``wfdb`` extra is installed, so the wfdb-backed decoding path is not + duplicated. When ``wfdb`` is not available, a built-in reader for the + 16-bit WFDB format used by PTB-XL is used so that this task's test suite + and the ``--synthetic`` example can run without the optional dependency — + neither path exists for PTB-XL elsewhere in the codebase, so the fallback + is kept here deliberately. Args: record_path: Path of the record without extension, e.g. @@ -105,9 +109,14 @@ def _aggregate_labels( ) -> list[str]: """Aggregate SCP codes to the chosen label set with likelihood filtering. - Extends :func:`~pyhealth.tasks.ptbxl.aggregate_diagnostic_superclasses` to - all six SCP-ECG label sets (diagnostic, subdiagnostic, superdiagnostic, - form, rhythm, all). + 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 diff --git a/tests/core/test_protoecgnet.py b/tests/core/test_protoecgnet.py index db6be5afb..4d18b816c 100644 --- a/tests/core/test_protoecgnet.py +++ b/tests/core/test_protoecgnet.py @@ -500,8 +500,13 @@ def test_project_prototypes_raises_on_shuffled_loader(self): 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 buffers - (class_weights, label_cooccurrence, fusion_mask). + 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"]} @@ -509,6 +514,12 @@ def test_state_dict_round_trip_preserves_forward_output(self): 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(): @@ -521,10 +532,13 @@ def test_state_dict_round_trip_preserves_forward_output(self): 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__": diff --git a/tests/core/test_ptbxl.py b/tests/core/test_ptbxl.py index af01117f4..da5460137 100644 --- a/tests/core/test_ptbxl.py +++ b/tests/core/test_ptbxl.py @@ -878,16 +878,120 @@ def test_diagnostic_samples_labels_nonempty(self): self.assertGreater(len(item["labels"]), 0) self.assertIn("strat_fold", item) - def test_standardize_none_keeps_raw_scale(self): + def test_standardize_record_versus_none(self): + """'record' forces per-lead mean≈0 and std≈1; None does not.""" from pyhealth.tasks import PTBXLMultilabelClassification - task = 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") - out = task._prepare_signal(rec, 100) - self.assertEqual(out.shape, (12, 1000)) - # Without standardization the mean is not forced to zero - self.assertIsNotNone(out) + + 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_fallback_reader_rejects_non_16bit_format(self): + """Fallback reader (no wfdb) raises ValueError on non-16-bit records.""" + import builtins + from pyhealth.tasks.ptbxl_multilabel_classification import read_wfdb_signal + + bad = Path(self.tmp) / "bad_fmt" + bad.parent.mkdir(parents=True, exist_ok=True) + name = bad.name + header = [f"{name} 12 100 1000"] + [ + # format 212 (not 16) — unsupported by the fallback. + f"{name}.dat 212 1000/mV 16 0 0 0 0 lead{i}" for i in range(12) + ] + Path(str(bad) + ".hea").write_text("\n".join(header) + "\n") + Path(str(bad) + ".dat").write_bytes(b"\x00" * 24000) + + real_import = builtins.__import__ + + def block_wfdb(name, *args, **kwargs): + if name == "wfdb" or name.startswith("wfdb."): + raise ImportError("blocked") + return real_import(name, *args, **kwargs) + + with patch("builtins.__import__", side_effect=block_wfdb): + with self.assertRaises(ValueError) as ctx: + read_wfdb_signal(str(bad)) + self.assertIn("Unsupported WFDB format", str(ctx.exception)) + + def test_fallback_reader_raises_on_missing_data_file(self): + """Fallback reader raises FileNotFoundError if the .dat is missing.""" + import builtins + from pyhealth.tasks.ptbxl_multilabel_classification import read_wfdb_signal + + hdr_only = Path(self.tmp) / "hdr_only" + hdr_only.parent.mkdir(parents=True, exist_ok=True) + name = hdr_only.name + header = [f"{name} 12 100 1000"] + [ + f"{name}.dat 16 1000/mV 16 0 0 0 0 lead{i}" for i in range(12) + ] + Path(str(hdr_only) + ".hea").write_text("\n".join(header) + "\n") + # Deliberately do NOT write the .dat file. + + real_import = builtins.__import__ + + def block_wfdb(name, *args, **kwargs): + if name == "wfdb" or name.startswith("wfdb."): + raise ImportError("blocked") + return real_import(name, *args, **kwargs) + + with patch("builtins.__import__", side_effect=block_wfdb): + with self.assertRaises(FileNotFoundError): + read_wfdb_signal(str(hdr_only)) + + def test_fallback_reader_roundtrip(self): + """Fallback reader (no wfdb) decodes the same shape as the delegated path.""" + import builtins + from pyhealth.tasks.ptbxl_multilabel_classification import read_wfdb_signal + + real_import = builtins.__import__ + + def block_wfdb(name, *args, **kwargs): + if name == "wfdb" or name.startswith("wfdb."): + raise ImportError("blocked") + return real_import(name, *args, **kwargs) + + rec = str(self.root / "records100" / "00000" / "00001_lr") + with patch("builtins.__import__", side_effect=block_wfdb): + sig = read_wfdb_signal(rec) + self.assertEqual(sig.shape, (1000, 12)) + self.assertEqual(sig.dtype, np.float32) if __name__ == "__main__": From 4dd101a205065457bc39ff2174183e424d3ccb6f Mon Sep 17 00:00:00 2001 From: Irvi Aini Date: Tue, 27 Oct 2026 22:06:49 +0800 Subject: [PATCH 10/14] fix: address second bot review on state_dict, interpretability claim, test rigor and docstring justifications MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The round-trip test added in the last commit passed only because class_weights and label_cooccurrence were never actually set — the underlying buffers were registered with register_buffer(name, None) and promoted to tensors only when supplied, so a checkpoint saved from a model with real label statistics carried keys a freshly constructed model lacked, and strict load_state_dict failed. Four fixes in response: - protoecgnet: register class_weights as torch.ones(num_labels) and label_cooccurrence as torch.zeros(num_labels, num_labels). ones is BCE's pos_weight identity, so the default is a no-op. The forward pass gates the contrastive term with bool(self.label_cooccurrence.any()), so an all-zeros buffer leaves the term off; the check is hoisted out of the per-branch loop so it fires once per forward pass. set_label_statistics docstring now spells out the buffer semantics. - test_state_dict_round_trip_preserves_forward_output: call set_label_statistics(class_weights=..., label_cooccurrence=...) before the save, so the test actually exercises the buffers its docstring claims to cover; also assert the two buffers match after load. - protoecgnet module docstring: correct the interpretability claim. Projected prototypes are feature-space representations, and in global mode (proto_time_len=None) the stored vector is a time-averaged feature rather than a raw ECG segment. Case-based explanations should display the matching window from the source sample, retrieved via the returned source-sample + window mapping. PR description updated to match. - ptbxl_multilabel_classification: strengthen the docstrings justifying read_wfdb_signal's fallback reader (needed so tests and --synthetic run without the optional wfdb extra; the wfdb-backed path still delegates to load_ptbxl_record) and _aggregate_labels (the public aggregate_diagnostic_superclasses doesn't take a label set or a likelihood threshold, so a call-through isn't viable). - tests: replace the weak test_standardize_none_keeps_raw_scale (asserted only assertIsNotNone) with test_standardize_record_versus_none that checks per-lead mean≈0 / std≈1 under 'record' and non-zero per-lead means under None. Add test_remove_baseline_changes_signal (filtered output differs, and its per-lead mean is closer to zero). Add three fallback-reader tests that block the wfdb import to force the fallback branches CI's wfdb install would otherwise skip (unsupported format, missing .dat, successful roundtrip). Task coverage expected to clear the 80% bar once CI reruns. Signed-off-by: Irvi Aini --- pyhealth/models/protoecgnet.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/pyhealth/models/protoecgnet.py b/pyhealth/models/protoecgnet.py index aea75ca31..5280dd3d7 100644 --- a/pyhealth/models/protoecgnet.py +++ b/pyhealth/models/protoecgnet.py @@ -704,6 +704,12 @@ def set_label_statistics( ) -> 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)``. @@ -828,11 +834,13 @@ def forward(self, **kwargs) -> dict[str, torch.Tensor]: if self.stage == "fusion": logits = self.fusion_classifier(sims) 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 = None - if self.lam_cnrst > 0 and self.label_cooccurrence.any(): - cooc = self.label_cooccurrence[index][:, index] + 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"] From f290e4879b8fac8d39d3e3091ad5878e417109e1 Mon Sep 17 00:00:00 2001 From: Irvi Aini Date: Tue, 27 Oct 2026 22:57:43 +0800 Subject: [PATCH 11/14] =?UTF-8?q?fix:=20address=20third=20bot=20review=20?= =?UTF-8?q?=E2=80=94=20single=20decoder,=20superclass=20reuse,=20=20=20hon?= =?UTF-8?q?or=20set=5Fpos=5Fweight,=20plus=20the=20four=20optional=20sugge?= =?UTF-8?q?stions?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three blocking issues: 1. Second WFDB decoder removed. read_wfdb_signal is deleted, the fallback 16-bit reader goes away, and _prepare_signal calls pyhealth.datasets.ptbxl.load_ptbxl_record directly — the same decoding path PTBXLSuperclassClassification uses. There is now one decoding path, so the decoded signal no longer depends on install state. test_ptbxl.py adds a module-level _WFDB_AVAILABLE check and @skipUnless on TestPTBXLMultilabelClassification; the three builtins.__import__-blocking fallback tests are removed because the code they covered is gone. A new test_load_signal_matches_superclass_task_signal asserts the multi-label task's _prepare_signal(standardize=None) matches load_ptbxl_record exactly. 2. Reuse PTBXLSuperclassClassification's aggregation. For label_set="superdiagnostic" with min_likelihood=0 the task's __call__ now routes through aggregate_diagnostic_superclasses + load_diagnostic_class_map, so the labels produced for the overlapping configuration match PTBXLSuperclassClassification record-for-record. Other label sets or non-zero likelihood thresholds keep using the generalising _aggregate_labels helper (which the public helper cannot replace — it doesn't take a label set or a likelihood threshold). Guarded by test_superdiagnostic_matches_superclass_task. 3. set_pos_weight no longer silently ignored. ProtoECGNet now overrides BaseModel.set_pos_weight to delegate to the base implementation (parses scalars, sequences, tensors and "balanced") and then mirrors self._pos_weight into self.class_weights. None resets class_weights to the ones default. The class docstring Note is corrected to describe the routing instead of a no-op. Guarded by test_set_pos_weight_routes_to_class_weights. Four optional suggestions: - docs/api/datasets.rst: revert the toctree title — master's titled entry (PTB-XL Dataset <...>) is restored and the untitled entry introduced on this branch is removed. - Example moved from examples/ to examples/ecg/ptbxl/ptbxl_multilabel_classification_protoecgnet.py (next to ptbxl_superclass_quickstart.py). PR description and run commands updated to the new path. - strat_fold now reads as int(float(event.strat_fold)) to match PTBXLSuperclassClassification and tolerate "9.0" strings from the CSV. - PR description now says results use one seed per cell and that the contrastive loss did not show a consistent gain over --no-contrastive at this compute budget; a multi-seed sweep would be needed to disentangle signal from seed noise. Signed-off-by: Irvi Aini --- docs/api/datasets.rst | 2 +- ...l_multilabel_classification_protoecgnet.py | 0 pyhealth/models/protoecgnet.py | 33 ++++- .../tasks/ptbxl_multilabel_classification.py | 125 ++++++---------- tests/core/test_protoecgnet.py | 25 ++++ tests/core/test_ptbxl.py | 134 ++++++++---------- 6 files changed, 156 insertions(+), 163 deletions(-) rename examples/{ => ecg/ptbxl}/ptbxl_multilabel_classification_protoecgnet.py (100%) diff --git a/docs/api/datasets.rst b/docs/api/datasets.rst index 41f795d22..a75b842ee 100644 --- a/docs/api/datasets.rst +++ b/docs/api/datasets.rst @@ -305,9 +305,9 @@ Available Datasets datasets/pyhealth.datasets.BMDHSDataset datasets/pyhealth.datasets.COVID19CXRDataset datasets/pyhealth.datasets.ChestXray14Dataset - datasets/pyhealth.datasets.PTBXLDataset datasets/pyhealth.datasets.PhysioNetDeIDDataset datasets/pyhealth.datasets.EEGBCIDataset + PTB-XL Dataset datasets/pyhealth.datasets.TUABDataset datasets/pyhealth.datasets.TUEVDataset datasets/pyhealth.datasets.ClinVarDataset diff --git a/examples/ptbxl_multilabel_classification_protoecgnet.py b/examples/ecg/ptbxl/ptbxl_multilabel_classification_protoecgnet.py similarity index 100% rename from examples/ptbxl_multilabel_classification_protoecgnet.py rename to examples/ecg/ptbxl/ptbxl_multilabel_classification_protoecgnet.py diff --git a/pyhealth/models/protoecgnet.py b/pyhealth/models/protoecgnet.py index 5280dd3d7..ed8c72732 100644 --- a/pyhealth/models/protoecgnet.py +++ b/pyhealth/models/protoecgnet.py @@ -547,9 +547,10 @@ class ProtoECGNet(BaseModel): Note: ``ProtoECGNet`` applies loss weighting via its own ``class_weights`` - parameter. The inherited - :meth:`~pyhealth.models.BaseModel.set_pos_weight` has no effect on - this model. + 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 @@ -697,6 +698,32 @@ 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. ``None`` resets + ``class_weights`` to ones. + + Args: + pos_weight: Same semantics as :meth:`BaseModel.set_pos_weight`. + dataset: Required when ``pos_weight="balanced"``. + """ + super().set_pos_weight(pos_weight, dataset) + if self._pos_weight is None: + self.class_weights = torch.ones(self.num_labels, device=self.device) + else: + self.class_weights = self._pos_weight.to(self.device).float() + def set_label_statistics( self, class_weights: torch.Tensor | None = None, diff --git a/pyhealth/tasks/ptbxl_multilabel_classification.py b/pyhealth/tasks/ptbxl_multilabel_classification.py index 8bab54711..caa13f4e7 100644 --- a/pyhealth/tasks/ptbxl_multilabel_classification.py +++ b/pyhealth/tasks/ptbxl_multilabel_classification.py @@ -2,7 +2,6 @@ import functools import os -import re from pathlib import Path from typing import Any, ClassVar @@ -19,80 +18,6 @@ "form", "rhythm", ) -_GAIN_PATTERN = re.compile(r"^([0-9.eE+-]+)(?:\(([-0-9]+)\))?(?:/.*)?$") - - -def read_wfdb_signal(record_path: str) -> np.ndarray: - """Reads a WFDB record and returns its physical signal. - - Delegates to :func:`~pyhealth.datasets.ptbxl.load_ptbxl_record` (the - canonical PTB-XL reader in ``pyhealth.datasets.ptbxl``) when the optional - ``wfdb`` extra is installed, so the wfdb-backed decoding path is not - duplicated. When ``wfdb`` is not available, a built-in reader for the - 16-bit WFDB format used by PTB-XL is used so that this task's test suite - and the ``--synthetic`` example can run without the optional dependency — - neither path exists for PTB-XL elsewhere in the codebase, so the fallback - is kept here deliberately. - - Args: - record_path: Path of the record without extension, e.g. - ``".../records100/00000/00001_lr"``. - - Returns: - A float32 array of shape ``(num_samples, num_leads)`` in millivolts. - - Raises: - FileNotFoundError: If the header or data file does not exist. - ValueError: If the record uses an unsupported WFDB storage format. - - Examples: - >>> signal = read_wfdb_signal("records100/00000/00001_lr") # doctest: +SKIP - >>> signal.shape # doctest: +SKIP - (1000, 12) - """ - # Prefer load_ptbxl_record (master's canonical reader) when wfdb is available. - try: - from pyhealth.datasets.ptbxl import load_ptbxl_record - - # load_ptbxl_record returns (n_leads, n_samples); transpose to compat shape. - return load_ptbxl_record(record_path).T - except ImportError: - pass - - # Fallback: read 16-bit PTB-XL WFDB records without the wfdb package. - header_path = record_path + ".hea" - if not os.path.isfile(header_path): - raise FileNotFoundError(f"WFDB header not found: {header_path}") - with open(header_path, "r", encoding="utf-8") as handle: - lines = [ln.strip() for ln in handle if ln.strip()] - lines = [ln for ln in lines if not ln.startswith("#")] - first = lines[0].split() - num_leads, num_samples = int(first[1]), int(first[3]) - gains, baselines, files = [], [], set() - for line in lines[1 : 1 + num_leads]: - fields = line.split() - files.add(fields[0]) - if fields[1].split("+")[0] != "16": - raise ValueError( - f"Unsupported WFDB format {fields[1]} in {header_path}; " - "install `wfdb` to read this record." - ) - match = _GAIN_PATTERN.match(fields[2]) - gain = float(match.group(1)) if match else 200.0 - gains.append(gain if gain != 0 else 200.0) - baselines.append( - int(match.group(2)) if match and match.group(2) else int(fields[4]) - ) - if len(files) != 1: - raise ValueError("Only records with a single signal file are supported.") - data_path = os.path.join(os.path.dirname(header_path), files.pop()) - if not os.path.isfile(data_path): - raise FileNotFoundError(f"WFDB data file not found: {data_path}") - digital = np.fromfile(data_path, dtype=" 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). @@ -244,7 +176,9 @@ def _prepare_signal(self, signal_file: str, sampling_rate: int) -> np.ndarray: Returns: Float32 array of shape ``(12, 10 * sampling_rate)``. """ - signal = read_wfdb_signal(signal_file).T # (leads, samples) + 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]))) @@ -263,6 +197,14 @@ def _prepare_signal(self, signal_file: str, sampling_rate: int) -> np.ndarray: 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`. @@ -274,20 +216,39 @@ def __call__(self, patient: Any) -> list[dict[str, Any]]: """ from pyhealth.datasets.ptbxl import parse_scp_codes - statements = _load_scp_statements(self.scp_statements_path) + 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"): - scp_codes = parse_scp_codes(event.scp_codes) - labels = _aggregate_labels( - scp_codes, statements, self.label_set, self.min_likelihood - ) + labels = aggregate(event.scp_codes) if not labels: continue sampling_rate = int(event.sampling_rate) samples.append( { "patient_id": patient.patient_id, - "strat_fold": int(event.strat_fold), + # 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, } diff --git a/tests/core/test_protoecgnet.py b/tests/core/test_protoecgnet.py index 4d18b816c..780974c2d 100644 --- a/tests/core/test_protoecgnet.py +++ b/tests/core/test_protoecgnet.py @@ -491,6 +491,31 @@ def test_prototype_contributions_sum_to_logit_fusion_stage(self): 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`` (regression: previously a silent no-op).""" + model = ProtoECGNet(self.dataset, prototypes_per_label=1, **TINY) + + # Scalar: broadcast to a (1,)-shape tensor, assigned verbatim. + model.set_pos_weight(4.0) + torch.testing.assert_close( + model.class_weights.detach().cpu(), torch.tensor([4.0]) + ) + + # Per-label sequence: length-6 vector covering self.dataset's labels. + 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) + ) + + # 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(model.num_labels), + ) + 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) diff --git a/tests/core/test_ptbxl.py b/tests/core/test_ptbxl.py index da5460137..dea5c8645 100644 --- a/tests/core/test_ptbxl.py +++ b/tests/core/test_ptbxl.py @@ -717,6 +717,17 @@ 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.""" @@ -773,13 +784,6 @@ def test_invalid_standardize_raises(self): scp_statements_path=self.scp_path, standardize="dataset" ) - def test_read_wfdb_signal_roundtrip(self): - from pyhealth.tasks.ptbxl_multilabel_classification import read_wfdb_signal - rec = str(self.root / "records100" / "00000" / "00001_lr") - sig = read_wfdb_signal(rec) - self.assertEqual(sig.shape[1], 12) - self.assertEqual(sig.dtype, np.float32) - def test_prepare_signal_shape_and_dtype(self): from pyhealth.tasks import PTBXLMultilabelClassification task = PTBXLMultilabelClassification( @@ -862,10 +866,39 @@ def test_form_label_set_empty_in_fixture(self): labels = _aggregate_labels({"NORM": 100.0, "IMI": 80.0}, stmts, "form", 0.0) self.assertEqual(labels, []) - def test_read_wfdb_signal_missing_header_raises(self): - from pyhealth.tasks.ptbxl_multilabel_classification import read_wfdb_signal - with self.assertRaises(FileNotFoundError): - read_wfdb_signal(str(self.root / "records100" / "00000" / "nonexistent")) + 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 @@ -923,75 +956,22 @@ def test_remove_baseline_changes_signal(self): float(np.abs(out_no_bl.mean(axis=1)).max()), ) - def test_fallback_reader_rejects_non_16bit_format(self): - """Fallback reader (no wfdb) raises ValueError on non-16-bit records.""" - import builtins - from pyhealth.tasks.ptbxl_multilabel_classification import read_wfdb_signal - - bad = Path(self.tmp) / "bad_fmt" - bad.parent.mkdir(parents=True, exist_ok=True) - name = bad.name - header = [f"{name} 12 100 1000"] + [ - # format 212 (not 16) — unsupported by the fallback. - f"{name}.dat 212 1000/mV 16 0 0 0 0 lead{i}" for i in range(12) - ] - Path(str(bad) + ".hea").write_text("\n".join(header) + "\n") - Path(str(bad) + ".dat").write_bytes(b"\x00" * 24000) - - real_import = builtins.__import__ - - def block_wfdb(name, *args, **kwargs): - if name == "wfdb" or name.startswith("wfdb."): - raise ImportError("blocked") - return real_import(name, *args, **kwargs) - - with patch("builtins.__import__", side_effect=block_wfdb): - with self.assertRaises(ValueError) as ctx: - read_wfdb_signal(str(bad)) - self.assertIn("Unsupported WFDB format", str(ctx.exception)) + def test_load_signal_matches_superclass_task_signal(self): + """Both tasks must load the exact same signal tensor from the same record. - def test_fallback_reader_raises_on_missing_data_file(self): - """Fallback reader raises FileNotFoundError if the .dat is missing.""" - import builtins - from pyhealth.tasks.ptbxl_multilabel_classification import read_wfdb_signal - - hdr_only = Path(self.tmp) / "hdr_only" - hdr_only.parent.mkdir(parents=True, exist_ok=True) - name = hdr_only.name - header = [f"{name} 12 100 1000"] + [ - f"{name}.dat 16 1000/mV 16 0 0 0 0 lead{i}" for i in range(12) - ] - Path(str(hdr_only) + ".hea").write_text("\n".join(header) + "\n") - # Deliberately do NOT write the .dat file. - - real_import = builtins.__import__ - - def block_wfdb(name, *args, **kwargs): - if name == "wfdb" or name.startswith("wfdb."): - raise ImportError("blocked") - return real_import(name, *args, **kwargs) - - with patch("builtins.__import__", side_effect=block_wfdb): - with self.assertRaises(FileNotFoundError): - read_wfdb_signal(str(hdr_only)) - - def test_fallback_reader_roundtrip(self): - """Fallback reader (no wfdb) decodes the same shape as the delegated path.""" - import builtins - from pyhealth.tasks.ptbxl_multilabel_classification import read_wfdb_signal - - real_import = builtins.__import__ - - def block_wfdb(name, *args, **kwargs): - if name == "wfdb" or name.startswith("wfdb."): - raise ImportError("blocked") - return real_import(name, *args, **kwargs) + 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") - with patch("builtins.__import__", side_effect=block_wfdb): - sig = read_wfdb_signal(rec) - self.assertEqual(sig.shape, (1000, 12)) - self.assertEqual(sig.dtype, np.float32) + prepped = task._prepare_signal(rec, 100) + canonical = load_ptbxl_record(rec)[:, :1000] + np.testing.assert_allclose(prepped, canonical, atol=1e-5) if __name__ == "__main__": From 2216e3a4e37cafc5b6d0e80c7919262f9328188d Mon Sep 17 00:00:00 2001 From: Irvi Aini Date: Wed, 28 Oct 2026 15:50:25 +0800 Subject: [PATCH 12/14] =?UTF-8?q?fix:=20address=20fourth=20bot=20review=20?= =?UTF-8?q?=E2=80=94=20stable=20class=5Fweights=20shape=20under=20=20=20se?= =?UTF-8?q?t=5Fpos=5Fweight,=20balanced=20test,=20example=20path,=20traine?= =?UTF-8?q?r-level=20test?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Blocker: set_pos_weight(scalar) left class_weights at shape (1,) because the base class's reshape(-1) produces a (1,)-tensor for a scalar. A subsequent state_dict() carried the (1,) shape, which failed strict load_state_dict into a freshly constructed model with the (num_labels,) default. Fix in protoecgnet.set_pos_weight: when the parsed weight has numel()==1, expand to (num_labels,) and clone; when it has numel()==num_labels, assign as-is; otherwise raise ValueError with a clear message. Docstring updated to describe the broadcast and the error case. Test coverage: - test_set_pos_weight_routes_to_class_weights now asserts class_weights.shape == (num_labels,) for both the scalar and the per-label paths, exercises the wrong-length ValueError, and ends with a state_dict round-trip after set_pos_weight(2.5) that verifies strict load_state_dict works. - Add test_set_pos_weight_balanced_uses_dataset: calls set_pos_weight("balanced", dataset), checks the resulting shape and that all weights are finite and positive, and checks that "balanced" without a dataset raises. - Add test_trainer_runs_through_staged_protocol: a real pyhealth.trainer.Trainer walks through feature_extractor → joint → project_prototypes → classifier, each for 1 epoch, and the final inference call must return finite logits across all labels. Optional: - examples/ecg/ptbxl/ptbxl_multilabel_classification_protoecgnet.py usage docstring: all three CLI examples now point at the file's current path (examples/ecg/ptbxl/...), not the old examples/ root. Signed-off-by: Irvi Aini --- ...l_multilabel_classification_protoecgnet.py | 6 +- pyhealth/models/protoecgnet.py | 22 ++++- tests/core/test_protoecgnet.py | 84 +++++++++++++++++-- 3 files changed, 99 insertions(+), 13 deletions(-) diff --git a/examples/ecg/ptbxl/ptbxl_multilabel_classification_protoecgnet.py b/examples/ecg/ptbxl/ptbxl_multilabel_classification_protoecgnet.py index 21c4d5e04..014444342 100644 --- a/examples/ecg/ptbxl/ptbxl_multilabel_classification_protoecgnet.py +++ b/examples/ecg/ptbxl/ptbxl_multilabel_classification_protoecgnet.py @@ -20,17 +20,17 @@ Usage: # Quick check on synthetic ECGs (no download needed, runs on CPU): - python examples/ptbxl_multilabel_classification_protoecgnet.py --synthetic + 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/ptbxl_multilabel_classification_protoecgnet.py \ + 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/ptbxl_multilabel_classification_protoecgnet.py \ + python examples/ecg/ptbxl/ptbxl_multilabel_classification_protoecgnet.py \ --root physionet.org/files/ptb-xl/1.0.3 --label-set all --paper-groups Note: diff --git a/pyhealth/models/protoecgnet.py b/pyhealth/models/protoecgnet.py index ed8c72732..63dae0e4e 100644 --- a/pyhealth/models/protoecgnet.py +++ b/pyhealth/models/protoecgnet.py @@ -711,18 +711,32 @@ def set_pos_weight( ``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. ``None`` resets - ``class_weights`` to ones. + ``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) - else: - self.class_weights = self._pos_weight.to(self.device).float() + 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, diff --git a/tests/core/test_protoecgnet.py b/tests/core/test_protoecgnet.py index 780974c2d..ca1c37b35 100644 --- a/tests/core/test_protoecgnet.py +++ b/tests/core/test_protoecgnet.py @@ -493,28 +493,57 @@ def test_prototype_contributions_sum_to_logit_fusion_stage(self): def test_set_pos_weight_routes_to_class_weights(self): """ProtoECGNet.set_pos_weight must mirror the parsed weight into - ``class_weights`` (regression: previously a silent no-op).""" + ``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 a (1,)-shape tensor, assigned verbatim. + # 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.tensor([4.0]) + model.class_weights.detach().cpu(), torch.full((n,), 4.0) ) - # Per-label sequence: length-6 vector covering self.dataset's labels. + # 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(model.num_labels), + 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_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) @@ -522,6 +551,49 @@ def test_project_prototypes_raises_on_shuffled_loader(self): 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. From 968f2cb8c887d5a6c295da34ce162f15dbe4cb7a Mon Sep 17 00:00:00 2001 From: Irvi Aini Date: Wed, 28 Oct 2026 17:06:06 +0800 Subject: [PATCH 13/14] fix: persist training stage in state_dict so learned-fusion checkpoints reload with the correct forward path Blocker: self.stage was a plain Python attribute, not part of state_dict. A ProtoECGNet(..., fusion="learned") trained and saved in the "fusion" stage reloaded into stage="all", so forward returned branch-scatter logits instead of fusion_classifier output and prototype_contributions returned branch-classifier contributions instead of fusion ones. Silent, no warning. Fix with the PyTorch-native hook pair: - ProtoECGNet.get_extra_state() returns {"stage": self.stage}; PyTorch stores this under "_extra_state" in state_dict. - ProtoECGNet.set_extra_state(state) calls self.set_stage(state["stage"]) on load, which also recomputes _frozen_modules so BN freezing stays correct after reload. Test coverage: The existing test_state_dict_round_trip_preserves_forward_output used fusion="learned" but never called set_stage("fusion"), so it couldn't catch this. Add test_state_dict_round_trip_preserves_fusion_stage: a model is moved into the fusion stage, parameters are nudged, then it is saved and loaded into a fresh model. Asserts clone.stage == "fusion" after load, logits match exactly, and prototype_contributions match per-branch (also pins the fusion-branch contribution path). Signed-off-by: Irvi Aini --- pyhealth/models/protoecgnet.py | 16 +++++++++++++ tests/core/test_protoecgnet.py | 43 ++++++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+) diff --git a/pyhealth/models/protoecgnet.py b/pyhealth/models/protoecgnet.py index 63dae0e4e..99042bb5e 100644 --- a/pyhealth/models/protoecgnet.py +++ b/pyhealth/models/protoecgnet.py @@ -781,6 +781,22 @@ def compute_label_statistics( 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. diff --git a/tests/core/test_protoecgnet.py b/tests/core/test_protoecgnet.py index ca1c37b35..ac587e699 100644 --- a/tests/core/test_protoecgnet.py +++ b/tests/core/test_protoecgnet.py @@ -532,6 +532,49 @@ def test_set_pos_weight_routes_to_class_weights(self): 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. From fee6fd0e01189725f4f3c25d25ffe28fdb0aa6fe Mon Sep 17 00:00:00 2001 From: Irvi Aini Date: Thu, 29 Oct 2026 20:39:38 +0800 Subject: [PATCH 14/14] fix: route logits through learned fusion head in the "all" stage; document class_weights difference from BaseModel.set_pos_weight Blocker 1 (medium): learned fusion head ignored outside "fusion" stage. With fusion="learned" and self.stage == "all" (the default, and the stage the class docstring describes as "trains everything with one loss"), forward() went through the branch-scatter path and never touched fusion_classifier. The fusion head stayed at random init and never received gradients, so the model silently behaved like fusion="average" with extra unused parameters. Same mismatch in prototype_contributions. Fix: add ProtoECGNet._uses_fusion_head() returning True iff fusion_classifier is not None and stage in {"all", "fusion"}. forward() and prototype_contributions() both route through fusion_classifier when this is True, keeping the paper's staged protocol intact for "joint" / "classifier" / "prototypes" (which still go through the per-branch classifiers so the branches get trained in isolation before the fusion head is fine-tuned). The fusion argument's docstring and prototype_contributions's Returns block are updated to point at _uses_fusion_head. Tests: - test_fusion_head_affects_logits_in_all_and_fusion_stages: perturb fusion_classifier.weight and check that logits change for stage in {"all", "fusion"} and are unchanged for stage in {"joint", "classifier"}. Catches regressions in both directions. - test_all_stage_trains_fusion_head_with_learned_fusion: backprop through model(**batch)["loss"] in the "all" stage must give fusion_classifier.weight a non-zero gradient. Blocker 2 (low): class_weights total differs from BaseModel.set_pos_weight. compute_label_statistics uses counts.sum() (sum of per-label positive counts) as the normalization total, following the paper's reference code. BaseModel.set_pos_weight("balanced") uses the sample count instead, and both end up in the same class_weights buffer. Docstring Note added spelling out the difference and recommending callers mix them deliberately rather than by accident. --- pyhealth/models/protoecgnet.py | 57 ++++++++++++++++++++++++++-------- tests/core/test_protoecgnet.py | 50 +++++++++++++++++++++++++++++ 2 files changed, 94 insertions(+), 13 deletions(-) diff --git a/pyhealth/models/protoecgnet.py b/pyhealth/models/protoecgnet.py index 99042bb5e..9971d223e 100644 --- a/pyhealth/models/protoecgnet.py +++ b/pyhealth/models/protoecgnet.py @@ -533,7 +533,9 @@ class ProtoECGNet(BaseModel): 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"`` - (train a linear layer over all prototype similarities in the + (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. @@ -770,9 +772,17 @@ def compute_label_statistics( labels: Multi-hot training labels ``(num_samples, num_labels)``. Returns: - ``(class_weights, cooccurrence)``. Class weights follow the + ``(class_weights, cooccurrence)``. Class weights follow the paper's reference code, ``(total - count) / count`` where ``total`` is the - sum of all label counts. + 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) @@ -841,6 +851,20 @@ def train(self, mode: bool = True) -> "ProtoECGNet": 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) @@ -887,9 +911,15 @@ def forward(self, **kwargs) -> dict[str, torch.Tensor]: 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) - logits = self._scatter({n: o[0] for n, o in outputs.items()}, batch) - if self.stage == "fusion": + # 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()) @@ -1017,18 +1047,19 @@ def prototype_contributions( Returns: For every branch a tensor equal to similarity times classifier weight. Shape is ``(batch, num_labels, num_branch_prototypes)`` - when ``stage="fusion"`` (fusion classifier weights; 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). + 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.stage == "fusion" and self.fusion_classifier is not None + use_fusion = self._uses_fusion_head() out = {} offset = 0 for name in self.branch_names: diff --git a/tests/core/test_protoecgnet.py b/tests/core/test_protoecgnet.py index ac587e699..f6eaf2688 100644 --- a/tests/core/test_protoecgnet.py +++ b/tests/core/test_protoecgnet.py @@ -362,6 +362,56 @@ def test_two_branches_with_learned_fusion(self): 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(