import os os.environ.setdefault("JAX_PLATFORMS", "cpu") # JAX is only used on CPU to deserialise the Equinox checkpoints import spaces # noqa: E402 (must precede torch) import torch # noqa: E402 import anndata as ad # noqa: E402 import gradio as gr # noqa: E402 import numpy as np # noqa: E402 import pandas as pd # noqa: E402 import plotly.graph_objects as go # noqa: E402 from huggingface_hub import hf_hub_download # noqa: E402 import mappfn_jax # noqa: E402 import mappfn_torch # noqa: E402 REPO = "marvinsxtr/MapPFN" DEVICE = os.environ.get("MAPPFN_DEVICE", "cuda") NUM_CELLS = 200 # cells per population, as in the paper's evaluation config MAX_SHOTS = 8 # num_shots used for fine-tuning / evaluation DATASETS = { "Frangieh et al. 2021 (melanoma, CRISPR KO)": { "key": "frangieh", "default_context": "IFNγ", "default_gene": "B2M", }, "Papalexi et al. 2021 (THP-1, ECCITE-seq)": { "key": "papalexi", "default_context": "stimulated", "default_gene": "IFNGR1", }, } # --------------------------------------------------------------------------- # Data # --------------------------------------------------------------------------- DATA = {} for label, meta in DATASETS.items(): path = hf_hub_download(REPO, f"{meta['key']}/split_10.h5ad", repo_type="dataset") a = ad.read_h5ad(path) X = np.asarray(a.X, dtype=np.float32) obs = a.obs[["context", "treatment", "split"]].astype(str).reset_index(drop=True) groups = obs.groupby(["context", "treatment", "split"]).indices by_ct = obs.groupby(["context", "treatment"]).indices tvec = {} for (ctx, tr), idx in by_ct.items(): if tr != "control": tvec[tr] = np.asarray(a.obsm["treatment"][idx[0]], dtype=np.float32) split_of = {(c, t): s for (c, t, s) in groups} contexts = sorted(obs["context"].unique()) perts = {c: sorted(t for (cc, t) in by_ct if cc == c and t != "control") for c in contexts} DATA[meta["key"]] = dict( X=X, obs=obs, groups=groups, by_ct=by_ct, tvec=tvec, split_of=split_of, contexts=contexts, perts=perts, genes=list(a.var_names), ) del a # --------------------------------------------------------------------------- # Model: Equinox checkpoint -> PyTorch (EMA weights, as used by the authors for inference). # The released *fine-tuned* checkpoints were trained on a different gene panel/ordering than the released # h5ad splits (their marginals only match the data after sorting), so we serve the pre-trained checkpoint, # which is gene-agnostic and reproduces the paper's "MapPFN (pre-trained)" numbers. # --------------------------------------------------------------------------- _eqx = mappfn_jax.load_ema_model(hf_hub_download(REPO, "model.ckpt"), in_dim=50) MODEL = mappfn_torch.from_equinox(_eqx).eval().to(DEVICE) del _eqx # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- SPLIT_NOTE = { "test": "held-out test condition of the paper's split", "val": "validation condition of the paper's split", "train": "training condition of the paper's split (excluded from the in-context examples when queried)", } def _gene_label(d, ctx, gene): s = d["split_of"].get((ctx, gene)) return SPLIT_NOTE.get(s, "no measured cells in this context") def _cells(d, ctx, tr, splits, n, rng): idx = np.concatenate([d["groups"][(ctx, tr, s)] for s in splits if (ctx, tr, s) in d["groups"]]) sel = rng.choice(idx, size=n, replace=len(idx) < n) return d["X"][sel] def update_choices(dataset, context=None, gene=None): d = DATA[DATASETS[dataset]["key"]] ctx = context if context in d["contexts"] else DATASETS[dataset]["default_context"] genes = d["perts"][ctx] g = gene if gene in genes else (DATASETS[dataset]["default_gene"] if DATASETS[dataset]["default_gene"] in genes else genes[0]) return gr.update(choices=d["contexts"], value=ctx), gr.update(choices=genes, value=g) def update_genes(dataset, context, gene=None): d = DATA[DATASETS[dataset]["key"]] if context not in d["perts"]: return gr.update() genes = d["perts"][context] return gr.update(choices=genes, value=gene if gene in genes else genes[0]) C_CTRL, C_OBS, C_PRED = "#8a8985", "#2a78d6", "#eb6834" def _layout(fig, title, xlab, ylab, height=460): fig.update_layout( title=dict(text=title, x=0.01, font=dict(size=15)), xaxis_title=xlab, yaxis_title=ylab, height=height, template="plotly_white", margin=dict(l=60, r=20, t=60, b=50), legend=dict(orientation="h", yanchor="bottom", y=1.0, xanchor="right", x=1), hoverlabel=dict(font_size=12), ) fig.update_xaxes(gridcolor="rgba(0,0,0,0.06)", zerolinecolor="rgba(0,0,0,0.25)") fig.update_yaxes(gridcolor="rgba(0,0,0,0.06)", zerolinecolor="rgba(0,0,0,0.25)") return fig def _delta_plot(genes, ctrl, pred, obs, gene, top_k=15): d_pred = pred.mean(0) - ctrl.mean(0) d_obs = obs.mean(0) - ctrl.mean(0) if obs is not None else None score = np.abs(d_obs) + np.abs(d_pred) if d_obs is not None else np.abs(d_pred) order = np.argsort(-score)[:top_k][::-1] names = [genes[i] for i in order] fig = go.Figure() if d_obs is not None: fig.add_bar(y=names, x=d_obs[order], orientation="h", name="Observed", marker_color=C_OBS, hovertemplate="%{y}
observed Δ = %{x:.3f}") fig.add_bar(y=names, x=d_pred[order], orientation="h", name="MapPFN prediction", marker_color=C_PRED, hovertemplate="%{y}
predicted Δ = %{x:.3f}") fig.update_layout(barmode="group", bargap=0.25, bargroupgap=0.08) return _layout(fig, f"Mean expression change vs. control after {gene} perturbation (top {top_k} genes)", "Δ mean log-expression (perturbed − control)", None, height=520) def _pca_plot(ctrl, pred, obs): ref = np.concatenate([ctrl] + ([obs] if obs is not None else [pred])) mu = ref.mean(0) _, _, vt = np.linalg.svd(ref - mu, full_matrices=False) proj = lambda x: (x - mu) @ vt[:2].T # noqa: E731 fig = go.Figure() series = [("Control cells", ctrl, C_CTRL)] if obs is not None: series.append(("Observed perturbed cells", obs, C_OBS)) series.append(("MapPFN predicted cells", pred, C_PRED)) for name, x, c in series: p = proj(x) fig.add_scatter(x=p[:, 0], y=p[:, 1], mode="markers", name=name, marker=dict(size=8, color=c, opacity=0.6, line=dict(width=1, color="white")), hovertemplate=name + "
PC1 %{x:.2f}, PC2 %{y:.2f}") return _layout(fig, "Cell-level distributions (PCA of the 50-gene panel)", "PC1", "PC2") def _energy(x, y): def md(a, b): return np.sqrt(((a[:, None, :] - b[None, :, :]) ** 2).sum(-1)).mean() return float(2 * md(x, y) - md(x, x) - md(y, y)) # --------------------------------------------------------------------------- # Inference # --------------------------------------------------------------------------- @spaces.GPU(duration=15) def _run_model(obs_q, intv, treat, guidance, seed): model = MODEL dev = DEVICE obs_t = torch.from_numpy(obs_q).to(dev)[None, None] intv_t = torch.from_numpy(intv).to(dev)[None] if intv is not None else None tr_t = torch.from_numpy(treat).to(dev)[None] gen = torch.Generator(device=dev).manual_seed(int(seed)) out = mappfn_torch.sample(model, obs_t, intv_t, tr_t, guidance=float(guidance), step_size=0.01, generator=gen) return out[0].float().cpu().numpy() def predict( dataset: str, context: str, gene: str, num_shots: int = 8, guidance: float = 2.0, seed: int = 0, ): """Predict the single-cell expression distribution after perturbing a gene with MapPFN. MapPFN conditions on control cells of the chosen biological context plus `num_shots` other perturbation experiments measured in the same context (in-context examples), and generates 200 post-perturbation cells over the 50-gene panel via flow matching. Args: dataset: Single-cell perturbation dataset (Frangieh or Papalexi). context: Biological context / condition (e.g. "IFNγ" or "stimulated"). gene: Gene whose perturbation (knockout) should be predicted. num_shots: Number of in-context perturbation experiments (0-8). guidance: Classifier-free guidance weight (paper default 2.0). seed: Random seed for cell sub-sampling, context selection and ODE noise. Returns: Summary markdown, mean-shift bar chart, PCA scatter, per-gene table, CSV of predicted cells. """ if dataset not in DATASETS: raise gr.Error("Unknown dataset.") ds_key = DATASETS[dataset]["key"] d = DATA[ds_key] if context not in d["contexts"]: raise gr.Error(f"Unknown context '{context}' for this dataset.") if gene not in d["tvec"]: raise gr.Error(f"'{gene}' is not a perturbation target in this dataset.") num_shots = int(max(0, min(MAX_SHOTS, int(num_shots)))) rng = np.random.default_rng(int(seed)) # Query: control cells of the context (the "control" split, as in the authors' test datasets) obs_q = _cells(d, context, "control", ["control"], NUM_CELLS, rng) # In-context experiments: training-split perturbations of the same context, excluding the query pool = [t for t in d["perts"][context] if t != gene and (context, t, "train") in d["groups"]] n = min(num_shots, len(pool)) shots = list(rng.choice(pool, size=n, replace=False)) if n else [] intv = np.stack([_cells(d, context, t, ["train"], NUM_CELLS, rng) for t in shots]) if shots else None treat = np.stack([d["tvec"][t][None] for t in [*shots, gene]]).astype(np.float32) # NB: context control cells for each shot are drawn from the same control pool; only the query's # control population enters the model (obs_data[:, -1:]), exactly as in `unpack_batch`. pred = _run_model(obs_q, intv, treat, guidance, seed) obs_cells = None if (context, gene) in d["by_ct"]: idx = d["by_ct"][(context, gene)] obs_cells = d["X"][rng.choice(idx, size=NUM_CELLS, replace=len(idx) < NUM_CELLS)] genes = d["genes"] ctrl_mean = obs_q.mean(0) table = pd.DataFrame({ "gene": genes, "control mean": ctrl_mean, "predicted mean": pred.mean(0), "predicted Δ": pred.mean(0) - ctrl_mean, }) lines = [ f"### {gene} perturbation · {context}", f"**Status of this perturbation:** {_gene_label(d, context, gene)} ", f"**In-context experiments ({len(shots)}):** {', '.join(shots) if shots else 'none (zero-shot)'}", ] if obs_cells is not None: table["observed mean"] = obs_cells.mean(0) table["observed Δ"] = obs_cells.mean(0) - ctrl_mean dp, do = table["predicted Δ"].to_numpy(), table["observed Δ"].to_numpy() r = float(np.corrcoef(dp, do)[0, 1]) e_pred = _energy(pred, obs_cells) e_ctrl = _energy(obs_q, obs_cells) lines += [ "", "| metric | MapPFN | no-change baseline (control cells) |", "|---|---|---|", f"| Pearson r of Δ vs observed Δ | **{r:.3f}** | – |", f"| MSE of mean expression | **{float(((pred.mean(0) - obs_cells.mean(0)) ** 2).mean()):.4f}** | " f"{float(((ctrl_mean - obs_cells.mean(0)) ** 2).mean()):.4f} |", f"| Energy distance to observed cells | **{e_pred:.3f}** | {e_ctrl:.3f} |", ] table = table.reindex(table["predicted Δ"].abs().sort_values(ascending=False).index).round(4) csv_path = f"/tmp/mappfn_{ds_key}_{gene}_{context}_{seed}.csv".replace(" ", "_") pd.DataFrame(pred, columns=genes).to_csv(csv_path, index_label="cell") return ( "\n".join(lines), _delta_plot(genes, obs_q, pred, obs_cells, gene), _pca_plot(obs_q, pred, obs_cells), table, csv_path, ) # --------------------------------------------------------------------------- # UI # --------------------------------------------------------------------------- CSS = """ #col-container { max-width: 1200px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ DEFAULT_DS = list(DATASETS)[0] _d0 = DATA[DATASETS[DEFAULT_DS]["key"]] with gr.Blocks() as demo: with gr.Column(elem_id="col-container"): gr.Markdown( """ # MapPFN: Learning Causal Perturbation Maps in Context A prior-data fitted network that predicts **single-cell post-perturbation expression distributions** in an unseen biological context by reading a handful of other perturbation experiments *in context*. Pick a dataset, a cell context and a gene to knock out: MapPFN generates 200 perturbed cells from the context's control cells and compares them with the measured cells (when available). [Paper](https://huggingface.co/papers/2601.21092) · [Code](https://github.com/marvinsxtr/MapPFN) · [Model](https://huggingface.co/marvinsxtr/MapPFN) · [Data](https://huggingface.co/datasets/marvinsxtr/MapPFN) Runs the **pre-trained** checkpoint, trained only on the synthetic causal prior, with no fine-tuning on real cells. As the paper reports for this setting, it gets the direction of effects right but often overshoots their size, so on mean error it can score worse than the "no change" baseline. The released fine-tuned checkpoints don't match the gene panel of the released data splits, so they are not used here. """ ) with gr.Row(): with gr.Column(scale=1): dataset = gr.Dropdown(list(DATASETS), value=DEFAULT_DS, label="Dataset") context = gr.Dropdown(_d0["contexts"], value=DATASETS[DEFAULT_DS]["default_context"], label="Biological context", allow_custom_value=True) gene = gr.Dropdown(_d0["perts"][DATASETS[DEFAULT_DS]["default_context"]], value=DATASETS[DEFAULT_DS]["default_gene"], label="Gene to perturb", allow_custom_value=True) with gr.Accordion("Advanced settings", open=False): num_shots = gr.Slider(0, MAX_SHOTS, value=8, step=1, label="In-context experiments (shots)") guidance = gr.Slider(1.0, 4.0, value=2.0, step=0.1, label="Classifier-free guidance") seed = gr.Number(value=0, precision=0, label="Seed") run = gr.Button("Predict perturbation", variant="primary") summary = gr.Markdown() with gr.Column(scale=2): delta_plot = gr.Plot(label="Mean shift per gene") pca_plot = gr.Plot(label="Cell distributions") with gr.Accordion("Per-gene table & predicted cells", open=False): table = gr.Dataframe(label="Per-gene means (log-normalised expression)", wrap=True) csv_out = gr.File(label="Predicted cells (CSV, 200 cells × 50 genes)") outputs = [summary, delta_plot, pca_plot, table, csv_out] inputs = [dataset, context, gene, num_shots, guidance, seed] gr.Examples( examples=[ ["Frangieh et al. 2021 (melanoma, CRISPR KO)", "IFNγ", "B2M"], ["Frangieh et al. 2021 (melanoma, CRISPR KO)", "IFNγ", "STAT3"], ["Frangieh et al. 2021 (melanoma, CRISPR KO)", "IFNγ", "CD58"], ["Papalexi et al. 2021 (THP-1, ECCITE-seq)", "stimulated", "IFNGR1"], ["Papalexi et al. 2021 (THP-1, ECCITE-seq)", "stimulated", "JAK2"], ["Papalexi et al. 2021 (THP-1, ECCITE-seq)", "stimulated", "STAT1"], ], inputs=[dataset, context, gene], outputs=outputs, fn=predict, cache_examples=True, cache_mode="lazy", ) dataset.change(update_choices, [dataset, context, gene], [context, gene], api_visibility="private") context.change(update_genes, [dataset, context, gene], [gene], api_visibility="private") run.click(predict, inputs, outputs) if __name__ == "__main__": demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)