| """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" |
|
|
|
|
| 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) |
| return X, notes |
|
|
|
|
| |
| 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) |
| |
| |
| 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() |
| if pool == "cls": |
| e = h[:, 0] |
| else: |
| gm = att.clone().float() |
| gm[:, 0] = 0 |
| lens = att.sum(1) |
| gm[torch.arange(len(b)), lens - 1] = 0 |
| empty = gm.sum(1) == 0 |
| 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 |
|
|
|
|
| |
| 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: |
| 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) |
| 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) |
| 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 |
|
|
|
|
| |
| 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() |
|
|
