| """Evaluate fixed populations and intervention catalogs with condition-level rows.""" |
|
|
| from __future__ import annotations |
| import json, time |
| from pathlib import Path |
| import numpy as np |
| import torch |
| from pivot.evaluation.metrics import ( |
| correlation, |
| mmd2, |
| sliced_wasserstein, |
| energy_distance, |
| retrieval_metrics, |
| bootstrap, |
| ) |
| from pivot.evaluation.rewards import Reward |
| from pivot.evaluation.inference import ( |
| encode_label, |
| forward_predict, |
| reward_guidance, |
| project_and_rerank, |
| ) |
|
|
|
|
| def save_json(path, value): |
| """Write strict JSON. Undefined numerical statistics become null.""" |
|
|
| def clean(x): |
| if isinstance(x, dict): |
| return {str(k): clean(v) for k, v in x.items()} |
| if isinstance(x, (list, tuple)): |
| return [clean(v) for v in x] |
| if isinstance(x, (float, np.floating)): |
| return float(x) if np.isfinite(x) else None |
| if isinstance(x, np.integer): |
| return int(x) |
| return x |
|
|
| p = Path(path) |
| p.parent.mkdir(parents=True, exist_ok=True) |
| p.write_text(json.dumps(clean(value), indent=2, allow_nan=False)) |
|
|
|
|
| def evaluate( |
| data, |
| predict, |
| output, |
| method="PIVOT", |
| partition="test", |
| catalog="single", |
| reward_kind="cosine", |
| n_cells=128, |
| seed=0, |
| model=None, |
| device="cpu", |
| guidance_steps=25, |
| step_size=0.5, |
| k_nearest=10, |
| initialization="best", |
| max_targets=None, |
| ): |
| """Forward and inverse evaluation on disjoint query and reference cells. |
| |
| Feature errors use normalized expression, with each target's observed top-20 |
| effect genes selected once. Population metrics use PCA coordinates. Inverse |
| query populations are disjoint from the reserved candidate-outcome cells. |
| The target condition is the aggregation and bootstrap unit. |
| """ |
| rng = np.random.default_rng(seed) |
| ci = data.indices(partition, True) |
| if not len(ci): |
| raise ValueError("Evaluation needs controls in the requested partition") |
| ci = rng.choice(ci, min(n_cells, len(ci)), replace=False) |
| c0 = data.emb[ci] |
| control_expr = np.asarray(data.Xhvg[ci].mean(0)).ravel() |
| control_latent = c0.mean(0) |
| candidates = ( |
| data.singles |
| if catalog == "single" |
| else data.combos if catalog == "combination" else data.perturbations |
| ) |
| if not candidates: |
| raise ValueError("Candidate catalog is empty") |
| labels = data.labels(partition) |
| if max_targets: |
| labels = sorted( |
| rng.choice(labels, min(max_targets, len(labels)), replace=False) |
| ) |
| |
| |
| t0 = time.perf_counter() |
| predictions = {p: np.asarray(predict(c0, p)) for p in candidates} |
| catalog_seconds = time.perf_counter() - t0 |
| query_ids = data.indices(partition, False) |
| refids = data.indices("reference", False) |
| reference = { |
| p: data.emb[np.intersect1d(refids, data.pert_to_idx[p])] for p in candidates |
| } |
| if any(len(v) == 0 for v in reference.values()): |
| raise ValueError("Every candidate needs independent reference outcomes") |
| forward = [] |
| inverse = [] |
| for p in labels: |
| ids = np.intersect1d(query_ids, data.pert_to_idx[p]) |
| ids = rng.choice(ids, min(n_cells, len(ids)), replace=False) |
| target = data.emb[ids] |
| pred = predictions.get(p) |
| if pred is None: |
| pred = np.asarray(predict(c0, p)) |
| expression = np.asarray(data.Xhvg[ids].mean(0)).ravel() |
| pred_expr = data.decode_to_genes(pred).mean(0) |
| effect = expression - control_expr |
| pred_effect = pred_expr - control_expr |
| de = np.argsort(-np.abs(effect))[: min(20, len(effect))] |
| train_genes = {g for q in data.labels("train") for g in data.parse(q)} |
| forward.append( |
| { |
| "target": p, |
| "query_cells": len(ids), |
| "genes_observed_in_training": all( |
| g in train_genes for g in data.parse(p) |
| ), |
| "mse_expression": float(np.mean((pred_expr - expression) ** 2)), |
| "effect_pearson": correlation(pred_effect, effect), |
| "effect_pearson_de20": correlation(pred_effect[de], effect[de]), |
| "mmd2": mmd2(pred, target, data.meta["mmd_gamma"]), |
| "sliced_wasserstein": sliced_wasserstein(pred, target, seed=seed), |
| "energy": energy_distance(pred, target), |
| "variance_ratio": float( |
| np.var(pred, axis=0).sum() |
| / max(np.var(target, axis=0).sum(), 1e-12) |
| ), |
| } |
| ) |
| if p not in candidates: |
| continue |
| reward = Reward( |
| reward_kind, |
| target_sample=target, |
| control_ref=control_latent, |
| gamma=data.meta["mmd_gamma"], |
| device=device, |
| ) |
| score = lambda pop: float( |
| reward(torch.as_tensor(pop, dtype=torch.float32, device=device)) |
| .mean() |
| .detach() |
| .cpu() |
| ) |
| measured = {q: score(v) for q, v in reference.items()} |
| scores = {q: score(v) for q, v in predictions.items()} |
| ranked = sorted(candidates, key=lambda q: (-scores[q], q)) |
| tied = sum(abs(scores[q] - scores[ranked[0]]) < 1e-10 for q in candidates) |
|
|
| def row(name, ranking, cost, seconds, continuous=None): |
| selected = ranking[0] |
| out = { |
| "target": p, |
| "search": name, |
| **retrieval_metrics(ranking, p), |
| "selected": selected, |
| "predicted_reward": scores.get(selected), |
| "measured_reward": measured[selected], |
| "measured_regret": max(measured.values()) - measured[selected], |
| "selected_gene_overlap": len( |
| set(data.parse(selected)) & set(data.parse(p)) |
| ) |
| / max(len(set(data.parse(p))), 1), |
| "catalog_size": len(candidates), |
| "candidate_cell_evaluations": cost, |
| "seconds": seconds, |
| "top_score_ties": tied, |
| } |
| if continuous is not None: |
| out["continuous_reward"] = continuous |
| return out |
|
|
| inverse.append( |
| row( |
| "exhaustive", |
| ranked, |
| len(candidates) * len(c0), |
| catalog_seconds / len(labels), |
| ) |
| ) |
| if model is not None and guidance_steps > 0: |
| |
| |
| init = ( |
| ranked[0] if initialization == "best" else str(rng.choice(candidates)) |
| ) |
| e = encode_label(model, data, init, device) |
| ct = torch.as_tensor(c0, device=device) |
| t1 = time.perf_counter() |
| es = reward_guidance(model, ct, reward, e, guidance_steps, step_size) |
| guided = project_and_rerank( |
| model, data, candidates, es, ct, reward, k_nearest, k_nearest, device |
| ) |
| continuous = score(forward_predict(model, ct, es).cpu().numpy()) |
| names = [q for q, _ in guided] |
| cost = ( |
| (len(candidates) if initialization == "best" else 0) |
| + guidance_steps |
| + min(k_nearest, len(candidates)) |
| ) |
| gr = row( |
| "guidance_" + initialization, |
| names, |
| cost * len(c0), |
| time.perf_counter() |
| - t1 |
| + (catalog_seconds / len(labels) if initialization == "best" else 0), |
| continuous, |
| ) |
| gr["initial_candidate"] = init |
| gr["predicted_reward"] = guided[0][1] |
| gr["projection_reward_change"] = guided[0][1] - continuous |
| inverse.append(gr) |
| summary = { |
| "forward": { |
| k: bootstrap([r[k] for r in forward], seed) |
| for k in ( |
| "mse_expression", |
| "effect_pearson", |
| "effect_pearson_de20", |
| "mmd2", |
| "sliced_wasserstein", |
| "energy", |
| "variance_ratio", |
| ) |
| }, |
| "inverse": {}, |
| } |
| for search in sorted({r["search"] for r in inverse}): |
| rows = [r for r in inverse if r["search"] == search] |
| summary["inverse"][search] = { |
| k: bootstrap([r[k] for r in rows], seed) |
| for k in ( |
| "top1", |
| "top5", |
| "ndcg10", |
| "measured_reward", |
| "measured_regret", |
| "selected_gene_overlap", |
| ) |
| } |
| result = { |
| "protocol": "split-first-v1", |
| "method": method, |
| "data_meta": data.meta, |
| "partition": partition, |
| "catalog": catalog, |
| "reward": reward_kind, |
| "seed": seed, |
| "n_controls": len(c0), |
| "query_cell_ids": data.obs.iloc[query_ids].cell_id.tolist(), |
| "guidance": { |
| "steps": guidance_steps, |
| "step_size": step_size, |
| "k_nearest": k_nearest, |
| "initialization": initialization, |
| }, |
| "summary": summary, |
| "forward": forward, |
| "inverse": inverse, |
| } |
| save_json(output, result) |
| return result |
|
|
