| """Per-holdout zeroshot baselines for the main_metrics figure. |
| |
| All baselines computed from zeroshot-eligible training data only (i.e., the |
| panel-free subset that drops held-out cells from train) so the figure can |
| honestly say "this is what's achievable without panel observations of the |
| test cell line." |
| |
| Three baselines: |
| - Global Mean Zeroshot: per-gene mean of Y/D/P across all zeroshot train rows; |
| broadcast to all test rows (no drug, no cell information used). |
| - Perturbation Mean Zeroshot: per-(drug, gene) mean over zeroshot train rows; |
| for each test row, look up the mean of its drug. |
| - Rhaister Zeroshot Shared: shared-γ-per-gene diagonal model |
| (HP_ZS_DIAG_SHARED=1) through the same pipeline as the headline zeroshot |
| model (so it gets the same calnet calibration). |
| |
| The headline drug-specific γ_{t,g} model values are already in |
| titration_cells_zeroshot_results.json at level L=45. |
| |
| Saved to: zeroshot_baselines_results.json |
| """ |
| import json |
| import multiprocessing as mp |
| import os |
| import sys |
| from concurrent.futures import ProcessPoolExecutor, as_completed |
|
|
|
|
| import numpy as np |
|
|
| HOLDOUTS = [5, 6, 7, 8, 9] |
| METRICS = [ |
| "pdex_static/pearson_delta_mean", |
| "pdex_static/auprc_p05", |
| "state/pearson_delta_mean", |
| "state/de_overlap_mean", |
| "state/de_spearman_sig", |
| "state/pr_auc_mean", |
| "state/spearman_lfc_sig_mean", |
| "state/discrimination_mean", |
| ] |
| OUTPUT = "zeroshot_baselines_results.json" |
|
|
|
|
| def _load_test_arrays(h): |
| cp = f"/tmp/tahoe_cache_combined/tahoe/{h}_holdout" |
| return ( |
| np.load(f"{cp}/Y_test.npy"), |
| np.load(f"{cp}/P_test.npy"), |
| np.load(f"{cp}/D_test.npy"), |
| np.load(f"{cp}/F_test.npy"), |
| ) |
|
|
|
|
| def _direct_baseline(h, kind): |
| """Direct (no model pipeline) baseline: kind in {'global_mean', 'pert_mean'}. |
| Predicts Y/D/P from train statistics and derives F via Benjamini-Hochberg.""" |
| from rhaister.prepare_combined import prepare_all, evaluate, pvalues_to_fdr_bh |
| data = prepare_all(f"tahoe_{h}_holdout", zeroshot=True) |
| Y_test, P_test, D_test, F_test = _load_test_arrays(h) |
|
|
| Y_train = np.asarray(data["Y_train"], dtype=np.float64) |
| D_train = np.asarray(data["D_train"], dtype=np.float64) |
| P_train = np.asarray(data["P_train"], dtype=np.float64) |
| train_treats = np.asarray(data["train_treatments"]) |
| test_treats = np.asarray(data["test_treatments"]) |
| test_cells = np.asarray(data["test_cells"]) |
| n_test = data["n_test"] |
| n_genes = Y_train.shape[1] |
|
|
| if kind == "global_mean": |
| Y_pred = np.broadcast_to(Y_train.mean(0), (n_test, n_genes)).copy() |
| D_pred = np.broadcast_to(D_train.mean(0), (n_test, n_genes)).copy() |
| P_pred = np.broadcast_to(P_train.mean(0), (n_test, n_genes)).copy() |
| elif kind == "pert_mean": |
| drugs = sorted(set(train_treats)) |
| d_to_idx = {d: i for i, d in enumerate(drugs)} |
| n_drugs = len(drugs) |
| Y_d = np.zeros((n_drugs, n_genes), dtype=np.float64) |
| D_d = np.zeros((n_drugs, n_genes), dtype=np.float64) |
| P_d = np.zeros((n_drugs, n_genes), dtype=np.float64) |
| counts = np.zeros(n_drugs) |
| for i, t in enumerate(train_treats): |
| j = d_to_idx[t] |
| Y_d[j] += Y_train[i] |
| D_d[j] += D_train[i] |
| P_d[j] += P_train[i] |
| counts[j] += 1 |
| Y_d /= counts[:, None].clip(min=1) |
| D_d /= counts[:, None].clip(min=1) |
| P_d /= counts[:, None].clip(min=1) |
| Y_global = Y_train.mean(0) |
| D_global = D_train.mean(0) |
| P_global = P_train.mean(0) |
| Y_pred = np.zeros((n_test, n_genes)) |
| D_pred = np.zeros((n_test, n_genes)) |
| P_pred = np.zeros((n_test, n_genes)) |
| for i, t in enumerate(test_treats): |
| j = d_to_idx.get(t) |
| if j is not None: |
| Y_pred[i] = Y_d[j] |
| D_pred[i] = D_d[j] |
| P_pred[i] = P_d[j] |
| else: |
| Y_pred[i] = Y_global |
| D_pred[i] = D_global |
| P_pred[i] = P_global |
| else: |
| raise ValueError(f"unknown direct baseline kind: {kind!r}") |
| F_pred = pvalues_to_fdr_bh(P_pred) |
|
|
| return evaluate( |
| Y_test, Y_pred, P_test, P_pred, D_test, D_pred, F_test, F_pred, |
| test_cells, test_treats, data["gene_cols"], compute_discrimination=False, |
| ) |
|
|
|
|
| def _pipeline_baseline(h, env_overrides): |
| """Run the full zeroshot pipeline with given env overrides (e.g. |
| HP_ZS_DIAG_SHARED=1).""" |
| from rhaister import train |
| from rhaister.prepare_combined import prepare_all |
| os.environ["HP_ZS_MODEL"] = "diagonal" |
| os.environ.setdefault("HP_ZS_DIAG_Z", "H") |
| for k, v in env_overrides.items(): |
| os.environ[k] = v |
| data = prepare_all(f"tahoe_{h}_holdout", zeroshot=True) |
| return train.train_and_evaluate( |
| experiment_name=f"zs_baseline_h{h}", |
| split_name=f"tahoe_{h}_holdout", |
| data=data, |
| log=False, |
| zeroshot=True, |
| compute_discrimination=False, |
| ) |
|
|
|
|
| def run_one(task): |
| h, name, mode, payload = task |
| print(f"[h={h} {name}] starting", flush=True) |
| if mode == "direct": |
| metrics = _direct_baseline(h, payload) |
| elif mode == "pipeline": |
| metrics = _pipeline_baseline(h, payload) |
| else: |
| raise ValueError(mode) |
| return h, name, {m: float(metrics[m]) for m in METRICS if m in metrics} |
|
|
|
|
| BASELINES = [ |
| |
| ("Global Mean Zeroshot", "direct", "global_mean"), |
| ("Perturbation Mean Zeroshot", "direct", "pert_mean"), |
| ("Rhaister Zeroshot Shared", "pipeline", {"HP_ZS_DIAG_SHARED": "1"}), |
| ] |
|
|
|
|
| def main(): |
| n_workers = 4 |
| for i, arg in enumerate(sys.argv): |
| if arg == "--workers" and i + 1 < len(sys.argv): |
| n_workers = int(sys.argv[i + 1]) |
|
|
| tasks = [(h, name, mode, payload) |
| for h in HOLDOUTS for (name, mode, payload) in BASELINES] |
| print(f"Dispatching {len(tasks)} (holdout, baseline) configs across {n_workers} workers") |
|
|
| results = {b[0]: {} for b in BASELINES} |
|
|
| ctx = mp.get_context("spawn") |
| n_done = 0 |
| with ProcessPoolExecutor(max_workers=n_workers, mp_context=ctx) as ex: |
| futures = {ex.submit(run_one, t): t for t in tasks} |
| for fut in as_completed(futures): |
| h_req, name_req, *_ = futures[fut] |
| try: |
| h, name, row = fut.result() |
| except Exception as exc: |
| n_done += 1 |
| print(f"[FAILED h={h_req} {name_req}] {exc} ({n_done}/{len(tasks)})", flush=True) |
| continue |
| results[name][str(h)] = row |
| n_done += 1 |
| short = " ".join(f"{m.split('/')[-1]}={row.get(m, float('nan')):.3f}" for m in METRICS) |
| print(f"[done h={h} {name:32s}] {short} ({n_done}/{len(tasks)})", flush=True) |
|
|
| out = { |
| "experiment": "zeroshot_baselines", |
| "description": ( |
| "Per-holdout panel-free zeroshot baselines. Global Mean and " |
| "Perturbation Mean are direct statistical baselines (no model " |
| "pipeline) computed from zeroshot-eligible train data; the Shared " |
| "diagonal goes through the model pipeline with HP_ZS_DIAG_SHARED=1 " |
| "(γ_g shared across drugs). The headline drug-specific γ_{t,g} " |
| "results live in titration_cells_zeroshot_results.json (L=45)." |
| ), |
| "holdouts": HOLDOUTS, |
| "baselines": [b[0] for b in BASELINES], |
| "metrics_per_baseline": results, |
| } |
| with open(OUTPUT, "w") as f: |
| json.dump(out, f, indent=2) |
| print(f"\nwrote {OUTPUT}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
