#!/usr/bin/env python3 """ Benchmark comparison: GIDFlow vs MLP baseline on LOCO drug repurposing. Trains a DrugReflector-style MLP classifier on the same data and compares it against GIDFlow's generative approach. Usage: python scripts/benchmark_drug_reflector.py \ --checkpoint outputs/norman2019/checkpoint.pt \ --dataset norman2019 \ --max-conditions 50 \ --output outputs/benchmark/ """ import argparse import os import sys import warnings warnings.filterwarnings("ignore") sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) import numpy as np import pandas as pd import torch import torch.nn.functional as F from torch.utils.data import DataLoader, random_split from gidflow.benchmark.baseline_mlp import ( DrugReflectorBaseline, build_baseline_dataset_from_scperturb, train_baseline, ) from gidflow.benchmark.comparison import compute_rank_metrics, compare_methods, permutation_test from gidflow.data import ScPerturbPopulationDataset, parse_perturbation_targets from gidflow.drug_library import DrugLibrary, DrugRecord from gidflow.predictor import GIDFlowPredictor DATASET_PATHS = { "norman2019": "/data/boom/Protein/regulatory_field/data/raw/scPerturb/rna_protein/NormanWeissman2019_filtered.h5ad", "replogle": "/data/boom/Protein/regulatory_field/data/raw/scPerturb/rna_protein/ReplogleWeissman2022_K562_essential.h5ad", } def run_gidflow_loco( predictor: GIDFlowPredictor, ds: ScPerturbPopulationDataset, max_conditions: int = None, n_bootstrap: int = 0, seed: int = 42, ) -> pd.DataFrame: """Run GIDFlow LOCO validation (reusing eval_loco logic).""" rng = np.random.default_rng(seed) conditions = ds._conditions n_cond = len(conditions) if max_conditions: n_cond = min(max_conditions, n_cond) conditions = conditions[:n_cond] results = [] for i, cond in enumerate(conditions): true_pert = cond["pert_name"] print(f"\r GIDFlow [{i+1}/{n_cond}] {true_pert:30s}", end="", flush=True) ns = min(32, len(cond["ctrl_idx"])) nt = min(32, len(cond["cell_idx"])) src_idx = rng.choice(cond["ctrl_idx"], size=ns, replace=False) tgt_idx = rng.choice(cond["cell_idx"], size=nt, replace=False) source_cells = torch.from_numpy(ds._X[src_idx]).float() target_cells = torch.from_numpy(ds._X[tgt_idx]).float() # Build library INCLUDING the true perturbation. # LOCO here means leaving the condition out of model fitting, not # removing the true perturbation from the retrieval candidates. all_records = [] for other_cond in ds._conditions: targets = parse_perturbation_targets(other_cond["pert_name"]) pert_vec = np.zeros(ds.num_genes, dtype=np.float32) for gene in targets: gene = gene.strip() if gene in ds._gene_to_idx: pert_vec[ds._gene_to_idx[gene]] = 1.0 all_records.append(DrugRecord( name=other_cond["pert_name"], gene_targets=targets, pert_vector=pert_vec, cell_type="lymphoblasts", source_dataset="Norman2019", metadata={"is_true": other_cond["pert_name"] == true_pert}, )) lib = DrugLibrary(all_records, ds.gene_names.tolist(), name="loco") predictor.drug_library = lib try: result = predictor.drug_repurposing( source_cells=source_cells, target_cells=target_cells, gene_names=ds.gene_names.tolist(), top_k=50, n_bootstrap=n_bootstrap, compute_pathways=False, ) except Exception as e: print(f"\n ERROR on {true_pert}: {e}") continue rank_row = result.ranked_drugs[result.ranked_drugs["drug_name"] == true_pert] if len(rank_row) > 0: rank = int(rank_row.iloc[0]["rank"]) score = float(rank_row.iloc[0]["composite_score"]) else: rank = lib.num_drugs + 1 score = 0.0 results.append({ "condition": true_pert, "true_perturbation": true_pert, "rank": rank, "in_top1": rank <= 1, "in_top5": rank <= 5, "in_top10": rank <= 10, "composite_score": score, "pearson_r": float(rank_row.iloc[0]["pearson_r"]) if len(rank_row) > 0 else 0.0, "de_overlap": float(rank_row.iloc[0]["de_overlap"]) if len(rank_row) > 0 else 0.0, "score_std": float(rank_row.iloc[0]["score_std"]) if len(rank_row) > 0 else 0.0, "n_source": ns, "n_target": nt, }) print() return pd.DataFrame(results) def run_baseline_loco( model: DrugReflectorBaseline, dataset, max_conditions: int = None, device: str = "cpu", seed: int = 42, ) -> pd.DataFrame: """Run MLP baseline LOCO validation. For each condition: 1. Compute signature = mean(target) - mean(source) 2. Run through MLP classifier 3. Record rank of true perturbation """ rng = np.random.default_rng(seed) conditions = dataset._conditions n_cond = len(conditions) if max_conditions: n_cond = min(max_conditions, n_cond) conditions = conditions[:n_cond] # Build pert name → index mapping all_pert_names = [c["pert_name"] for c in dataset._conditions] pert_to_idx = {name: i for i, name in enumerate(all_pert_names)} results = [] model.eval() model = model.to(device) with torch.no_grad(): for i, cond in enumerate(conditions): true_pert = cond["pert_name"] print(f"\r MLP [{i+1}/{n_cond}] {true_pert:30s}", end="", flush=True) ns = min(32, len(cond["ctrl_idx"])) nt = min(32, len(cond["cell_idx"])) src_idx = rng.choice(cond["ctrl_idx"], size=ns, replace=False) tgt_idx = rng.choice(cond["cell_idx"], size=nt, replace=False) src = dataset._X[src_idx].mean(axis=0) tgt = dataset._X[tgt_idx].mean(axis=0) sig = torch.from_numpy(tgt - src).float().unsqueeze(0).to(device) logits = model(sig).squeeze(0) probs = F.softmax(logits, dim=-1) sorted_indices = torch.argsort(probs, descending=True) true_idx = pert_to_idx.get(true_pert, -1) if true_idx >= 0: matches = (sorted_indices == true_idx).nonzero(as_tuple=True)[0] rank = int(matches[0].item()) + 1 if len(matches) > 0 else len(all_pert_names) + 1 else: rank = len(all_pert_names) + 1 results.append({ "condition": true_pert, "true_perturbation": true_pert, "rank": rank, "in_top1": rank <= 1, "in_top5": rank <= 5, "in_top10": rank <= 10, "composite_score": float(probs[true_idx]) if true_idx >= 0 else 0.0, "n_source": ns, "n_target": nt, }) print() return pd.DataFrame(results) def main(): parser = argparse.ArgumentParser() parser.add_argument("--checkpoint", required=True) parser.add_argument("--dataset", required=True, choices=["norman2019", "replogle"]) parser.add_argument("--max-conditions", type=int, default=None) parser.add_argument("--n-bootstrap", type=int, default=0) parser.add_argument("--device", default="auto") parser.add_argument("--output", default="outputs/benchmark") parser.add_argument("--skip-training", action="store_true", help="Skip MLP training (use existing model)") args = parser.parse_args() os.makedirs(args.output, exist_ok=True) device = "cpu" if args.device == "auto" else args.device # Load dataset h5ad_path = DATASET_PATHS[args.dataset] print(f"Loading dataset: {h5ad_path}") ds = ScPerturbPopulationDataset( h5ad_path=h5ad_path, n_hvg=2000, min_cells_per_cond=20, max_source_cells=32, max_target_cells=32, control_label="control", use_single_pert_only=False, seed=42, ) print(f" {len(ds)} conditions, {ds.num_genes} genes") # ---- GIDFlow ---- print("\n" + "=" * 60) print(" Running GIDFlow LOCO...") print("=" * 60) predictor = GIDFlowPredictor( checkpoint_path=args.checkpoint, drug_library="norman2019", device=device, n_bootstrap=args.n_bootstrap, seed=42, ) gidflow_results = run_gidflow_loco(predictor, ds, args.max_conditions, args.n_bootstrap) gidflow_csv = os.path.join(args.output, "gidflow_loco.csv") gidflow_results.to_csv(gidflow_csv, index=False) print(f"GIDFlow results saved to {gidflow_csv}") # ---- MLP Baseline ---- print("\n" + "=" * 60) print(" Training MLP Baseline...") print("=" * 60) baseline_path = os.path.join(args.output, "mlp_baseline.pt") if args.skip_training and os.path.exists(baseline_path): print(f" Loading existing baseline from {baseline_path}") state = torch.load(baseline_path, map_location=device, weights_only=False) model = DrugReflectorBaseline( input_dim=state["input_dim"], n_classes=state["n_classes"], hidden_dims=state.get("hidden_dims", [1024, 1024]), dropout=state.get("dropout", 0.2), ) model.load_state_dict(state["model_state_dict"]) else: # Build training data sig_dataset, gene_names = build_baseline_dataset_from_scperturb( h5ad_path, n_hvg=2000, min_cells_per_cond=20, max_source_cells=32, max_target_cells=32, seed=42, force_include_pert_genes=True, ) print(f" Training data: {len(sig_dataset)} signatures, {sig_dataset.n_classes} classes") # Train/val split n_train = int(0.8 * len(sig_dataset)) n_val = len(sig_dataset) - n_train train_ds, val_ds = random_split( sig_dataset, [n_train, n_val], generator=torch.Generator().manual_seed(42), ) train_loader = DataLoader(train_ds, batch_size=32, shuffle=True) val_loader = DataLoader(val_ds, batch_size=32, shuffle=False) model = DrugReflectorBaseline( input_dim=sig_dataset.signatures.shape[1], n_classes=sig_dataset.n_classes, ) history = train_baseline( model, train_loader, val_loader, n_epochs=30, device=device, verbose=True, ) # Save torch.save({ "model_state_dict": model.state_dict(), "input_dim": model.mlp[0].in_features, "n_classes": model.mlp[-1].out_features, "history": history, }, baseline_path) print(f" Baseline saved to {baseline_path}") # Run baseline LOCO print("\n" + "=" * 60) print(" Running MLP Baseline LOCO...") print("=" * 60) baseline_results = run_baseline_loco(model, ds, args.max_conditions, device) baseline_csv = os.path.join(args.output, "baseline_loco.csv") baseline_results.to_csv(baseline_csv, index=False) print(f"Baseline results saved to {baseline_csv}") # ---- Comparison ---- print("\n" + "=" * 60) print(" Comparison Results") print("=" * 60) comparison = compare_methods({ "GIDFlow": gidflow_results, "MLP_Baseline": baseline_results, }) print(comparison.to_string(index=False)) comparison_path = os.path.join(args.output, "comparison.csv") comparison.to_csv(comparison_path, index=False) print(f"\nComparison saved to {comparison_path}") # Permutation test on MRR if len(gidflow_results) == len(baseline_results): gidflow_mrr = 1.0 / gidflow_results["rank"].clip(lower=1) baseline_mrr = 1.0 / baseline_results["rank"].clip(lower=1) perm_result = permutation_test(gidflow_mrr.tolist(), baseline_mrr.tolist()) print(f"\nPermutation test (MRR):") print(f" GIDFlow MRR: {gidflow_mrr.mean():.4f}") print(f" MLP MRR: {baseline_mrr.mean():.4f}") print(f" Observed diff: {perm_result['observed_diff']:.4f}") print(f" p-value: {perm_result['p_value']:.4f}") print(f" Significant (p<0.05): {perm_result['significant']}") if __name__ == "__main__": main()