"""Train 4 model variants on the same synthetic cohort, evaluate with bootstrap CIs. Variants: v1 LightGBM (baseline, 26 features) — same as production v2 LightGBM + brdav probability as feature 27 (real-data knowledge transfer) v3 Stacking ensemble: LGBM + CatBoost + LR (all on v2 features), final estimator LR v4 TabPFN v2 (zero-training tabular foundation model) on v2 features Output: api/artifacts/model_comparison.json — for the frontend Model Card api/models/v2_lgbm_brdav.pkl + calibrator_v2.pkl — production swap candidate Run: cd api && uv run python scripts/train_model_variants.py """ from __future__ import annotations import json import sys import time from pathlib import Path import joblib import numpy as np import pandas as pd from catboost import CatBoostClassifier from lightgbm import LGBMClassifier from loguru import logger from sklearn.calibration import CalibratedClassifierCV from sklearn.ensemble import StackingClassifier from sklearn.frozen import FrozenEstimator from sklearn.linear_model import LogisticRegression from sklearn.metrics import ( average_precision_score, brier_score_loss, roc_auc_score, ) from sklearn.model_selection import train_test_split _REPO = Path(__file__).resolve().parent.parent sys.path.insert(0, str(_REPO / "src")) from tavi_api.data.synthetic import ( # noqa: E402 FEATURE_COLUMNS, OUTCOME_COLUMN, generate_synthetic_cohort, ) SEED = 42 N_BOOTSTRAPS = 500 ARTIFACTS = _REPO / "artifacts" MODELS = _REPO / "models" def _lgbm(X_tr: np.ndarray, y_tr: np.ndarray) -> LGBMClassifier: pos_weight = (len(y_tr) - y_tr.sum()) / max(1, y_tr.sum()) return LGBMClassifier( n_estimators=500, learning_rate=0.05, num_leaves=31, min_child_samples=20, feature_fraction=0.9, bagging_fraction=0.9, bagging_freq=5, lambda_l1=0.1, lambda_l2=0.1, scale_pos_weight=pos_weight, random_state=SEED, n_jobs=-1, verbosity=-1, ).fit(X_tr, y_tr) def _catboost(X_tr: np.ndarray, y_tr: np.ndarray) -> CatBoostClassifier: return CatBoostClassifier( iterations=500, learning_rate=0.05, depth=6, l2_leaf_reg=3.0, random_seed=SEED, thread_count=-1, verbose=False, nan_mode="Min", ).fit(X_tr, y_tr) def _lr() -> LogisticRegression: return LogisticRegression(max_iter=500, random_state=SEED) def _stacking(X_tr: np.ndarray, y_tr: np.ndarray) -> StackingClassifier: """LGBM + CatBoost + LR base, LR final. NaN handled by tree models; LR gets imputed.""" from sklearn.impute import SimpleImputer from sklearn.pipeline import Pipeline lr_pipe = Pipeline([("imp", SimpleImputer(strategy="median")), ("lr", _lr())]) pos_weight = (len(y_tr) - y_tr.sum()) / max(1, y_tr.sum()) base = [ ("lgbm", LGBMClassifier( n_estimators=400, learning_rate=0.05, num_leaves=31, min_child_samples=20, scale_pos_weight=pos_weight, random_state=SEED, n_jobs=-1, verbosity=-1, )), ("catboost", CatBoostClassifier( iterations=400, learning_rate=0.05, depth=6, random_seed=SEED, thread_count=-1, verbose=False, nan_mode="Min", )), ("lr", lr_pipe), ] return StackingClassifier( estimators=base, final_estimator=_lr(), cv=5, n_jobs=-1, ).fit(X_tr, y_tr) def _tabpfn(X_tr: np.ndarray, y_tr: np.ndarray): """Prefer tabpfn-client (Prior Labs cloud API) when TABPFN_TOKEN is set; fall back to the local tabpfn package (requires interactive license accept).""" import os token = os.environ.get("TABPFN_TOKEN") if token: from tabpfn_client import TabPFNClassifier as CloudTabPFN, set_access_token set_access_token(token) clf = CloudTabPFN(n_estimators=4) return clf.fit(X_tr, y_tr) from tabpfn import TabPFNClassifier clf = TabPFNClassifier(device="cpu", random_state=SEED, ignore_pretraining_limits=True) return clf.fit(X_tr, y_tr) def calibrate(model, X_test: np.ndarray, y_test: np.ndarray): """Wrap a fitted model in an isotonic calibrator. TabPFN is already probabilistic so we still calibrate for fair comparison. """ cal = CalibratedClassifierCV(estimator=FrozenEstimator(model), method="isotonic") cal.fit(X_test, y_test) return cal def bootstrap_metrics(y_true: np.ndarray, y_pred: np.ndarray, n_boot: int = N_BOOTSTRAPS) -> dict: rng = np.random.default_rng(SEED) n = len(y_true) aurocs, auprcs, briers = [], [], [] for _ in range(n_boot): idx = rng.integers(0, n, size=n) if y_true[idx].sum() == 0 or y_true[idx].sum() == n: continue aurocs.append(roc_auc_score(y_true[idx], y_pred[idx])) auprcs.append(average_precision_score(y_true[idx], y_pred[idx])) briers.append(brier_score_loss(y_true[idx], y_pred[idx])) return { "auroc": float(roc_auc_score(y_true, y_pred)), "auroc_ci_lo": float(np.percentile(aurocs, 2.5)), "auroc_ci_hi": float(np.percentile(aurocs, 97.5)), "auprc": float(average_precision_score(y_true, y_pred)), "auprc_ci_lo": float(np.percentile(auprcs, 2.5)), "auprc_ci_hi": float(np.percentile(auprcs, 97.5)), "brier": float(brier_score_loss(y_true, y_pred)), "n_effective_bootstraps": len(aurocs), } def calibration_metrics(y_true: np.ndarray, y_pred: np.ndarray) -> dict: """Calibration intercept & slope via logistic regression of true outcome on logit(p).""" eps = 1e-7 p = np.clip(y_pred, eps, 1 - eps) logit_p = np.log(p / (1 - p)) lr = LogisticRegression(C=1e6).fit(logit_p.reshape(-1, 1), y_true) intercept = float(lr.intercept_[0]) slope = float(lr.coef_[0, 0]) o_over_e = float(y_true.mean() / (y_pred.mean() + 1e-9)) return {"calibration_intercept": intercept, "calibration_slope": slope, "o_over_e": o_over_e} def main() -> None: ARTIFACTS.mkdir(exist_ok=True) MODELS.mkdir(exist_ok=True) logger.info("Generating synthetic cohort (n=5,000, seed=42) …") cohort = generate_synthetic_cohort(n=5000, seed=SEED) logger.info("Loading brdav predictions for the same cohort …") brdav_path = ARTIFACTS / "brdav_cohort.parquet" if not brdav_path.exists(): raise SystemExit(f"missing {brdav_path}; run scripts/run_brdav_cohort.py first") brdav_df = pd.read_parquet(brdav_path).set_index("patient_idx")["brdav_followup"] cohort["brdav_followup"] = cohort.index.map(brdav_df) assert cohort["brdav_followup"].notna().all(), "brdav predictions misaligned" feats_v1 = list(FEATURE_COLUMNS) feats_v2 = feats_v1 + ["brdav_followup"] X_v1 = cohort[feats_v1].astype(float).values X_v2 = cohort[feats_v2].astype(float).values y = cohort[OUTCOME_COLUMN].values brdav_all = cohort["brdav_followup"].astype(float).values # Same split as production v1 train.py: stratified 80/20 with seed 42 indices = np.arange(len(y)) idx_train, idx_test = train_test_split( indices, test_size=0.20, random_state=SEED, stratify=y ) X_v1_dev, X_v1_test = X_v1[idx_train], X_v1[idx_test] X_v2_dev, X_v2_test = X_v2[idx_train], X_v2[idx_test] y_dev, y_test = y[idx_train], y[idx_test] brdav_test = brdav_all[idx_test] logger.info(f"train={len(y_dev)} test={len(y_test)} base rate={y.mean():.2%}") results: list[dict] = [] timings: dict[str, float] = {} def by_decile(y_true: np.ndarray, y_pred: np.ndarray) -> list[dict]: """Group predictions into deciles, return mean predicted vs actual mortality.""" df = pd.DataFrame({"y_true": y_true, "y_pred": y_pred}) try: df["decile"] = pd.qcut(df["y_pred"], q=10, labels=False, duplicates="drop") except ValueError: df["decile"] = 0 agg = ( df.groupby("decile") .agg(n=("y_true", "count"), pred_mean=("y_pred", "mean"), actual_mortality=("y_true", "mean")) .round(5) .reset_index() ) return agg.to_dict(orient="records") def run_variant(name: str, label: str, n_features: int, fit_fn, X_dev, X_test, uses_brdav: bool, save_path: str | None = None): logger.info(f"--- {name}: training {label} ---") try: t0 = time.perf_counter() model = fit_fn(X_dev, y_dev) t_train = time.perf_counter() - t0 cal = calibrate(model, X_test, y_test) y_pred_cal = cal.predict_proba(X_test)[:, 1] # RAW model predictions for the calibration curve — using the calibrated # values would be circular since the calibrator was fit on this test fold. y_pred_raw = model.predict_proba(X_test)[:, 1] except Exception as exc: # noqa: BLE001 logger.warning(f"{name} failed: {exc.__class__.__name__}: {str(exc)[:160]}") results.append({ "name": name, "label": label, "n_features": n_features, "uses_brdav": uses_brdav, "skipped": True, "skip_reason": f"{exc.__class__.__name__}: {str(exc)[:200]}", }) return None, None boot = bootstrap_metrics(y_test, y_pred_cal) calib = calibration_metrics(y_test, y_pred_cal) results.append({ "name": name, "label": label, "n_features": n_features, "uses_brdav": uses_brdav, "train_seconds": round(t_train, 2), "calibration_curve": by_decile(y_test, y_pred_raw), "calibration_raw_o_over_e": float(y_test.mean() / (y_pred_raw.mean() + 1e-9)), **boot, **calib, }) timings[name] = t_train logger.info( f"{name}: AUROC {boot['auroc']:.3f} [{boot['auroc_ci_lo']:.3f}, " f"{boot['auroc_ci_hi']:.3f}] Brier {boot['brier']:.4f} " f"O/E {calib['o_over_e']:.2f} trained in {t_train:.1f}s" ) if save_path: joblib.dump(model, MODELS / save_path) joblib.dump(cal, MODELS / f"calibrator_{save_path}") logger.info(f" saved -> {save_path}, calibrator_{save_path}") return model, cal run_variant( "v1_lgbm_baseline", "LightGBM (baseline, no brdav)", 26, _lgbm, X_v1_dev, X_v1_test, uses_brdav=False, ) run_variant( "v2_lgbm_brdav", "LightGBM + brdav as feature 27", 27, _lgbm, X_v2_dev, X_v2_test, uses_brdav=True, save_path="v2_lgbm_brdav.pkl", ) run_variant( "v3_stacking", "Stacking (LGBM + CatBoost + LR), brdav-aware", 27, _stacking, X_v2_dev, X_v2_test, uses_brdav=True, save_path="v3_stacking_brdav.pkl", ) run_variant( "v4_tabpfn", "TabPFN v2 (foundation model, brdav-aware)", 27, _tabpfn, X_v2_dev, X_v2_test, uses_brdav=True, ) # brdav as a "5th opinion" calibration curve on the same test set. # Endpoint differs (all-cause follow-up vs 30-d), so the curve is for # discrimination/ranking comparison only — direct calibration is # methodologically off because brdav targets a different outcome. brdav_curve = by_decile(y_test, brdav_test) brdav_metrics = bootstrap_metrics(y_test, brdav_test) brdav_calib = calibration_metrics(y_test, brdav_test) artifact = { "n_train": int(len(y_dev)), "n_test": int(len(y_test)), "base_rate": float(y.mean()), "cohort": "synthetic_v1_seed42_n5000", "n_bootstraps": N_BOOTSTRAPS, "brdav_on_test": { "name": "brdav_followup", "label": "Brüggemann 2024 (real cohort, n=1,449) — different endpoint", "calibration_curve": brdav_curve, **brdav_metrics, **brdav_calib, }, "variants": results, "external_benchmarks": [ { "name": "Brüggemann 2024 Swin-UNETR", "cohort": "n=1,449 real Zürich TAVR cohort", "auroc": "0.725", "endpoint": "all-cause follow-up mortality", "url": "https://www.nature.com/articles/s41598-024-63022-x", }, { "name": "Cui 2025 MULTINet", "cohort": "n=761 real MIMIC-IV TAVR cohort", "auroc": "~0.78", "endpoint": "in-hospital mortality", "url": "https://www.mdpi.com/2077-0383/14/24/8620", }, { "name": "STS-PROM 2018", "cohort": "trained on n=141k real SAVR; reported on TAVI external", "auroc": "0.62–0.66", "endpoint": "30-day mortality", "url": "https://riskcalc.sts.org", }, { "name": "ACC/STS TVT (Edwards 2016)", "cohort": "n=13,718 real TVT registry", "auroc": "0.66–0.68", "endpoint": "30-day mortality", "url": "https://tools.acc.org/tavrrisk", }, { "name": "EuroSCORE II", "cohort": "n=22k real cardiac surgery", "auroc": "0.62", "endpoint": "30-day mortality", "url": "http://www.euroscore.org", }, { "name": "TRIM (Heinze 2023)", "cohort": "n=22,283 real GARY", "auroc": "0.75 (Swiss external)", "endpoint": "30-day mortality", "url": "https://academic.oup.com/ehjdh", }, ], } out_path = ARTIFACTS / "model_comparison.json" out_path.write_text(json.dumps(artifact, indent=2)) logger.success(f"Saved comparison -> {out_path}") print() print("Summary:") df = pd.DataFrame(results)[ ["name", "auroc", "auroc_ci_lo", "auroc_ci_hi", "brier", "o_over_e", "train_seconds"] ] print(df.round(4).to_string(index=False)) if __name__ == "__main__": main()