TxOptimus

Zero-shot drug repurposing via TxGNN on PrimeKG and OptimusKG.

TxOptimus is a unified Python package that wraps TxGNN (Huang et al., Nature Medicine, 2024) with two knowledge graph backends:

• TxGNN Prime — the original baseline trained on Harvard PrimeKG
• TxGNN Optimus — an enriched model trained on the Zitnik Lab's OptimusKG, a multimodal knowledge graph validated by LLM extraction (PaperQA3)

OptimusKG achieves a ~34% improvement in Macro AUROC over PrimeKG in zero-shot drug repurposing for oncology under the cell_proliferation split.

Installation

# Create a dedicated environment (recommended)
conda create -n txoptimus python=3.8
conda activate txoptimus

# Install DGL first (not available on PyPI, requires separate wheel)
# For CUDA 12.1:
pip install https://data.dgl.ai/wheels/torch-2.4/cu121/dgl-2.4.0%2Bcu121-cp38-cp38-manylinux1_x86_64.whl
# For CPU only:
# pip install https://data.dgl.ai/wheels/torch-2.4/cpu/dgl-2.4.0-cp38-cp38-manylinux1_x86_64.whl

# Install TxOptimus
pip install txoptimus

⚠ System Requirements:

• RAM: ≥32 GB (OptimusKG loads a 21.8M-edge graph)
• Storage: ≥8 GB free disk space
• Python: 3.8+ (tested on 3.8.20)
• GPU: Optional (CPU inference supported)

Quick Start

1. Download data files

txoptimus --setup              # downloads both PrimeKG and OptimusKG data
txoptimus --setup prime        # only PrimeKG
txoptimus --setup optimus      # only OptimusKG

2. Run zero-shot drug repurposing

# Using OptimusKG (recommended)
txoptimus --engine optimus --diseases "oral cavity cancer" "glioblastoma"

# Using PrimeKG baseline
txoptimus --engine prime --diseases "oral cavity cancer"

# Benchmark mode — runs both and generates a comparison report
txoptimus --engine benchmark --diseases "oral cavity cancer"

3. Fuzzy disease matching

TxOptimus uses PubMedBERT embeddings to fuzzy-match your disease terms against the knowledge graph vocabulary. If your input doesn't match exactly, the tool will show the top-10 closest matches and let you pick:

Disease term: "breast cancer"

Matched terms in OPTIMUS:
  [1] invasive breast carcinoma (similarity: 0.94)
  [2] breast carcinoma (similarity: 0.91)
  [3] triple-negative breast cancer (similarity: 0.89)

Select terms (e.g., '1 3' or 'all'): 1 3

4. Enable explainability

# Add --graphmask to enable GraphMask pathway attribution
txoptimus --engine optimus --diseases "oral cavity cancer" --graphmask

Note: GraphMask requires ≥30 GB RAM and takes 2-4 hours on CPU.

Full CLI Reference

Output Files

Citation

If you use TxOptimus in your research, please cite the original TxGNN paper:

@article{huang2024txgnn,
  title={Zero-shot prediction of therapeutic use with geometric deep learning and clinician centered design},
  author={Huang, Kexin and Chandak, Payal and Wang, Qianwen and Haber, Shreyas and Zitnik, Marinka},
  journal={Nature Medicine},
  year={2024}
}

License

MIT License — see LICENSE for details.

Author

arsalanriaz