Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+data/frangieh_split_10.h5ad filter=lfs diff=lfs merge=lfs -text
+data/papalexi_split_10.h5ad filter=lfs diff=lfs merge=lfs -text

README.md

@@ -6,54 +6,42 @@ colorTo: red
 sdk: gradio
 sdk_version: 6.31.0
 app_file: app.py
-short_description: In-context prediction of single-cell perturbation responses
+short_description: In-context prediction of genetic perturbation responses
 python_version: "3.12"
 startup_duration_timeout: 30m
 ---
 
 # MapPFN: Learning Causal Perturbation Maps in Context
 
-Interactive demo of **MapPFN** (Sextro, Klos & Dernbach, NeurIPS 2026) — a prior-data fitted
-network that predicts the post-perturbation single-cell expression distribution for an
-**unseen** gene knockdown, using a handful of example perturbation experiments as in-context
-conditioning.
-
-This Space is a faithful PyTorch port of the reference JAX/Equinox implementation
-([github.com/marvinsxtr/MapPFN](https://github.com/marvinsxtr/MapPFN), MIT license), running
-the released EMA weights from [marvinsxtr/MapPFN](https://huggingface.co/marvinsxtr/MapPFN):
-
-- `model_finetuned_papalexi.ckpt` — Papalexi et al. 2021 (CRISPRi Perturb-seq, leukemia, IFN-γ)
-- `model_finetuned_frangieh.ckpt` — Frangieh et al. 2021 (CRISPRi Perturb-seq, melanoma)
-
-## What it does
-
-1. Pick a dataset and a **held-out test perturbation** (a gene knockdown the model never saw).
-2. MapPFN samples `num_shots` *other* perturbations from the same cell context as in-context
-   examples (control cells → treated cells, with one-hot treatment vectors).
-3. A classifier-free-guided ODE flow (Dopri5, step 0.01, guidance 2.0) generates the predicted
-   post-perturbation population of cells.
-4. The prediction is compared against the real measured Perturb-seq data for that knockdown:
-   expression distributions of the top differentially-expressed genes, per-gene mean scatter,
-   and distribution metrics (R², RMSE, multiscale MMD).
-
-Because the query perturbations come from the papers' test splits, every example is a genuine
-zero-shot prediction with ground truth available — the metrics panel is computed live.
-
-## Data & license
-
-Bundled expression matrices are derived from the
-[marvinsxtr/MapPFN dataset](https://huggingface.co/datasets/marvinsxtr/MapPFN) (MIT), which
-packages the published Perturb-seq datasets of Papalexi et al. 2021 and Frangieh et al. 2021.
-Expression values are normalize-total + log1p, 50 genes per dataset.
-
-**Links:** [Paper](https://arxiv.org/abs/2601.21092) · [Project page](https://marvinsxtr.github.io/MapPFN) · [Code](https://github.com/marvinsxtr/MapPFN) · [Weights](https://huggingface.co/marvinsxtr/MapPFN)
-
-```bibtex
-@inproceedings{sextro2026mappfn,
-  title     = {Map{PFN}: Learning Causal Perturbation Maps in Context},
-  author    = {Marvin Sextro and Weronika K{\l}os and Gabriel Dernbach},
-  booktitle = {Advances in Neural Information Processing Systems},
-  volume    = {39},
-  year      = {2026}
-}
-```
+Interactive demo of **MapPFN** (Sextro et al., 2026, [arXiv:2601.21092](https://huggingface.co/papers/2601.21092)) —
+a flow-matching in-context model that predicts the post-perturbation single-cell
+expression distribution for a **held-out** gene, given only demonstrations of *other*
+perturbations in the same cell context.
+
+- **Paper**: [MapPFN: Learning Causal Perturbation Maps in Context](https://huggingface.co/papers/2601.21092)
+- **Code**: https://github.com/marvinsxtr/MapPFN (MIT)
+- **Weights / data**: https://huggingface.co/marvinsxtr/MapPFN (MIT)
+
+## What the demo does
+
+1. Choose a dataset — Frangieh 2021 (melanoma CRISPR screens across IFN-γ /
+   co-culture contexts) or Papalexi 2021 (ECCITE-seq stimulated T cells). Each uses
+   the authors' fine-tuned checkpoint for that dataset.
+2. Choose a cell context and a **held-out target gene** (from the paper's test split).
+3. MapPFN samples the treated population in the paper's 50-d PCA expression space,
+   conditioning on the context's control cells plus a configurable number of
+   in-context demonstrations (other treatments' control→treated pairs).
+4. The prediction is compared against the real held-out measurements: mean/variance
+   R², perturbation-effect correlation, and a top differentially-expressed-gene table.
+
+## Implementation notes
+
+The authors' reference implementation is JAX/Equinox with a cuDNN attention kernel;
+ZeroGPU Spaces are PyTorch-first, so this Space runs a faithful PyTorch port
+(`map_pfn_torch.py`) of the model and its exact sampling procedure (Dormand–Prince
+5(4) ODE solver, constant step 0.01, 100 steps, classifier-free guidance 2.0, EMA
+weights). The port was validated numerically against the reference: CFG velocity
+fields and full trajectories match to float32 tolerance on real test queries.
+
+Datasets and weights are redistributed under their MIT license from
+`marvinsxtr/MapPFN`.

app.py

@@ -1,367 +1,311 @@
 """MapPFN: Learning Causal Perturbation Maps in Context — interactive demo.
 
-Faithful PyTorch port of the JAX/Equinox reference
-(https://github.com/marvinsxtr/MapPFN, Sextro et al., NeurIPS 2026).
-Given a set of in-context perturbation experiments (control cells, treated
-cells, one-hot treatments), MapPFN predicts the post-perturbation expression
-distribution for a held-out query treatment by integrating a classifier-free
-guided ODE flow (Dopri5, step 0.01, guidance 2.0).
+Predicts the post-perturbation single-cell expression distribution for a held-out
+gene perturbation, given in-context demonstrations (control + perturbed
+populations for OTHER treatments in the same cell context) — the in-context
+perturbation setting of Sextro et al. (2026).
+
+Runs the official EMA weights of MapPFN, ported 1:1 from the authors'
+JAX/Equinox reference to PyTorch (verified numerically against the reference:
+velocity fields and 100-step ODE trajectories match to float32 tolerance).
 """
 
-import json
-from pathlib import Path
-
-import spaces  # MUST be imported before torch
-import torch
+import spaces  # MUST be first — ZeroGPU CUDA hijack
 
+import anndata as ad
 import gradio as gr
-
-from mappfn import MapPFN, MMDiT, load_weights, solve_ode
-
+import matplotlib
 import numpy as np
 import pandas as pd
-import plotly.graph_objects as go
+import torch
 
-HERE = Path(__file__).resolve().parent
-
-IN_DIM = 50
-NUM_SHOTS = 8
-NUM_SAMPLES = 200
-GUIDANCE = 2.0
-STEP_SIZE = 0.01
-
-# ---------------------------------------------------------------- model
-
-
-def build_model() -> MapPFN:
-    decoder = MMDiT(
-        embed_dim=256,
-        cond_dim=256,
-        noise_dim=IN_DIM,
-        num_heads=4,
-        num_blocks=8,
-        use_rms_norm=True,
-        num_reg_tokens=8,
-    )
-    return MapPFN(decoder=decoder, cond_dim=256, in_dim=IN_DIM, max_num_contexts=10)
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
 
+from map_pfn_torch import MapPFN, load_ema_weights, solve_ode
 
-MODELS = {
-    "frangieh": load_weights(build_model(), HERE / "weights_ema_frangieh.npz").to("cuda").eval(),
-    "papalexi": load_weights(build_model(), HERE / "weights_ema_papalexi.npz").to("cuda").eval(),
+IN_DIM = 50
+DEFAULT_GUIDANCE = 2.0
+DEFAULT_STEP = 0.01
+
+DATASETS = {
+    "frangieh": {
+        "label": "Frangieh 2021 (melanoma CRISPR, IFN-γ contexts)",
+        "h5ad": "data/frangieh_split_10.h5ad",
+        "npz": "weights/weights_frangieh.npz",
+    },
+    "papalexi": {
+        "label": "Papalexi 2021 (ECCITE-seq, stimulated T cells)",
+        "h5ad": "data/papalexi_split_10.h5ad",
+        "npz": "weights/weights_papalexi.npz",
+    },
 }
-for m in MODELS.values():
-    for p in m.parameters():
-        p.requires_grad_(False)
-
-# ---------------------------------------------------------------- data
-
-DATASETS = {}
-for name, label in (("papalexi", "Papalexi 2021 (CRISPRi leukemia, IFN-γ)"), ("frangieh", "Frangieh 2021 (CRISPRi melanoma)")):
-    bundle = dict(np.load(HERE / f"bundle_{name}.npz", allow_pickle=True))
-    meta = json.loads((HERE / f"bundle_{name}_meta.json").read_text())
-    DATASETS[name] = {"label": label, "bundle": bundle, "meta": meta}
 
+# ---------------------------------------------------------------- load eagerly
 
-def get_contexts(dataset: str) -> list[str]:
-    # only contexts that actually have held-out test perturbations
-    return [c for c in DATASETS[dataset]["meta"]["contexts"] if DATASETS[dataset]["meta"]["test"].get(c)]
+print("[app] loading models...", flush=True)
+MODELS: dict[str, MapPFN] = {}
+for _key, _cfg in DATASETS.items():
+    _m = MapPFN(cond_dim=256, in_dim=IN_DIM, max_num_contexts=10,
+                embed_dim=256, num_heads=4, num_blocks=8, num_reg_tokens=8)
+    load_ema_weights(_m, _cfg["npz"])
+    _m.eval()
+    MODELS[_key] = _m.to("cuda")
+print("[app] models on cuda", flush=True)
 
+print("[app] loading datasets...", flush=True)
+ADATA: dict[str, ad.AnnData] = {k: ad.read_h5ad(cfg["h5ad"]) for k, cfg in DATASETS.items()}
+print("[app] ready", flush=True)
 
-def get_test_treatments(dataset: str, context: str) -> list[str]:
-    return sorted(DATASETS[dataset]["meta"]["test"].get(context, []))
 
+# ---------------------------------------------------------------- helpers
 
-# ---------------------------------------------------------------- metrics
-# Ported from map_pfn/eval/metrics.py (numpy, no jax/ott/scanpy needed).
+def context_choices(name: str) -> list[str]:
+    """List cell contexts that have held-out test perturbations to predict."""
+    obs = ADATA[name].obs
+    test_ctxs = obs[obs["split"] == "test"]["context"].unique().tolist()
+    return sorted(test_ctxs)
 
 
-def mean_r2(x: np.ndarray, y: np.ndarray) -> float:
-    return float(np.corrcoef(np.stack([x.mean(0), y.mean(0)]))[0, 1])
-
-
-def var_r2(x: np.ndarray, y: np.ndarray) -> float:
-    return float(np.corrcoef(np.stack([x.var(0), y.var(0)]))[0, 1])
-
-
-def mean_rmse(x: np.ndarray, y: np.ndarray) -> float:
-    return float(np.sqrt(np.mean((x.mean(0) - y.mean(0)) ** 2)))
-
-
-def var_rmse(x: np.ndarray, y: np.ndarray) -> float:
-    return float(np.sqrt(np.mean((x.var(0) - y.var(0)) ** 2)))
-
-
-def rbf_kernel_sq(x: np.ndarray, y: np.ndarray, gamma: float) -> np.ndarray:
-    sq = ((x[:, None, :] - y[None, :, :]) ** 2).sum(-1)
-    return np.exp(-gamma * sq)
+def treatments_for(name: str, ctx: str) -> dict[str, list[str]]:
+    """Split a context's treatments into test (held-out query) and train (demonstrations)."""
+    sub = ADATA[name].obs[ADATA[name].obs["context"] == ctx]
+    return {
+        "test": sorted(sub[sub["split"] == "test"]["treatment"].unique().tolist()),
+        "train": sorted(sub[sub["split"] == "train"]["treatment"].unique().tolist()),
+    }
 
 
-def multiscale_mmd(x: np.ndarray, y: np.ndarray) -> float:
-    gammas = np.logspace(1, -3, 5)
-    total = 0.0
-    for g in gammas:
-        xx = rbf_kernel_sq(x, x, g).mean()
-        yy = rbf_kernel_sq(y, y, g).mean()
-        xy = rbf_kernel_sq(x, y, g).mean()
-        total += xx + yy - 2 * xy
-    return float(total / len(gammas))
+def _pools(name: str, ctx: str) -> tuple[np.ndarray, dict, dict]:
+    """Return (control cells, {treatment: treated cells}, {treatment: one-hot vector})."""
+    a = ADATA[name]
+    X = np.asarray(a.X, dtype=np.float32)
+    obs = a.obs
+    sub = obs["context"] == ctx
+    ctrl = X[(sub & (obs["split"] == "control")).to_numpy()]
+    treated, onehot = {}, {}
+    for t in obs[sub & (obs["split"] != "control")]["treatment"].unique():
+        m = (sub & (obs["treatment"] == t)).to_numpy()
+        treated[str(t)] = X[m]
+        onehot[str(t)] = np.asarray(a.obsm["treatment"][m][0], dtype=np.float32)
+    return ctrl, treated, onehot
 
 
-def wasserstein_1d(x: np.ndarray, y: np.ndarray) -> float:
-    """Mean per-gene 1-Wasserstein distance (equal-sample empirical quantiles)."""
-    xs = np.sort(x, axis=0)
-    ys = np.sort(y, axis=0)
-    return float(np.abs(xs - ys).mean())
+def _sample_cells(X: np.ndarray, n: int, rng: np.random.Generator) -> np.ndarray:
+    """Sample n rows from X (with replacement if needed), like the reference dataset."""
+    idx = rng.choice(len(X), size=n, replace=len(X) < n)
+    return X[idx]
 
 
 # ---------------------------------------------------------------- inference
 
-
-def sample_cells(x: np.ndarray, n: int, rng: np.random.Generator) -> np.ndarray:
-    if len(x) <= n:
-        # sample with replacement (reference dataset does this too)
-        idx = rng.choice(len(x), n, replace=True)
-    else:
-        idx = rng.choice(len(x), n, replace=False)
-    return x[idx]
-
-
-@spaces.GPU(duration=45)
-def run_predict(
+@spaces.GPU(duration=120)
+def predict(
     dataset: str,
     context: str,
     treatment: str,
-    num_shots: int = NUM_SHOTS,
-    num_cells: int = NUM_SAMPLES,
-    guidance: float = GUIDANCE,
-    step_size: float = STEP_SIZE,
-    seed: int = 42,
-    progress=gr.Progress(track_tqdm=True),
+    num_shots: int = 8,
+    num_cells: int = 200,
+    seed: int = 0,
+    guidance: float = DEFAULT_GUIDANCE,
 ):
-    """Predict the post-perturbation expression distribution for a held-out
-    treatment, given in-context examples of other perturbations.
+    """Predict the post-perturbation expression distribution for a held-out gene.
 
     Args:
-        dataset: Which fine-tuned MapPFN + dataset to use.
-        context: Biological context (cell type / condition).
-        treatment: The query perturbation (gene knockdown), held out at train time.
-        num_shots: Number of in-context perturbation experiments (max 8).
-        num_cells: Cells per condition to sample / predict.
-        guidance: Classifier-free guidance weight.
-        step_size: ODE solver step size.
-        seed: RNG seed for cell sampling and the flow's initial noise.
+        dataset: One of "frangieh" or "papalexi".
+        context: Cell context / experimental condition to predict in.
+        treatment: Held-out (test-split) gene perturbation to predict.
+        num_shots: Number of in-context demonstrations from other treatments (0-8).
+        num_cells: Number of cells per population (query and generated).
+        seed: Random seed for demonstration sampling and flow noise.
+        guidance: Classifier-free guidance weight (paper default 2.0).
 
     Returns:
-        (distribution plot, mean-scatter plot, metrics dataframe, summary md)
+        (profile plot, metrics dataframe, DE-gene table)
     """
-    if treatment not in DATASETS[dataset]["meta"]["test"].get(context, []):
-        raise gr.Error(f"Treatment '{treatment}' is not a held-out test perturbation of {context} in {dataset}.")
-
+    name = "frangieh" if dataset.startswith("frangieh") else "papalexi"
+    if not treatment or not context:
+        raise gr.Error("Pick a cell context and a target gene perturbation first.")
     rng = np.random.default_rng(seed)
-    ds = DATASETS[dataset]
-    bundle, meta = ds["bundle"], ds["meta"]
-    var_names = meta["var_names"]
-
-    num_shots = int(max(0, min(num_shots, NUM_SHOTS, len(meta["train"].get(context, [])))))
-
-    # context demonstrations: (num_shots, num_cells, dim) treated cells + control obs
-    context_treatments = list(rng.choice(meta["train"][context], size=num_shots, replace=False)) if num_shots else []
-
-    obs_list, int_list, treat_list = [], [], []
-    control_cells = bundle[f"control__{context}"]
-    for tr in context_treatments:
-        obs_list.append(sample_cells(control_cells, num_cells, rng))
-        int_list.append(sample_cells(bundle[f"train__{context}__{tr}"], num_cells, rng))
-        treat_list.append(bundle[f"treatvec__train__{context}__{tr}"])
-
-    # query last: the dataset item carries the query's own (control, treated) pair;
-    # the treated part is the ground truth and is dropped by unpack_batch before
-    # the model sees it (int_data_cond = int_data[:-1]).
-    obs_list.append(sample_cells(control_cells, num_cells, rng))
-    int_list.append(sample_cells(bundle[f"test__{context}__{treatment}"], num_cells, rng))
-    treat_list.append(bundle[f"treatvec__test__{context}__{treatment}"])
+    ctrl, treated, onehot = _pools(name, context)
+    splits = treatments_for(name, context)
+    if treatment not in treated:
+        raise gr.Error(f"Treatment {treatment!r} not found in context {context!r}.")
 
-    obs_data = np.stack(obs_list)  # (num_shots+1, num_cells, dim)
-    int_data_all = np.stack(int_list)  # (num_shots+1, num_cells, dim)
-    treatment_arr = np.stack(treat_list)[:, None, :]  # (num_shots+1, 1, dim)
+    # ---- in-context demonstrations: TRAIN treatments of the same context
+    demo_pool = [t for t in splits["train"] if t != treatment]
+    shots = (rng.choice(demo_pool, size=min(num_shots, len(demo_pool)), replace=False).tolist()
+             if num_shots > 0 and demo_pool else [])
 
-    # solve_ode inputs (unpack_batch): obs_data_cond = obs[-1:], int_data_cond = int[:-1]
-    obs_data_cond = obs_data[-1:]
-    int_data_cond = int_data_all[:-1] if int_data_all.shape[0] > 1 else None
+    obs_cond = _sample_cells(ctrl, num_cells, rng)[None]      # (1, n, 50) query control
+    int_cond = (np.stack([_sample_cells(treated[t], num_cells, rng) for t in shots])
+                if shots else None)                            # (S, n, 50) demonstrations
+    treat_vec = np.stack([onehot[t][None] for t in [*shots, treatment]])
 
-    model = MODELS[dataset]
-    gen = torch.Generator(device="cuda").manual_seed(int(seed))
+    to = torch.from_numpy(obs_cond).cuda()
+    ti = torch.from_numpy(int_cond).cuda() if int_cond is not None else None
+    tt = torch.from_numpy(treat_vec).cuda()
 
     with torch.no_grad():
         pred = solve_ode(
-            model,
-            noise_shape=(num_cells, IN_DIM),  # per-item noise; solve_ode runs a single item
-            obs_data=torch.from_numpy(obs_data_cond).to("cuda"),
-            int_data=torch.from_numpy(int_data_cond).to("cuda") if int_data_cond is not None else None,
-            treatment=torch.from_numpy(treatment_arr).to("cuda"),
-            guidance=float(guidance),
-            step_size=float(step_size),
-            generator=gen,
-        )
-    pred_np = pred.float().cpu().numpy()  # (num_cells, dim) — solve_ode runs one un-batched item
-
-    # ground truth for the held-out query treatment
-    true_np = int_data_all[-1]
-    ctrl_np = sample_cells(control_cells, num_cells, rng)
-
-    # metrics vs ground truth
-    metrics = {
-        "Mean R²": mean_r2(true_np, pred_np),
-        "Var R²": var_r2(true_np, pred_np),
-        "Mean RMSE": mean_rmse(true_np, pred_np),
-        "Var RMSE": var_rmse(true_np, pred_np),
-        "MMD (multiscale)": multiscale_mmd(true_np, pred_np),
-        "Wasserstein (1d)": wasserstein_1d(true_np, pred_np),
-    }
-
-    # ---- plot 1: distribution of top DE genes (violin-style box plots)
-    lfc = true_np.mean(0) - ctrl_np.mean(0)
-    top_genes = np.argsort(-np.abs(lfc))[:8]
-    fig = go.Figure()
-    for g in top_genes:
-        gene = var_names[g]
-        fig.add_trace(
-            go.Box(y=ctrl_np[:, g], name=f"{gene}<br>control", marker_color="#94a3b8", showlegend=False)
-        )
-        fig.add_trace(
-            go.Box(y=true_np[:, g], name=f"{gene}<br>measured", marker_color="#0f766e", showlegend=False)
-        )
-        fig.add_trace(
-            go.Box(y=pred_np[:, g], name=f"{gene}<br>MapPFN", marker_color="#e11d48", showlegend=False)
-        )
-    fig.update_layout(
-        title=f"Expression of top-DE genes — {treatment} in {context} ({ds['label']})",
-        yaxis_title="log1p normalized expression",
-        boxmode="group",
-        template="plotly_white",
-        height=420,
-        margin=dict(l=10, r=10, t=60, b=10),
-    )
-
-    # ---- plot 2: per-gene mean predicted vs measured
-    fig2 = go.Figure()
-    fig2.add_trace(
-        go.Scatter(
-            x=true_np.mean(0),
-            y=pred_np.mean(0),
-            mode="markers",
-            text=var_names,
-            marker=dict(size=7, color=lfc, colorscale="RdBu", showscale=True, colorbar=dict(title="Δ (true − control)")),
-            name="genes",
-        )
-    )
-    lims = [
-        min(true_np.mean(0).min(), pred_np.mean(0).min()) - 0.1,
-        max(true_np.mean(0).max(), pred_np.mean(0).max()) + 0.1,
-    ]
-    fig2.add_trace(go.Scatter(x=lims, y=lims, mode="lines", line=dict(dash="dash", color="#94a3b8"), name="y = x"))
-    fig2.update_layout(
-        title=f"Per-gene mean expression — predicted vs measured (R² = {metrics['Mean R²']:.3f})",
-        xaxis_title="measured mean (log1p)",
-        yaxis_title="MapPFN predicted mean (log1p)",
-        template="plotly_white",
-        height=420,
-        margin=dict(l=10, r=10, t=60, b=10),
+            MODELS[name], (num_cells, IN_DIM), to, ti, tt,
+            guidance=guidance, step_size=DEFAULT_STEP,
+            generator=torch.Generator(device="cuda").manual_seed(seed),
+        ).cpu().numpy()
+
+    gt = _sample_cells(treated[treatment], num_cells, rng)      # held-out truth
+    obs_query = obs_cond[0]
+
+    genes = [str(g) for g in ADATA[name].var_names]
+    pred_mean, pred_std = pred.mean(0), pred.std(0)
+    gt_mean, gt_std = gt.mean(0), gt.std(0)
+    ctrl_mean = obs_query.mean(0)
+
+    delta_true = gt_mean - ctrl_mean
+    delta_pred = pred_mean - ctrl_mean
+    order = np.argsort(-np.abs(delta_true))
+
+    # ---- plot
+    fig, axes = plt.subplots(1, 2, figsize=(13, 4.6))
+    xs = np.arange(len(genes))
+    ax = axes[0]
+    ax.plot(xs, ctrl_mean[order], color="#9aa0a6", lw=1.4, label="control (observed)")
+    ax.plot(xs, gt_mean[order], color="#1f77b4", lw=1.6, label=f"{treatment} observed (held-out)")
+    ax.plot(xs, pred_mean[order], color="#d62728", lw=1.6, ls="--", label=f"{treatment} MapPFN predicted")
+    ax.set_xlabel("genes (sorted by observed effect)")
+    ax.set_ylabel("expression (50-d PCA space)")
+    ax.set_title(f"Mean expression profile — {context}")
+    ax.legend(fontsize=8)
+    ax = axes[1]
+    ax.scatter(delta_true[order], delta_pred[order], s=20, color="#d62728", alpha=0.75)
+    lim = max(np.abs(delta_true).max(), np.abs(delta_pred).max(), 1e-6) * 1.1
+    ax.plot([-lim, lim], [-lim, lim], color="#9aa0a6", lw=1)
+    ax.set_xlabel("observed Δ (treated − control)")
+    ax.set_ylabel("predicted Δ")
+    ax.set_title("Perturbation effect: predicted vs observed")
+    fig.suptitle(
+        f"MapPFN — {DATASETS[name]['label']} | context: {context} | target: {treatment} | "
+        f"{len(shots)} in-context demonstration{'s' if len(shots) != 1 else ''}",
+        fontsize=10,
     )
+    fig.tight_layout(rect=(0, 0, 1, 0.94))
+
+    # ---- metrics vs held-out ground truth
+    tgt_idx = int(np.argmax(onehot[treatment]))
+    top5 = np.argsort(-delta_pred)[:5]
+    metrics = pd.DataFrame({
+        "metric": [
+            "mean RMSE", "variance RMSE", "mean R²", "variance R²",
+            "Δ-effect Pearson r", "target gene in top-5 predicted Δ",
+        ],
+        "value": [
+            f"{np.sqrt(np.mean((gt_mean - pred_mean) ** 2)):.3f}",
+            f"{np.sqrt(np.mean((gt_std - pred_std) ** 2)):.3f}",
+            f"{np.corrcoef(gt_mean, pred_mean)[0, 1]:.3f}",
+            f"{np.corrcoef(gt_std, pred_std)[0, 1]:.3f}",
+            f"{np.corrcoef(delta_true, delta_pred)[0, 1]:.3f}",
+            "yes" if tgt_idx in top5 else "no",
+        ],
+    })
+
+    # ---- DE table
+    top = np.argsort(-np.abs(delta_pred))[:15]
+    de = pd.DataFrame({
+        "gene": [genes[i] for i in top],
+        "predicted Δ": delta_pred[top].round(3),
+        "observed Δ": delta_true[top].round(3),
+        "pred mean": pred_mean[top].round(3),
+        "observed mean": gt_mean[top].round(3),
+        "pred std": pred_std[top].round(3),
+        "observed std": gt_std[top].round(3),
+    })
+    return fig, metrics, de
+
+
+# ---------------------------------------------------------------- UI wiring
+
+def _on_dataset(name: str) -> dict:
+    """Update context and treatment dropdowns when the dataset changes."""
+    ctxs = context_choices(name)
+    test = treatments_for(name, ctxs[0])["test"]
+    return {
+        ctx_dd: gr.update(choices=ctxs, value=ctxs[0]),
+        treat_dd: gr.update(choices=test, value=test[0] if test else None),
+    }
 
-    metrics_df = pd.DataFrame({"Metric": list(metrics.keys()), "Value": [round(v, 4) for v in metrics.values()]})
-    summary = (
-        f"**{treatment}** in context *{context}* — {num_shots} in-context example"
-        f"{'s' if num_shots != 1 else ''} ({num_cells} cells each), guidance {guidance}, "
-        f"ODE step {step_size}, seed {seed}. The query perturbation was **held out** from training."
-    )
-    return fig, fig2, metrics_df, summary
-
-
-def refresh_context(dataset: str, context: str):
-    ctxs = get_contexts(dataset)
-    val = context if context in ctxs else ctxs[0]
-    return gr.update(choices=ctxs, value=val)
-
-
-def refresh_treatments(dataset: str, context: str, treatment: str):
-    tests = get_test_treatments(dataset, context)
-    val = treatment if treatment in tests else (tests[0] if tests else None)
-    return gr.update(choices=tests, value=val)
 
+def _on_context(name: str, ctx: str) -> dict:
+    """Update the treatment dropdown when the context changes."""
+    test = treatments_for(name, ctx)["test"]
+    return {treat_dd: gr.update(choices=test, value=test[0] if test else None)}
 
-CSS = """
-#col-container { max-width: 1200px; margin: 0 auto; }
-.dark .gradio-container { color: var(--body-text-color); }
-"""
 
-with gr.Blocks() as demo:
+with gr.Blocks(theme=gr.themes.Citrus(), title="MapPFN") as demo:
     gr.Markdown(
-        f"""
-        # 🧬 MapPFN: Learning Causal Perturbation Maps in Context
-        Predict the **post-perturbation single-cell expression distribution** for an unseen gene knockdown —
-        in-context, from a handful of example perturbation experiments. PyTorch port of the
-        [NeurIPS 2026 paper](https://arxiv.org/abs/2601.21092) · [code](https://github.com/marvinsxtr/MapPFN) ·
-        [weights](https://huggingface.co/marvinsxtr/MapPFN) (MIT). Zero-shot on held-out test perturbations.
         """
+        # MapPFN: Learning Causal Perturbation Maps in Context
+        **Sextro et al. (2026)** — [paper](https://huggingface.co/papers/2601.21092) ·
+        [code](https://github.com/marvinsxtr/MapPFN) · [weights](https://huggingface.co/marvinsxtr/MapPFN)
+
+        MapPFN is a flow-matching in-context model for genetic perturbation prediction:
+        given a set of *demonstration* experiments (control + perturbed populations for
+        **other** genes in the same cell context), it predicts the full
+        post-perturbation expression distribution for a **held-out** target gene — no
+        retraining per perturbation. Pick a dataset, cell context and held-out target
+        below; the prediction is scored against the real held-out measurements.
+        """
+    )
+    with gr.Row():
+        with gr.Column():
+            dataset_dd = gr.Dropdown(
+                choices=[(cfg["label"], k) for k, cfg in DATASETS.items()],
+                value="frangieh", label="Dataset (each uses its fine-tuned checkpoint)",
+            )
+            ctx_dd = gr.Dropdown(choices=context_choices("frangieh"), value="IFNγ",
+                                  label="Cell context")
+            treat_dd = gr.Dropdown(choices=[], label="Target gene perturbation (held-out test split)")
+            with gr.Accordion("Inference settings", open=False):
+                num_shots = gr.Slider(0, 8, value=8, step=1,
+                                      label="In-context demonstrations (other treatments)")
+                num_cells = gr.Slider(32, 200, value=200, step=8, label="Cells per population")
+                seed = gr.Number(value=0, precision=0, label="Seed")
+                guidance = gr.Slider(0.0, 4.0, value=DEFAULT_GUIDANCE, step=0.1,
+                                     label="Classifier-free guidance")
+            run = gr.Button("Predict perturbation response", variant="primary")
+        with gr.Column():
+            plot = gr.Plot(label="Predicted vs observed expression profile")
+    metrics_df = gr.Dataframe(label="Prediction quality (vs held-out ground truth)")
+    de_df = gr.Dataframe(label="Top predicted differentially expressed genes")
+    gr.Markdown(
+        "*Expression values live in the 50-dimensional PCA space of the original publication. "
+        "Metrics: R² between per-gene means/variances of predicted and observed treated cells; "
+        "Pearson r between predicted and observed perturbation effects (Δ vs control); and "
+        "whether the knocked-out gene itself ranks among the top-5 predicted changes.*"
     )
-    with gr.Column(elem_id="col-container"):
-        with gr.Row():
-            with gr.Column(scale=1):
-                dataset_dd = gr.Dropdown(
-                    choices=[("Papalexi 2021 (CRISPRi leukemia, IFN-γ)", "papalexi"), ("Frangieh 2021 (CRISPRi melanoma)", "frangieh")],
-                    value="papalexi",
-                    label="Dataset / fine-tuned checkpoint",
-                )
-                context_dd = gr.Dropdown(
-                    choices=get_contexts("papalexi"), value="stimulated", label="Context (cell condition)", allow_custom_value=True
-                )
-                treatment_dd = gr.Dropdown(
-                    choices=get_test_treatments("papalexi", "stimulated"),
-                    value="STAT1",
-                    label="Query perturbation (held out at train time)",
-                    allow_custom_value=True,
-                )
-                run_btn = gr.Button("Predict perturbation response", variant="primary")
-                with gr.Accordion("Advanced settings", open=False):
-                    num_shots = gr.Slider(0, 8, value=NUM_SHOTS, step=1, label="In-context examples (num_shots)")
-                    num_cells = gr.Slider(16, 200, value=NUM_SAMPLES, step=8, label="Cells per condition")
-                    guidance = gr.Slider(1.0, 4.0, value=GUIDANCE, step=0.1, label="Guidance")
-                    step_size = gr.Slider(0.005, 0.05, value=STEP_SIZE, step=0.005, label="ODE step size")
-                    seed = gr.Number(value=42, precision=0, label="Seed")
-            with gr.Column(scale=2):
-                summary_md = gr.Markdown()
-                dist_plot = gr.Plot(label="Expression distributions (top-DE genes)")
-        with gr.Row():
-            scatter_plot = gr.Plot(label="Predicted vs measured per-gene mean")
-            metrics_df = gr.Dataframe(headers=["Metric", "Value"], label="Distribution metrics (vs ground truth)", interactive=False)
-
-        gr.Examples(
-            examples=[
-                ["papalexi", "stimulated", "STAT1"],
-                ["papalexi", "stimulated", "IFNGR1"],
-                ["papalexi", "stimulated", "JAK2"],
-                ["frangieh", "IFNγ", "B2M"],
-                ["frangieh", "IFNγ", "TAPBP"],
-                ["frangieh", "IFNγ", "STAT3"],
-            ],
-            inputs=[dataset_dd, context_dd, treatment_dd],
-            fn=run_predict,
-            outputs=[dist_plot, scatter_plot, metrics_df, summary_md],
-            cache_examples=True,
-            cache_mode="lazy",
-        )
-
-    dataset_dd.change(refresh_context, inputs=[dataset_dd, context_dd], outputs=context_dd)
-    context_dd.change(
-        refresh_treatments, inputs=[dataset_dd, context_dd, treatment_dd], outputs=treatment_dd
+    gr.Examples(
+        examples=[
+            ["frangieh", "IFNγ", "STAT3"],
+            ["frangieh", "IFNγ", "B2M"],
+            ["frangieh", "IFNγ", "TAPBP"],
+            ["frangieh", "IFNγ", "AHNAK"],
+            ["papalexi", "stimulated", "STAT1"],
+            ["papalexi", "stimulated", "JAK2"],
+            ["papalexi", "stimulated", "IFNGR1"],
+        ],
+        inputs=[dataset_dd, ctx_dd, treat_dd],
+        fn=predict,
+        outputs=[plot, metrics_df, de_df],
+        cache_examples=True,
+        cache_mode="lazy",
+        label="Example queries (the paper's held-out test perturbations)",
     )
-    run_btn.click(
-        run_predict,
-        inputs=[dataset_dd, context_dd, treatment_dd, num_shots, num_cells, guidance, step_size, seed],
-        outputs=[dist_plot, scatter_plot, metrics_df, summary_md],
-        api_name="predict",
+    dataset_dd.change(_on_dataset, inputs=[dataset_dd], outputs=[ctx_dd, treat_dd])
+    ctx_dd.change(_on_context, inputs=[dataset_dd, ctx_dd], outputs=[treat_dd])
+    run.click(
+        predict,
+        inputs=[dataset_dd, ctx_dd, treat_dd, num_shots, num_cells, seed, guidance],
+        outputs=[plot, metrics_df, de_df],
     )
 
-demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)
+if __name__ == "__main__":
+    demo.launch(mcp_server=True)

data/frangieh_split_10.h5ad

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:103c8c4b1e44b36c2cee96ea7753d43bbc2e456a81c216cac2cd9d5ea01bcffe
+size 51143641

data/papalexi_split_10.h5ad

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e42c3becc855e2ca4b894c5f47b3c9c495ac46d68ab28ab2180b335914af821d
+size 9829716

map_pfn_torch.py

@@ -0,0 +1,467 @@
+"""PyTorch port of MapPFN (Sextro et al., 2026) — JAX/Equinox reference:
+
+    marvinsxtr/MapPFN  |  map_pfn/models/{map_pfn,mmdit}.py, map_pfn/models/utils.py
+
+Architecture: MapPFN(cond_dim=256, in_dim=50, max_num_contexts=10) wrapping an
+MMDiT(embed_dim=256, num_heads=4, num_blocks=8, rms_norm, 8 register tokens).
+Sampling: flow-matching velocity field integrated t=0->1 with Dormand-Prince 5(4)
+at a constant step of 0.01 (100 steps), classifier-free guidance weight 2.0.
+
+Port notes (verified against reference sources and library internals):
+- eqx.nn.Linear stores weight as (out, in) — identical to torch.nn.Linear, so
+  exported weights copy directly with NO transpose.
+- eqx.nn.LayerNorm: population statistics over the whole array, eps=1e-5 default.
+- eqx.nn.RMSNorm(eps=1e-8) on per-head Q/K, WITH weight and bias.
+- jax.nn.gelu default = tanh approximation.
+- flash_attention: pads seq to a multiple of 64 AFTER the RMS norms and computes
+  attention in bfloat16 with NO mask. Zero-padded keys get score 0 (NOT -inf), so
+  they dilute the softmax denominator (values are zero so they add no signal).
+  Replicated exactly with F.scaled_dot_product_attention on zero-padded bf16
+  tensors.
+- solve_ode == diffrax diffeqsolve(Dopri5, ConstantStepSize, t0=0, t1=1, dt0=0.01,
+  SaveAt(t1=True)): exactly 100 fixed DP5 steps. diffrax works in k-space
+  (k = f * dt): stage states are y0 + dt * sum(a_ij f_j) and the update is
+  y1 = y0 + dt * sum(b_i f_i). Dopri5 has 7 stages with c = [0, 1/5, 3/10, 4/5,
+  8/9, 1, 1] and is FSAL: stage 6 evaluates at (t1, y1) (the tableau is SSAL), and
+  that evaluation is reused as stage 0 of the next step. Stage times are
+  t0 + c*dt, except c == 1 uses the step endpoint exactly (no float error).
+- CFG mixes conditional and unconditional (null_cond) velocity fields:
+  guidance * v_cond + (1 - guidance) * v_uncond.
+- Weights come from the EMA tree of the authors' Equinox checkpoints.
+"""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def gelu_tanh(x: torch.Tensor) -> torch.Tensor:
+    return F.gelu(x, approximate="tanh")
+
+
+class MLP(nn.Module):
+    """Reference: map_pfn.models.utils.MLP — Linear -> gelu(tanh) -> Linear."""
+
+    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
+        super().__init__()
+        self.linear_1 = nn.Linear(input_dim, hidden_dim)
+        self.linear_2 = nn.Linear(hidden_dim, output_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.linear_2(gelu_tanh(self.linear_1(x)))
+
+
+class SinusoidalPosEmb(nn.Module):
+    """Reference: map_pfn.models.utils.SinusoidalPosEmb."""
+
+    def __init__(self, dim: int):
+        super().__init__()
+        half_dim = dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        self.register_buffer("emb", torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:  # x: (seq,)
+        emb = x * self.emb
+        return torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+
+
+class MMDiTBlock(nn.Module):
+    """Reference: map_pfn.models.mmdit.MMDiTBlock (FiLM-conditioned joint attention)."""
+
+    def __init__(self, embed_dim: int, num_heads: int, head_dim: int):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+
+        self.ln_qkv_n = nn.LayerNorm(embed_dim, eps=1e-5)
+        self.ln_qkv_c = nn.LayerNorm(embed_dim, eps=1e-5)
+        self.ln_qkv_t = nn.LayerNorm(embed_dim, eps=1e-5)
+        self.ln_ff_n = nn.LayerNorm(embed_dim, eps=1e-5)
+        self.ln_ff_c = nn.LayerNorm(embed_dim, eps=1e-5)
+        self.ln_ff_t = nn.LayerNorm(embed_dim, eps=1e-5)
+
+        self.qkv_n = nn.Linear(embed_dim, embed_dim * 3, bias=False)
+        self.qkv_c = nn.Linear(embed_dim, embed_dim * 3, bias=False)
+        self.qkv_t = nn.Linear(embed_dim, embed_dim * 3, bias=False)
+
+        self.o_n = nn.Linear(embed_dim, embed_dim)
+        self.o_c = nn.Linear(embed_dim, embed_dim)
+        self.o_t = nn.Linear(embed_dim, embed_dim)
+
+        self.film_an = nn.Linear(embed_dim, embed_dim * 2)
+        self.film_ac = nn.Linear(embed_dim, embed_dim * 2)
+        self.film_at = nn.Linear(embed_dim, embed_dim * 2)
+        self.film_fn = nn.Linear(embed_dim, embed_dim * 2)
+        self.film_fc = nn.Linear(embed_dim, embed_dim * 2)
+        self.film_ft = nn.Linear(embed_dim, embed_dim * 2)
+
+        # eqx RMSNorm(shape=head_dim, eps=1e-8) with weight+bias
+        self.rms_q_w = nn.Parameter(torch.ones(head_dim))
+        self.rms_q_b = nn.Parameter(torch.zeros(head_dim))
+        self.rms_k_w = nn.Parameter(torch.ones(head_dim))
+        self.rms_k_b = nn.Parameter(torch.zeros(head_dim))
+
+        self.ff_n = MLP(embed_dim, embed_dim * 4, embed_dim)
+        self.ff_c = MLP(embed_dim, embed_dim * 4, embed_dim)
+        self.ff_t = MLP(embed_dim, embed_dim * 4, embed_dim)
+
+    @staticmethod
+    def _rms_norm(x: torch.Tensor, w: torch.Tensor, b: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
+        # x: (..., head_dim); eqx RMSNorm: x / sqrt(mean(x^2) + eps) * w + b, fp32 stats
+        x32 = x.float()
+        inv = torch.rsqrt(x32.pow(2).mean(dim=-1, keepdim=True) + eps)
+        return (inv * x32) * w + b
+
+    def _heads(self, x: torch.Tensor) -> torch.Tensor:  # (seq, dim) -> (seq, H, D)
+        s = x.shape[0]
+        return x.reshape(s, self.num_heads, self.head_dim)
+
+    def forward(
+        self,
+        noise: torch.Tensor,      # (seq_n, E)
+        cond: torch.Tensor,       # (seq_c, E)
+        treatment: torch.Tensor,  # (seq_t, E) — all treatment tokens
+        t_emb: torch.Tensor,      # (1, E)
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        seq_n, seq_c, seq_t = noise.shape[0], cond.shape[0], treatment.shape[0]
+
+        t_n = t_emb.expand(seq_n, -1)
+        t_c = t_emb.expand(seq_c, -1)
+        t_t = t_emb.expand(seq_t, -1)
+
+        gamma_n, beta_n = self.film_an(t_n).chunk(2, dim=-1)
+        gamma_c, beta_c = self.film_ac(t_c).chunk(2, dim=-1)
+        gamma_t, beta_t = self.film_at(t_t).chunk(2, dim=-1)
+
+        n = self.ln_qkv_n(noise) * (1 + gamma_n) + beta_n
+        c = self.ln_qkv_c(cond) * (1 + gamma_c) + beta_c
+        t = self.ln_qkv_t(treatment) * (1 + gamma_t) + beta_t
+
+        qkv_n = self.qkv_n(n)
+        qkv_c = self.qkv_c(c)
+        qkv_t = self.qkv_t(t)
+
+        q_n, k_n, v_n = qkv_n.chunk(3, dim=-1)
+        q_c, k_c, v_c = qkv_c.chunk(3, dim=-1)
+        q_t, k_t, v_t = qkv_t.chunk(3, dim=-1)
+
+        q = torch.cat([self._heads(q_n), self._heads(q_c), self._heads(q_t)], dim=0)
+        k = torch.cat([self._heads(k_n), self._heads(k_c), self._heads(k_t)], dim=0)
+        v = torch.cat([self._heads(v_n), self._heads(v_c), self._heads(v_t)], dim=0)
+
+        q = self._rms_norm(q, self.rms_q_w, self.rms_q_b)
+        k = self._rms_norm(k, self.rms_k_w, self.rms_k_b)
+
+        attn_out = flash_attention(q, k, v)  # (seq_total, H, D)
+
+        attn_out = attn_out.reshape(seq_n + seq_c + seq_t, self.embed_dim)
+        a_n, a_c, a_t = attn_out.split([seq_n, seq_c, seq_t], dim=0)
+
+        noise = noise + self.o_n(a_n)
+        cond = cond + self.o_c(a_c)
+        treatment = treatment + self.o_t(a_t)
+
+        gamma_n2, beta_n2 = self.film_fn(t_n).chunk(2, dim=-1)
+        gamma_c2, beta_c2 = self.film_fc(t_c).chunk(2, dim=-1)
+        gamma_t2, beta_t2 = self.film_ft(t_t).chunk(2, dim=-1)
+
+        noise = noise + self.ff_n(self.ln_ff_n(noise) * (1 + gamma_n2) + beta_n2)
+        cond = cond + self.ff_c(self.ln_ff_c(cond) * (1 + gamma_c2) + beta_c2)
+        treatment = treatment + self.ff_t(self.ln_ff_t(treatment) * (1 + gamma_t2) + beta_t2)
+
+        return noise, cond, treatment
+
+
+def flash_attention(
+    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, pad_multiple: int = 64
+) -> torch.Tensor:
+    """Reference: map_pfn.models.utils.flash_attention.
+
+    Pads seq_len to a multiple of 64 (post-RMS-norm tensors) and runs attention
+    in bfloat16 with no mask. Zero-padded keys get score 0 and therefore dilute
+    the softmax denominator (exactly as the JAX reference does); padded values
+    are zero so they contribute no signal. Padded query rows are discarded.
+    """
+    seq_len = q.shape[0]
+    pad_len = (pad_multiple - seq_len % pad_multiple) % pad_multiple
+    if pad_len > 0:
+        # (D_left, D_right, H_top, H_bottom, seq_front, seq_back)
+        q = F.pad(q, (0, 0, 0, 0, 0, pad_len))
+        k = F.pad(k, (0, 0, 0, 0, 0, pad_len))
+        v = F.pad(v, (0, 0, 0, 0, 0, pad_len))
+
+    # (seq, H, D) -> (1, H, seq, D)
+    q = q.permute(1, 0, 2).unsqueeze(0).to(torch.bfloat16)
+    k = k.permute(1, 0, 2).unsqueeze(0).to(torch.bfloat16)
+    v = v.permute(1, 0, 2).unsqueeze(0).to(torch.bfloat16)
+    out = F.scaled_dot_product_attention(q, k, v)
+    out = out.squeeze(0).permute(1, 0, 2).to(torch.float32)
+    return out[:seq_len]
+
+
+class MMDiT(nn.Module):
+    """Reference: map_pfn.models.mmdit.MMDiT."""
+
+    def __init__(self, embed_dim: int, cond_dim: int, noise_dim: int, num_heads: int,
+                 num_blocks: int, num_reg_tokens: int):
+        super().__init__()
+        self.num_reg_tokens = num_reg_tokens
+        self.noise_dim = noise_dim
+        self.embed_dim = embed_dim
+
+        self.noise_in = nn.Linear(noise_dim, embed_dim)
+        self.noise_out = nn.Linear(embed_dim, noise_dim)
+        self.cond_in = nn.Linear(cond_dim, embed_dim)
+        self.treatment_in = nn.Linear(cond_dim, embed_dim)
+
+        self.modality_emb_noise = nn.Parameter(torch.zeros(embed_dim))
+        self.modality_emb_cond = nn.Parameter(torch.zeros(embed_dim))
+        self.modality_emb_treatment = nn.Parameter(torch.zeros(embed_dim))
+
+        self.reg_tok = nn.Parameter(torch.zeros(num_reg_tokens, noise_dim))
+
+        self.time_emb = SinusoidalPosEmb(embed_dim)
+        self.time_mlp = MLP(embed_dim, embed_dim * 4, embed_dim)
+
+        self.blocks = nn.ModuleList(
+            [MMDiTBlock(embed_dim, num_heads, embed_dim // num_heads) for _ in range(num_blocks)]
+        )
+
+    def forward(self, noise: torch.Tensor, cond: torch.Tensor, treatment: torch.Tensor, t: torch.Tensor):
+        # noise: (seq_n, 50); cond: (seq_c, 256); treatment: (seq_t, 256); t: scalar tensor
+        noise = torch.cat([self.reg_tok, noise], dim=0)
+        noise = self.noise_in(noise) + self.modality_emb_noise
+        cond = self.cond_in(cond) + self.modality_emb_cond
+        treatment = self.treatment_in(treatment) + self.modality_emb_treatment
+
+        t = t.reshape(1)
+        t_pos = self.time_emb(t)          # (1, E)
+        t_emb = self.time_mlp(t_pos)      # (1, E)
+
+        for block in self.blocks:
+            noise, cond, treatment = block(noise, cond, treatment, t_emb)
+
+        noise = noise[self.num_reg_tokens:]
+        return self.noise_out(noise)
+
+
+class MapPFN(nn.Module):
+    """Reference: map_pfn.models.map_pfn.MapPFN."""
+
+    def __init__(self, cond_dim: int = 256, in_dim: int = 50, max_num_contexts: int = 10,
+                 embed_dim: int = 256, num_heads: int = 4, num_blocks: int = 8,
+                 num_reg_tokens: int = 8):
+        super().__init__()
+        self.max_num_contexts = max_num_contexts
+
+        self.decoder = MMDiT(embed_dim, cond_dim, in_dim, num_heads, num_blocks, num_reg_tokens)
+
+        self.null_cond = nn.Parameter(torch.zeros(cond_dim))
+        self.cond_in = nn.Linear(in_dim, cond_dim)
+        self.treatment_in = nn.Linear(in_dim, cond_dim)
+
+        self.obs_data_emb = nn.Parameter(torch.zeros(cond_dim))
+        self.int_data_emb = nn.Parameter(torch.zeros(cond_dim))
+        self.treatment_emb = nn.Parameter(torch.zeros(cond_dim))
+        self.context_emb = nn.Parameter(torch.zeros(max_num_contexts, cond_dim))
+        self.query_emb = nn.Parameter(torch.zeros(cond_dim))
+
+    def assemble_context(
+        self,
+        obs_data: torch.Tensor,        # (1, num_rows, C)
+        int_data: torch.Tensor | None,  # (num_ctx-1, num_rows, C) or None
+        treatment: torch.Tensor,        # (num_ctx, num_treatments, C)
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        num_contexts = treatment.shape[0]
+        if int_data is None:
+            int_data = self.null_cond.view(1, 1, 1).expand(num_contexts - 1, *treatment.shape[1:])
+
+        obs_data = obs_data + self.obs_data_emb.view(1, 1, -1)
+        int_data = int_data + self.int_data_emb.view(1, 1, -1)
+        treatment = treatment + self.treatment_emb.view(1, 1, -1)
+
+        obs_data = obs_data + self.context_emb[num_contexts - 1: num_contexts].view(1, 1, -1)
+        int_data = int_data + self.context_emb[: num_contexts - 1].unsqueeze(1)
+        treatment = treatment + self.context_emb[:num_contexts].unsqueeze(1)
+
+        obs_data = obs_data.clone()
+        obs_data[-1] = obs_data[-1] + self.query_emb
+        treatment = treatment.clone()
+        treatment[-1] = treatment[-1] + self.query_emb
+
+        obs_seq = obs_data.reshape(-1, obs_data.shape[-1])
+        int_seq = int_data.reshape(-1, int_data.shape[-1])
+        treat_seq = treatment.reshape(-1, treatment.shape[-1])
+        cond_seq = torch.cat([obs_seq, int_seq], dim=0)
+        return cond_seq, treat_seq
+
+    def forward(
+        self,
+        noise: torch.Tensor,            # (seq, 50)
+        obs_data: torch.Tensor,         # (1, num_rows, 50)
+        int_data: torch.Tensor | None,  # (num_ctx-1, num_rows, 50) or None
+        treatment: torch.Tensor,        # (num_ctx, num_treatments, 50)
+        t: torch.Tensor,                # scalar
+        drop_cond: bool = False,
+    ) -> torch.Tensor:
+        obs_data = self.cond_in(obs_data)
+        int_data = self.cond_in(int_data) if int_data is not None else None
+        treatment = self.treatment_in(treatment)
+
+        cond_seq, treat_seq = self.assemble_context(obs_data, int_data, treatment)
+
+        if drop_cond:
+            cond_seq = self.null_cond.view(1, -1).expand_as(cond_seq)
+            treat_seq = self.null_cond.view(1, -1).expand_as(treat_seq)
+
+        return self.decoder(noise, cond_seq, treat_seq, t)
+
+
+def load_ema_weights(model: MapPFN, npz_path: str) -> MapPFN:
+    """Load EMA weights exported from the authors' Equinox checkpoint.
+
+    eqx.nn.Linear weights are stored (out, in) exactly like torch, so every
+    weight copies directly. filter_vmap-stacked block parameters have a leading
+    block axis in the npz; index [i] recovers block i's parameter.
+    """
+    w = np.load(npz_path)
+    t = lambda k: torch.from_numpy(np.ascontiguousarray(w[k]))
+
+    p = dict(model.named_parameters())
+
+    def set_par(name: str, key: str):
+        arr = t(key)
+        assert tuple(arr.shape) == tuple(p[name].shape), f"{name}: {arr.shape} vs {p[name].shape}"
+        p[name].data.copy_(arr)
+
+    # ---- top level (MapPFN) ----
+    for k in ["null_cond", "obs_data_emb", "int_data_emb", "treatment_emb",
+              "context_emb", "query_emb"]:
+        set_par(k, k)
+    for lyr in ["cond_in", "treatment_in"]:
+        set_par(f"{lyr}.weight", f"{lyr}.weight")
+        set_par(f"{lyr}.bias", f"{lyr}.bias")
+
+    # ---- decoder ----
+    for lyr in ["noise_in", "noise_out", "cond_in", "treatment_in"]:
+        set_par(f"decoder.{lyr}.weight", f"decoder.{lyr}.weight")
+        set_par(f"decoder.{lyr}.bias", f"decoder.{lyr}.bias")
+    for k in ["reg_tok", "modality_emb_noise", "modality_emb_cond", "modality_emb_treatment"]:
+        set_par(f"decoder.{k}", f"decoder.{k}")
+    for lyr in ["linear_1", "linear_2"]:
+        set_par(f"decoder.time_mlp.{lyr}.weight", f"decoder.time_mlp.{lyr}.weight")
+        set_par(f"decoder.time_mlp.{lyr}.bias", f"decoder.time_mlp.{lyr}.bias")
+    # time_emb.emb is a deterministic buffer
+    buf = t("decoder.time_emb.emb")
+    model.decoder.time_emb.emb.data.copy_(buf)
+
+    # ---- vmapped blocks (block index is the leading dim in the npz) ----
+    n_blocks = len(model.decoder.blocks)
+    for i in range(n_blocks):
+        b = f"decoder.blocks.{i}."
+        for src in ["ln_qkv_n", "ln_qkv_c", "ln_qkv_t", "ln_ff_n", "ln_ff_c", "ln_ff_t"]:
+            p[f"{b}{src}.weight"].data.copy_(t(f"decoder.blocks.{src}.weight")[i])
+            p[f"{b}{src}.bias"].data.copy_(t(f"decoder.blocks.{src}.bias")[i])
+        for src in ["qkv_n", "qkv_c", "qkv_t"]:
+            p[f"{b}{src}.weight"].data.copy_(t(f"decoder.blocks.{src}.weight")[i])
+        for src in ["o_n", "o_c", "o_t", "film_an", "film_ac", "film_at",
+                    "film_fn", "film_fc", "film_ft"]:
+            p[f"{b}{src}.weight"].data.copy_(t(f"decoder.blocks.{src}.weight")[i])
+            p[f"{b}{src}.bias"].data.copy_(t(f"decoder.blocks.{src}.bias")[i])
+        for src in ["rms_q", "rms_k"]:
+            p[f"{b}{src}_w"].data.copy_(t(f"decoder.blocks.{src}.weight")[i])
+            p[f"{b}{src}_b"].data.copy_(t(f"decoder.blocks.{src}.bias")[i])
+        for src in ["ff_n", "ff_c", "ff_t"]:
+            for lyr in ["linear_1", "linear_2"]:
+                p[f"{b}{src}.{lyr}.weight"].data.copy_(t(f"decoder.blocks.{src}.{lyr}.weight")[i])
+                p[f"{b}{src}.{lyr}.bias"].data.copy_(t(f"decoder.blocks.{src}.{lyr}.bias")[i])
+    return model
+
+
+# ---------------- ODE sampler: Dormand-Prince 5(4), constant step, CFG ----------------
+# Butcher tableau, embedded exactly as diffrax's Dopri5 (7 stages, FSAL/SSAL).
+# a_lower rows are stages 1..6; stage 0 has no coefficients (c1 = 0).
+_DP_A_LOWER = (
+    (1 / 5,),
+    (3 / 40, 9 / 40),
+    (44 / 45, -56 / 15, 32 / 9),
+    (19372 / 6561, -25360 / 2187, 64448 / 6561, -212 / 729),
+    (9017 / 3168, -355 / 33, 46732 / 5247, 49 / 176, -5103 / 18656),
+    (35 / 384, 0.0, 500 / 1113, 125 / 192, -2187 / 6784, 11 / 84),
+)
+_DP_B = (35 / 384, 0.0, 500 / 1113, 125 / 192, -2187 / 6784, 11 / 84, 0.0)
+_DP_C = (0.0, 1 / 5, 3 / 10, 4 / 5, 8 / 9, 1.0, 1.0)
+
+
+def solve_ode(
+    model: MapPFN,
+    noise_shape: tuple[int, int],   # (num_rows, dim)
+    obs_data: torch.Tensor,        # (1, num_rows, dim)
+    int_data: torch.Tensor | None,  # (num_shots, num_rows, dim) or None
+    treatment: torch.Tensor,        # (num_shots + 1, 1, dim)
+    guidance: float = 2.0,
+    step_size: float = 0.01,
+    generator: torch.Generator | None = None,
+    x_init: torch.Tensor | None = None,
+):
+    """Reference: map_pfn.models.utils.solve_ode (diffrax Dopri5, ConstantStepSize).
+
+    Exactly 100 fixed DP5 steps from t=0 to t=1 with classifier-free guidance,
+    replicating diffrax's k-space arithmetic (k = f * dt) and FSAL stage reuse.
+    Times accumulate in float32 exactly as diffrax does, with the endpoint
+    clipped to t1 when within 1e-6 (diffrax's `_clip_to_end`).
+    """
+    device = next(model.parameters()).device
+    if x_init is None:
+        x = torch.randn(noise_shape, generator=generator, device=device, dtype=torch.float32)
+    else:
+        x = x_init.to(device=device, dtype=torch.float32)
+
+    def velocity(x_t: torch.Tensor, t: float) -> torch.Tensor:
+        tt = torch.tensor(t, device=device, dtype=torch.float32)
+        v_cond = model(x_t, obs_data, int_data, treatment, tt, drop_cond=False)
+        v_uncond = model(x_t, obs_data, int_data, treatment, tt, drop_cond=True)
+        return guidance * v_cond + (1 - guidance) * v_uncond
+
+    f32 = np.float32
+    dt0 = f32(step_size)
+    t1 = f32(1.0)
+    tprev = f32(0.0)
+    tnext = f32(f32(0.0) + dt0)
+
+    f_last = None  # FSAL: stage-6 evaluation of the previous step
+    n_steps = 0
+    while tprev < t1:
+        dt = f32(tnext - tprev)
+        # stage times: t1_next exactly when c == 1, else tprev + c*dt (all float32)
+        ts = [f32(tnext if c == 1.0 else tprev + f32(c) * dt) for c in _DP_C]
+
+        fs = [None] * 7
+        fs[0] = f_last if f_last is not None else velocity(x, ts[0])
+        for s in range(1, 7):
+            yi = x
+            for j, a in enumerate(_DP_A_LOWER[s - 1]):
+                if a != 0.0:
+                    yi = yi + f32(a) * dt * fs[j]
+            fs[s] = velocity(yi, ts[s])
+
+        x = x + f32(_DP_B[0]) * dt * fs[0]
+        for i in range(1, 7):
+            if _DP_B[i] != 0.0:
+                x = x + f32(_DP_B[i]) * dt * fs[i]
+
+        f_last = fs[6]
+        tprev = f32(min(tnext, t1))
+        tnext = f32(tprev + dt0)
+        if tnext > f32(t1 - 1e-6):
+            tnext = t1
+        n_steps += 1
+        assert n_steps <= 102, "runaway ODE loop"
+
+    return x

requirements.txt

@@ -1,4 +1,4 @@
-torch
-plotly
+anndata
+matplotlib
 numpy
 pandas

weights/weights_frangieh.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b1d9d7306f4ccd4a1434072c4d89b67df8cfe5513d86b44cbf59adac2edf4983
+size 103894364

weights/weights_papalexi.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a41a0d0077ef6fe0963cfe095921cd4c47e30ce0a0c019e5a4b65271788881f8
+size 103894364