scripts/analysis/72_emergent_axes.py

"""per-class-residual PCA on the 128d projections; skin only (only cached)."""
from pathlib import Path
import json, warnings, numpy as np, pandas as pd, anndata as ad
warnings.filterwarnings("ignore")
from sklearn.decomposition import PCA

import os as _os
from pathlib import Path as _Path
PANDA_ROOT = _Path(_os.environ.get("PANDA_ROOT", str(_Path(__file__).resolve().parents[2])))
OUT = Path(str(PANDA_ROOT / "discovery")); OUT.mkdir(exist_ok=True)

P = ad.read_h5ad(str(PANDA_ROOT / "discovery/pan_skin/marker/50_aldrich_projections.h5ad"))
Z = np.asarray(P.obsm["Z_projection"])
pred = (P.obs["pred_bbse_label"] if "pred_bbse_label" in P.obs
        else P.obs["pred_label"]).astype(str).values
print(f"[skin] Z shape {Z.shape}, n_pred_classes={len(np.unique(pred))}", flush=True)

Zres = np.zeros_like(Z)
for cls in np.unique(pred):
    m = pred == cls
    if m.sum() < 2: continue
    Zres[m] = Z[m] - Z[m].mean(0, keepdims=True)

pca = PCA(n_components=20, random_state=0).fit(Zres)
ev = pca.explained_variance_ratio_
print(f"[skin] top-10 residual PC EV: {[f'{e:.4f}' for e in ev[:10]]}", flush=True)
print(f"[skin] cumulative top-10: {ev[:10].cumsum()[-1]:.3f}", flush=True)
print(f"[skin] cumulative top-20: {ev.cumsum()[-1]:.3f}", flush=True)

scores = pca.transform(Zres)  # (N, 20)
aux_cols = {}
if "log10_counts" in P.obs:
    aux_cols["log10_counts"] = P.obs["log10_counts"].astype(float).values
elif "n_counts" in P.obs:
    aux_cols["log10_counts"] = np.log10(P.obs["n_counts"].astype(float).values + 1)
if "missing_hvg_frac" in P.obs:
    aux_cols["missing_hvg_frac"] = P.obs["missing_hvg_frac"].astype(float).values
if "max_cos" in P.obs:
    aux_cols["max_cos"] = P.obs["max_cos"].astype(float).values
else:
    import torch
    ck = torch.load(str(PANDA_ROOT / "checkpoints/pan_skin/panda_final.pt"),
                    map_location="cpu", weights_only=False)
    protos = ck["prototypes"]
    protos = protos / (np.linalg.norm(protos, axis=1, keepdims=True) + 1e-8)
    Zn = Z / (np.linalg.norm(Z, axis=1, keepdims=True) + 1e-8)
    aux_cols["max_cos"] = (Zn @ protos.T).max(axis=1)

if "genotype" in P.obs:
    genotype_bin = (P.obs["genotype"] == "En1-cKO").astype(int).values
    aux_cols["genotype_cKO"] = genotype_bin.astype(float)

print(f"[skin] auxiliaries: {list(aux_cols.keys())}", flush=True)

rows = []
for pc in range(10):
    row = {"PC": f"PC{pc+1}", "EV": float(ev[pc])}
    for aux_name, aux_vals in aux_cols.items():
        r = float(np.corrcoef(scores[:, pc], aux_vals)[0, 1])
        row[f"corr_{aux_name}"] = r
    rows.append(row)
axes_df = pd.DataFrame(rows)
axes_df.to_csv(OUT / "72_emergent_axes_skin.csv", index=False)
print(axes_df.to_string(index=False), flush=True)

loadings = pca.components_[:5]  # (5, 128)
np.save(OUT / "72_emergent_axes_skin_pc_loadings.npy", loadings)

summary = {
    "system": "pan_skin",
    "total_ev_top10": float(ev[:10].sum()),
    "total_ev_top20": float(ev.sum()),
    "top10_ev": [float(e) for e in ev[:10]],
    "auxiliaries": list(aux_cols.keys()),
}
json.dump(summary, open(OUT / "72_emergent_axes_summary.json", "w"), indent=2)
print(f"\nwrote {OUT}/72_emergent_axes_summary.json", flush=True)