"""Bootstrap 95 % confidence intervals for ReMAP-PET key metrics. Stage-1 metrics (153 test subjects): - SUVR MAE - Pearson r (voxel-level across all subjects x regions) - PET->SUVR Recall@1 (retrieval) Clinical probe metrics: - AD vs CN AUROC (logistic regression on PET embeddings) - 3-way (CN/MCI/AD) AUROC Usage (from /data/Albus/Brain): CUDA_VISIBLE_DEVICES=1 python scripts/bootstrap_ci.py """ from __future__ import annotations import argparse import sys from pathlib import Path import numpy as np import pandas as pd import torch import torch.nn.functional as F from torch.utils.data import DataLoader from sklearn.linear_model import LogisticRegression from sklearn.metrics import balanced_accuracy_score, roc_auc_score from sklearn.pipeline import make_pipeline from sklearn.preprocessing import LabelEncoder, StandardScaler, label_binarize # -- project imports (scripts/ is the working dir's sibling) ----------------- sys.path.insert(0, str(Path(__file__).resolve().parent)) from pet_vlm_dataset import PETSUVRDataset, collate_pet_suvr from train_pet_foundation import PETSUVRFoundationModel, build_encoder # --------------------------------------------------------------------------- # helpers copied from evaluate_pet_foundation.py # --------------------------------------------------------------------------- def _pearson_flat(pred: np.ndarray, target: np.ndarray) -> float: p = pred.reshape(-1) t = target.reshape(-1) if p.std() < 1e-8 or t.std() < 1e-8: return float("nan") return float(np.corrcoef(p, t)[0, 1]) def _retrieval_recall_at_1(logits: np.ndarray) -> float: ranks = [] for i in range(logits.shape[0]): order = np.argsort(-logits[i]) rank = int(np.where(order == i)[0][0]) + 1 ranks.append(rank) return float(np.mean(np.asarray(ranks) <= 1)) # --------------------------------------------------------------------------- # Stage-1: forward pass -> per-subject arrays # --------------------------------------------------------------------------- @torch.no_grad() def collect_stage1( model: PETSUVRFoundationModel, loader: DataLoader, device: torch.device, ) -> dict[str, np.ndarray]: """Return pred_suvr, target_suvr, pet_z, suvr_z (all numpy, N-first).""" model.eval() pred_chunks, target_chunks = [], [] pet_z_chunks, suvr_z_chunks = [], [] for batch in loader: image = batch["image"].to(device, non_blocking=True) suvr = batch["suvr"].to(device, non_blocking=True) outputs = model(image, suvr) pred_chunks.append(outputs["pred_suvr"].cpu().numpy()) target_chunks.append(suvr.cpu().numpy()) pet_feat = model.pet_encoder(image) pet_z = F.normalize(model.pet_projector(pet_feat), dim=-1) suvr_z = F.normalize(model.suvr_encoder(suvr), dim=-1) pet_z_chunks.append(pet_z.cpu().numpy()) suvr_z_chunks.append(suvr_z.cpu().numpy()) return { "pred": np.concatenate(pred_chunks, axis=0), "target": np.concatenate(target_chunks, axis=0), "pet_z": np.concatenate(pet_z_chunks, axis=0), "suvr_z": np.concatenate(suvr_z_chunks, axis=0), } def stage1_metrics(d: dict[str, np.ndarray], idx: np.ndarray) -> dict[str, float]: """Compute stage-1 metrics on a subset given by *idx*. MAE and Pearson work fine with duplicate indices (bootstrap). For retrieval R@1 we need unique subjects (duplicates would make the diagonal ground-truth ambiguous), so we deduplicate *idx* first. """ pred = d["pred"][idx] target = d["target"][idx] # retrieval: use unique indices only uid = np.unique(idx) pet_z = d["pet_z"][uid] suvr_z = d["suvr_z"][uid] logits = pet_z @ suvr_z.T return { "mae": float(np.mean(np.abs(pred - target))), "pearson": _pearson_flat(pred, target), "pet_suvr_r1": _retrieval_recall_at_1(logits), } # --------------------------------------------------------------------------- # Clinical: extract embeddings, train probe, evaluate # --------------------------------------------------------------------------- @torch.no_grad() def extract_embeddings( model: PETSUVRFoundationModel, manifest: Path, output_size: tuple[int, int, int], batch_size: int, num_workers: int, device: torch.device, ) -> tuple[pd.DataFrame, np.ndarray]: dataset = PETSUVRDataset(manifest, output_size=output_size) loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, collate_fn=collate_pet_suvr) feats = [] model.eval() for batch in loader: image = batch["image"].to(device, non_blocking=True) pet_feat = model.pet_encoder(image) pet_z = F.normalize(model.pet_projector(pet_feat), dim=-1) feats.append(pet_z.cpu().numpy()) return pd.read_csv(manifest), np.concatenate(feats, axis=0) def _subset_cls(df, x, column, labels): mask = df[column].isin(labels).to_numpy() return x[mask], df.loc[mask, column].astype(str).to_numpy() def train_probe(x_train, y_train, x_val, y_val): """Train logistic probe with C sweep; return best model + encoder.""" enc = LabelEncoder() enc.fit(np.concatenate([y_train, y_val])) y_tr = enc.transform(y_train) y_v = enc.transform(y_val) best_m, best_s = None, -np.inf for c in [0.01, 0.03, 0.1, 0.3, 1.0, 3.0, 10.0]: m = make_pipeline(StandardScaler(), LogisticRegression(C=c, max_iter=5000, class_weight="balanced")) m.fit(x_train, y_tr) s = balanced_accuracy_score(y_v, m.predict(x_val)) if s > best_s: best_m, best_s = m, s return best_m, enc def clinical_auroc(model_probe, encoder, x_test, y_test): """Return AUROC (binary or macro-OVR).""" y_int = encoder.transform(y_test) proba = model_probe.predict_proba(x_test) if len(encoder.classes_) == 2: return roc_auc_score(y_int, proba[:, 1]) else: y_bin = label_binarize(y_int, classes=np.arange(len(encoder.classes_))) return roc_auc_score(y_bin, proba, average="macro", multi_class="ovr") # --------------------------------------------------------------------------- # Bootstrap # --------------------------------------------------------------------------- def bootstrap_ci( metric_fn, n: int, B: int = 1000, seed: int = 42, alpha: float = 0.05, ) -> tuple[float, float, float]: """ metric_fn(idx) -> float where idx is array of resampled indices. Returns (point_estimate, lo, hi) for the (1-alpha) CI. """ rng = np.random.RandomState(seed) all_idx = np.arange(n) point = metric_fn(all_idx) boots = np.empty(B) for b in range(B): idx = rng.choice(n, size=n, replace=True) boots[b] = metric_fn(idx) lo = float(np.percentile(boots, 100 * alpha / 2)) hi = float(np.percentile(boots, 100 * (1 - alpha / 2))) return point, lo, hi def bootstrap_clinical_auroc( probe, encoder, x_train, y_train_raw, x_val, y_val_raw, x_test, y_test_raw, B: int = 1000, seed: int = 42, alpha: float = 0.05, ) -> tuple[float, float, float]: """ Bootstrap over the *test* set only (probe is fixed). """ rng = np.random.RandomState(seed) n = len(y_test_raw) all_idx = np.arange(n) point = clinical_auroc(probe, encoder, x_test, y_test_raw) boots = np.empty(B) for b in range(B): idx = rng.choice(n, size=n, replace=True) try: boots[b] = clinical_auroc(probe, encoder, x_test[idx], y_test_raw[idx]) except ValueError: # can happen if a resample has only one class boots[b] = np.nan boots = boots[~np.isnan(boots)] lo = float(np.percentile(boots, 100 * alpha / 2)) hi = float(np.percentile(boots, 100 * (1 - alpha / 2))) return point, lo, hi # --------------------------------------------------------------------------- # main # --------------------------------------------------------------------------- def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--checkpoint", type=Path, default=Path("runs/foundation/medicalnet_layer4_regalign_best.pt")) parser.add_argument("--test-manifest", type=Path, default=Path("metadata/splits/test.csv")) parser.add_argument("--train-clinical", type=Path, default=Path("data/metadata/splits/train_clinical_server.csv")) parser.add_argument("--val-clinical", type=Path, default=Path("data/metadata/splits/val_clinical_server.csv")) parser.add_argument("--test-clinical", type=Path, default=Path("data/metadata/splits/test_clinical_server.csv")) parser.add_argument("--batch-size", type=int, default=4) parser.add_argument("--num-workers", type=int, default=2) parser.add_argument("--B", type=int, default=1000, help="bootstrap resamples") parser.add_argument("--seed", type=int, default=42) args = parser.parse_args() # ---- load model ------------------------------------------------------- device = torch.device("cuda" if torch.cuda.is_available() else "cpu") ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False) saved = ckpt.get("args", {}) class _Args: pass margs = _Args() margs.backbone = saved.get("backbone", "medicalnet") margs.medicalnet_weights = Path(saved.get("medicalnet_weights", "pretrained/medicalnet/resnet_50_23dataset.pth")) margs.brainiac_weights = Path(saved.get("brainiac_weights", "pretrained/brainiac/backbone.safetensors")) margs.brainfm_weights = Path("pretrained/brainfm/assets/brainfm_pretrained.pth") margs.brainfm_code_root = Path("pretrained/brainfm") margs.swinunetr_weights = Path("pretrained/swinunetr/model_swinvit.pt") margs.sam_med3d_weights = Path("pretrained/sam-med3d/sam_med3d_turbo.pth") margs.output_size = tuple(saved.get("output_size", (96, 96, 96))) embed_dim = saved.get("embed_dim", 256) freeze_encoder = bool(saved.get("freeze_encoder", False)) output_size = margs.output_size # build model dataset_tmp = PETSUVRDataset(args.test_manifest, output_size=output_size) n_regions = int(dataset_tmp[0]["suvr"].numel()) encoder = build_encoder(margs) model = PETSUVRFoundationModel(encoder, n_regions, embed_dim, freeze_encoder).to(device) model.load_state_dict(ckpt["model"], strict=True) model.eval() print(f"Loaded checkpoint: {args.checkpoint}", flush=True) print(f"backbone={margs.backbone} embed_dim={embed_dim} " f"freeze={freeze_encoder} output_size={output_size}", flush=True) # ===== STAGE 1 ========================================================= print("\n===== Stage-1 evaluation (test set) =====", flush=True) test_ds = PETSUVRDataset(args.test_manifest, output_size=output_size) test_loader = DataLoader(test_ds, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, collate_fn=collate_pet_suvr) d = collect_stage1(model, test_loader, device) N = d["pred"].shape[0] print(f" N = {N}", flush=True) for name in ("mae", "pearson", "pet_suvr_r1"): fn = lambda idx, _n=name: stage1_metrics(d, idx)[_n] pt, lo, hi = bootstrap_ci(fn, N, B=args.B, seed=args.seed) print(f" {name:20s} {pt:.4f} 95% CI [{lo:.4f}, {hi:.4f}]", flush=True) # ===== CLINICAL ======================================================== print("\n===== Clinical downstream probes =====", flush=True) train_df, x_train_all = extract_embeddings( model, args.train_clinical, output_size, args.batch_size, args.num_workers, device) val_df, x_val_all = extract_embeddings( model, args.val_clinical, output_size, args.batch_size, args.num_workers, device) test_df, x_test_all = extract_embeddings( model, args.test_clinical, output_size, args.batch_size, args.num_workers, device) # ---- AD vs CN --------------------------------------------------------- print("\n -- AD vs CN --", flush=True) x_tr, y_tr = _subset_cls(train_df, x_train_all, "clinical_label", ["CN", "AD"]) x_v, y_v = _subset_cls(val_df, x_val_all, "clinical_label", ["CN", "AD"]) x_te, y_te = _subset_cls(test_df, x_test_all, "clinical_label", ["CN", "AD"]) print(f" train={len(y_tr)} val={len(y_v)} test={len(y_te)}", flush=True) probe_ad, enc_ad = train_probe(x_tr, y_tr, x_v, y_v) pt, lo, hi = bootstrap_clinical_auroc( probe_ad, enc_ad, x_tr, y_tr, x_v, y_v, x_te, y_te, B=args.B, seed=args.seed) print(f" {'ad_vs_cn_auroc':20s} {pt:.4f} 95% CI [{lo:.4f}, {hi:.4f}]", flush=True) # ---- 3-way CN / MCI / AD --------------------------------------------- print("\n -- 3-way (CN / MCI / AD) --", flush=True) x_tr3, y_tr3 = _subset_cls(train_df, x_train_all, "clinical_label", ["CN", "MCI", "AD"]) x_v3, y_v3 = _subset_cls(val_df, x_val_all, "clinical_label", ["CN", "MCI", "AD"]) x_te3, y_te3 = _subset_cls(test_df, x_test_all, "clinical_label", ["CN", "MCI", "AD"]) print(f" train={len(y_tr3)} val={len(y_v3)} test={len(y_te3)}", flush=True) probe_3w, enc_3w = train_probe(x_tr3, y_tr3, x_v3, y_v3) pt, lo, hi = bootstrap_clinical_auroc( probe_3w, enc_3w, x_tr3, y_tr3, x_v3, y_v3, x_te3, y_te3, B=args.B, seed=args.seed) print(f" {'3way_auroc':20s} {pt:.4f} 95% CI [{lo:.4f}, {hi:.4f}]", flush=True) print("\nDone.", flush=True) if __name__ == "__main__": main()