DOCSEARCHbyAlgolia
Submit Docs

TxGNN

repository·main·Indexed 43 days ago

https://github.com/mims-harvard/txgnn

A graph neural network for zero-shot prediction of therapeutic uses for diseases. TxGNN leverages geometric deep learning on large knowledge graphs to identify drug repurposing opportunities. It features a multi-stage pipeline including pre-training for link prediction, fine-tuning with metric learning, and Graph XAI training via Graphmask for interpretability. The framework includes TxData for knowledge graph partitioning and TxEval for disease-centric evaluation.

Tokens
4.6K
Snippets
15
Records
17
Agent score
54%

What's inside TxGNN

  1. Configure TxGNN data splits

    main

    The TxData.prepare_split() method allows you to choose different data partitioning strategies:

    • complex_disease: The systematic split used in the paper. Samples a set of diseases and moves all their treatments to the test set to ensure zero-shot capability.
    • disease_area: Splits based on disease ontology. Moves all treatments for a specific disease area to the test set. Available areas: cell_proliferation, mental_health, cardiovascular, anemia, adrenal_gland, autoimmune, metabolic_disorder, diabetes, neurodigenerative.
    • random: Randomly shuffles drug-disease pairs.
    • disease_eval: Used for deploying/evaluating a specific disease. Set disease_eval_idx to the index of the disease of interest to mask it and use others for context.
    • full_graph: Trains on the entire network without masking. Uses 95% for training and 5% for validation (no test set).
    # Example: Disease area split
TxData.prepare_split(split = 'cell_proliferation', seed = 42)

# Example: Evaluating a specific disease index
TxData.prepare_split(split = 'disease_eval', disease_eval_idx = 'XX')

  2. Reproduce TxGNN training and results

    main

    To reproduce the results presented in the paper, use the following files located in the reproduce/ directory:

    • Training: Run run_txgnn.sh or train.py to initiate the training process.
    • Visualization: Use gather result-txgnn.ipynb to generate the plots used in the paper.
    • Raw Data: result_more_metrics.csv contains the raw result files.

    Model checkpoints can be downloaded from Google Drive.

  3. Install TxGNN

    main

    To install TxGNN, create a Python 3.8 environment using conda, install PyTorch compatible with your CUDA version, and then install the DGL (Deep Graph Library) version 0.5.2. Finally, install the TxGNN package via pip.

    Note: If you intend to use the disease_area split, you must also install PyG (PyTorch Geometric) following the official PyG installation instructions, as legacy data processing code relies on PyG utility functions.

    conda create --name txgnn_env python=3.8
conda activate txgnn_env
# Install PyTorch via https://pytorch.org/ with your CUDA versions
conda install -c dglteam dgl-cuda{$CUDA_VERSION}==0.5.2
pip install TxGNN

  4. Train and fine-tune the TxGNN model

    main

    The TxGNN workflow typically involves two stages:

    1. Pretraining: Use pretrain to train the model on the dataset.
    2. Finetuning: Use finetune to refine the model.

    Note: For actual use, n_epoch in finetune should typically be set to a higher value (e.g., 500) than the demo value of 30.

    # Pretraining
TxGNN.pretrain(n_epoch = 2, 
               learning_rate = 1e-3,
               batch_size = 1024, 
               train_print_per_n = 20)

# Finetuning
TxGNN.finetune(n_epoch = 30, 
               learning_rate = 5e-4,
               train_print_per_n = 5,
               valid_per_n = 20)

  5. Identify drug negative indices for evaluation

    main

    When evaluating model performance, you may need to identify the indices of drugs that do not have a specific relationship (indication or contraindication) with a disease in the training set. This is done by preparing the split and filtering the training dataframe for the reverse relation (prefixed with rev_).

    # Example: Getting negative drug indices for a specific relation and split
txdata.prepare_split(split = 'autoimmune', seed = 1)

# 'rel' would be 'indication' or 'contraindication'
rel = 'indication'
drug_neg_index = txdata.df_train[txdata.df_train.relation == 'rev_' + rel].y_idx.unique().astype(int)

  6. Evaluate TxGNN models with TxEval

    main

    Use the TxEval class to perform disease-centric evaluations. You can evaluate the entire test set or specific disease indices.

    To retrieve and save graph XAI gates (edge importance) into a .pkl file, use TxGNN.retrieve_save_gates(SAVED_PATH).

    from txgnn import TxEval
TxEval = TxEval(model = TxGNN)

# Evaluate entire test set
result = TxEval.eval_disease_centric(disease_idxs = 'test_set', 
                                     show_plot = False, 
                                     verbose = True, 
                                     save_result = True,
                                     return_raw = False,
                                     save_name = 'SAVE_PATH')

# Evaluate specific diseases
result = TxEval.eval_disease_centric(disease_idxs = [9907.0, 12787.0], 
                                     relation = 'indication', 
                                     save_result = False)

# Retrieve graph XAI gates
gates = TxGNN.retrieve_save_gates('SAVED_PATH')

  7. Initialize TxGNN with TxData and TxGNN

    main

    To use the TxGNN API, first initialize TxData with a path to your data folder, then prepare a data split. Use the resulting TxData object to initialize the TxGNN model class. You can specify the device (e.g., 'cuda:0') and Weights & Biases (wandb) configuration via proj_name and exp_name.

    from txgnn import TxData, TxGNN, TxEval

# Download/load knowledge graph dataset
TxData = TxData(data_folder_path = './data')
TxData.prepare_split(split = 'complex_disease', seed = 42)
TxGNN = TxGNN(data = TxData, 
              weight_bias_track = False,
              proj_name = 'TxGNN', # wandb project name
              exp_name = 'TxGNN', # wandb experiment name
              device = 'cuda:0' # define your cuda device
              )

  8. Configure TxGNN model parameters

    main

    Use TxGNN.model_initialize() to set the architecture and similarity measures for the model.

    Key parameters include:

    • n_hid, n_inp, n_out: Hidden, input, and output dimensions.
    • proto: Boolean to enable the metric learning module.
    • proto_num: Number of similar diseases to retrieve for augmentation.
    • attention: Boolean for using an attention layer (set to False if using graph XAI).
    • sim_measure: Disease signature method. Options: ['all_nodes_profile', 'protein_profile', 'protein_random_walk'].
    • agg_measure: Aggregation method for similar disease embeddings. Options: ['rarity', 'avg'].
    • num_walks: Number of sampled walks (required if sim_measure is 'protein_random_walk').
    • path_length: Path length for walks (required if sim_measure is 'protein_random_walk').
    TxGNN.model_initialize(n_hid = 100, 
                      n_inp = 100, 
                      n_out = 100, 
                      proto = True, 
                      proto_num = 3, 
                      attention = False, 
                      sim_measure = 'all_nodes_profile', 
                      agg_measure = 'rarity', 
                      num_walks = 200, 
                      path_length = 2 
                      )

  9. Train and fine-tune TxGNN models

    main

    TxGNN supports several training stages:

    1. Pre-training: Use TxGNN.pretrain() for link prediction across all edge types.
    2. Fine-tuning: Use TxGNN.finetune() on drug-disease relations using metric learning.
    3. Graph XAI Training: After link prediction training, use TxGNN.train_graphmask() to train a graph explainability model.

    To manage models, use TxGNN.save_model(path) to save weights and TxGNN.load_pretrained(path) to load them. For Graphmask models, use TxGNN.save_graphmask_model(path) and TxGNN.load_pretrained_graphmask(path).

    # Pre-training
TxGNN.pretrain(n_epoch = 2, 
               learning_rate = 1e-3,
               batch_size = 1024, 
               train_print_per_n = 20)

# Fine-tuning
TxGNN.finetune(n_epoch = 500, 
               learning_rate = 5e-4,
               train_print_per_n = 5,
               valid_per_n = 20,
               save_name = finetune_result_path)

# Graph XAI training
TxGNN.train_graphmask(relation = 'indication',
                      learning_rate = 3e-4,
                      allowance = 0.005,
                      epochs_per_layer = 3,
                      penalty_scaling = 1,
                      valid_per_n = 20)

  10. Load pretrained TxGNN checkpoints

    main

    You can load pre-trained model checkpoints for evaluation using the TxGNN.load_pretrained method. This allows you to bypass the training phase and immediately perform zero-shot predictions or evaluations using the weights provided in the repository's distribution.

    # Example usage pattern
model = TxGNN.load_pretrained(checkpoint_path)

  11. Prepare knowledge graph data with TxData

    main

    Use TxData to load your knowledge graph dataset from a specified directory and prepare data splits. You can define splits using the prepare_split method, specifying the split type (e.g., 'complex_disease') and a random seed.

    from txgnn import TxData

TxData = TxData(data_folder_path = '/n/scratch3/users/k/kh278/kg/')
TxData.prepare_split(split = 'complex_disease', seed = 42, no_kg = False)

  12. Evaluate model predictions with TxEval

    main

    Use TxEval to evaluate the performance of a TxGNN model. You can evaluate specific disease indices or the entire test set using eval_disease_centric.

    from txgnn import TxEval

TxEval = TxEval(model = TxGNN)

# Evaluate individual diseases
result = TxEval.eval_disease_centric(disease_idxs = [12661.0, 11318.0], 
                                     relation = 'indication', 
                                     save_result = False)

# Evaluate the entire test set
result = TxEval.eval_disease_centric(disease_idxs = 'test_set', 
                                     show_plot = False, 
                                     verbose = True, 
                                     save_result = True,
                                     return_raw = False)

# Retrieve disease indices for the test set
TxEval.retrieve_disease_idxs_test_set('indication')