| """true zero-shot on Nestorowa GSE81682 (1920 smart-seq2 FACS-labeled 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")
|
| NEST = ROOT / "data/raw/nestorowa_combined.h5ad"
|
|
|
|
|
| COARSE = {
|
| "LT-HSC": "LT-HSC",
|
| "MPP": "HSPC",
|
| "GMP": "HSPC",
|
| "myeloid": "HSPC",
|
| "erythroid": "HSPC",
|
| "megakaryocyte":"HSPC",
|
| "basophil-mast":"HSPC",
|
| "lymphoid": "HSPC",
|
| "unassigned": "HSPC",
|
| "UNK": "HSPC",
|
| }
|
|
|
|
|
| def infer(a, variant):
|
| ck = torch.load(ROOT / f"checkpoints/hematopoiesis/{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/hematopoiesis/harmonized/corpus_stats.npz", allow_pickle=True)
|
| pca = pickle.load(open(ROOT / "data/corpus/hematopoiesis/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():
|
| a = ad.read_h5ad(NEST)
|
| print(f"[nest] {a.shape} facs gates: {a.obs['cell_type'].value_counts().to_dict()}", flush=True)
|
|
|
| for variant in ("pca", "marker"):
|
| print(f"\n=== {variant.upper()} ===", flush=True)
|
| pred, probs, classes = infer(a, variant)
|
| pred_coarse = np.array([COARSE.get(p, "HSPC") for p in pred])
|
| y_true = a.obs["cell_type"].astype(str).values
|
| mask = y_true != "unknown"
|
| acc = accuracy_score(y_true[mask], pred_coarse[mask])
|
| f1 = f1_score(y_true[mask], pred_coarse[mask], average="macro", zero_division=0)
|
| rep = classification_report(y_true[mask], pred_coarse[mask], zero_division=0, output_dict=True)
|
| print(f"[eval-{variant}] n_labeled={mask.sum()} coarse-acc={acc:.4f} macro-f1={f1:.4f}", flush=True)
|
| fine_by_gate = pd.crosstab(a.obs["cell_type"].astype(str), pd.Series(pred))
|
| print(fine_by_gate.to_string(), flush=True)
|
|
|
| out = ROOT / f"discovery/hematopoiesis/{variant}"
|
| out.mkdir(parents=True, exist_ok=True)
|
| (out / "94_nestorowa_zero_shot.json").write_text(json.dumps({
|
| "variant": variant, "n_cells_total": int(a.n_obs), "n_cells_labeled": int(mask.sum()),
|
| "coarse_acc": float(acc), "coarse_f1": float(f1),
|
| "coarse_per_class": rep,
|
| "fine_by_gate": fine_by_gate.to_dict(),
|
| "max_cos_p50": float(np.median(probs.max(axis=1))),
|
| }, indent=2, default=str))
|
| pd.DataFrame({"cell_id": a.obs_names, "facs_gate": y_true,
|
| "pred_fine": pred, "pred_coarse": pred_coarse,
|
| "max_cos": probs.max(axis=1)}).to_csv(out / "94_nestorowa_predictions.csv", index=False)
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
