Spaces:
Running on Zero
Running on Zero
| 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}<br>observed Δ = %{x:.3f}<extra></extra>") | |
| fig.add_bar(y=names, x=d_pred[order], orientation="h", name="MapPFN prediction", marker_color=C_PRED, | |
| hovertemplate="%{y}<br>predicted Δ = %{x:.3f}<extra></extra>") | |
| 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 + "<br>PC1 %{x:.2f}, PC2 %{y:.2f}<extra></extra>") | |
| 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 | |
| # --------------------------------------------------------------------------- | |
| 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) | |
| <sub>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.</sub> | |
| """ | |
| ) | |
| 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) | |