Blame for experiments/43_rnafm_nicheformer.py

"""Frozen Nicheformer spot embeddings of the 5 RNA sets (each Visium spot / virtual spot = one pseudo-cell).

    python 43_rnafm_nicheformer.py [set ...]          (default: all of rna_sets.SETS)

Model     : theislab/Nicheformer (HF, MIT), revision 0ba4ba162f9ad3d6c8b713761954c31be2fa69ed, file model.safetensors
            (12-layer post-norm nn.TransformerEncoder, d=512, 16 heads, FFN 1024, learnable PE, context 1500).
            Downloaded to ext/nicheformer/hf/; the HF remote code (modeling_/tokenization_nicheformer.py) was read but is
            NOT executed: the forward pass is re-implemented below (same nn.TransformerEncoderLayer stack) and loaded
            with a strict state-dict load (the unused template `encoder_layer.*`, the MLM head and the `pos` buffer
            are the only checkpoint keys skipped, and that set is asserted).
            Technology means: ext/nicheformer/repo/data/model_means/dissociated_mean_script.npy from the official
            GitHub repo theislab/nicheformer (commit 485cadbc5caa15119adfd54228f8a8af835fcabc); gene order =
            model.h5ad var (20,310 Ensembl ids, token = column index + 30, == vocab.json).
Inputs    : rna_sets.load(set) raw(-like) counts + HGNC symbols. Symbols -> Ensembl via cache/external/
            dlpfc_genes_full.csv (10x GRCh38 reference), fallback Geneformer gene_name_id_dict_gc104M.pkl; genes not in
            the Nicheformer vocabulary are dropped (several symbols hitting one Ensembl id are summed), i.e. the spot is
            aligned to model.h5ad like the official `ad.concat([model, adata])` step.
Tokenising: official NicheformerDataset / tokenization notebooks: counts -> size-factor normalise to 1e4 (over the
            aligned genes) -> divide by the technology mean (nan->0, rounded half-up, 0 -> 1, as in the notebooks) ->
            rank the non-zero genes by descending value (stable sort; official uses numba quicksort, so only ties may
            be ordered differently) -> token = gene index + 30, zero-padded to 1500. Context tokens prepended in the
            official order [specie, assay, modality] and the sequence truncated to context_length 1500 (=> at most
            1497 gene tokens). Pad 0 -> 1 and attention padding mask = (token == 1), as in models/_utils.complete_masking
            with p = 0.
Context   : Visium was not in the pretraining corpus (SpatialCorpus-110M = dissociated 10x/Smart-seq + MERFISH/CosMx/
            Xenium/ISS) and has no technology token (the "visium": 9 entry in the get_embeddings notebook is a typo
            for Xenium = 9 in every tokenisation notebook and the HF vocab). Chosen: specie = human (5),
            modality = spatial (4), assay = 10x 3' v3 (11) -- Visium v1/v2 is a whole-transcriptome poly(A) 3'
            sequencing assay, so the closest documented technology is 10x 3' v3, and its matching technology mean is
            the dissociated (10x) one (the image-based means only cover their 300-1000-gene panels). A sensitivity
            run with modality = dissociated (3) is logged for hest only (not saved).
Readout   : official Nicheformer.get_embeddings(layer=-1, with_context=False): last-layer output, drop the 3 context
            positions, mean over the remaining 1497 positions (this INCLUDES padding positions, exactly as the
            official code; only spots with < 1497 detected vocabulary genes are affected and the fraction is logged).
            D = 512, not L2-normalised. bf16 autocast inference.
Output    : rna_sets.save("nicheformer", set, E) -> cache/rnafm/nicheformer_<set>.npy (n_spot, 512) float16.
Sanity    : finite, mean pairwise cosine of 200 random spots, and (hest) organ agreement of cosine 10-NN vs base rate.
"""
import hashlib
import json
import os
import pickle
import sys
import time

import anndata as ad
import numpy as np
import pandas as pd
import scipy.sparse as sp
import torch
import torch.nn as nn
from safetensors.torch import load_file

from common import CACHE, ROOT
from rna_sets import SETS, load, save

MODEL = "nicheformer"
REPO, REV = "theislab/Nicheformer", "0ba4ba162f9ad3d6c8b713761954c31be2fa69ed"
GH_COMMIT = "485cadbc5caa15119adfd54228f8a8af835fcabc"
EXT = os.path.join(ROOT, "ext", "nicheformer")
HF_DIR, GH_DIR = os.path.join(EXT, "hf"), os.path.join(EXT, "repo")
PAD, AUX, CTX_LEN = 1, 30, 1500
SPECIE_HUMAN, ASSAY_10X3V3, MOD_SPATIAL, MOD_DISSOC = 5, 11, 4, 3
dev = "cuda"
torch.manual_seed(0)


def sha(p):
    return hashlib.sha256(open(p, "rb").read()).hexdigest()[:16]


def ensure_download():
    if not os.path.exists(os.path.join(HF_DIR, "model.safetensors")):
        from huggingface_hub import snapshot_download
        snapshot_download(REPO, revision=REV, local_dir=HF_DIR)
    assert os.path.exists(os.path.join(GH_DIR, "data", "model_means", "dissociated_mean_script.npy")), \
        f"git clone https://github.com/theislab/nicheformer {GH_DIR} (commit {GH_COMMIT})"


class Nicheformer(nn.Module):
    """Re-implementation of NicheformerModel (HF modeling_nicheformer.py / official models/_nicheformer.py)."""

    def __init__(self, cfg):
        super().__init__()
        self.embeddings = nn.Embedding(cfg["n_tokens"] + 5, cfg["dim_model"], padding_idx=1)
        self.positional_embedding = nn.Embedding(cfg["context_length"], cfg["dim_model"])
        layer = nn.TransformerEncoderLayer(cfg["dim_model"], cfg["nheads"], cfg["dim_feedforward"], dropout=0.0,
                                           batch_first=True, layer_norm_eps=1e-12)
        self.encoder = nn.TransformerEncoder(layer, cfg["nlayers"], enable_nested_tensor=False)

    def get_embeddings(self, ids):
        pad = ids.eq(PAD)
        x = self.embeddings(ids) + self.positional_embedding(torch.arange(ids.shape[1], device=ids.device))
        for blk in self.encoder.layers:
            x = blk(x, src_key_padding_mask=pad, is_causal=False)
        return x[:, 3:].float().mean(1)        # with_context=False, mean incl. pad positions (official)


def load_model():
    cfg = json.load(open(os.path.join(HF_DIR, "config.json")))
    assert cfg["learnable_pe"] and cfg["specie"] and cfg["assay"] and cfg["modality"]
    sd = load_file(os.path.join(HF_DIR, "model.safetensors"))
    keep = {k[len("nicheformer."):]: v for k, v in sd.items()
            if k.startswith("nicheformer.") and not k.startswith("nicheformer.encoder_layer.") and k != "nicheformer.pos"}
    skipped = sorted(set(sd) - {"nicheformer." + k for k in keep})
    assert all(k.startswith(("nicheformer.encoder_layer.", "classifier_head.")) or k == "nicheformer.pos"
               for k in skipped), skipped
    m = Nicheformer(cfg)
    m.load_state_dict(keep, strict=True)
    print(f"weights {sha(os.path.join(HF_DIR, 'model.safetensors'))}: loaded {len(keep)} tensors strictly; "
          f"skipped {len(skipped)} (template encoder_layer / MLM head / pos buffer)", flush=True)
    return m.eval().to(dev), cfg


def gene_reference():
    ref = ad.read_h5ad(os.path.join(HF_DIR, "model.h5ad"), backed=None)
    ens = np.asarray(ref.var_names, dtype=str)
    vocab = json.load(open(os.path.join(HF_DIR, "vocab.json")))
    assert all(vocab[g] == i + AUX for i, g in enumerate(ens))
    mean = np.load(os.path.join(GH_DIR, "data", "model_means", "dissociated_mean_script.npy"))
    assert mean.shape == ens.shape
    mean = np.nan_to_num(mean)
    r = np.where((mean % 1) >= 0.5, np.ceil(mean), np.floor(mean))
    mean = np.where(r == 0, 1, r)
    sym2ens = dict(pd.read_csv(os.path.join(CACHE, "external", "dlpfc_genes_full.csv"))[["gene", "ensembl_id"]].values)
    from huggingface_hub import hf_hub_download
    gf = pickle.load(open(hf_hub_download("ctheodoris/Geneformer", "geneformer/gene_name_id_dict_gc104M.pkl"), "rb"))
    for k, v in gf.items():
        sym2ens.setdefault(k, v)
    return ens, mean, sym2ens


def tokenize(X, genes, ens, mean, sym2ens, modality):
    e2i = {e: i for i, e in enumerate(ens)}
    src = [(j, e2i[sym2ens[g]]) for j, g in enumerate(genes) if g in sym2ens and sym2ens[g] in e2i]
    rows, cols = zip(*src)
    M = sp.csr_matrix((np.ones(len(src), np.float32), (rows, cols)), shape=(len(genes), len(ens)))
    A = (X @ M).tocsr()                                          # spot x model genes (aligned like concat with model)
    lib = np.asarray(A.sum(1)).ravel()
    cov = dict(n_in=len(genes), n_mapped=len(src), n_model_genes=int((np.asarray(M.sum(0)).ravel() > 0).sum()),
               umi_frac=float(np.median(lib / np.maximum(np.asarray(X.sum(1)).ravel(), 1))))
    lib = lib + (lib == 0)
    A = (sp.diags(1e4 / lib) @ A).tocsr()
    ids = np.full((A.shape[0], CTX_LEN), PAD, np.int64)
    ids[:, :3] = [SPECIE_HUMAN, ASSAY_10X3V3, modality]
    ngen = np.zeros(A.shape[0], int)
    for r in range(A.shape[0]):
        s, e = A.indptr[r], A.indptr[r + 1]
        nz, v = A.indices[s:e], A.data[s:e] / mean[A.indices[s:e]]
        keep = v > 0
        nz, v = nz[keep], v[keep]
        order = nz[np.argsort(-v, kind="stable")][:CTX_LEN - 3]
        ids[r, 3:3 + len(order)] = order + AUX
        ngen[r] = len(order)
    cov["expr_genes_median"] = float(np.median(np.diff(A.indptr)))
    cov["tokens_median"] = float(np.median(ngen))
    cov["frac_spots_padded"] = float((ngen < CTX_LEN - 3).mean())
    return ids, cov


@torch.no_grad()
def embed(model, ids, bs=32):
    out = np.zeros((len(ids), 512), np.float32)
    i = 0
    while i < len(ids):
        try:
            b = torch.from_numpy(ids[i:i + bs]).to(dev)
            with torch.autocast("cuda", dtype=torch.bfloat16):
                out[i:i + bs] = model.get_embeddings(b).cpu().numpy()
            i += bs
        except torch.OutOfMemoryError:
            torch.cuda.empty_cache()
            bs = max(1, bs // 2)
            print(f"  OOM -> batch {bs}", flush=True)
    return out


def sanity(E, name):
    rng = np.random.default_rng(0)
    S = E[rng.choice(len(E), min(200, len(E)), replace=False)]
    S = S / np.linalg.norm(S, axis=1, keepdims=True)
    C = S @ S.T
    mc = float(C[np.triu_indices(len(S), 1)].mean())
    msg = f"[{MODEL}] {name}: finite={np.isfinite(E).all()} std={E.std(0).mean():.4f} mean-pair-cos(200)={mc:.4f}"
    if name == "hest":
        org = pd.read_csv(os.path.join(CACHE, "hest", "meta.csv")).organ.values
        assert len(org) == len(E)
        Z = torch.from_numpy(E / np.linalg.norm(E, axis=1, keepdims=True)).to(dev)
        nn_idx = []
        for k in range(0, len(Z), 2048):
            s = Z[k:k + 2048] @ Z.T
            s[torch.arange(s.shape[0]), torch.arange(k, k + s.shape[0])] = -2
            nn_idx.append(s.topk(10, dim=1).indices.cpu().numpy())
        nn_idx = np.concatenate(nn_idx)
        knn = float((org[nn_idx] == org[:, None]).mean())
        p = pd.Series(org).value_counts(normalize=True).values
        msg += f"  10-NN same-organ={knn:.3f} (base rate {float((p ** 2).sum()):.3f})"
    print(msg, flush=True)
    return mc


def main():
    sets = sys.argv[1:] or SETS
    t0 = time.time()
    ensure_download()
    model, cfg = load_model()
    ens, mean, sym2ens = gene_reference()
    print(f"context tokens: specie=human({SPECIE_HUMAN}) assay=10x 3' v3({ASSAY_10X3V3}) modality=spatial({MOD_SPATIAL}); "
          f"technology mean = dissociated_mean_script.npy (GitHub {GH_COMMIT[:7]})", flush=True)
    for name in sets:
        t1 = time.time()
        X, genes = load(name)
        ids, cov = tokenize(X, genes, ens, mean, sym2ens, MOD_SPATIAL)
        print(f"[{MODEL}] {name}: {X.shape[0]} spots; genes in {cov['n_in']}, mapped to vocab {cov['n_mapped']} "
              f"({cov['n_mapped'] / cov['n_in']:.1%}) -> {cov['n_model_genes']} model genes; median UMI fraction kept "
              f"{cov['umi_frac']:.3f}; expressed vocab genes/spot median {cov['expr_genes_median']:.0f}; gene tokens "
              f"median {cov['tokens_median']:.0f}; spots with padding {cov['frac_spots_padded']:.1%}", flush=True)
        E = embed(model, ids)
        sanity(E, name)
        save(MODEL, name, E)
        if name == "hest":                                 # sensitivity: modality token = dissociated
            ids[:, 2] = MOD_DISSOC
            E2 = embed(model, ids)
            c = (E * E2).sum(1) / np.linalg.norm(E, axis=1) / np.linalg.norm(E2, axis=1)
            print(f"  sensitivity modality=dissociated(3): cosine to saved per spot median {np.median(c):.4f}", flush=True)
            sanity(E2, "hest")
        print(f"[{MODEL}] {name}: {time.time() - t1:.0f}s", flush=True)
    print(f"[{MODEL}] done {time.time() - t0:.0f}s; peak GPU {torch.cuda.max_memory_allocated() / 2**30:.2f} GB", flush=True)


if __name__ == "__main__":
    main()