PDGrapher/scripts/predict_drugs.py

#!/usr/bin/env python3
"""
Command-line interface for GIDFlow drug repurposing predictions.

Usage:
    # From cell populations
    python scripts/predict_drugs.py \
        --checkpoint outputs/norman2019/checkpoint.pt \
        --source data/control.h5ad \
        --target data/disease.h5ad \
        --drug-library norman2019 \
        --cell-type K562 \
        --top-k 50 \
        --output results/

    # From v-score signature
    python scripts/predict_drugs.py \
        --checkpoint outputs/norman2019/checkpoint.pt \
        --signature vscores.csv \
        --drug-library norman2019 \
        --top-k 50 \
        --output results/

    # From config file
    python scripts/predict_drugs.py --config configs/drug_repurposing.yaml
"""
import argparse
import os
import sys
import warnings

warnings.filterwarnings("ignore")
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))

import pandas as pd
import torch

from gidflow.predictor import GIDFlowPredictor


def load_signature(path: str) -> pd.Series:
    """Load a v-score signature from CSV or TSV.

    Expected format: gene_name, value (one gene per row)
    or: first column = gene names, second = values.

    Parameters
    ----------
    path : str

    Returns
    -------
    pd.Series indexed by gene names
    """
    # Try CSV first
    for sep in [",", "\t"]:
        try:
            df = pd.read_csv(path, sep=sep, header=None, index_col=0)
            if df.shape[1] >= 1:
                series = df.iloc[:, 0]
                series.name = os.path.basename(path)
                return series
        except Exception:
            continue

    raise ValueError(f"Could not load signature from {path}")


def main():
    parser = argparse.ArgumentParser(
        description="GIDFlow Drug Repurposing Predictions",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    # Required
    parser.add_argument("--checkpoint", required=True, help="Path to model checkpoint (.pt)")
    parser.add_argument("--drug-library", default="norman2019",
                        help="Drug library name (norman2019, replogle) or path to h5ad")

    # Input options (mutually exclusive)
    input_group = parser.add_mutually_exclusive_group(required=True)
    input_group.add_argument("--source", help="Path to source (control) h5ad")
    input_group.add_argument("--signature", help="Path to pre-computed v-score signature CSV")

    parser.add_argument("--target", help="Path to target (disease) h5ad (required if --source given)")
    parser.add_argument("--cell-type", default=None, help="Filter drug library to cell type")
    parser.add_argument("--top-k", type=int, default=50, help="Number of top drugs to return")
    parser.add_argument("--n-bootstrap", type=int, default=10, help="Bootstrap samples for CI")
    parser.add_argument("--seed", type=int, default=42, help="Random seed")
    parser.add_argument("--output", default="results", help="Output directory")
    parser.add_argument("--compute-pathways", action="store_true", default=True,
                        help="Run pathway enrichment")
    parser.add_argument("--no-pathways", action="store_true",
                        help="Skip pathway enrichment")
    parser.add_argument("--config", help="Config file (yaml)")

    args = parser.parse_args()

    # Handle config
    if args.config:
        try:
            import yaml
            with open(args.config) as f:
                cfg = yaml.safe_load(f)
            pred_cfg = cfg.get("predictor", {})
            drug_cfg = cfg.get("drug_library", {})
            args.checkpoint = pred_cfg.get("checkpoint", args.checkpoint)
            args.drug_library = drug_cfg.get("h5ad_path", args.drug_library)
            args.cell_type = drug_cfg.get("cell_type", args.cell_type)
            args.top_k = pred_cfg.get("top_k", args.top_k)
            args.n_bootstrap = pred_cfg.get("n_bootstrap", args.n_bootstrap)
            args.compute_pathways = not pred_cfg.get("skip_pathways", False)
        except ImportError:
            print("Warning: pyyaml not installed, ignoring config file")

    # Validate input
    if args.source and not args.target:
        parser.error("--target is required when --source is given")
    if not args.source and not args.signature:
        parser.error("Must provide either --source+--target or --signature")

    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Device: {device}")

    # Load data
    source_cells = None
    target_cells = None
    signature = None

    if args.source:
        try:
            import anndata as ad
            source_cells = ad.read_h5ad(args.source)
            target_cells = ad.read_h5ad(args.target)
            print(f"Loaded source: {source_cells.n_obs} cells, {source_cells.n_vars} genes")
            print(f"Loaded target: {target_cells.n_obs} cells, {target_cells.n_vars} genes")
        except ImportError:
            raise ImportError("pip install anndata")
    elif args.signature:
        signature = load_signature(args.signature)
        print(f"Loaded signature: {len(signature)} genes")

    # Initialize predictor
    predictor = GIDFlowPredictor(
        checkpoint_path=args.checkpoint,
        drug_library=args.drug_library,
        device=device,
        n_bootstrap=args.n_bootstrap,
        seed=args.seed,
    )

    # Run prediction
    compute_pathways = args.compute_pathways and not args.no_pathways

    # Check gene dimensions
    if source_cells is not None:
        input_genes = source_cells.n_vars
        model_genes = predictor.num_genes
        if input_genes != model_genes:
            print(f"\nERROR: Input data has {input_genes} genes, "
                  f"but model expects {model_genes} genes.", file=sys.stderr)
            print(f"\nPlease pre-process your data to match the model's gene space:", file=sys.stderr)
            print(f"  1. Load the h5ad used during training", file=sys.stderr)
            print(f"  2. Select the same HVG set ({model_genes} genes)", file=sys.stderr)
            print(f"  3. Use those genes as input", file=sys.stderr)
            sys.exit(1)

    result = predictor.drug_repurposing(
        source_cells=source_cells,
        target_cells=target_cells,
        signature=signature,
        cell_type=args.cell_type,
        top_k=args.top_k,
        compute_pathways=compute_pathways,
    )

    # Print results
    print(f"\n{'=' * 60}")
    print(f"  Drug Repurposing Results")
    print(f"{'=' * 60}")
    print(f"  Drugs screened: {result.metadata.get('n_drugs_screened', '?')}")
    print()
    for _, row in result.ranked_drugs.head(20).iterrows():
        print(f"  #{int(row['rank']):3d} {row['drug_name']:30s}  "
              f"score={row['composite_score']:.3f}  "
              f"pearson_r={row['pearson_r']:.3f}  "
              f"de_overlap={row['de_overlap']:.3f}")

    # Save results
    os.makedirs(args.output, exist_ok=True)

    # Rankings CSV
    csv_path = os.path.join(args.output, "drug_rankings.csv")
    result.to_csv(csv_path)
    print(f"\nRankings saved to {csv_path}")

    # Top predicted populations
    pop_dir = os.path.join(args.output, "predicted_populations")
    result.save_predicted_populations(pop_dir)
    print(f"Predicted populations saved to {pop_dir}/")

    # Pathway enrichment
    if compute_pathways and result.pathway_enrichment:
        for drug_name, enr_df in result.pathway_enrichment.items():
            if enr_df is not None and len(enr_df) > 0:
                safe_name = drug_name.replace("+", "_").replace("/", "_")
                enr_path = os.path.join(args.output, f"pathways_{safe_name}.csv")
                enr_df.to_csv(enr_path, index=False)
        print(f"Pathway enrichment saved to {args.output}/pathways_*.csv")


if __name__ == "__main__":
    main()