| """true zero-shot on baron test-half (943 mouse islet cells held out of corpus) under pca+marker variants."""
|
| from pathlib import Path
|
| import warnings, json, sys, pickle, numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp, torch
|
| warnings.filterwarnings("ignore"); sc.settings.verbosity = 0
|
| 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 import PANDAEncoder
|
| from sklearn.metrics import accuracy_score, f1_score, classification_report
|
|
|
| ROOT = Path(str(PANDA_ROOT))
|
| DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| BARON = ROOT / "data/corpus/pancreas/held_out_labeled/baron_GSE84133_mouse_test.h5ad"
|
|
|
|
|
| def infer(a, variant):
|
| ck = torch.load(ROOT / f"checkpoints/pancreas/{variant}/panda_final.pt",
|
| map_location=DEVICE, weights_only=False)
|
| classes = ck["classes"]; marker_genes = ck.get("marker_genes", [])
|
| stats = np.load(ROOT / "data/corpus/pancreas/harmonized/corpus_stats.npz", allow_pickle=True)
|
| pca = pickle.load(open(ROOT / "data/corpus/pancreas/harmonized/pca_basis.pkl", "rb"))
|
| hvgs = [str(g) for g in stats["shared_hvgs"]]
|
| hvg2i = {g: i for i, g in enumerate(hvgs)}
|
| common = [g for g in a.var_names.astype(str) if g in hvg2i]
|
| a_c = a[:, common].copy()
|
| sc.pp.normalize_total(a_c, target_sum=1e4); sc.pp.log1p(a_c)
|
| X = a_c.X.toarray().astype(np.float32) if sp.issparse(a_c.X) else a_c.X.astype(np.float32)
|
| Xf = np.zeros((a.n_obs, len(hvgs)), dtype=np.float32)
|
| Xf[:, np.array([hvg2i[g] for g in common])] = X
|
| Xz = np.clip((Xf - stats["mean"].astype(np.float32)) / stats["std"].astype(np.float32), -10, 10)
|
| Xpca = pca.transform(Xz).astype(np.float32)
|
|
|
| Xmark = None
|
| if variant == "marker":
|
| mvals = np.zeros((a.n_obs, len(marker_genes)), dtype=np.float32)
|
| for j, g in enumerate(marker_genes):
|
| if g in a.var_names:
|
| col = a[:, g].X
|
| if sp.issparse(col): col = col.toarray()
|
| mvals[:, j] = col.flatten().astype(np.float32)
|
| mmu = mvals.mean(axis=0, keepdims=True); msig = mvals.std(axis=0, keepdims=True) + 1e-6
|
| Xmark = np.clip((mvals - mmu) / msig, -5, 5).astype(np.float32)
|
|
|
| model = PANDAEncoder(variant=variant, n_pca=50,
|
| n_markers=len(marker_genes) if variant == "marker" else 0,
|
| n_classes=len(classes), n_sub=3,
|
| n_datasets=len(ck["datasets"])).to(DEVICE).eval()
|
| model.load_state_dict(ck["model"])
|
|
|
| preds, probs = [], []
|
| with torch.no_grad():
|
| for i in range(0, a.n_obs, 4096):
|
| xb = torch.from_numpy(Xpca[i:i+4096]).to(DEVICE)
|
| xmb = torch.from_numpy(Xmark[i:i+4096]).to(DEVICE) if Xmark is not None else None
|
| aux = torch.zeros(len(xb), 2, device=DEVICE)
|
| out = model(xb, aux, x_markers=xmb, lam_dann=0.0)
|
| mc = model.max_sub_cos(out["z"])
|
| preds.append(mc.argmax(dim=1).cpu().numpy())
|
| probs.append(torch.softmax(mc / 0.07, dim=1).cpu().numpy())
|
| return np.array([classes[i] for i in np.concatenate(preds)]), np.concatenate(probs), classes
|
|
|
|
|
| def main():
|
| print(f"[baron] loading {BARON}", flush=True)
|
| a = ad.read_h5ad(BARON)
|
| y_true = a.obs["canonical_label"].astype(str).values
|
| print(f"[baron] {a.shape} true labels: {pd.Series(y_true).value_counts().to_dict()}", flush=True)
|
|
|
| for variant in ("pca", "marker"):
|
| print(f"\n=== {variant.upper()} ===", flush=True)
|
| pred, probs, classes = infer(a, variant)
|
|
|
| mask = np.isin(y_true, classes)
|
| acc = accuracy_score(y_true[mask], pred[mask])
|
| f1 = f1_score(y_true[mask], pred[mask], average="macro", zero_division=0)
|
| rep = classification_report(y_true[mask], pred[mask], zero_division=0, output_dict=True)
|
| print(f"[eval-{variant}] n={mask.sum()} acc={acc:.4f} macro-f1={f1:.4f}", flush=True)
|
| out = ROOT / f"discovery/pancreas/{variant}"
|
| out.mkdir(parents=True, exist_ok=True)
|
| (out / "93_baron_zero_shot.json").write_text(json.dumps({
|
| "variant": variant, "n_cells": int(mask.sum()), "n_classes_eval": int(len(set(y_true[mask]))),
|
| "acc": float(acc), "macro_f1": float(f1), "per_class": rep,
|
| "predicted_dist": pd.Series(pred).value_counts().to_dict(),
|
| "true_dist": pd.Series(y_true).value_counts().to_dict(),
|
| }, indent=2, default=str))
|
| pd.DataFrame({"cell_id": a.obs_names, "true_label": y_true, "pred_label": pred,
|
| "max_cos": probs.max(axis=1)}).to_csv(
|
| out / "93_baron_predictions.csv", index=False)
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
