Blame for app.py

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
# ---------------------------------------------------------------------------
@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)

<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)