txoptimus 0.1.0__py3-none-any.whl

txoptimus/__init__.py

@@ -0,0 +1,8 @@
+"""
+TxOptimus — Zero-shot drug repurposing via TxGNN on PrimeKG and OptimusKG.
+"""
+
+__version__ = "0.1.0"
+__author__ = "arsalanriaz"
+
+from .core import TxData, TxGNN, TxEval

txoptimus/core/TxData.py

@@ -0,0 +1,160 @@
+import os
+import math
+import copy
+import pickle
+
+import numpy as np
+import pandas as pd
+from tqdm.auto import tqdm
+import dgl
+
+from .utils import preprocess_kg, create_split, process_disease_area_split, create_dgl_graph, evaluate_graph_construct, convert2str, data_download_wrapper, reverse_rel_generation
+
+import warnings
+warnings.filterwarnings("ignore")
+
+
+class TxData:
+    
+    def __init__(self, data_folder_path):
+        if not os.path.exists(data_folder_path):
+            os.mkdir(data_folder_path)
+            
+        self.data_folder = data_folder_path # the data folder, contains the kg.csv
+        data_download_wrapper('https://dataverse.harvard.edu/api/access/datafile/7144484', os.path.join(self.data_folder, 'kg.csv'))
+        data_download_wrapper('https://dataverse.harvard.edu/api/access/datafile/7144482', os.path.join(self.data_folder, 'node.csv'))
+        data_download_wrapper('https://dataverse.harvard.edu/api/access/datafile/7144483', os.path.join(self.data_folder, 'edges.csv'))
+        
+        
+    def prepare_split(self, split = 'complex_disease',
+                     disease_eval_idx = None,
+                     seed = 42,
+                     no_kg = False,
+                     test_size = 0.05,
+                     mask_ratio = 0.1, 
+                     one_hop = False):
+        
+        if split not in ['random', 'complex_disease', 'complex_disease_cv', 'disease_eval', 'cell_proliferation', 'mental_health', 'cardiovascular', 'anemia', 'adrenal_gland','autoimmune', 'metabolic_disorder', 'diabetes', 'neurodigenerative', 'full_graph', 'downstream_pred', 'few_edeges_to_kg', 'few_edeges_to_indications']:
+            raise ValueError("Please select one of the following supported splits: 'random', 'complex_disease', 'disease_eval', 'cell_proliferation', 'mental_health', 'cardiovascular', 'anemia', 'adrenal_gland'")
+            
+        if disease_eval_idx is not None:
+            split = 'disease_eval'
+            print('disease eval index is not none, use the individual disease split...')
+        self.split = split
+        
+        if split in ['cell_proliferation', 'mental_health', 'cardiovascular', 'anemia', 'adrenal_gland', 'autoimmune', 'metabolic_disorder', 'diabetes', 'neurodigenerative']:
+            
+            if test_size != 0.05:
+                folder_name = split + '_kg' + '_frac' + str(test_size)
+            elif one_hop:
+                folder_name = split + '_kg' + '_one_hop_ratio' + str(mask_ratio)
+            else:
+                folder_name = split + '_kg'
+           
+            if not os.path.exists(os.path.join(self.data_folder, folder_name)):
+                os.mkdir(os.path.join(self.data_folder, folder_name))
+            kg_path = os.path.join(self.data_folder, folder_name, 'kg_directed.csv')
+        else:
+            kg_path = os.path.join(self.data_folder, 'kg_directed.csv')
+            
+        if os.path.exists(kg_path):
+            print('Found saved processed KG... Loading...')
+            df = pd.read_csv(kg_path)
+        else:
+            if os.path.exists(os.path.join(self.data_folder, 'kg.csv')):
+                print('First time usage... Mapping TxData raw KG to directed csv... it takes several minutes...')
+                preprocess_kg(self.data_folder, split, test_size,one_hop, mask_ratio)
+                df = pd.read_csv(kg_path)
+            else:
+                raise ValueError("KG file path does not exist...")
+        
+        if split == 'disease_eval':
+            split_data_path = os.path.join(self.data_folder, self.split + '_' + str(disease_eval_idx))
+        elif split == 'downstream_pred':
+            split_data_path = os.path.join(self.data_folder, self.split + '_downstream_pred')
+            disease_eval_idx = [11394.,  6353., 12696., 14183., 12895.,  9128., 12623., 15129.,
+                                   12897., 12860.,  7611., 13113.,  4029., 14906., 13438., 13177.,
+                                   13335., 12896., 12879., 12909.,  4815., 12766., 12653.]
+        elif no_kg:
+            split_data_path = os.path.join(self.data_folder, self.split + '_no_kg_' + str(seed))
+        elif test_size != 0.05:
+            split_data_path = os.path.join(self.data_folder, self.split + '_' + str(seed)) + '_frac' + str(test_size)
+        elif one_hop:
+            split_data_path = os.path.join(self.data_folder, self.split + '_' + str(seed)) + '_one_hop_ratio' + str(mask_ratio)
+        else:
+            split_data_path = os.path.join(self.data_folder, self.split + '_' + str(seed))
+        
+        if no_kg:
+            sub_kg = ['off-label use', 'indication', 'contraindication']
+            df = df[df.relation.isin(sub_kg)].reset_index(drop = True)        
+        
+        if not os.path.exists(os.path.join(split_data_path, 'train.csv')):
+            if not os.path.exists(split_data_path):
+                os.mkdir(split_data_path)           
+            print('Creating splits... it takes several minutes...')
+            df_train, df_valid, df_test = create_split(df, split, disease_eval_idx, split_data_path, seed)
+        else:
+            print('Splits detected... Loading splits....')
+            df_train = pd.read_csv(os.path.join(split_data_path, 'train.csv'))
+            df_valid = pd.read_csv(os.path.join(split_data_path, 'valid.csv'))
+            df_test = pd.read_csv(os.path.join(split_data_path, 'test.csv'))
+            
+            # Safeguard: if loaded splits are missing rev_ edges, generate them
+            if not any(df_train.relation.str.startswith('rev_')):
+                print('Safeguard: generating missing rev_ relations for loaded splits...')
+                unique_rel = df[['x_type', 'relation', 'y_type']].drop_duplicates()
+                df_train = reverse_rel_generation(df, df_train, unique_rel)
+                df_valid = reverse_rel_generation(df, df_valid, unique_rel)
+                df_test = reverse_rel_generation(df, df_test, unique_rel)
+
+        if split not in ['random', 'complex_disease', 'complex_disease_cv', 'disease_eval', 'full_graph', 'downstream_pred', 'few_edeges_to_indications', 'few_edeges_to_kg']:
+            # in disease area split
+            df_test = process_disease_area_split(self.data_folder, df, df_test, split)
+        
+        print('Creating DGL graph....')
+        # create dgl graph
+        g = create_dgl_graph(df_train, df)
+         
+        self.G = g         
+        self.df, self.df_train, self.df_valid, self.df_test = df, df_train, df_valid, df_test
+        self.disease_eval_idx = disease_eval_idx
+        self.no_kg = no_kg
+        self.seed = seed
+        print('Done!')
+        
+        
+    def retrieve_id_mapping(self):
+        df = self.df
+        df['x_id'] = df['x_id'].apply(lambda x: convert2str(x))
+        df['y_id'] = df['y_id'].apply(lambda x: convert2str(x))
+
+        idx2id_drug = dict(df[df.x_type == 'drug'][['x_idx', 'x_id']].drop_duplicates().values)
+        idx2id_drug.update(dict(df[df.y_type == 'drug'][['y_idx', 'y_id']].drop_duplicates().values))
+
+        idx2id_disease = dict(df[df.x_type == 'disease'][['x_idx', 'x_id']].drop_duplicates().values)
+        idx2id_disease.update(dict(df[df.y_type == 'disease'][['y_idx', 'y_id']].drop_duplicates().values))
+
+        # Only read columns needed for name lookups, and filter to drug/disease rows
+        # This avoids loading all 21.8M rows into memory (~6GB)
+        df_ = pd.read_csv(os.path.join(self.data_folder, 'kg.csv'),
+                          usecols=['x_id', 'x_type', 'x_name', 'y_id', 'y_type', 'y_name'])
+        df_dd = df_[df_.x_type.isin(['drug', 'disease']) | df_.y_type.isin(['drug', 'disease'])].copy()
+        
+        df_dd['x_id'] = df_dd['x_id'].apply(lambda x: convert2str(x))
+        df_dd['y_id'] = df_dd['y_id'].apply(lambda x: convert2str(x))
+
+        id2name_disease = dict(zip(df_dd.loc[df_dd.x_type == 'disease', 'x_id'], 
+                                    df_dd.loc[df_dd.x_type == 'disease', 'x_name']))
+        id2name_disease.update(dict(zip(df_dd.loc[df_dd.y_type == 'disease', 'y_id'],
+                                         df_dd.loc[df_dd.y_type == 'disease', 'y_name'])))
+
+        id2name_drug = dict(zip(df_dd.loc[df_dd.x_type == 'drug', 'x_id'],
+                                 df_dd.loc[df_dd.x_type == 'drug', 'x_name']))
+        id2name_drug.update(dict(zip(df_dd.loc[df_dd.y_type == 'drug', 'y_id'],
+                                       df_dd.loc[df_dd.y_type == 'drug', 'y_name'])))
+        
+        return {'id2name_drug': id2name_drug,
+                'id2name_disease': id2name_disease,
+                'idx2id_disease': idx2id_disease,
+                'idx2id_drug': idx2id_drug
+               }

txoptimus/core/TxEval.py

@@ -0,0 +1,42 @@
+from .utils import *
+
+class TxEval:
+    
+    def __init__(self, model):
+        self.df, self.df_train, self.df_valid, self.df_test, self.data_folder, self.G, self.best_model, self.weight_bias_track, self.wandb = model.df, model.df_train, model.df_valid, model.df_test, model.data_folder, model.G, model.best_model, model.weight_bias_track, model.wandb
+        self.device = model.device
+        self.disease_rel_types = ['rev_contraindication', 'rev_indication', 'rev_off-label use']
+        self.split = model.split
+        
+    def eval_disease_centric(self, disease_idxs, relation = None, save_result = False, show_plot = False, verbose = False, save_name = None, return_raw = False, simulate_random = True):
+        if self.split == 'full_graph':
+            # set only_prediction to True during full graph training
+            only_prediction = True
+        else:
+            only_prediction = False
+            
+        if disease_idxs == 'test_set':
+            disease_idxs = None
+        
+        self.out = disease_centric_evaluation(self.df, self.df_train, self.df_valid, self.df_test, self.data_folder, self.G, self.best_model,self.device, disease_idxs, relation, self.weight_bias_track, self.wandb, show_plot, verbose, return_raw, simulate_random, only_prediction)
+        
+        if save_result:
+            import pickle, os
+            if save_name is None:
+                save_name = os.path.join(self.data_folder, 'disease_centric_eval.pkl')
+            with open(save_name, 'wb') as f:
+                pickle.dump(self.out, f)
+        return self.out
+    
+    def retrieve_disease_idxs_test_set(self, relation):
+        relation = 'rev_' + relation
+        df_train_valid = pd.concat([self.df_train, self.df_valid])
+        df_dd = self.df_test[self.df_test.relation.isin(self.disease_rel_types)]
+        df_dd_train = df_train_valid[df_train_valid.relation.isin(self.disease_rel_types)]
+
+        df_rel_dd = df_dd[df_dd.relation == relation]        
+        return df_rel_dd.x_idx.unique()
+    
+    
+    def retrieve_all_disease_idxs(self):
+        return np.unique(self.df[self.df.x_type == 'disease'].x_idx.unique().tolist() + self.df[self.df.y_type == 'disease'].y_idx.unique().tolist())