scripts/bootstrap_all_baselines.py

"""Bootstrap 95% CI for ALL frozen baseline encoders (Stage-1 only: MAE + R@1).

Runs 1000-resample bootstrap on the 153-subject test set for each baseline.

Usage (from /data/Albus/Brain):
    CUDA_VISIBLE_DEVICES=2 python scripts/bootstrap_all_baselines.py
"""
from __future__ import annotations

import sys
import time
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader

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


# -- baselines ---------------------------------------------------------------
BASELINES = [
    ("MedicalNet frozen",  "runs/foundation/medicalnet_frozen_mlp.pt"),
    ("BrainIAC frozen",    "runs/foundation/brainiac_frozen_mlp.pt"),
    ("BrainFM frozen",     "runs/foundation/brainfm_frozen_mlp_b4_best.pt"),
    ("SAM-Med3D frozen",   "runs/foundation/sam_med3d_frozen_mlp_best.pt"),
    ("SwinUNETR frozen",   "runs/foundation/swinunetr_frozen_mlp_best.pt"),
]

TEST_MANIFEST = Path("metadata/splits/test.csv")
B = 1000
SEED = 42
BATCH_SIZE = 4


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))


@torch.no_grad()
def collect_stage1(model, loader, device):
    model.eval()
    pred_c, tgt_c, pz_c, sz_c = [], [], [], []
    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_c.append(outputs["pred_suvr"].cpu().numpy())
        tgt_c.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)
        pz_c.append(pet_z.cpu().numpy())
        sz_c.append(suvr_z.cpu().numpy())
    return {
        "pred":   np.concatenate(pred_c),
        "target": np.concatenate(tgt_c),
        "pet_z":  np.concatenate(pz_c),
        "suvr_z": np.concatenate(sz_c),
    }


def stage1_metrics(d, idx):
    pred   = d["pred"][idx]
    target = d["target"][idx]
    uid    = np.unique(idx)
    logits = d["pet_z"][uid] @ d["suvr_z"][uid].T
    return {
        "mae":         float(np.mean(np.abs(pred - target))),
        "pet_suvr_r1": _retrieval_recall_at_1(logits),
    }


def bootstrap_ci(metric_fn, n, B=1000, seed=42):
    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, 2.5))
    hi = float(np.percentile(boots, 97.5))
    return point, lo, hi


def load_model(ckpt_path, device):
    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    saved = ckpt.get("args", {})

    class _A:
        pass

    a = _A()
    a.backbone           = saved.get("backbone", "medicalnet")
    a.medicalnet_weights = Path(saved.get("medicalnet_weights",
                                          "pretrained/medicalnet/resnet_50_23dataset.pth"))
    a.brainiac_weights   = Path(saved.get("brainiac_weights",
                                          "pretrained/brainiac/backbone.safetensors"))
    a.brainfm_weights    = Path(saved.get("brainfm_weights",
                                          "pretrained/brainfm/assets/brainfm_pretrained.pth"))
    a.brainfm_code_root  = Path(saved.get("brainfm_code_root", "pretrained/brainfm"))
    a.swinunetr_weights  = Path(saved.get("swinunetr_weights",
                                          "pretrained/swinunetr/model_swinvit.pt"))
    a.sam_med3d_weights  = Path(saved.get("sam_med3d_weights",
                                          "pretrained/sam-med3d/sam_med3d_turbo.pth"))
    a.output_size        = tuple(saved.get("output_size", (96, 96, 96)))
    embed_dim            = saved.get("embed_dim", 256)
    freeze               = bool(saved.get("freeze_encoder", False))

    ds_tmp    = PETSUVRDataset(TEST_MANIFEST, output_size=a.output_size)
    n_regions = int(ds_tmp[0]["suvr"].numel())
    encoder   = build_encoder(a)
    model     = PETSUVRFoundationModel(encoder, n_regions, embed_dim, freeze).to(device)
    model.load_state_dict(ckpt["model"], strict=True)
    model.eval()
    return model, a.output_size


def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}", flush=True)

    results = []
    for name, ckpt_path in BASELINES:
        t0 = time.time()
        print(f"\n{'='*60}", flush=True)
        print(f"  {name}  ({ckpt_path})", flush=True)
        print(f"{'='*60}", flush=True)

        model, output_size = load_model(ckpt_path, device)
        bs = 2 if "sam_med3d" in ckpt_path else BATCH_SIZE
        ds     = PETSUVRDataset(TEST_MANIFEST, output_size=output_size)
        loader = DataLoader(ds, batch_size=bs, shuffle=False,
                            num_workers=2, collate_fn=collate_pet_suvr)
        d = collect_stage1(model, loader, device)
        N = d["pred"].shape[0]
        print(f"  N = {N}", flush=True)

        for metric_name in ("mae", "pet_suvr_r1"):
            fn = lambda idx, _m=metric_name: stage1_metrics(d, idx)[_m]
            pt, lo, hi = bootstrap_ci(fn, N, B=B, seed=SEED)
            print(f"  {metric_name:20s}  {pt:.4f}   95% CI [{lo:.4f}, {hi:.4f}]", flush=True)
            results.append((name, metric_name, pt, lo, hi))

        # free GPU memory
        del model
        torch.cuda.empty_cache()
        print(f"  elapsed: {time.time()-t0:.1f}s", flush=True)

    # ---- summary table ----
    print(f"\n\n{'='*70}", flush=True)
    print(f"SUMMARY: Bootstrap 95% CI  (B={B}, seed={SEED})", flush=True)
    print(f"{'='*70}", flush=True)
    print(f"{'Model':<22s}  {'MAE':>8s}  {'MAE 95% CI':>18s}  {'R@1':>8s}  {'R@1 95% CI':>18s}", flush=True)
    print("-"*70, flush=True)
    for i in range(0, len(results), 2):
        nm = results[i][0]
        mae_pt, mae_lo, mae_hi = results[i][2], results[i][3], results[i][4]
        r1_pt,  r1_lo,  r1_hi  = results[i+1][2], results[i+1][3], results[i+1][4]
        print(f"{nm:<22s}  {mae_pt:8.4f}  [{mae_lo:.4f}, {mae_hi:.4f}]  {r1_pt:8.4f}  [{r1_lo:.4f}, {r1_hi:.4f}]", flush=True)
    print(f"{'='*70}", flush=True)


if __name__ == "__main__":
    main()