PDGrapher/scripts/benchmark_drug_reflector.py

#!/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()