Blame for experiments/15_extract_marker_fm.py

"""Frozen single-cell foundation-model embeddings of the per-tile marker vector.

    python 15_extract_marker_fm.py geneformer [--value prop|count] [--pool mean|cls]
    python 15_extract_marker_fm.py scgpt      [--value prop|count]

Each ORION tile is treated as one pseudo-cell: the "expression" of each mapped gene is that
marker's positive-cell fraction (default, `*_prop`) or positive-cell count (`--value count`) in
the tile; every other gene is unexpressed.  Both models only see the *within-cell ordering* of
expressed genes (Geneformer: rank-value encoding; scGPT: per-cell quantile binning), so global
scale (how many cells / how much signal the tile has) is not representable — only composition.

Marker -> gene (protein panel -> HGNC symbol -> Ensembl via Geneformer gc104M dictionary):
  CD31->PECAM1, CD45->PTPRC, CD45RO->PTPRC (isoform of the same gene: merged, value=max(CD45,CD45RO)),
  CD68->CD68, CD4->CD4, FOXP3->FOXP3, CD8a->CD8A, CD20->MS4A1, PD-L1->CD274, CD3e->CD3E,
  CD163->CD163, E-cadherin->CDH1, PD-1->PDCD1, Ki67->MKI67, Pan-CK->KRT8+KRT18+KRT19 (same value
  for all three; the pan-cytokeratin antibody cocktail targets these simple-epithelium keratins),
  SMA->ACTA2.  Hoechst (DNA stain) has no gene -> dropped.
  PD-1_prop is degenerate in the ORION dataframes (>=0.998 in every tile, Spearman 0.02 with
  PD-1_count), so with --value prop the PDCD1 value is PD-1_count / n_nuclei (results/tile_stats.csv),
  clipped to [0, 1].

Outputs  cache/markers_<name>.npy  (3187, 1, D) float16, row order = tile_id
         cache/markers_<name>_meta.json  (mapping, choices, sanity numbers)
<name> = geneformer | scgpt, with "-count" / "-cls" suffixes for the non-default options.
"""
import argparse
import glob
import json
import os
import pickle
import time

import numpy as np
import pandas as pd
import torch
import torch.nn as nn

from common import CACHE, ROOT, load_all_tiles, marker_columns, marker_matrix

MARKER2GENES = {
    "CD31": ["PECAM1"], "CD45": ["PTPRC"], "CD45RO": ["PTPRC"], "CD68": ["CD68"], "CD4": ["CD4"],
    "FOXP3": ["FOXP3"], "CD8a": ["CD8A"], "CD20": ["MS4A1"], "PD-L1": ["CD274"], "CD3e": ["CD3E"],
    "CD163": ["CD163"], "E-cadherin": ["CDH1"], "PD-1": ["PDCD1"], "Ki67": ["MKI67"],
    "Pan-CK": ["KRT8", "KRT18", "KRT19"], "SMA": ["ACTA2"],
}
DROPPED = {"Hoechst": "nuclear DNA stain, no gene product"}
GF_REPO, GF_REV = "ctheodoris/Geneformer", None
SCGPT_REPO = "wanglab/scGPT-human"   # official bowang-lab whole-human pretrained checkpoint + vocab


def gene_expression(df, value):
    """(N, G) pseudo-cell expression over the mapped genes (unique, ordered)."""
    genes = list(dict.fromkeys(g for gs in MARKER2GENES.values() for g in gs))
    X = pd.DataFrame(0.0, index=df.index, columns=genes)
    notes = {}
    for m, gs in MARKER2GENES.items():
        if value == "prop":
            if m == "PD-1":
                ts = pd.read_csv(os.path.join(ROOT, "results", "tile_stats.csv"), index_col="tile_id")
                nn_ = ts.loc[df.tile_id, "n_nuclei"].values.astype(float)
                v = np.clip(df["PD-1_count"].values / np.maximum(nn_, 1.0), 0, 1)
                notes["PD-1"] = "PD-1_count / max(n_nuclei,1), clipped to [0,1] (PD-1_prop degenerate ~1.0)"
            else:
                v = df[f"{m}_prop"].values.astype(float)
        else:
            v = df[f"{m}_count"].values.astype(float)
        for g in gs:
            X[g] = np.maximum(X[g].values, v)   # PTPRC = max(CD45, CD45RO)
    return X, notes


# ----------------------------------------------------------------------------- Geneformer
def hf_file(repo, f):
    from huggingface_hub import hf_hub_download
    return hf_hub_download(repo, f)


def run_geneformer(X, pool, device, bs=256):
    from transformers import BertModel
    tokd = pickle.load(open(hf_file(GF_REPO, "geneformer/token_dictionary_gc104M.pkl"), "rb"))
    medd = pickle.load(open(hf_file(GF_REPO, "geneformer/gene_median_dictionary_gc104M.pkl"), "rb"))
    n2id = pickle.load(open(hf_file(GF_REPO, "geneformer/gene_name_id_dict_gc104M.pkl"), "rb"))
    mdir = os.path.dirname(hf_file(GF_REPO, "Geneformer-V2-104M/model.safetensors"))
    hf_file(GF_REPO, "Geneformer-V2-104M/config.json")
    ens = [n2id[g] for g in X.columns]
    tok = np.array([tokd[e] for e in ens])
    med = np.array([medd[e] for e in ens])
    CLS, EOS, PAD = tokd["<cls>"], tokd["<eos>"], tokd["<pad>"]

    model, info = BertModel.from_pretrained(mdir, output_loading_info=True, add_pooling_layer=False)
    model.eval().to(device)
    # Geneformer tokenizer: X / n_counts * 1e4 / gene_median, rank descending, drop zeros.
    # The per-cell n_counts factor does not change the within-cell order, so it is omitted.
    V = X.values / med[None]
    seqs = []
    for row in V:
        nz = np.nonzero(row > 0)[0]
        order = nz[np.argsort(-row[nz], kind="stable")]
        seqs.append([CLS] + tok[order].tolist() + [EOS])
    L = max(map(len, seqs))
    embs = []
    with torch.no_grad():
        for i in range(0, len(seqs), bs):
            b = seqs[i:i + bs]
            ids = torch.full((len(b), L), PAD, dtype=torch.long)
            for j, s in enumerate(b):
                ids[j, :len(s)] = torch.tensor(s)
            ids = ids.to(device)
            att = (ids != PAD).long()
            hs = model(input_ids=ids, attention_mask=att, output_hidden_states=True).hidden_states
            h = hs[-2].float()          # penultimate layer (Geneformer EmbExtractor default emb_layer=-1)
            if pool == "cls":
                e = h[:, 0]
            else:                       # mean over gene tokens (exclude <cls>, <eos>, pads)
                gm = att.clone().float()
                gm[:, 0] = 0
                lens = att.sum(1)
                gm[torch.arange(len(b)), lens - 1] = 0
                empty = gm.sum(1) == 0  # no expressed gene -> mean over <cls>,<eos>
                gm[empty] = att[empty].float()
                e = (h * gm[..., None]).sum(1) / gm.sum(1, keepdim=True)
            embs.append(e.cpu())
    E = torch.cat(embs).numpy()
    n_empty = int(sum(len(s) == 2 for s in seqs))
    meta = {"model": f"{GF_REPO}/Geneformer-V2-104M", "hidden": int(E.shape[1]),
            "missing_keys": info["missing_keys"], "unexpected_keys": info["unexpected_keys"],
            "ensembl": dict(zip(X.columns, ens)), "token_ids": dict(zip(X.columns, tok.tolist())),
            "gene_median_gc104M": dict(zip(X.columns, med.round(4).tolist())),
            "input": "<cls> + genes ranked by value/gene_median (desc, nonzero only) + <eos>",
            "layer": "penultimate (hidden_states[-2])",
            "pool": "CLS token" if pool == "cls" else "mean over gene tokens (excl. <cls>/<eos>/pad); "
                    "tiles with no expressed gene: mean over <cls>,<eos>",
            "n_tiles_no_expressed_gene": n_empty,
            "seq_len_mean": float(np.mean([len(s) - 2 for s in seqs]))}
    return E, meta


# ----------------------------------------------------------------------------- scGPT
class GeneEncoder(nn.Module):
    def __init__(self, n, d, pad):
        super().__init__()
        self.embedding = nn.Embedding(n, d, padding_idx=pad)
        self.enc_norm = nn.LayerNorm(d)

    def forward(self, x):
        return self.enc_norm(self.embedding(x))


class ContinuousValueEncoder(nn.Module):
    def __init__(self, d, max_value=512):
        super().__init__()
        self.linear1, self.linear2, self.norm = nn.Linear(1, d), nn.Linear(d, d), nn.LayerNorm(d)
        self.max_value = max_value

    def forward(self, x):
        x = torch.clamp(x.unsqueeze(-1), max=self.max_value)
        return self.norm(self.linear2(torch.relu(self.linear1(x))))


class ExprDecoder(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.fc = nn.Sequential(nn.Linear(d, d), nn.LeakyReLU(), nn.Linear(d, d), nn.LeakyReLU(), nn.Linear(d, 1))

    def forward(self, x):
        return self.fc(x).squeeze(-1)


class MVCDecoder(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.gene2query, self.W = nn.Linear(d, d), nn.Linear(d, d, bias=False)


class ScGPT(nn.Module):
    """Re-implementation of scgpt.model.TransformerModel for the whole-human checkpoint
    (continuous value encoder, post-norm encoder, no DSBN/batch tokens); the flash-attn packed
    Wqkv weights map 1:1 onto nn.MultiheadAttention.in_proj_*."""

    def __init__(self, ntoken, d=512, nhead=8, d_hid=512, nlayers=12, pad=0):
        super().__init__()
        self.encoder = GeneEncoder(ntoken, d, pad)
        self.flag_encoder = nn.Embedding(2, d)
        self.value_encoder = ContinuousValueEncoder(d)
        layer = nn.TransformerEncoderLayer(d, nhead, d_hid, dropout=0.0, batch_first=True, norm_first=False)
        self.transformer_encoder = nn.TransformerEncoder(layer, nlayers, enable_nested_tensor=False)
        self.decoder = ExprDecoder(d)
        self.mvc_decoder = MVCDecoder(d)

    def encode(self, genes, values, pad_mask):
        x = self.encoder(genes) + self.value_encoder(values)
        return self.transformer_encoder(x, src_key_padding_mask=pad_mask)


def scgpt_binning(row, n_bins, rng):
    """scgpt.preprocess.binning for a row of strictly positive values (random tie-splitting _digitize)."""
    bins = np.quantile(row, np.linspace(0, 1, n_bins - 1))
    left, right = np.digitize(row, bins), np.digitize(row, bins, right=True)
    return np.ceil(rng.random(len(row)) * (right - left) + left).astype(np.int64)


def run_scgpt(X, device, bs=256, seed=0):
    from huggingface_hub import hf_hub_download
    args = json.load(open(hf_hub_download(SCGPT_REPO, "args.json")))
    vocab = json.load(open(hf_hub_download(SCGPT_REPO, "vocab.json")))
    sd = torch.load(hf_hub_download(SCGPT_REPO, "best_model.pt"), map_location="cpu", weights_only=True)
    sd = {k.replace("self_attn.Wqkv.", "self_attn.in_proj_"): v for k, v in sd.items()}
    PAD, CLS = vocab[args["pad_token"]], vocab["<cls>"]
    model = ScGPT(len(vocab), args["embsize"], args["nheads"], args["d_hid"], args["nlayers"], PAD)
    model.load_state_dict(sd, strict=True)
    model.eval().to(device)
    n_bins, pad_value = args["n_bins"], args["pad_value"]
    gid = np.array([vocab[g] for g in X.columns])
    rng = np.random.default_rng(seed)
    seqs = []
    for row in X.values:   # scgpt.tasks.cell_emb: nonzero genes only, <cls> first with value=pad_value, then binning
        nz = np.nonzero(row > 0)[0]
        vals = scgpt_binning(row[nz], n_bins, rng).astype(np.float32) if len(nz) else np.zeros(0, np.float32)
        seqs.append((np.r_[CLS, gid[nz]], np.r_[np.float32(pad_value), vals]))
    L = max(len(g) for g, _ in seqs)
    embs, rec = [], []
    with torch.no_grad():
        for i in range(0, len(seqs), bs):
            b = seqs[i:i + bs]
            g = torch.full((len(b), L), PAD, dtype=torch.long)
            v = torch.full((len(b), L), float(pad_value))
            for j, (gg, vv) in enumerate(b):
                g[j, :len(gg)], v[j, :len(vv)] = torch.from_numpy(gg), torch.from_numpy(vv)
            g, v = g.to(device), v.to(device)
            mask = g.eq(PAD)
            h = model.encode(g, v, mask)
            embs.append(h[:, 0].float().cpu())
            pred = model.decoder(h)        # expression decoder: should track the visible input bins
            for j, (gg, vv) in enumerate(b):
                if len(gg) > 3:
                    rec.append(np.corrcoef(pred[j, 1:len(gg)].cpu().numpy(), vv[1:])[0, 1])
    E = torch.cat(embs).numpy()
    E = E / np.linalg.norm(E, axis=1, keepdims=True)   # as scgpt.tasks.embed_data
    meta = {"model": SCGPT_REPO + " (best_model.pt, args.json, vocab.json; official whole-human pretrained)",
            "hidden": int(E.shape[1]), "strict_load": True,
            "note_tdc": "tdc/scGPT weights differ from the official checkpoint and its original_model.bin "
                        "contains a cls_decoder (fine-tuned) -> not used",
            "vocab_ids": dict(zip(X.columns, gid.tolist())),
            "input": "<cls>(value=pad_value=-2) + expressed genes, values = scGPT per-cell quantile binning "
                     f"(n_bins={n_bins}, random tie-split, seed={seed}); padding masked",
            "pool": "<cls> token output of last layer, L2-normalised (scgpt.tasks.embed_data)",
            "n_tiles_no_expressed_gene": int(sum(len(gg) == 1 for gg, _ in seqs)),
            "decoder_recon_pearson_mean": float(np.nanmean(rec))}
    return E, meta


# ----------------------------------------------------------------------------- sanity
def ridge_cv_spearman(Z, Y, seed=0):
    from scipy.stats import spearmanr
    from sklearn.linear_model import RidgeCV
    from sklearn.model_selection import KFold, cross_val_predict
    from sklearn.pipeline import make_pipeline
    from sklearn.preprocessing import StandardScaler
    m = make_pipeline(StandardScaler(), RidgeCV(alphas=np.logspace(-2, 5, 15), alpha_per_target=True))
    P = cross_val_predict(m, Z, Y, cv=KFold(5, shuffle=True, random_state=seed))
    P, Y = P.reshape(len(Y), -1), Y.reshape(len(Y), -1)
    return np.array([spearmanr(P[:, k], Y[:, k])[0] for k in range(Y.shape[1])])


def sanity(E, df):
    Z = E.astype(np.float32)
    out = {"finite": bool(np.isfinite(Z).all()), "n_unique_rows": int(len(np.unique(E, axis=0)))}
    Zn = Z / np.linalg.norm(Z, axis=1, keepdims=True)
    S = Zn @ Zn.T
    iu = np.triu_indices(len(Z), 1)
    out["mean_pairwise_cos"] = float(S[iu].mean())
    Zc = Z - Z.mean(0)
    Zc /= np.linalg.norm(Zc, axis=1, keepdims=True) + 1e-8
    out["mean_pairwise_cos_centered"] = float((Zc @ Zc.T)[iu].mean())
    ev = np.linalg.svd(Z - Z.mean(0), compute_uv=False) ** 2
    out["pcs_for_90pct_var"] = int(np.searchsorted(np.cumsum(ev) / ev.sum(), 0.9) + 1)
    prop, _ = marker_columns(df)
    Y = df[prop].values.astype(np.float64)
    informative = [i for i, c in enumerate(prop) if c != "PD-1_prop"]
    raw = np.load(os.path.join(CACHE, "markers.npy")) if os.path.exists(os.path.join(CACHE, "markers.npy")) \
        else marker_matrix(df)
    for tag, X in [("emb", Z), ("raw_marker_baseline", raw)]:
        r = ridge_cv_spearman(X, Y)
        out[f"{tag}_prop_spearman_mean17"] = float(np.nanmean(r))
        out[f"{tag}_prop_spearman_mean16_noPD1"] = float(np.nanmean(r[informative]))
        out[f"{tag}_prop_spearman_per_marker"] = {c: round(float(x), 3) for c, x in zip(prop, r)}
        ts_p = os.path.join(ROOT, "results", "tile_stats.csv")
        if os.path.exists(ts_p):
            nn_ = pd.read_csv(ts_p, index_col="tile_id").loc[df.tile_id, "n_nuclei"].values
            out[f"{tag}_log1p_nnuclei_spearman"] = float(ridge_cv_spearman(X, np.log1p(nn_)[:, None])[0])
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("model", choices=["geneformer", "scgpt"])
    ap.add_argument("--value", choices=["prop", "count"], default="prop")
    ap.add_argument("--pool", choices=["mean", "cls"], default="mean", help="geneformer only")
    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    a = ap.parse_args()
    t0 = time.time()
    df = load_all_tiles()
    X, notes = gene_expression(df, a.value)
    name = a.model + ("-count" if a.value == "count" else "") + ("-cls" if a.model == "geneformer" and a.pool == "cls" else "")
    if a.model == "geneformer":
        E, meta = run_geneformer(X, a.pool, a.device)
    else:
        E, meta = run_scgpt(X, a.device)
    t_emb = time.time() - t0
    if a.device.startswith("cuda"):
        meta["peak_gpu_mem_GB"] = round(torch.cuda.max_memory_allocated() / 1e9, 3)
    E16 = E.astype(np.float16)[:, None, :]
    assert E16.shape[0] == len(df) == 3187
    np.save(os.path.join(CACHE, f"markers_{name}.npy"), E16)
    meta.update({"name": name, "shape": list(E16.shape), "dtype": "float16", "row_order": "tile_id",
                 "value_source": "*_prop (positive fraction)" if a.value == "prop" else "*_count (positive cells)",
                 "value_notes": notes, "marker_to_genes": MARKER2GENES, "dropped_markers": DROPPED,
                 "merges": {"PTPRC": "max(CD45, CD45RO) — CD45RO is an isoform of PTPRC",
                            "KRT8/KRT18/KRT19": "all = Pan-CK"},
                 "scale_note": "both models encode only within-cell gene order -> tile-level abundance/cellularity "
                               "is not representable",
                 "runtime_embed_s": round(t_emb, 1)})
    s = sanity(E16[:, 0], df)
    meta["sanity"] = s
    meta["runtime_total_s"] = round(time.time() - t0, 1)
    json.dump(meta, open(os.path.join(CACHE, f"markers_{name}_meta.json"), "w"), indent=1)
    print(f"[{name}] saved {E16.shape} float16  embed {t_emb:.1f}s  total {meta['runtime_total_s']}s")
    for k in ["n_tiles_no_expressed_gene", "decoder_recon_pearson_mean", "missing_keys", "unexpected_keys", "peak_gpu_mem_GB"]:
        if k in meta:
            print(f"  {k}: {meta[k]}")
    for k, v in s.items():
        if not k.endswith("per_marker"):
            print(f"  {k}: {v:.4f}" if isinstance(v, float) else f"  {k}: {v}")
    print("  per-marker Spearman (emb):", s["emb_prop_spearman_per_marker"])


if __name__ == "__main__":
    main()