scripts/sulic/sulic_panda_heldout.py

"""sulic-in-panda held-out 5-fold cv. Test A: binary facs (placode vs epi). Test C: 4-way placode subtype."""
from __future__ import annotations
from pathlib import Path
import warnings, json, sys, time
warnings.filterwarnings("ignore")

import numpy as np
import anndata as ad
import scanpy as sc
import scipy.sparse as sp
import torch
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from sklearn.decomposition import PCA
from sklearn.model_selection import StratifiedKFold
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import roc_auc_score, accuracy_score

import os as _os
from pathlib import Path as _Path
PANDA_ROOT = _Path(_os.environ.get("PANDA_ROOT", str(_Path(__file__).resolve().parents[2])))
sys.path.insert(0, str(PANDA_ROOT))
from panda.pan_skin.model import (
    PANDAEncoder, supcon_loss, vicreg_loss, prototype_infonce
)

SULIC_H5AD = Path(str(PANDA_ROOT / "data/processed/sulic/adata_sulic_clustered.h5ad"))
OUT = Path(str(PANDA_ROOT / "scripts/sulic"))
OUT.mkdir(parents=True, exist_ok=True)

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
N_FOLDS = 5


class SulicDataset(Dataset):
    def __init__(self, X, y):
        self.X = X.astype(np.float32); self.y = y.astype(np.int64)
    def __len__(self): return self.X.shape[0]
    def __getitem__(self, i):
        return (torch.from_numpy(self.X[i]),
                torch.tensor(self.y[i]),
                torch.zeros(1, dtype=torch.int64),                # single-dataset
                torch.zeros(2, dtype=torch.float32))              # no aux


class PxKSampler:
    def __init__(self, y, P=None, K=16, n_batches=80, seed=0):
        self.y = np.asarray(y)
        self.classes = np.unique(self.y)
        self.P = P or len(self.classes)
        self.K = K
        self.n_batches = n_batches
        self.rng = np.random.default_rng(seed)
        self.by_cls = {c: np.where(self.y == c)[0] for c in self.classes}
    def __iter__(self):
        for _ in range(self.n_batches):
            classes_p = self.rng.choice(self.classes,
                                        size=min(self.P, len(self.classes)),
                                        replace=False)
            batch = []
            for c in classes_p:
                idx = self.by_cls[c]
                take = self.K
                pick = self.rng.choice(idx, size=take, replace=(len(idx) < take))
                batch.extend(pick.tolist())
            yield batch
    def __len__(self): return self.n_batches


def prepare_pca(a, n_pca=50):
    if a.raw is not None:
        a = a.raw.to_adata()
    sc.pp.normalize_total(a, target_sum=1e4)
    sc.pp.log1p(a)
    sc.pp.highly_variable_genes(a, n_top_genes=2000, flavor="seurat", subset=False)
    a = a[:, a.var["highly_variable"]].copy()
    X = a.X.toarray() if sp.issparse(a.X) else a.X
    scaler = StandardScaler().fit(X)
    Xz = np.clip(scaler.transform(X), -10, 10)
    pca = PCA(n_components=n_pca, random_state=42).fit(Xz)
    Xp = pca.transform(Xz).astype(np.float32)
    return a, Xp


def train_fold(Xp, y, classes, tr, te, fold_id, ensemble_seeds=5):
    K = len(classes)
    all_probs = []
    for seed in range(ensemble_seeds):
        torch.manual_seed(fold_id * 100 + seed)
        np.random.seed(fold_id * 100 + seed)
        torch.cuda.empty_cache()
        model = PANDAEncoder(n_pca=Xp.shape[1], n_classes=K,
                             n_datasets=1).to(DEVICE)
        opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4)
        ds = SulicDataset(Xp[tr], y[tr])
        sampler = PxKSampler(y[tr], K=16, n_batches=80, seed=seed)
        loader = DataLoader(ds, batch_sampler=sampler, num_workers=0)
        for stage, ne in enumerate([15, 20, 25]):
            for e in range(ne):
                for X_b, y_b, _, aux_b in loader:
                    X_b, y_b = X_b.to(DEVICE), y_b.to(DEVICE)
                    aux_b = aux_b.to(DEVICE)
                    out = model(X_b, aux_b, lam_dann=0.0)
                    L_sup = supcon_loss(out["z"], y_b)
                    L_vic = vicreg_loss(out["z"])
                    L_ce  = F.cross_entropy(out["logits"], y_b, label_smoothing=0.05)
                    total = L_sup + 1.0 * L_vic + 0.4 * L_ce
                    if stage >= 1:
                        proto_ref = model.prototypes.detach().clone()
                        L_p = prototype_infonce(out["z"], y_b, proto_ref)
                        total = total + 0.6 * L_p
                    opt.zero_grad(); total.backward()
                    torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
                    opt.step()
                    if stage >= 1:
                        model.update_prototypes(out["z"].detach(), y_b)
        model.eval()
        with torch.no_grad():
            Xt = torch.from_numpy(Xp[te]).to(DEVICE)
            aux = torch.zeros(len(te), 2, device=DEVICE)
            out = model(Xt, aux, lam_dann=0.0)
            cos = out["z"] @ model.prototypes.T
            probs = torch.softmax(cos / 0.07, dim=1).cpu().numpy()
        all_probs.append(probs)
    ensemble_probs = np.mean(all_probs, axis=0)
    pred = ensemble_probs.argmax(axis=1)
    yte = y[te]
    return pred, ensemble_probs, yte


def evaluate_test(name, X, y_bin_or_multi, class_list):
    print(f"\n=== {name} ===", flush=True)
    skf = StratifiedKFold(n_splits=N_FOLDS, shuffle=True, random_state=42)
    aurocs, accs = [], []
    for fold, (tr, te) in enumerate(skf.split(X, y_bin_or_multi)):
        t0 = time.time()
        pred, probs, yte = train_fold(X, y_bin_or_multi, class_list, tr, te, fold)
        acc = accuracy_score(yte, pred)
        if len(class_list) == 2:
            auc = roc_auc_score(yte, probs[:, 1])
        else:
            try:
                auc = roc_auc_score(np.eye(len(class_list))[yte], probs,
                                    average="macro", multi_class="ovr")
            except Exception:
                auc = float("nan")
        print(f"[{name} fold {fold}] acc={acc:.4f} AUROC={auc:.4f} "
              f"wall={time.time()-t0:.0f}s", flush=True)
        aurocs.append(auc); accs.append(acc)
    print(f"[{name}] MEAN AUROC = {np.mean(aurocs):.4f} +- {np.std(aurocs):.4f}", flush=True)
    print(f"[{name}] MEAN ACC   = {np.mean(accs):.4f} +- {np.std(accs):.4f}", flush=True)
    return {"aurocs": aurocs, "accs": accs,
            "mean_auroc": float(np.mean(aurocs)),
            "std_auroc":  float(np.std(aurocs)),
            "mean_acc":   float(np.mean(accs)),
            "std_acc":    float(np.std(accs))}


def main():
    print(f"[sulic-panda] loading {SULIC_H5AD}", flush=True)
    a = ad.read_h5ad(SULIC_H5AD)
    print(f"[sulic-panda] shape {a.shape}, samples: {a.obs['sample'].value_counts().to_dict()}",
          flush=True)

    a_p, Xp = prepare_pca(a, n_pca=50)
    print(f"[sulic-panda] Xp {Xp.shape}", flush=True)

    y_A = (a.obs["sample"].isin(["Placode1", "Placode2"])).astype(int).values
    print(f"[sulic-panda] Test A class balance: {np.bincount(y_A).tolist()}", flush=True)
    resA = evaluate_test("TestA", Xp, y_A, ["Epithelium", "Placode"])

    if "placode_enriched" in a.obs.columns:
        mask_p = (a.obs["placode_enriched"] == 1).values
        sub = a[mask_p].copy()
        if "paper_subtype" not in sub.obs.columns:
            # fallback: kmeans on placode-cell Xp gives 4 pseudo-subtypes
            print("[sulic-panda] paper_subtype missing — deriving 4-way clustering on Xp", flush=True)
            from sklearn.cluster import KMeans
            Xp_sub = Xp[mask_p]
            km = KMeans(n_clusters=4, random_state=42, n_init=10).fit(Xp_sub)
            paper_subtype = np.array([f"PlacodeK{i}" for i in km.labels_])
        else:
            paper_subtype = sub.obs["paper_subtype"].astype(str).values
        cls = sorted(np.unique(paper_subtype))
        y_C = np.array([cls.index(v) for v in paper_subtype], dtype=np.int64)
        Xp_C = Xp[mask_p]
        print(f"[sulic-panda] Test C n={len(y_C)}, classes={cls}, "
              f"counts={np.bincount(y_C).tolist()}", flush=True)
        resC = evaluate_test("TestC", Xp_C, y_C, cls)
    else:
        resC = None

    # save
    result = {"testA": resA, "testC": resC}
    with open(OUT / "sulic_panda_heldout_results.json", "w") as f:
        json.dump(result, f, indent=2)
    print(f"\n[sulic-panda] wrote {OUT}/sulic_panda_heldout_results.json", flush=True)


if __name__ == "__main__":
    main()