"""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_.npy (3187, 1, D) float16, row order = tile_id cache/markers__meta.json (mapping, choices, sanity numbers) = 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[""], tokd[""], tokd[""] 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 , , 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 , 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": " + genes ranked by value/gene_median (desc, nonzero only) + ", "layer": "penultimate (hidden_states[-2])", "pool": "CLS token" if pool == "cls" else "mean over gene tokens (excl. //pad); " "tiles with no expressed gene: mean over ,", "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[""] 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, 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": "(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": " 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()