From a968b86b880e95aa1998845a95dfe4b96fd2eebb Mon Sep 17 00:00:00 2001 From: lehendo Date: Fri, 26 Jun 2026 01:17:10 -0700 Subject: [PATCH 1/3] New benchmarking wrapper --- benchmarks/__init__.py | 0 benchmarks/cpbench/__init__.py | 0 benchmarks/cpbench/benchmark.py | 547 ++++++++++++++++++++++++++++++++ 3 files changed, 547 insertions(+) create mode 100644 benchmarks/__init__.py create mode 100644 benchmarks/cpbench/__init__.py create mode 100644 benchmarks/cpbench/benchmark.py diff --git a/benchmarks/__init__.py b/benchmarks/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/benchmarks/cpbench/__init__.py b/benchmarks/cpbench/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/benchmarks/cpbench/benchmark.py b/benchmarks/cpbench/benchmark.py new file mode 100644 index 000000000..1f1b21b2a --- /dev/null +++ b/benchmarks/cpbench/benchmark.py @@ -0,0 +1,547 @@ +import argparse +import importlib +import json +import os +import sys +import time +from datetime import datetime +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np +import torch + +PROTOCOL = { + "version": "v1", + "split_ratios": [0.60, 0.15, 0.10, 0.15], + "split_seed": 42, + "batch_size": 32, + "standard_alphas": [0.05, 0.10, 0.20], +} + +TASK_REGISTRY: Dict[str, Dict[str, Any]] = { + "sleep_staging_isruc": { + "description": "5-class sleep staging on ISRUC-I (polysomnography)", + "dataset_class": "ISRUCDataset", + "dataset_module": "pyhealth.datasets", + "task_fn": "sleep_staging_isruc_fn", + "task_module": "pyhealth.tasks", + "model_class": "SparcNet", + "model_module": "pyhealth.models", + "feature_keys": ["signal"], + "label_key": "label", + "mode": "multiclass", + "split_strategy": "by_patient", + "monitor": "accuracy", + "monitor_criterion": "max", + "default_epochs": 30, + }, + "sleep_staging_sleepedf": { + "description": "5-class sleep staging on Sleep-EDF", + "dataset_class": "SleepEDFDataset", + "dataset_module": "pyhealth.datasets", + "task_fn": "sleep_staging_sleepedf_fn", + "task_module": "pyhealth.tasks", + "model_class": "SparcNet", + "model_module": "pyhealth.models", + "feature_keys": ["signal"], + "label_key": "label", + "mode": "multiclass", + "split_strategy": "by_patient", + "monitor": "accuracy", + "monitor_criterion": "max", + "default_epochs": 30, + }, + "sleep_staging_shhs": { + "description": "5-class sleep staging on SHHS", + "dataset_class": "SHHSDataset", + "dataset_module": "pyhealth.datasets", + "task_fn": "sleep_staging_shhs_fn", + "task_module": "pyhealth.tasks", + "model_class": "SparcNet", + "model_module": "pyhealth.models", + "feature_keys": ["signal"], + "label_key": "label", + "mode": "multiclass", + "split_strategy": "by_patient", + "monitor": "accuracy", + "monitor_criterion": "max", + "default_epochs": 30, + }, + "eeg_tuev": { + "description": "6-class EEG event detection on TUEV", + "dataset_class": "TUEVDataset", + "dataset_module": "pyhealth.datasets", + "task_fn": "EEGEventsTUEV", + "task_module": "pyhealth.tasks", + "model_class": "ContraWR", + "model_module": "pyhealth.models", + "feature_keys": ["signal"], + "label_key": "label", + "mode": "multiclass", + "split_strategy": "by_patient", + "monitor": "accuracy", + "monitor_criterion": "max", + "default_epochs": 30, + }, + "mortality_mimic4": { + "description": "In-hospital mortality prediction on MIMIC-IV (binary)", + "dataset_class": "MIMIC4Dataset", + "dataset_module": "pyhealth.datasets", + "task_fn": "mortality_prediction_mimic4_fn", + "task_module": "pyhealth.tasks", + "model_class": "Transformer", + "model_module": "pyhealth.models", + "feature_keys": ["conditions", "procedures", "drugs"], + "label_key": "label", + "mode": "binary", + "split_strategy": "by_patient", + "monitor": "roc_auc", + "monitor_criterion": "max", + "default_epochs": 20, + }, + "readmission_mimic4": { + "description": "30-day readmission prediction on MIMIC-IV (binary)", + "dataset_class": "MIMIC4Dataset", + "dataset_module": "pyhealth.datasets", + "task_fn": "readmission_prediction_mimic4_fn", + "task_module": "pyhealth.tasks", + "model_class": "Transformer", + "model_module": "pyhealth.models", + "feature_keys": ["conditions", "procedures", "drugs"], + "label_key": "label", + "mode": "binary", + "split_strategy": "by_patient", + "monitor": "roc_auc", + "monitor_criterion": "max", + "default_epochs": 20, + }, +} + +METHOD_REGISTRY: Dict[str, Dict[str, Any]] = { + "base": { + "description": "Split conformal prediction with APS score (Vovk et al. 2005)", + "class": "BaseConformal", + "module": "pyhealth.calib.predictionset", + "needs_embeddings": False, + "extra_kwargs": {"score_type": "aps"}, + "paper": "Vovk, Gammerman, Shafer. Algorithmic Learning in a Random World (2005)", + }, + "label": { + "description": "LABEL: Least Ambiguous set-valued classifiers with bounded error (Sadinle et al. 2019)", + "class": "LABEL", + "module": "pyhealth.calib.predictionset", + "needs_embeddings": False, + "extra_kwargs": {}, + "paper": "Sadinle, Lei, Wasserman. JASA (2019)", + }, + "cluster": { + "description": "Cluster-based conformal prediction with per-cluster thresholds", + "class": "ClusterLabel", + "module": "pyhealth.calib.predictionset", + "needs_embeddings": True, + "extra_kwargs": {"n_clusters": 5}, + "paper": "Cluster-based CP baseline", + }, + "neighborhood": { + "description": "Neighborhood-based conformal prediction using patient similarity", + "class": "NeighborhoodLabel", + "module": "pyhealth.calib.predictionset", + "needs_embeddings": True, + "extra_kwargs": {}, + "paper": "Neighborhood-based CP baseline", + }, + "covariate": { + "description": "Covariate shift-corrected conformal prediction via KDE (Tibshirani et al. 2019)", + "class": "CovariateLabel", + "module": "pyhealth.calib.predictionset", + "needs_embeddings": True, + "extra_kwargs": {}, + "paper": "Tibshirani, Barber, Candes, Ramdas. NeurIPS (2019)", + }, +} + + +def _import(module: str, name: str) -> Any: + mod = importlib.import_module(module) + return getattr(mod, name) + + +def _device() -> str: + return "cuda" if torch.cuda.is_available() else "cpu" + + +def _print_header(title: str) -> None: + width = 72 + print("\n" + "=" * width) + print(f" {title}") + print("=" * width) + + +def _print_section(title: str) -> None: + print(f"\n── {title} " + "─" * max(0, 68 - len(title))) + + +def load_dataset(task_cfg: Dict[str, Any], data_path: str, dev: bool) -> Any: + DatasetClass = _import(task_cfg["dataset_module"], task_cfg["dataset_class"]) + dataset = DatasetClass(root=data_path, dev=dev) + try: + task_fn = _import(task_cfg["task_module"], task_cfg["task_fn"]) + except AttributeError: + raise ValueError( + f"Task function '{task_cfg['task_fn']}' not found in " + f"'{task_cfg['task_module']}'. Check TASK_REGISTRY." + ) + dataset = dataset.set_task(task_fn) + return dataset + + +def split_dataset( + dataset: Any, + task_cfg: Dict[str, Any], + seed: int, + ratios: List[float], +) -> Tuple[Any, Any, Any, Any]: + from pyhealth.datasets import split_by_patient_conformal, split_by_sample_conformal + + strategy = task_cfg.get("split_strategy", "by_patient") + if strategy == "by_patient": + return split_by_patient_conformal(dataset, ratios=ratios, seed=seed) + elif strategy == "by_sample": + return split_by_sample_conformal(dataset, ratios=ratios, seed=seed) + else: + raise ValueError(f"Unknown split strategy: {strategy!r}") + + +def build_model(task_cfg: Dict[str, Any], dataset: Any, checkpoint: Optional[str]) -> Any: + ModelClass = _import(task_cfg["model_module"], task_cfg["model_class"]) + model = ModelClass( + dataset=dataset, + feature_keys=task_cfg["feature_keys"], + label_key=task_cfg["label_key"], + mode=task_cfg["mode"], + ) + if checkpoint: + state = torch.load(checkpoint, map_location="cpu") + if isinstance(state, dict) and "model_state_dict" in state: + state = state["model_state_dict"] + model.load_state_dict(state) + print(f" Loaded checkpoint: {checkpoint}") + return model + + +def train_model( + model: Any, + train_data: Any, + val_data: Any, + task_cfg: Dict[str, Any], + epochs: int, + batch_size: int, + output_dir: str, +) -> Any: + from pyhealth.datasets import get_dataloader + from pyhealth.trainer import Trainer + + train_dl = get_dataloader(train_data, batch_size=batch_size, shuffle=True) + val_dl = get_dataloader(val_data, batch_size=batch_size, shuffle=False) + + trainer = Trainer( + model=model, + device=_device(), + enable_logging=True, + output_path=output_dir, + exp_name="base_model", + ) + trainer.train( + train_dataloader=train_dl, + val_dataloader=val_dl, + epochs=epochs, + monitor=task_cfg.get("monitor"), + monitor_criterion=task_cfg.get("monitor_criterion", "max"), + load_best_model_at_last=True, + ) + return model + + +def build_cp_model( + method_cfg: Dict[str, Any], + base_model: Any, + alpha: float, + cal_data: Any, + test_data: Any, + batch_size: int, + dev: bool, +) -> Any: + CPClass = _import(method_cfg["module"], method_cfg["class"]) + extra = method_cfg.get("extra_kwargs", {}) + + cp_model = CPClass(model=base_model, alpha=alpha, debug=dev, **extra) + + if method_cfg["needs_embeddings"]: + cp_model.calibrate(cal_dataset=cal_data, test_dataset=test_data) + else: + cp_model.calibrate(cal_dataset=cal_data) + + return cp_model + + +def evaluate( + cp_model: Any, + test_data: Any, + alpha: float, + batch_size: int, + mode: str, +) -> Dict[str, float]: + from pyhealth.datasets import get_dataloader + from pyhealth.trainer import Trainer, get_metrics_fn + import pyhealth.metrics.prediction_set as pset_metrics + + test_dl = get_dataloader(test_data, batch_size=batch_size, shuffle=False) + trainer = Trainer(model=cp_model, device=_device(), enable_logging=False) + y_true_all, y_prob_all, loss, extra = trainer.inference( + test_dl, additional_outputs=["y_predset"] + ) + + y_predset = extra["y_predset"] + miscoverage = pset_metrics.miscoverage_overall_ps(y_predset, y_true_all) + empirical_coverage = 1.0 - miscoverage + coverage_gap = abs(alpha - miscoverage) + avg_set_size = float(pset_metrics.size(y_predset)) + rejection_rate = float(pset_metrics.rejection_rate(y_predset)) + + base_metrics = get_metrics_fn(mode)(y_true_all, y_prob_all, metrics=["accuracy"]) + accuracy = base_metrics.get("accuracy", float("nan")) + + return { + "alpha": alpha, + "empirical_coverage": round(float(empirical_coverage), 4), + "target_coverage": round(1.0 - alpha, 4), + "coverage_gap": round(float(coverage_gap), 4), + "avg_set_size": round(avg_set_size, 4), + "rejection_rate": round(rejection_rate, 4), + "accuracy": round(float(accuracy), 4), + "n_test": int(len(y_true_all)), + "n_cal": None, + } + + +def print_results_table(all_results: List[Dict[str, Any]], task: str, method: str) -> None: + _print_section("Results") + header = ( + f" {'alpha':>6} {'target_cov':>10} {'emp_cov':>8} " + f"{'cov_gap':>8} {'set_size':>9} {'reject%':>8} {'accuracy':>9}" + ) + print(f"\n Task: {task} Method: {method}") + print(f" Protocol: splits={PROTOCOL['split_ratios']} seed={PROTOCOL['split_seed']}") + print() + print(header) + print(" " + "-" * (len(header) - 2)) + for r in all_results: + flag = " ✓" if r["coverage_gap"] <= 0.01 else (" ✗" if r["coverage_gap"] > 0.05 else " ") + print( + f" {r['alpha']:>6.2f} {r['target_coverage']:>10.4f} " + f"{r['empirical_coverage']:>8.4f} {r['coverage_gap']:>8.4f} " + f"{r['avg_set_size']:>9.4f} {r['rejection_rate']:>8.4f} " + f"{r['accuracy']:>9.4f}{flag}" + ) + print() + print(" ✓ = coverage_gap ≤ 0.01 ✗ = coverage_gap > 0.05") + print(f" n_test = {all_results[0]['n_test']} n_cal = {all_results[0]['n_cal']}") + + +def save_results(all_results: List[Dict[str, Any]], task: str, method: str, output_path: str) -> None: + record = { + "cpbench_protocol": PROTOCOL, + "task": task, + "method": method, + "method_paper": METHOD_REGISTRY[method]["paper"], + "timestamp": datetime.utcnow().isoformat() + "Z", + "results": all_results, + } + with open(output_path, "w") as f: + json.dump(record, f, indent=2) + print(f"\n Results saved → {output_path}") + + +def run_benchmark( + task: str, + method: str, + data_path: str, + alphas: List[float], + seed: int, + epochs: int, + checkpoint: Optional[str], + output_path: Optional[str], + dev: bool, + output_dir: str, +) -> List[Dict[str, Any]]: + task_cfg = TASK_REGISTRY[task] + method_cfg = METHOD_REGISTRY[method] + batch_size = PROTOCOL["batch_size"] + ratios = PROTOCOL["split_ratios"] + + _print_header(f"CPBench | task: {task} | method: {method}") + + _print_section("Loading dataset") + print(f" {task_cfg['description']}") + print(f" data_path = {data_path}") + dataset = load_dataset(task_cfg, data_path, dev) + print(f" Dataset size: {len(dataset)} samples") + + _print_section("Splitting dataset (fixed protocol)") + print(f" ratios={ratios} seed={seed}") + train_data, val_data, cal_data, test_data = split_dataset(dataset, task_cfg, seed, ratios) + print( + f" train={len(train_data)} val={len(val_data)} " + f"cal={len(cal_data)} test={len(test_data)}" + ) + n_cal = len(cal_data) + + _print_section("Base model") + print(f" {task_cfg['model_class']} (mode={task_cfg['mode']})") + model = build_model(task_cfg, dataset, checkpoint) + + if checkpoint is None: + _print_section(f"Training base model ({epochs} epochs)") + t0 = time.time() + model = train_model(model, train_data, val_data, task_cfg, epochs, batch_size, output_dir) + print(f" Training time: {time.time() - t0:.1f}s") + else: + print(f" Skipping training (checkpoint provided)") + + model.to(_device()) + model.eval() + + all_results: List[Dict[str, Any]] = [] + + for alpha in alphas: + _print_section(f"Conformal method: {method} (α={alpha})") + print(f" {method_cfg['description']}") + print(f" Calibrating on {n_cal} samples...") + + t0 = time.time() + try: + cp_model = build_cp_model( + method_cfg, model, alpha, cal_data, test_data, batch_size, dev + ) + except TypeError: + cp_model = _import(method_cfg["module"], method_cfg["class"])( + model=model, + alpha=alpha, + debug=dev, + **method_cfg.get("extra_kwargs", {}), + ) + cp_model.calibrate(cal_dataset=cal_data) + + cp_model.to(_device()) + cal_time = time.time() - t0 + print(f" Calibration time: {cal_time:.1f}s") + + print(f" Evaluating on {len(test_data)} test samples...") + result = evaluate(cp_model, test_data, alpha, batch_size, task_cfg["mode"]) + result["n_cal"] = n_cal + result["cal_time_s"] = round(cal_time, 2) + all_results.append(result) + + print_results_table(all_results, task, method) + + if output_path: + save_results(all_results, task, method, output_path) + + return all_results + + +def list_registry() -> None: + print("\nRegistered Tasks:") + print("-" * 60) + for name, cfg in TASK_REGISTRY.items(): + print(f" {name:<30} {cfg['description']}") + + print("\nRegistered CP Methods:") + print("-" * 60) + for name, cfg in METHOD_REGISTRY.items(): + print(f" {name:<14} {cfg['description']}") + print(f" {'':14} Paper: {cfg['paper']}") + print() + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="CPBench: Standardized conformal prediction benchmarking for PyHealth", + ) + + parser.add_argument("--list", action="store_true", help="List all registered tasks and methods, then exit.") + parser.add_argument("--task", type=str, choices=list(TASK_REGISTRY.keys()), help="Task name (see --list).") + parser.add_argument("--method", type=str, choices=list(METHOD_REGISTRY.keys()), help="CP method name (see --list).") + parser.add_argument("--data-path", type=str, help="Path to the raw dataset on disk.") + parser.add_argument( + "--alpha", + type=float, + nargs="+", + default=None, + help=f"Miscoverage rate(s). Defaults to standard set {PROTOCOL['standard_alphas']}.", + ) + parser.add_argument( + "--seed", + type=int, + default=PROTOCOL["split_seed"], + help=f"Split seed. Default: {PROTOCOL['split_seed']}. Override only for sensitivity analysis.", + ) + parser.add_argument( + "--epochs", + type=int, + default=None, + help="Training epochs. Defaults to per-task default. Ignored when --checkpoint is given.", + ) + parser.add_argument("--checkpoint", type=str, default=None, help="Pre-trained model checkpoint (.pth). Skips training.") + parser.add_argument("--output", type=str, default=None, help="Path to write results JSON.") + parser.add_argument("--output-dir", type=str, default="./cpbench_runs", help="Directory for Trainer logs.") + parser.add_argument("--dev", action="store_true", help="Dev mode: tiny dataset subset, single alpha (0.10).") + + return parser.parse_args() + + +def main() -> None: + args = parse_args() + + if args.list: + list_registry() + sys.exit(0) + + missing = [f for f in ("task", "method", "data_path") if not getattr(args, f, None)] + if missing: + print(f"Error: --{', --'.join(m.replace('_', '-') for m in missing)} required.") + sys.exit(1) + + if args.alpha is not None: + alphas = args.alpha + if alphas != PROTOCOL["standard_alphas"]: + print( + f" [cpbench] Note: using non-standard alpha={alphas}. " + f"Published results must include all of {PROTOCOL['standard_alphas']}." + ) + elif args.dev: + alphas = [0.10] + else: + alphas = PROTOCOL["standard_alphas"] + + task_cfg = TASK_REGISTRY[args.task] + epochs = args.epochs if args.epochs is not None else task_cfg.get("default_epochs", 20) + + os.makedirs(args.output_dir, exist_ok=True) + + run_benchmark( + task=args.task, + method=args.method, + data_path=args.data_path, + alphas=alphas, + seed=args.seed, + epochs=epochs, + checkpoint=args.checkpoint, + output_path=args.output, + dev=args.dev, + output_dir=args.output_dir, + ) + + +if __name__ == "__main__": + main() From 4db5e3a51d5d61b2434180f2268e198ee543b641 Mon Sep 17 00:00:00 2001 From: lehendo Date: Fri, 26 Jun 2026 01:27:37 -0700 Subject: [PATCH 2/3] Add multiple seeds and 0.01 alpha --- benchmarks/cpbench/benchmark.py | 284 ++++++++++++++++---------------- 1 file changed, 142 insertions(+), 142 deletions(-) diff --git a/benchmarks/cpbench/benchmark.py b/benchmarks/cpbench/benchmark.py index 1f1b21b2a..daf08e41d 100644 --- a/benchmarks/cpbench/benchmark.py +++ b/benchmarks/cpbench/benchmark.py @@ -13,9 +13,9 @@ PROTOCOL = { "version": "v1", "split_ratios": [0.60, 0.15, 0.10, 0.15], - "split_seed": 42, + "seeds": [42, 43, 44, 45, 46], "batch_size": 32, - "standard_alphas": [0.05, 0.10, 0.20], + "standard_alphas": [0.01, 0.05, 0.10, 0.20], } TASK_REGISTRY: Dict[str, Dict[str, Any]] = { @@ -83,38 +83,6 @@ "monitor_criterion": "max", "default_epochs": 30, }, - "mortality_mimic4": { - "description": "In-hospital mortality prediction on MIMIC-IV (binary)", - "dataset_class": "MIMIC4Dataset", - "dataset_module": "pyhealth.datasets", - "task_fn": "mortality_prediction_mimic4_fn", - "task_module": "pyhealth.tasks", - "model_class": "Transformer", - "model_module": "pyhealth.models", - "feature_keys": ["conditions", "procedures", "drugs"], - "label_key": "label", - "mode": "binary", - "split_strategy": "by_patient", - "monitor": "roc_auc", - "monitor_criterion": "max", - "default_epochs": 20, - }, - "readmission_mimic4": { - "description": "30-day readmission prediction on MIMIC-IV (binary)", - "dataset_class": "MIMIC4Dataset", - "dataset_module": "pyhealth.datasets", - "task_fn": "readmission_prediction_mimic4_fn", - "task_module": "pyhealth.tasks", - "model_class": "Transformer", - "model_module": "pyhealth.models", - "feature_keys": ["conditions", "procedures", "drugs"], - "label_key": "label", - "mode": "binary", - "split_strategy": "by_patient", - "monitor": "roc_auc", - "monitor_criterion": "max", - "default_epochs": 20, - }, } METHOD_REGISTRY: Dict[str, Dict[str, Any]] = { @@ -160,6 +128,14 @@ }, } +AGGREGATE_METRICS = [ + "empirical_coverage", + "coverage_gap", + "avg_set_size", + "rejection_rate", + "accuracy", +] + def _import(module: str, name: str) -> Any: mod = importlib.import_module(module) @@ -212,21 +188,14 @@ def split_dataset( raise ValueError(f"Unknown split strategy: {strategy!r}") -def build_model(task_cfg: Dict[str, Any], dataset: Any, checkpoint: Optional[str]) -> Any: +def build_model(task_cfg: Dict[str, Any], dataset: Any) -> Any: ModelClass = _import(task_cfg["model_module"], task_cfg["model_class"]) - model = ModelClass( + return ModelClass( dataset=dataset, feature_keys=task_cfg["feature_keys"], label_key=task_cfg["label_key"], mode=task_cfg["mode"], ) - if checkpoint: - state = torch.load(checkpoint, map_location="cpu") - if isinstance(state, dict) and "model_state_dict" in state: - state = state["model_state_dict"] - model.load_state_dict(state) - print(f" Loaded checkpoint: {checkpoint}") - return model def train_model( @@ -237,6 +206,7 @@ def train_model( epochs: int, batch_size: int, output_dir: str, + seed: int, ) -> Any: from pyhealth.datasets import get_dataloader from pyhealth.trainer import Trainer @@ -249,7 +219,7 @@ def train_model( device=_device(), enable_logging=True, output_path=output_dir, - exp_name="base_model", + exp_name=f"seed_{seed}", ) trainer.train( train_dataloader=train_dl, @@ -268,12 +238,10 @@ def build_cp_model( alpha: float, cal_data: Any, test_data: Any, - batch_size: int, dev: bool, ) -> Any: CPClass = _import(method_cfg["module"], method_cfg["class"]) extra = method_cfg.get("extra_kwargs", {}) - cp_model = CPClass(model=base_model, alpha=alpha, debug=dev, **extra) if method_cfg["needs_embeddings"]: @@ -290,72 +258,115 @@ def evaluate( alpha: float, batch_size: int, mode: str, -) -> Dict[str, float]: +) -> Dict[str, Any]: from pyhealth.datasets import get_dataloader from pyhealth.trainer import Trainer, get_metrics_fn import pyhealth.metrics.prediction_set as pset_metrics test_dl = get_dataloader(test_data, batch_size=batch_size, shuffle=False) trainer = Trainer(model=cp_model, device=_device(), enable_logging=False) - y_true_all, y_prob_all, loss, extra = trainer.inference( + y_true_all, y_prob_all, _, extra = trainer.inference( test_dl, additional_outputs=["y_predset"] ) y_predset = extra["y_predset"] miscoverage = pset_metrics.miscoverage_overall_ps(y_predset, y_true_all) - empirical_coverage = 1.0 - miscoverage - coverage_gap = abs(alpha - miscoverage) - avg_set_size = float(pset_metrics.size(y_predset)) - rejection_rate = float(pset_metrics.rejection_rate(y_predset)) base_metrics = get_metrics_fn(mode)(y_true_all, y_prob_all, metrics=["accuracy"]) - accuracy = base_metrics.get("accuracy", float("nan")) return { "alpha": alpha, - "empirical_coverage": round(float(empirical_coverage), 4), "target_coverage": round(1.0 - alpha, 4), - "coverage_gap": round(float(coverage_gap), 4), - "avg_set_size": round(avg_set_size, 4), - "rejection_rate": round(rejection_rate, 4), - "accuracy": round(float(accuracy), 4), + "empirical_coverage": round(float(1.0 - miscoverage), 4), + "coverage_gap": round(float(abs(alpha - miscoverage)), 4), + "avg_set_size": round(float(pset_metrics.size(y_predset)), 4), + "rejection_rate": round(float(pset_metrics.rejection_rate(y_predset)), 4), + "accuracy": round(float(base_metrics.get("accuracy", float("nan"))), 4), "n_test": int(len(y_true_all)), "n_cal": None, } -def print_results_table(all_results: List[Dict[str, Any]], task: str, method: str) -> None: +def aggregate_results( + per_seed_results: List[List[Dict[str, Any]]], + alphas: List[float], +) -> List[Dict[str, Any]]: + aggregated = [] + for i, alpha in enumerate(alphas): + alpha_rows = [seed_rows[i] for seed_rows in per_seed_results] + row: Dict[str, Any] = { + "alpha": alpha, + "target_coverage": round(1.0 - alpha, 4), + "n_seeds": len(alpha_rows), + "n_test": alpha_rows[0]["n_test"], + "n_cal": alpha_rows[0]["n_cal"], + } + for metric in AGGREGATE_METRICS: + vals = [r[metric] for r in alpha_rows] + row[f"{metric}_mean"] = round(float(np.mean(vals)), 4) + row[f"{metric}_std"] = round(float(np.std(vals)), 4) + aggregated.append(row) + return aggregated + + +def print_aggregated_table( + aggregated: List[Dict[str, Any]], + task: str, + method: str, + seeds: List[int], +) -> None: _print_section("Results") + print(f"\n Task: {task} Method: {method}") + print(f" Protocol: splits={PROTOCOL['split_ratios']} seeds={seeds}") + print(f" Reporting mean ± std across {len(seeds)} seeds\n") + + col = 16 header = ( - f" {'alpha':>6} {'target_cov':>10} {'emp_cov':>8} " - f"{'cov_gap':>8} {'set_size':>9} {'reject%':>8} {'accuracy':>9}" + f" {'alpha':>6} {'target_cov':>10} " + f"{'emp_cov':<{col}} {'cov_gap':<{col}} " + f"{'set_size':<{col}} {'reject%':<{col}} {'accuracy':<{col}}" ) - print(f"\n Task: {task} Method: {method}") - print(f" Protocol: splits={PROTOCOL['split_ratios']} seed={PROTOCOL['split_seed']}") - print() print(header) print(" " + "-" * (len(header) - 2)) - for r in all_results: - flag = " ✓" if r["coverage_gap"] <= 0.01 else (" ✗" if r["coverage_gap"] > 0.05 else " ") + + for r in aggregated: + def fmt(m: str) -> str: + return f"{r[f'{m}_mean']:.4f} ± {r[f'{m}_std']:.4f}" + + gap_mean = r["coverage_gap_mean"] + flag = " ✓" if gap_mean <= 0.01 else (" ✗" if gap_mean > 0.05 else " ") + print( f" {r['alpha']:>6.2f} {r['target_coverage']:>10.4f} " - f"{r['empirical_coverage']:>8.4f} {r['coverage_gap']:>8.4f} " - f"{r['avg_set_size']:>9.4f} {r['rejection_rate']:>8.4f} " - f"{r['accuracy']:>9.4f}{flag}" + f"{fmt('empirical_coverage'):<{col}} {fmt('coverage_gap'):<{col}} " + f"{fmt('avg_set_size'):<{col}} {fmt('rejection_rate'):<{col}} " + f"{fmt('accuracy'):<{col}}{flag}" ) + print() - print(" ✓ = coverage_gap ≤ 0.01 ✗ = coverage_gap > 0.05") - print(f" n_test = {all_results[0]['n_test']} n_cal = {all_results[0]['n_cal']}") + print(" ✓ = mean coverage_gap ≤ 0.01 ✗ = mean coverage_gap > 0.05") + print(f" n_test ≈ {aggregated[0]['n_test']} n_cal ≈ {aggregated[0]['n_cal']}") -def save_results(all_results: List[Dict[str, Any]], task: str, method: str, output_path: str) -> None: +def save_results( + aggregated: List[Dict[str, Any]], + per_seed_results: List[List[Dict[str, Any]]], + seeds: List[int], + task: str, + method: str, + output_path: str, +) -> None: record = { "cpbench_protocol": PROTOCOL, "task": task, "method": method, "method_paper": METHOD_REGISTRY[method]["paper"], "timestamp": datetime.utcnow().isoformat() + "Z", - "results": all_results, + "seeds": seeds, + "aggregated": aggregated, + "per_seed": { + str(seed): per_seed_results[i] for i, seed in enumerate(seeds) + }, } with open(output_path, "w") as f: json.dump(record, f, indent=2) @@ -367,9 +378,8 @@ def run_benchmark( method: str, data_path: str, alphas: List[float], - seed: int, + seeds: List[int], epochs: int, - checkpoint: Optional[str], output_path: Optional[str], dev: bool, output_dir: str, @@ -380,6 +390,7 @@ def run_benchmark( ratios = PROTOCOL["split_ratios"] _print_header(f"CPBench | task: {task} | method: {method}") + print(f" Seeds: {seeds} Alphas: {alphas}") _print_section("Loading dataset") print(f" {task_cfg['description']}") @@ -387,67 +398,62 @@ def run_benchmark( dataset = load_dataset(task_cfg, data_path, dev) print(f" Dataset size: {len(dataset)} samples") - _print_section("Splitting dataset (fixed protocol)") - print(f" ratios={ratios} seed={seed}") - train_data, val_data, cal_data, test_data = split_dataset(dataset, task_cfg, seed, ratios) - print( - f" train={len(train_data)} val={len(val_data)} " - f"cal={len(cal_data)} test={len(test_data)}" - ) - n_cal = len(cal_data) + per_seed_results: List[List[Dict[str, Any]]] = [] + + for seed in seeds: + _print_header(f"Seed {seed} ({seeds.index(seed) + 1}/{len(seeds)})") - _print_section("Base model") - print(f" {task_cfg['model_class']} (mode={task_cfg['mode']})") - model = build_model(task_cfg, dataset, checkpoint) + _print_section("Splitting") + print(f" ratios={ratios} seed={seed}") + train_data, val_data, cal_data, test_data = split_dataset(dataset, task_cfg, seed, ratios) + print( + f" train={len(train_data)} val={len(val_data)} " + f"cal={len(cal_data)} test={len(test_data)}" + ) + n_cal = len(cal_data) - if checkpoint is None: - _print_section(f"Training base model ({epochs} epochs)") + _print_section(f"Training {task_cfg['model_class']} ({epochs} epochs)") + model = build_model(task_cfg, dataset) t0 = time.time() - model = train_model(model, train_data, val_data, task_cfg, epochs, batch_size, output_dir) + model = train_model(model, train_data, val_data, task_cfg, epochs, batch_size, output_dir, seed) print(f" Training time: {time.time() - t0:.1f}s") - else: - print(f" Skipping training (checkpoint provided)") + model.to(_device()) + model.eval() - model.to(_device()) - model.eval() + seed_rows: List[Dict[str, Any]] = [] - all_results: List[Dict[str, Any]] = [] + for alpha in alphas: + _print_section(f"α={alpha} method={method}") + print(f" Calibrating on {n_cal} samples...") - for alpha in alphas: - _print_section(f"Conformal method: {method} (α={alpha})") - print(f" {method_cfg['description']}") - print(f" Calibrating on {n_cal} samples...") + t0 = time.time() + try: + cp_model = build_cp_model(method_cfg, model, alpha, cal_data, test_data, dev) + except TypeError: + cp_model = _import(method_cfg["module"], method_cfg["class"])( + model=model, alpha=alpha, debug=dev, **method_cfg.get("extra_kwargs", {}) + ) + cp_model.calibrate(cal_dataset=cal_data) - t0 = time.time() - try: - cp_model = build_cp_model( - method_cfg, model, alpha, cal_data, test_data, batch_size, dev - ) - except TypeError: - cp_model = _import(method_cfg["module"], method_cfg["class"])( - model=model, - alpha=alpha, - debug=dev, - **method_cfg.get("extra_kwargs", {}), - ) - cp_model.calibrate(cal_dataset=cal_data) - - cp_model.to(_device()) - cal_time = time.time() - t0 - print(f" Calibration time: {cal_time:.1f}s") - - print(f" Evaluating on {len(test_data)} test samples...") - result = evaluate(cp_model, test_data, alpha, batch_size, task_cfg["mode"]) - result["n_cal"] = n_cal - result["cal_time_s"] = round(cal_time, 2) - all_results.append(result) - - print_results_table(all_results, task, method) + cp_model.to(_device()) + cal_time = time.time() - t0 + print(f" Calibration: {cal_time:.1f}s | Evaluating on {len(test_data)} samples...") + + result = evaluate(cp_model, test_data, alpha, batch_size, task_cfg["mode"]) + result["n_cal"] = n_cal + result["cal_time_s"] = round(cal_time, 2) + result["seed"] = seed + seed_rows.append(result) + + per_seed_results.append(seed_rows) + + aggregated = aggregate_results(per_seed_results, alphas) + print_aggregated_table(aggregated, task, method, seeds) if output_path: - save_results(all_results, task, method, output_path) + save_results(aggregated, per_seed_results, seeds, task, method, output_path) - return all_results + return aggregated def list_registry() -> None: @@ -481,21 +487,21 @@ def parse_args() -> argparse.Namespace: help=f"Miscoverage rate(s). Defaults to standard set {PROTOCOL['standard_alphas']}.", ) parser.add_argument( - "--seed", + "--seeds", type=int, - default=PROTOCOL["split_seed"], - help=f"Split seed. Default: {PROTOCOL['split_seed']}. Override only for sensitivity analysis.", + nargs="+", + default=None, + help=f"Seeds for splitting. Defaults to {PROTOCOL['seeds']}.", ) parser.add_argument( "--epochs", type=int, default=None, - help="Training epochs. Defaults to per-task default. Ignored when --checkpoint is given.", + help="Training epochs. Defaults to per-task default.", ) - parser.add_argument("--checkpoint", type=str, default=None, help="Pre-trained model checkpoint (.pth). Skips training.") parser.add_argument("--output", type=str, default=None, help="Path to write results JSON.") parser.add_argument("--output-dir", type=str, default="./cpbench_runs", help="Directory for Trainer logs.") - parser.add_argument("--dev", action="store_true", help="Dev mode: tiny dataset subset, single alpha (0.10).") + parser.add_argument("--dev", action="store_true", help="Dev mode: tiny subset, single seed (42), single alpha (0.10).") return parser.parse_args() @@ -512,17 +518,12 @@ def main() -> None: print(f"Error: --{', --'.join(m.replace('_', '-') for m in missing)} required.") sys.exit(1) - if args.alpha is not None: - alphas = args.alpha - if alphas != PROTOCOL["standard_alphas"]: - print( - f" [cpbench] Note: using non-standard alpha={alphas}. " - f"Published results must include all of {PROTOCOL['standard_alphas']}." - ) - elif args.dev: + if args.dev: alphas = [0.10] + seeds = [42] else: - alphas = PROTOCOL["standard_alphas"] + alphas = args.alpha if args.alpha is not None else PROTOCOL["standard_alphas"] + seeds = args.seeds if args.seeds is not None else PROTOCOL["seeds"] task_cfg = TASK_REGISTRY[args.task] epochs = args.epochs if args.epochs is not None else task_cfg.get("default_epochs", 20) @@ -534,9 +535,8 @@ def main() -> None: method=args.method, data_path=args.data_path, alphas=alphas, - seed=args.seed, + seeds=seeds, epochs=epochs, - checkpoint=args.checkpoint, output_path=args.output, dev=args.dev, output_dir=args.output_dir, From c091fe37dd5f44c7a45a62c96c20bf8ec2588855 Mon Sep 17 00:00:00 2001 From: lehendo Date: Mon, 20 Jul 2026 01:20:20 -0700 Subject: [PATCH 3/3] Complete refactor of benchmarking wrapper --- benchmarks/cpbench/benchmark.py | 786 +++++++++++------- pyhealth/calib/base_classes.py | 25 +- .../predictionset/base_conformal/__init__.py | 15 +- .../predictionset/cluster/cluster_label.py | 30 +- .../cluster/neighborhood_label.py | 8 + .../covariate/covariate_label.py | 95 ++- .../calib/predictionset/favmac/__init__.py | 8 +- pyhealth/calib/predictionset/label.py | 15 +- .../calib/predictionset/scrib/__init__.py | 8 +- 9 files changed, 635 insertions(+), 355 deletions(-) diff --git a/benchmarks/cpbench/benchmark.py b/benchmarks/cpbench/benchmark.py index daf08e41d..f287a4b0f 100644 --- a/benchmarks/cpbench/benchmark.py +++ b/benchmarks/cpbench/benchmark.py @@ -1,146 +1,128 @@ +""" +CPBench: a generic conformal-prediction benchmark harness for PyHealth. + +Rather than a curated, hand-maintained list of pre-wired (dataset, task, +model, CP method) combinations, CPBench lets you name ANY BaseDataset +subclass, ANY BaseTask subclass, ANY BaseModel subclass, and ANY +SetPredictor (conformal prediction-set method) subclass already available +in pyhealth, and it trains + calibrates + evaluates that exact combination. +Adding a new dataset, task, model, or CP method to pyhealth automatically +makes it available here too -- nothing in this file needs to change, as +long as the new piece follows pyhealth's existing BaseDataset/BaseTask/ +BaseModel/SetPredictor conventions. + +Usage: + List everything available: + python benchmark.py --list + + Full benchmark run: + python benchmark.py \\ + --dataset TUEVDataset --dataset-root /path/to/tuev \\ + --task EEGEventsTUEV \\ + --model ContraWR \\ + --cp-method LABEL \\ + --alpha 0.05 0.1 --seeds 42 43 --epochs 30 + + Quick smoke-test (tiny data subset, single seed/alpha, skips training): + python benchmark.py \\ + --dataset TUEVDataset --dataset-root /path/to/tuev \\ + --task EEGEventsTUEV --model ContraWR --cp-method LABEL --dev + + Dry run (resolve and validate the combination without touching data): + python benchmark.py \\ + --dataset TUEVDataset --task EEGEventsTUEV \\ + --model ContraWR --cp-method LABEL --dry-run + + A dataset/task/model needing extra constructor args (beyond `root=`, + no-args, or `dataset=`) can be given them via JSON passthrough, e.g.: + --dataset-kwargs '{"subset": "cassette"}' + --task-kwargs '{"chunk_duration": 30.0}' + --model-kwargs '{"embedding_dim": 256}' + --cp-kwargs '{"n_clusters": 5}' +""" + import argparse import importlib +import inspect import json import os import sys import time -from datetime import datetime +from datetime import datetime, timezone from typing import Any, Dict, List, Optional, Tuple import numpy as np import torch PROTOCOL = { - "version": "v1", + "version": "v2", "split_ratios": [0.60, 0.15, 0.10, 0.15], "seeds": [42, 43, 44, 45, 46], "batch_size": 32, "standard_alphas": [0.01, 0.05, 0.10, 0.20], } -TASK_REGISTRY: Dict[str, Dict[str, Any]] = { - "sleep_staging_isruc": { - "description": "5-class sleep staging on ISRUC-I (polysomnography)", - "dataset_class": "ISRUCDataset", - "dataset_module": "pyhealth.datasets", - "task_fn": "sleep_staging_isruc_fn", - "task_module": "pyhealth.tasks", - "model_class": "SparcNet", - "model_module": "pyhealth.models", - "feature_keys": ["signal"], - "label_key": "label", - "mode": "multiclass", - "split_strategy": "by_patient", - "monitor": "accuracy", - "monitor_criterion": "max", - "default_epochs": 30, - }, - "sleep_staging_sleepedf": { - "description": "5-class sleep staging on Sleep-EDF", - "dataset_class": "SleepEDFDataset", - "dataset_module": "pyhealth.datasets", - "task_fn": "sleep_staging_sleepedf_fn", - "task_module": "pyhealth.tasks", - "model_class": "SparcNet", - "model_module": "pyhealth.models", - "feature_keys": ["signal"], - "label_key": "label", - "mode": "multiclass", - "split_strategy": "by_patient", - "monitor": "accuracy", - "monitor_criterion": "max", - "default_epochs": 30, - }, - "sleep_staging_shhs": { - "description": "5-class sleep staging on SHHS", - "dataset_class": "SHHSDataset", - "dataset_module": "pyhealth.datasets", - "task_fn": "sleep_staging_shhs_fn", - "task_module": "pyhealth.tasks", - "model_class": "SparcNet", - "model_module": "pyhealth.models", - "feature_keys": ["signal"], - "label_key": "label", - "mode": "multiclass", - "split_strategy": "by_patient", - "monitor": "accuracy", - "monitor_criterion": "max", - "default_epochs": 30, - }, - "eeg_tuev": { - "description": "6-class EEG event detection on TUEV", - "dataset_class": "TUEVDataset", - "dataset_module": "pyhealth.datasets", - "task_fn": "EEGEventsTUEV", - "task_module": "pyhealth.tasks", - "model_class": "ContraWR", - "model_module": "pyhealth.models", - "feature_keys": ["signal"], - "label_key": "label", - "mode": "multiclass", - "split_strategy": "by_patient", - "monitor": "accuracy", - "monitor_criterion": "max", - "default_epochs": 30, - }, -} -METHOD_REGISTRY: Dict[str, Dict[str, Any]] = { - "base": { - "description": "Split conformal prediction with APS score (Vovk et al. 2005)", - "class": "BaseConformal", - "module": "pyhealth.calib.predictionset", - "needs_embeddings": False, - "extra_kwargs": {"score_type": "aps"}, - "paper": "Vovk, Gammerman, Shafer. Algorithmic Learning in a Random World (2005)", - }, - "label": { - "description": "LABEL: Least Ambiguous set-valued classifiers with bounded error (Sadinle et al. 2019)", - "class": "LABEL", - "module": "pyhealth.calib.predictionset", - "needs_embeddings": False, - "extra_kwargs": {}, - "paper": "Sadinle, Lei, Wasserman. JASA (2019)", - }, - "cluster": { - "description": "Cluster-based conformal prediction with per-cluster thresholds", - "class": "ClusterLabel", - "module": "pyhealth.calib.predictionset", - "needs_embeddings": True, - "extra_kwargs": {"n_clusters": 5}, - "paper": "Cluster-based CP baseline", - }, - "neighborhood": { - "description": "Neighborhood-based conformal prediction using patient similarity", - "class": "NeighborhoodLabel", - "module": "pyhealth.calib.predictionset", - "needs_embeddings": True, - "extra_kwargs": {}, - "paper": "Neighborhood-based CP baseline", - }, - "covariate": { - "description": "Covariate shift-corrected conformal prediction via KDE (Tibshirani et al. 2019)", - "class": "CovariateLabel", - "module": "pyhealth.calib.predictionset", - "needs_embeddings": True, - "extra_kwargs": {}, - "paper": "Tibshirani, Barber, Candes, Ramdas. NeurIPS (2019)", - }, -} +# -------------------------------------------------------------------------- +# Discovery: enumerate everything CPBench can plug together, straight from +# pyhealth's own public API. Nothing here is a hand-maintained registry -- +# if a class is exported from the relevant pyhealth package and subclasses +# the right base class, it is automatically available. +# -------------------------------------------------------------------------- + +def _discover(package_name: str, base_class_path: str) -> Dict[str, type]: + """Return {class_name: class} for every public class in `package_name` + that subclasses the class at `base_class_path` ("module.ClassName"), + excluding the base class itself.""" + module = importlib.import_module(package_name) + base_module_name, base_class_name = base_class_path.rsplit(".", 1) + base_class = getattr(importlib.import_module(base_module_name), base_class_name) + + found = {} + for name in dir(module): + if name.startswith("_"): + continue + obj = getattr(module, name) + if ( + inspect.isclass(obj) + and issubclass(obj, base_class) + and obj is not base_class + ): + found[name] = obj + return found + + +def discover_datasets() -> Dict[str, type]: + return _discover("pyhealth.datasets", "pyhealth.datasets.BaseDataset") + + +def discover_tasks() -> Dict[str, type]: + return _discover("pyhealth.tasks", "pyhealth.tasks.BaseTask") + + +def discover_models() -> Dict[str, type]: + return _discover("pyhealth.models", "pyhealth.models.BaseModel") -AGGREGATE_METRICS = [ - "empirical_coverage", - "coverage_gap", - "avg_set_size", - "rejection_rate", - "accuracy", -] +def discover_cp_methods() -> Dict[str, type]: + return _discover( + "pyhealth.calib.predictionset", "pyhealth.calib.base_classes.SetPredictor" + ) + + +def _resolve(registry: Dict[str, type], name: str, kind: str) -> type: + if name not in registry: + choices = ", ".join(sorted(registry.keys())) or "(none found)" + raise ValueError( + f"Unknown {kind} '{name}'.\nAvailable {kind}s: {choices}\n" + "(Run with --list to see this with descriptions.)" + ) + return registry[name] -def _import(module: str, name: str) -> Any: - mod = importlib.import_module(module) - return getattr(mod, name) +# -------------------------------------------------------------------------- +# Small utilities +# -------------------------------------------------------------------------- def _device() -> str: return "cuda" if torch.cuda.is_available() else "cpu" @@ -157,56 +139,105 @@ def _print_section(title: str) -> None: print(f"\n── {title} " + "─" * max(0, 68 - len(title))) -def load_dataset(task_cfg: Dict[str, Any], data_path: str, dev: bool) -> Any: - DatasetClass = _import(task_cfg["dataset_module"], task_cfg["dataset_class"]) - dataset = DatasetClass(root=data_path, dev=dev) +def _parse_kwargs_json(raw: Optional[str], flag_name: str) -> Dict[str, Any]: + if raw is None: + return {} try: - task_fn = _import(task_cfg["task_module"], task_cfg["task_fn"]) - except AttributeError: - raise ValueError( - f"Task function '{task_cfg['task_fn']}' not found in " - f"'{task_cfg['task_module']}'. Check TASK_REGISTRY." - ) - dataset = dataset.set_task(task_fn) - return dataset + parsed = json.loads(raw) + except json.JSONDecodeError as e: + raise ValueError(f"{flag_name} must be a valid JSON object: {e}") from e + if not isinstance(parsed, dict): + raise ValueError(f"{flag_name} must be a JSON object, got {type(parsed).__name__}") + return parsed + + +# -------------------------------------------------------------------------- +# Generic construction helpers -- each one introspects the target class's +# constructor rather than assuming a fixed argument set, and raises a clear, +# actionable error (pointing at the relevant --*-kwargs flag) if the given +# combination genuinely doesn't fit together, instead of guessing. +# -------------------------------------------------------------------------- + +def build_dataset( + dataset_cls: type, + root: Optional[str], + dev: bool, + dataset_kwargs: Dict[str, Any], +) -> Any: + ctor_params = inspect.signature(dataset_cls.__init__).parameters + accepts_catchall_kwargs = any( + p.kind == inspect.Parameter.VAR_KEYWORD for p in ctor_params.values() + ) + + kwargs = dict(dataset_kwargs) + if root is not None: + kwargs.setdefault("root", root) + if "dev" not in kwargs: + if "dev" in ctor_params or accepts_catchall_kwargs: + kwargs["dev"] = dev + elif dev: + print( + f" Note: {dataset_cls.__name__} does not support a 'dev' " + "flag; loading the full dataset." + ) + + try: + return dataset_cls(**kwargs) + except TypeError as e: + raise TypeError( + f"Failed to construct {dataset_cls.__name__}(**{kwargs}): {e}\n" + "Pass extra/renamed constructor args via --dataset-kwargs " + "(a JSON object)." + ) from e + + +def build_task(task_cls: type, task_kwargs: Dict[str, Any]) -> Any: + try: + return task_cls(**task_kwargs) + except TypeError as e: + raise TypeError( + f"Failed to construct {task_cls.__name__}(**{task_kwargs}): {e}\n" + "Pass its constructor args via --task-kwargs (a JSON object)." + ) from e + + +def build_model(model_cls: type, dataset: Any, model_kwargs: Dict[str, Any]) -> Any: + try: + return model_cls(dataset=dataset, **model_kwargs) + except TypeError as e: + raise TypeError( + f"Failed to construct {model_cls.__name__}(dataset=..., " + f"**{model_kwargs}): {e}\n" + "Pass extra constructor args via --model-kwargs (a JSON object)." + ) from e def split_dataset( - dataset: Any, - task_cfg: Dict[str, Any], - seed: int, - ratios: List[float], + dataset: Any, split_strategy: str, seed: int, ratios: List[float] ) -> Tuple[Any, Any, Any, Any]: from pyhealth.datasets import split_by_patient_conformal, split_by_sample_conformal - strategy = task_cfg.get("split_strategy", "by_patient") - if strategy == "by_patient": + if split_strategy == "by_patient": return split_by_patient_conformal(dataset, ratios=ratios, seed=seed) - elif strategy == "by_sample": + elif split_strategy == "by_sample": return split_by_sample_conformal(dataset, ratios=ratios, seed=seed) else: - raise ValueError(f"Unknown split strategy: {strategy!r}") - - -def build_model(task_cfg: Dict[str, Any], dataset: Any) -> Any: - ModelClass = _import(task_cfg["model_module"], task_cfg["model_class"]) - return ModelClass( - dataset=dataset, - feature_keys=task_cfg["feature_keys"], - label_key=task_cfg["label_key"], - mode=task_cfg["mode"], - ) + raise ValueError( + f"Unknown split strategy: {split_strategy!r}. Use 'by_patient' " + "or 'by_sample' (--split-strategy)." + ) def train_model( model: Any, train_data: Any, val_data: Any, - task_cfg: Dict[str, Any], epochs: int, batch_size: int, output_dir: str, seed: int, + monitor: Optional[str], + monitor_criterion: str, ) -> Any: from pyhealth.datasets import get_dataloader from pyhealth.trainer import Trainer @@ -225,37 +256,91 @@ def train_model( train_dataloader=train_dl, val_dataloader=val_dl, epochs=epochs, - monitor=task_cfg.get("monitor"), - monitor_criterion=task_cfg.get("monitor_criterion", "max"), - load_best_model_at_last=True, + monitor=monitor, + monitor_criterion=monitor_criterion, + load_best_model_at_last=(monitor is not None), ) return model +def construct_cp_predictor( + cp_cls: type, + base_model: Any, + alpha: Optional[float], + cp_kwargs: Dict[str, Any], + dev: bool, +) -> Tuple[Any, bool]: + """Construct (but do not calibrate) any SetPredictor generically, with + clear error wrapping. + + Construction only needs `base_model.mode` (set synchronously in + BaseModel.__init__ from the dataset's schema) -- it does not require the + model to be trained. This lets callers validate a (model, CP method) + pairing cheaply, before paying for training. + + Not every CP method is parameterized by a scalar/array `alpha` (e.g. + SCRIB uses `risk`, FavMac uses `target_cost`/`delta`). Returns + (cp_model, used_alpha): used_alpha is False when the class has no + `alpha` constructor parameter, in which case its real parameters must be + supplied via --cp-kwargs. + """ + ctor_params = inspect.signature(cp_cls.__init__).parameters + kwargs = dict(cp_kwargs) + used_alpha = "alpha" in ctor_params + if used_alpha and "alpha" not in kwargs: + kwargs["alpha"] = alpha + if "debug" in ctor_params and "debug" not in kwargs: + kwargs["debug"] = dev + + try: + cp_model = cp_cls(model=base_model, **kwargs) + except NotImplementedError as e: + raise NotImplementedError( + f"{cp_cls.__name__}(model=..., **{kwargs}) failed to construct " + f"for a model with mode='{base_model.mode}': {e}\n" + "This usually means the CP method's required classification " + "mode (see its docstring) doesn't match the task's output mode." + ) from e + except TypeError as e: + raise TypeError( + f"Failed to construct {cp_cls.__name__}(model=..., **{kwargs}): " + f"{e}\nCheck its constructor signature (see its docstring) and " + "pass required args via --cp-kwargs (a JSON object)." + ) from e + + return cp_model, used_alpha + + def build_cp_model( - method_cfg: Dict[str, Any], + cp_cls: type, base_model: Any, - alpha: float, + alpha: Optional[float], + cp_kwargs: Dict[str, Any], + train_data: Any, cal_data: Any, test_data: Any, dev: bool, -) -> Any: - CPClass = _import(method_cfg["module"], method_cfg["class"]) - extra = method_cfg.get("extra_kwargs", {}) - cp_model = CPClass(model=base_model, alpha=alpha, debug=dev, **extra) - - if method_cfg["needs_embeddings"]: - cp_model.calibrate(cal_dataset=cal_data, test_dataset=test_data) - else: - cp_model.calibrate(cal_dataset=cal_data) - - return cp_model +) -> Tuple[Any, bool]: + """Construct and calibrate any SetPredictor generically. + + Every SetPredictor in pyhealth.calib.predictionset is expected to accept + calibrate(cal_dataset, train_dataset=None, test_dataset=None) -- see + pyhealth.calib.base_classes.SetPredictor -- so this function never needs + to know which extra data a particular CP method needs; each method + extracts whatever it needs internally. + """ + cp_model, used_alpha = construct_cp_predictor(cp_cls, base_model, alpha, cp_kwargs, dev) + cp_model.calibrate( + cal_dataset=cal_data, train_dataset=train_data, test_dataset=test_data + ) + return cp_model, used_alpha def evaluate( cp_model: Any, test_data: Any, - alpha: float, + alpha: Optional[float], + used_alpha: bool, batch_size: int, mode: str, ) -> Dict[str, Any]: @@ -268,17 +353,12 @@ def evaluate( y_true_all, y_prob_all, _, extra = trainer.inference( test_dl, additional_outputs=["y_predset"] ) - y_predset = extra["y_predset"] - miscoverage = pset_metrics.miscoverage_overall_ps(y_predset, y_true_all) + # These are computed purely from the output prediction sets, so they're + # meaningful for every SetPredictor regardless of its calibration target. base_metrics = get_metrics_fn(mode)(y_true_all, y_prob_all, metrics=["accuracy"]) - - return { - "alpha": alpha, - "target_coverage": round(1.0 - alpha, 4), - "empirical_coverage": round(float(1.0 - miscoverage), 4), - "coverage_gap": round(float(abs(alpha - miscoverage)), 4), + result: Dict[str, Any] = { "avg_set_size": round(float(pset_metrics.size(y_predset)), 4), "rejection_rate": round(float(pset_metrics.rejection_rate(y_predset)), 4), "accuracy": round(float(base_metrics.get("accuracy", float("nan"))), 4), @@ -286,23 +366,43 @@ def evaluate( "n_cal": None, } + # Coverage-vs-target reporting only makes sense for methods parameterized + # by a scalar miscoverage rate alpha. + if used_alpha and alpha is not None: + miscoverage = pset_metrics.miscoverage_overall_ps(y_predset, y_true_all) + result["alpha"] = alpha + result["target_coverage"] = round(1.0 - alpha, 4) + result["empirical_coverage"] = round(float(1.0 - miscoverage), 4) + result["coverage_gap"] = round(float(abs(alpha - miscoverage)), 4) -def aggregate_results( - per_seed_results: List[List[Dict[str, Any]]], - alphas: List[float], -) -> List[Dict[str, Any]]: + return result + + +def aggregate_results(per_seed_results: List[List[Dict[str, Any]]]) -> List[Dict[str, Any]]: + """Aggregate mean/std across seeds for each alpha (or the single + no-alpha run, for methods without a scalar alpha).""" + n_rows = len(per_seed_results[0]) aggregated = [] - for i, alpha in enumerate(alphas): - alpha_rows = [seed_rows[i] for seed_rows in per_seed_results] + for i in range(n_rows): + rows = [seed_rows[i] for seed_rows in per_seed_results] row: Dict[str, Any] = { - "alpha": alpha, - "target_coverage": round(1.0 - alpha, 4), - "n_seeds": len(alpha_rows), - "n_test": alpha_rows[0]["n_test"], - "n_cal": alpha_rows[0]["n_cal"], + "n_seeds": len(rows), + "n_test": rows[0]["n_test"], + "n_cal": rows[0]["n_cal"], } - for metric in AGGREGATE_METRICS: - vals = [r[metric] for r in alpha_rows] + if "alpha" in rows[0]: + row["alpha"] = rows[0]["alpha"] + row["target_coverage"] = rows[0]["target_coverage"] + + metric_names = [ + m for m in ( + "empirical_coverage", "coverage_gap", + "avg_set_size", "rejection_rate", "accuracy", + ) + if m in rows[0] + ] + for metric in metric_names: + vals = [r[metric] for r in rows] row[f"{metric}_mean"] = round(float(np.mean(vals)), 4) row[f"{metric}_std"] = round(float(np.std(vals)), 4) aggregated.append(row) @@ -311,21 +411,30 @@ def aggregate_results( def print_aggregated_table( aggregated: List[Dict[str, Any]], - task: str, - method: str, + dataset_name: str, + task_name: str, + model_name: str, + cp_method_name: str, seeds: List[int], ) -> None: _print_section("Results") - print(f"\n Task: {task} Method: {method}") + print(f"\n Dataset={dataset_name} Task={task_name} Model={model_name} CP method={cp_method_name}") print(f" Protocol: splits={PROTOCOL['split_ratios']} seeds={seeds}") print(f" Reporting mean ± std across {len(seeds)} seeds\n") + has_alpha = "alpha" in aggregated[0] col = 16 - header = ( - f" {'alpha':>6} {'target_cov':>10} " - f"{'emp_cov':<{col}} {'cov_gap':<{col}} " - f"{'set_size':<{col}} {'reject%':<{col}} {'accuracy':<{col}}" - ) + + if has_alpha: + header = ( + f" {'alpha':>6} {'target_cov':>10} " + f"{'emp_cov':<{col}} {'cov_gap':<{col}} " + f"{'set_size':<{col}} {'reject%':<{col}} {'accuracy':<{col}}" + ) + else: + header = ( + f" {'set_size':<{col}} {'reject%':<{col}} {'accuracy':<{col}}" + ) print(header) print(" " + "-" * (len(header) - 2)) @@ -333,18 +442,24 @@ def print_aggregated_table( def fmt(m: str) -> str: return f"{r[f'{m}_mean']:.4f} ± {r[f'{m}_std']:.4f}" - gap_mean = r["coverage_gap_mean"] - flag = " ✓" if gap_mean <= 0.01 else (" ✗" if gap_mean > 0.05 else " ") - - print( - f" {r['alpha']:>6.2f} {r['target_coverage']:>10.4f} " - f"{fmt('empirical_coverage'):<{col}} {fmt('coverage_gap'):<{col}} " - f"{fmt('avg_set_size'):<{col}} {fmt('rejection_rate'):<{col}} " - f"{fmt('accuracy'):<{col}}{flag}" - ) + if has_alpha: + gap_mean = r["coverage_gap_mean"] + flag = " ✓" if gap_mean <= 0.01 else (" ✗" if gap_mean > 0.05 else " ") + print( + f" {r['alpha']:>6.2f} {r['target_coverage']:>10.4f} " + f"{fmt('empirical_coverage'):<{col}} {fmt('coverage_gap'):<{col}} " + f"{fmt('avg_set_size'):<{col}} {fmt('rejection_rate'):<{col}} " + f"{fmt('accuracy'):<{col}}{flag}" + ) + else: + print( + f" {fmt('avg_set_size'):<{col}} {fmt('rejection_rate'):<{col}} " + f"{fmt('accuracy'):<{col}}" + ) print() - print(" ✓ = mean coverage_gap ≤ 0.01 ✗ = mean coverage_gap > 0.05") + if has_alpha: + print(" ✓ = mean coverage_gap ≤ 0.01 ✗ = mean coverage_gap > 0.05") print(f" n_test ≈ {aggregated[0]['n_test']} n_cal ≈ {aggregated[0]['n_cal']}") @@ -352,156 +467,226 @@ def save_results( aggregated: List[Dict[str, Any]], per_seed_results: List[List[Dict[str, Any]]], seeds: List[int], - task: str, - method: str, + dataset_name: str, + task_name: str, + model_name: str, + cp_method_name: str, output_path: str, ) -> None: record = { "cpbench_protocol": PROTOCOL, - "task": task, - "method": method, - "method_paper": METHOD_REGISTRY[method]["paper"], - "timestamp": datetime.utcnow().isoformat() + "Z", + "dataset": dataset_name, + "task": task_name, + "model": model_name, + "cp_method": cp_method_name, + "timestamp": datetime.now(timezone.utc).isoformat(), "seeds": seeds, "aggregated": aggregated, - "per_seed": { - str(seed): per_seed_results[i] for i, seed in enumerate(seeds) - }, + "per_seed": {str(seed): per_seed_results[i] for i, seed in enumerate(seeds)}, } with open(output_path, "w") as f: json.dump(record, f, indent=2) print(f"\n Results saved → {output_path}") -def run_benchmark( - task: str, - method: str, - data_path: str, - alphas: List[float], - seeds: List[int], - epochs: int, - output_path: Optional[str], - dev: bool, - output_dir: str, -) -> List[Dict[str, Any]]: - task_cfg = TASK_REGISTRY[task] - method_cfg = METHOD_REGISTRY[method] +def run_benchmark(args: argparse.Namespace) -> List[Dict[str, Any]]: + datasets = discover_datasets() + tasks = discover_tasks() + models = discover_models() + cp_methods = discover_cp_methods() + + DatasetClass = _resolve(datasets, args.dataset, "dataset") + TaskClass = _resolve(tasks, args.task, "task") + ModelClass = _resolve(models, args.model, "model") + CPClass = _resolve(cp_methods, args.cp_method, "CP method") + + dataset_kwargs = _parse_kwargs_json(args.dataset_kwargs, "--dataset-kwargs") + task_kwargs = _parse_kwargs_json(args.task_kwargs, "--task-kwargs") + model_kwargs = _parse_kwargs_json(args.model_kwargs, "--model-kwargs") + cp_kwargs = _parse_kwargs_json(args.cp_kwargs, "--cp-kwargs") + + if args.dry_run: + _print_header( + f"CPBench dry-run | {args.dataset} + {args.task} + {args.model} " + f"+ {args.cp_method}" + ) + print(f" dataset kwargs: {dataset_kwargs}") + print(f" task kwargs: {task_kwargs}") + print(f" model kwargs: {model_kwargs}") + print(f" cp kwargs: {cp_kwargs}") + print("\n All four classes resolved successfully. Re-run without " + "--dry-run to actually train/calibrate/evaluate.") + return [] + batch_size = PROTOCOL["batch_size"] ratios = PROTOCOL["split_ratios"] - _print_header(f"CPBench | task: {task} | method: {method}") - print(f" Seeds: {seeds} Alphas: {alphas}") + _print_header( + f"CPBench | dataset: {args.dataset} | task: {args.task} | " + f"model: {args.model} | cp-method: {args.cp_method}" + ) + print(f" Seeds: {args.seeds} Alphas: {args.alpha}") - _print_section("Loading dataset") - print(f" {task_cfg['description']}") - print(f" data_path = {data_path}") - dataset = load_dataset(task_cfg, data_path, dev) + _print_section("Loading dataset + task") + print(f" data_path = {args.dataset_root}") + dataset = build_dataset(DatasetClass, args.dataset_root, args.dev, dataset_kwargs) + task = build_task(TaskClass, task_kwargs) + dataset = dataset.set_task(task) print(f" Dataset size: {len(dataset)} samples") + _print_section("Validating model / cp-method compatibility") + # Construction only needs model.mode (set from the dataset's schema, not + # from trained weights) -- so an incompatible (model, cp-method) pairing + # (e.g. a multilabel-only CP method against a multiclass task) is caught + # here, before spending time training, rather than only surfacing after + # every seed has already finished training. + probe_model = build_model(ModelClass, dataset, model_kwargs) + probe_alpha = args.alpha[0] if args.alpha else 0.1 + construct_cp_predictor(CPClass, probe_model, probe_alpha, cp_kwargs, args.dev) + print(f" model.mode='{probe_model.mode}' is compatible with cp-method '{args.cp_method}'") + del probe_model + per_seed_results: List[List[Dict[str, Any]]] = [] - for seed in seeds: - _print_header(f"Seed {seed} ({seeds.index(seed) + 1}/{len(seeds)})") + for seed in args.seeds: + _print_header(f"Seed {seed} ({args.seeds.index(seed) + 1}/{len(args.seeds)})") _print_section("Splitting") - print(f" ratios={ratios} seed={seed}") - train_data, val_data, cal_data, test_data = split_dataset(dataset, task_cfg, seed, ratios) + print(f" ratios={ratios} seed={seed} strategy={args.split_strategy}") + train_data, val_data, cal_data, test_data = split_dataset( + dataset, args.split_strategy, seed, ratios + ) print( f" train={len(train_data)} val={len(val_data)} " f"cal={len(cal_data)} test={len(test_data)}" ) n_cal = len(cal_data) - _print_section(f"Training {task_cfg['model_class']} ({epochs} epochs)") - model = build_model(task_cfg, dataset) + _print_section(f"Training {args.model} ({args.epochs} epochs)") + model = build_model(ModelClass, dataset, model_kwargs) t0 = time.time() - model = train_model(model, train_data, val_data, task_cfg, epochs, batch_size, output_dir, seed) + model = train_model( + model, train_data, val_data, args.epochs, batch_size, args.output_dir, + seed, args.monitor, args.monitor_criterion, + ) print(f" Training time: {time.time() - t0:.1f}s") model.to(_device()) model.eval() seed_rows: List[Dict[str, Any]] = [] + alphas = args.alpha for alpha in alphas: - _print_section(f"α={alpha} method={method}") + _print_section(f"α={alpha} cp-method={args.cp_method}") print(f" Calibrating on {n_cal} samples...") t0 = time.time() - try: - cp_model = build_cp_model(method_cfg, model, alpha, cal_data, test_data, dev) - except TypeError: - cp_model = _import(method_cfg["module"], method_cfg["class"])( - model=model, alpha=alpha, debug=dev, **method_cfg.get("extra_kwargs", {}) - ) - cp_model.calibrate(cal_dataset=cal_data) - + cp_model, used_alpha = build_cp_model( + CPClass, model, alpha, cp_kwargs, train_data, cal_data, test_data, + args.dev, + ) cp_model.to(_device()) cal_time = time.time() - t0 print(f" Calibration: {cal_time:.1f}s | Evaluating on {len(test_data)} samples...") - result = evaluate(cp_model, test_data, alpha, batch_size, task_cfg["mode"]) + result = evaluate(cp_model, test_data, alpha, used_alpha, batch_size, model.mode) result["n_cal"] = n_cal result["cal_time_s"] = round(cal_time, 2) result["seed"] = seed seed_rows.append(result) + if not used_alpha: + # This CP method has no scalar alpha to sweep over -- one run + # per seed is all there is; skip the rest of the alpha list. + if len(alphas) > 1: + print( + f" Note: {args.cp_method} has no 'alpha' constructor " + "parameter; ignoring the remaining --alpha values " + "for this method." + ) + break + per_seed_results.append(seed_rows) - aggregated = aggregate_results(per_seed_results, alphas) - print_aggregated_table(aggregated, task, method, seeds) + aggregated = aggregate_results(per_seed_results) + print_aggregated_table( + aggregated, args.dataset, args.task, args.model, args.cp_method, args.seeds + ) - if output_path: - save_results(aggregated, per_seed_results, seeds, task, method, output_path) + if args.output: + save_results( + aggregated, per_seed_results, args.seeds, + args.dataset, args.task, args.model, args.cp_method, args.output, + ) return aggregated def list_registry() -> None: - print("\nRegistered Tasks:") - print("-" * 60) - for name, cfg in TASK_REGISTRY.items(): - print(f" {name:<30} {cfg['description']}") - - print("\nRegistered CP Methods:") - print("-" * 60) - for name, cfg in METHOD_REGISTRY.items(): - print(f" {name:<14} {cfg['description']}") - print(f" {'':14} Paper: {cfg['paper']}") - print() + datasets = discover_datasets() + tasks = discover_tasks() + models = discover_models() + cp_methods = discover_cp_methods() + + def _first_line(cls: type) -> str: + doc = (cls.__doc__ or "").strip() + return doc.splitlines()[0].strip() if doc else "" + + for title, registry in ( + ("Datasets (pyhealth.datasets, BaseDataset subclasses)", datasets), + ("Tasks (pyhealth.tasks, BaseTask subclasses)", tasks), + ("Models (pyhealth.models, BaseModel subclasses)", models), + ("CP methods (pyhealth.calib.predictionset, SetPredictor subclasses)", cp_methods), + ): + print(f"\n{title}:") + print("-" * 70) + for name in sorted(registry): + print(f" {name:<30} {_first_line(registry[name])}") + print( + "\nNote: only classes that already subclass BaseDataset/BaseTask/" + "BaseModel/SetPredictor show up here. Legacy pre-2.0 datasets/tasks " + "not yet migrated to this API (e.g. classes still built on the " + "deprecated BaseSignalDataset/BaseEHRDataset stubs) are intentionally " + "excluded -- they are not usable with the current set_task()/model " + "API regardless of CPBench.\n" + "Not every dataset/task pairing is guaranteed to be compatible " + "(pyhealth has no dataset<->task compatibility registry yet); an " + "incompatible pairing will fail with an error when you run it.\n" + ) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( - description="CPBench: Standardized conformal prediction benchmarking for PyHealth", + description="CPBench: generic conformal-prediction benchmarking for PyHealth. " + "Name any dataset/task/model/CP-method already in pyhealth.", ) - parser.add_argument("--list", action="store_true", help="List all registered tasks and methods, then exit.") - parser.add_argument("--task", type=str, choices=list(TASK_REGISTRY.keys()), help="Task name (see --list).") - parser.add_argument("--method", type=str, choices=list(METHOD_REGISTRY.keys()), help="CP method name (see --list).") - parser.add_argument("--data-path", type=str, help="Path to the raw dataset on disk.") + parser.add_argument("--list", action="store_true", help="List all available datasets, tasks, models, and CP methods, then exit.") + parser.add_argument("--dataset", type=str, help="Dataset class name (see --list), e.g. TUEVDataset.") + parser.add_argument("--dataset-root", type=str, default=None, help="Path to the raw dataset on disk (passed as root=).") + parser.add_argument("--dataset-kwargs", type=str, default=None, help="JSON object of extra dataset constructor kwargs.") + parser.add_argument("--task", type=str, help="Task class name (see --list), e.g. EEGEventsTUEV.") + parser.add_argument("--task-kwargs", type=str, default=None, help="JSON object of task constructor kwargs.") + parser.add_argument("--model", type=str, help="Model class name (see --list), e.g. ContraWR.") + parser.add_argument("--model-kwargs", type=str, default=None, help="JSON object of extra model constructor kwargs.") + parser.add_argument("--cp-method", type=str, help="CP method class name (see --list), e.g. LABEL.") + parser.add_argument("--cp-kwargs", type=str, default=None, help="JSON object of extra CP-method constructor kwargs (required for methods with no 'alpha' param, e.g. SCRIB's risk=, FavMac's target_cost=/delta=).") parser.add_argument( - "--alpha", - type=float, - nargs="+", - default=None, - help=f"Miscoverage rate(s). Defaults to standard set {PROTOCOL['standard_alphas']}.", + "--alpha", type=float, nargs="+", default=None, + help=f"Miscoverage rate(s), for CP methods parameterized by alpha. Defaults to standard set {PROTOCOL['standard_alphas']}.", ) parser.add_argument( - "--seeds", - type=int, - nargs="+", - default=None, + "--seeds", type=int, nargs="+", default=None, help=f"Seeds for splitting. Defaults to {PROTOCOL['seeds']}.", ) - parser.add_argument( - "--epochs", - type=int, - default=None, - help="Training epochs. Defaults to per-task default.", - ) + parser.add_argument("--epochs", type=int, default=20, help="Training epochs. Default 20.") + parser.add_argument("--split-strategy", type=str, default="by_patient", choices=["by_patient", "by_sample"], help="Split strategy. Default by_patient.") + parser.add_argument("--monitor", type=str, default="accuracy", help="Validation metric to monitor for checkpoint selection. Default accuracy.") + parser.add_argument("--monitor-criterion", type=str, default="max", choices=["max", "min"], help="Whether higher or lower --monitor is better. Default max.") parser.add_argument("--output", type=str, default=None, help="Path to write results JSON.") parser.add_argument("--output-dir", type=str, default="./cpbench_runs", help="Directory for Trainer logs.") - parser.add_argument("--dev", action="store_true", help="Dev mode: tiny subset, single seed (42), single alpha (0.10).") + parser.add_argument("--dev", action="store_true", help="Dev mode: tiny data subset (if the dataset supports it), single seed (42), single alpha (0.10).") + parser.add_argument("--dry-run", action="store_true", help="Resolve and validate the dataset/task/model/cp-method combination without loading data or training.") return parser.parse_args() @@ -513,34 +698,21 @@ def main() -> None: list_registry() sys.exit(0) - missing = [f for f in ("task", "method", "data_path") if not getattr(args, f, None)] + missing = [f for f in ("dataset", "task", "model", "cp_method") if not getattr(args, f, None)] if missing: - print(f"Error: --{', --'.join(m.replace('_', '-') for m in missing)} required.") + print(f"Error: --{', --'.join(m.replace('_', '-') for m in missing)} required. Run --list to see options.") sys.exit(1) if args.dev: - alphas = [0.10] - seeds = [42] + args.alpha = [0.10] + args.seeds = [42] else: - alphas = args.alpha if args.alpha is not None else PROTOCOL["standard_alphas"] - seeds = args.seeds if args.seeds is not None else PROTOCOL["seeds"] - - task_cfg = TASK_REGISTRY[args.task] - epochs = args.epochs if args.epochs is not None else task_cfg.get("default_epochs", 20) + args.alpha = args.alpha if args.alpha is not None else PROTOCOL["standard_alphas"] + args.seeds = args.seeds if args.seeds is not None else PROTOCOL["seeds"] os.makedirs(args.output_dir, exist_ok=True) - run_benchmark( - task=args.task, - method=args.method, - data_path=args.data_path, - alphas=alphas, - seeds=seeds, - epochs=epochs, - output_path=args.output, - dev=args.dev, - output_dir=args.output_dir, - ) + run_benchmark(args) if __name__ == "__main__": diff --git a/pyhealth/calib/base_classes.py b/pyhealth/calib/base_classes.py index b68897b9e..78a425c7e 100644 --- a/pyhealth/calib/base_classes.py +++ b/pyhealth/calib/base_classes.py @@ -23,11 +23,34 @@ def to(self, device): class SetPredictor(ABC, torch.nn.Module): + """Base class for conformal-prediction-style prediction-set constructors. + + Interface convention for `calibrate()`: every concrete subclass should + accept `calibrate(self, cal_dataset, train_dataset=None, test_dataset=None, + **kwargs)`. `cal_dataset` is required; `train_dataset`/`test_dataset` are + optional and only meaningful to methods that need extra data beyond the + calibration set (e.g. embeddings-based methods like ClusterLabel, + CovariateLabel, NeighborhoodLabel) -- methods that don't need them should + still accept and ignore the two keyword arguments rather than omitting + them from the signature. + + This lets a generic caller (e.g. a benchmark harness) calibrate *any* + SetPredictor -- including ones not yet written -- with the same call: + `cp_model.calibrate(cal_dataset=cal, train_dataset=train, test_dataset=test)`, + without needing to special-case which extra data each method requires. + Methods that need embeddings should extract them internally (e.g. via + `pyhealth.calib.utils.extract_embeddings`) from whichever of + train_dataset/test_dataset they're given, rather than requiring the + caller to pre-compute and pass embeddings under method-specific kwarg + names -- pre-computed embeddings may still be accepted as additional + optional kwargs for callers who want to reuse/precompute them. + """ + def __init__(self, model, **kwargs) -> None: super().__init__() self.model = model - def calibrate(self, cal_dataset): + def calibrate(self, cal_dataset, train_dataset=None, test_dataset=None): ... def forward(self, **kwargs) -> Dict[str, torch.Tensor]: diff --git a/pyhealth/calib/predictionset/base_conformal/__init__.py b/pyhealth/calib/predictionset/base_conformal/__init__.py index 3451f9062..a6e44114c 100644 --- a/pyhealth/calib/predictionset/base_conformal/__init__.py +++ b/pyhealth/calib/predictionset/base_conformal/__init__.py @@ -15,7 +15,7 @@ "Inductive confidence machines for regression." ECML 2002. """ -from typing import Dict, Union +from typing import Dict, Optional, Union import numpy as np import torch @@ -209,11 +209,22 @@ def _compute_nc_scores( return scores - def calibrate(self, cal_dataset: IterableDataset): + def calibrate( + self, + cal_dataset: IterableDataset, + train_dataset: Optional[IterableDataset] = None, + test_dataset: Optional[IterableDataset] = None, + ): """Calibrate the thresholds for prediction set construction. Args: cal_dataset: Calibration set (held-out validation data) + train_dataset: Unused by this method. Accepted only so callers + that generically calibrate any SetPredictor (e.g. CPBench) + can pass the same (cal_dataset, train_dataset, test_dataset) + arguments to every CP method without branching on which + extra data each one needs. + test_dataset: Unused by this method (see train_dataset). """ # Get predictions and true labels cal_dataset_dict = prepare_numpy_dataset( diff --git a/pyhealth/calib/predictionset/cluster/cluster_label.py b/pyhealth/calib/predictionset/cluster/cluster_label.py index f56a325fa..5d342d890 100644 --- a/pyhealth/calib/predictionset/cluster/cluster_label.py +++ b/pyhealth/calib/predictionset/cluster/cluster_label.py @@ -137,6 +137,8 @@ def __init__( def calibrate( self, cal_dataset: IterableDataset, + train_dataset: Optional[IterableDataset] = None, + test_dataset: Optional[IterableDataset] = None, train_embeddings: Optional[np.ndarray] = None, cal_embeddings: Optional[np.ndarray] = None, batch_size: int = 32, @@ -151,14 +153,23 @@ def calibrate( Args: cal_dataset: Calibration set + train_dataset: Training set to extract embeddings from, used + only if train_embeddings is not already provided. ClusterLabel + needs a pool of training-set embeddings to fit clusters on -- + one of train_dataset or train_embeddings is required. + test_dataset: Unused by this method. Accepted only so callers + that generically calibrate any SetPredictor (e.g. CPBench) + can pass the same (cal_dataset, train_dataset, test_dataset) + arguments to every CP method without branching on which + extra data each one needs. train_embeddings: Optional pre-computed training embeddings - of shape (n_train, embedding_dim). If not provided, will - be extracted from the model (requires train_dataset parameter). + of shape (n_train, embedding_dim). If not provided, will be + extracted from train_dataset. cal_embeddings: Optional pre-computed calibration embeddings of shape (n_cal, embedding_dim). If not provided, will be extracted from cal_dataset. batch_size: Batch size for embedding extraction when - cal_embeddings is not provided. Default is 32. + cal_embeddings/train_embeddings are not provided. Default 32. Note: Either provide embeddings directly or ensure the model supports @@ -186,9 +197,16 @@ def calibrate( cal_embeddings = np.asarray(cal_embeddings) if train_embeddings is None: - raise ValueError( - "train_embeddings must be provided. " - "Extract embeddings from training set using extract_embeddings()." + if train_dataset is None: + raise ValueError( + "ClusterLabel needs a pool of training-set embeddings to " + "fit clusters on. Provide either train_embeddings " + "directly, or train_dataset so they can be extracted " + "automatically." + ) + print("Extracting embeddings from training set...") + train_embeddings = extract_embeddings( + self.model, train_dataset, batch_size=batch_size, device=self.device ) else: train_embeddings = np.asarray(train_embeddings) diff --git a/pyhealth/calib/predictionset/cluster/neighborhood_label.py b/pyhealth/calib/predictionset/cluster/neighborhood_label.py index 2d9f2dc6d..14fadd858 100644 --- a/pyhealth/calib/predictionset/cluster/neighborhood_label.py +++ b/pyhealth/calib/predictionset/cluster/neighborhood_label.py @@ -107,6 +107,8 @@ def __init__( def calibrate( self, cal_dataset: IterableDataset, + train_dataset: Optional[IterableDataset] = None, + test_dataset: Optional[IterableDataset] = None, cal_embeddings: Optional[np.ndarray] = None, batch_size: int = 32, ) -> None: @@ -120,6 +122,12 @@ def calibrate( Args: cal_dataset: Calibration dataset (for labels and predictions if cal_embeddings not provided). + train_dataset: Unused by this method. Accepted only so callers + that generically calibrate any SetPredictor (e.g. CPBench) + can pass the same (cal_dataset, train_dataset, test_dataset) + arguments to every CP method without branching on which + extra data each one needs. + test_dataset: Unused by this method (see train_dataset). cal_embeddings: Optional precomputed calibration embeddings (n_cal, embedding_dim). If None, extracted from cal_dataset. batch_size: Batch size for embedding extraction when cal_embeddings diff --git a/pyhealth/calib/predictionset/covariate/covariate_label.py b/pyhealth/calib/predictionset/covariate/covariate_label.py index f90430be8..bf2394c5d 100644 --- a/pyhealth/calib/predictionset/covariate/covariate_label.py +++ b/pyhealth/calib/predictionset/covariate/covariate_label.py @@ -30,7 +30,7 @@ from pyhealth.calib.base_classes import SetPredictor from pyhealth.calib.calibration.kcal.kde import RBFKernelMean -from pyhealth.calib.utils import prepare_numpy_dataset +from pyhealth.calib.utils import extract_embeddings, prepare_numpy_dataset from pyhealth.datasets import get_dataloader from pyhealth.models import BaseModel @@ -354,50 +354,71 @@ def __init__( def calibrate( self, cal_dataset: IterableDataset, + train_dataset: Optional[IterableDataset] = None, + test_dataset: Optional[IterableDataset] = None, cal_embeddings: Optional[np.ndarray] = None, test_embeddings: Optional[np.ndarray] = None, cal_weights: Optional[np.ndarray] = None, + batch_size: int = 32, ): """Calibrate the thresholds with covariate shift correction. This method supports three approaches for handling covariate shift: - - 1. **KDE-based (CoDrug approach)**: Provide cal_embeddings and - test_embeddings (and optionally kde_test/kde_cal). The method will + + 1. **KDE-based (CoDrug approach)**: Provide cal_embeddings/test_embeddings + (or cal_dataset/test_dataset, from which they'll be extracted + automatically) and, optionally, kde_test/kde_cal. The method will use kernel density estimation to compute likelihood ratios. - + 2. **Custom weights**: Directly provide cal_weights computed from your own covariate shift correction method (e.g., importance sampling, propensity scores, discriminator-based methods, etc.). - + 3. **Pre-fitted KDEs**: Provide kde_test and kde_cal during initialization - along with cal_embeddings here. + along with cal_embeddings (or cal_dataset) here. Args: - cal_dataset: Calibration set + cal_dataset: Calibration set. Also used to extract cal_embeddings + automatically when cal_embeddings is not explicitly given. + train_dataset: Unused by this method. Accepted only so callers + that generically calibrate any SetPredictor (e.g. CPBench) + can pass the same (cal_dataset, train_dataset, test_dataset) + arguments to every CP method without branching on which + extra data each one needs. + test_dataset: Test set to extract test_embeddings from, used only + if test_embeddings is not already provided and no KDEs are + set. CovariateLabel needs test-side density information to + correct for covariate shift -- one of test_dataset, + test_embeddings, or cal_weights is required (unless kde_test/ + kde_cal were already provided at __init__). cal_embeddings: Optional pre-computed calibration embeddings - of shape (n_cal, embedding_dim). If provided along with - test_embeddings and KDEs are not set, will be used to - compute likelihood ratios via KDE (CoDrug approach). + of shape (n_cal, embedding_dim). If not provided, extracted + from cal_dataset. test_embeddings: Optional pre-computed test embeddings - of shape (n_test, embedding_dim). Used with cal_embeddings - for KDE-based likelihood ratio computation. + of shape (n_test, embedding_dim). If not provided, extracted + from test_dataset when given. cal_weights: Optional custom weights for calibration samples of shape (n_cal,). If provided, these weights will be used directly instead of computing likelihood ratios via KDE. Weights should represent importance weights or likelihood ratios p_test(x) / p_cal(x). These will be normalized internally. + batch_size: Batch size for embedding extraction when + cal_embeddings/test_embeddings are not provided. Default 32. Note: You must provide ONE of: 1. cal_weights (custom weights), OR 2. kde_test and kde_cal during initialization, OR - 3. cal_embeddings and test_embeddings here + 3. test_embeddings or test_dataset here Examples: - >>> # Approach 1: KDE-based (CoDrug) - >>> model.calibrate(cal_dataset, cal_embeddings, test_embeddings) - >>> + >>> # Approach 1: KDE-based (CoDrug), embeddings extracted for you + >>> model.calibrate(cal_dataset, test_dataset=test_dataset) + >>> + >>> # Approach 1b: KDE-based (CoDrug), pre-computed embeddings + >>> model.calibrate(cal_dataset, cal_embeddings=cal_embeddings, + ... test_embeddings=test_embeddings) + >>> >>> # Approach 2: Custom weights (e.g., from importance sampling) >>> custom_weights = compute_importance_weights(cal_data, test_data) >>> model.calibrate(cal_dataset, cal_weights=custom_weights) @@ -425,37 +446,41 @@ def calibrate( likelihood_ratios = np.asarray(cal_weights, dtype=np.float64) print("Using custom calibration weights") else: - # Use KDE-based approach (CoDrug method) + # Use KDE-based approach (CoDrug method). cal_dataset is always + # available, so cal-side embeddings can always be obtained. + if cal_embeddings is None: + print("Extracting embeddings from calibration set...") + cal_embeddings = extract_embeddings( + self.model, cal_dataset, batch_size=batch_size, device=self.device + ) + else: + cal_embeddings = np.asarray(cal_embeddings) + # Check if we have KDEs if self.kde_test is None or self.kde_cal is None: - if cal_embeddings is None or test_embeddings is None: + if test_embeddings is None and test_dataset is not None: + print("Extracting embeddings from test set...") + test_embeddings = extract_embeddings( + self.model, test_dataset, batch_size=batch_size, + device=self.device, + ) + if test_embeddings is None: raise ValueError( - "Must provide ONE of:\n" + "CovariateLabel needs test-side density information " + "to correct for covariate shift. Provide ONE of:\n" " 1. cal_weights (custom weights), OR\n" " 2. kde_test and kde_cal during __init__, OR\n" - " 3. cal_embeddings and test_embeddings during calibrate()" + " 3. test_embeddings or test_dataset during calibrate()" ) - # Fit KDEs if embeddings provided + # Fit KDEs on the (now available) embeddings print("Fitting KDEs on provided embeddings (CoDrug approach)...") self.kde_cal, self.kde_test = fit_kde(cal_embeddings, test_embeddings) - # Use provided embeddings or extract from calibration data - if cal_embeddings is not None: - X = cal_embeddings - else: - # KDEs should already be provided in this case - # We just need to get the embeddings for likelihood ratio - # This assumes the model outputs embeddings - raise NotImplementedError( - "Automatic embedding extraction not yet supported. " - "Please provide cal_embeddings and test_embeddings." - ) - # Compute likelihood ratios using KDE print("Computing likelihood ratios via KDE...") likelihood_ratios = _compute_likelihood_ratio( - self.kde_test, self.kde_cal, X + self.kde_test, self.kde_cal, cal_embeddings ) # Normalize weights diff --git a/pyhealth/calib/predictionset/favmac/__init__.py b/pyhealth/calib/predictionset/favmac/__init__.py index bd57b644c..0c12122bc 100644 --- a/pyhealth/calib/predictionset/favmac/__init__.py +++ b/pyhealth/calib/predictionset/favmac/__init__.py @@ -197,11 +197,17 @@ def __init__( self.target_cost = target_cost self.delta = delta - def calibrate(self, cal_dataset): + def calibrate(self, cal_dataset, train_dataset=None, test_dataset=None): """Calibrate the cost-control procedure. :param cal_dataset: Calibration set. :type cal_dataset: Subset + :param train_dataset: Unused by this method. Accepted only so + callers that generically calibrate any SetPredictor (e.g. + CPBench) can pass the same (cal_dataset, train_dataset, + test_dataset) arguments to every CP method without branching on + which extra data each one needs. + :param test_dataset: Unused by this method (see train_dataset). """ _cal_data = prepare_numpy_dataset( self.model, cal_dataset, ["logit", "y_true"], debug=self.debug diff --git a/pyhealth/calib/predictionset/label.py b/pyhealth/calib/predictionset/label.py index 8ff87d070..6424eb4fc 100644 --- a/pyhealth/calib/predictionset/label.py +++ b/pyhealth/calib/predictionset/label.py @@ -9,7 +9,7 @@ """ -from typing import Dict, Union +from typing import Dict, Optional, Union import numpy as np import torch @@ -91,11 +91,22 @@ def __init__( self.t = None - def calibrate(self, cal_dataset: Subset): + def calibrate( + self, + cal_dataset: Subset, + train_dataset: Optional[Subset] = None, + test_dataset: Optional[Subset] = None, + ): """Calibrate the thresholds used to construct the prediction set. :param cal_dataset: Calibration set. :type cal_dataset: Subset + :param train_dataset: Unused by this method. Accepted only so + callers that generically calibrate any SetPredictor (e.g. + CPBench) can pass the same (cal_dataset, train_dataset, + test_dataset) arguments to every CP method without branching on + which extra data each one needs. + :param test_dataset: Unused by this method (see train_dataset). """ cal_dataset = prepare_numpy_dataset( self.model, cal_dataset, ["y_prob", "y_true"], debug=self.debug diff --git a/pyhealth/calib/predictionset/scrib/__init__.py b/pyhealth/calib/predictionset/scrib/__init__.py index 3406bef3b..999c697f1 100644 --- a/pyhealth/calib/predictionset/scrib/__init__.py +++ b/pyhealth/calib/predictionset/scrib/__init__.py @@ -267,11 +267,17 @@ def __init__( self.t = None - def calibrate(self, cal_dataset): + def calibrate(self, cal_dataset, train_dataset=None, test_dataset=None): """Calibrate/Search for the thresholds used to construct the prediction set. :param cal_dataset: Calibration set. :type cal_dataset: Subset + :param train_dataset: Unused by this method. Accepted only so + callers that generically calibrate any SetPredictor (e.g. + CPBench) can pass the same (cal_dataset, train_dataset, + test_dataset) arguments to every CP method without branching on + which extra data each one needs. + :param test_dataset: Unused by this method (see train_dataset). """ cal_dataset = prepare_numpy_dataset( self.model, cal_dataset, ["y_prob", "y_true"], debug=self.debug