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