""" Drug Repositioning Prediction Script for REDDA Model This script performs drug repurposing prediction using the trained REDDA model. It takes an input CSV with drug SMILES and calculates similarity to existing drugs in the knowledge graph, then predicts drug-disease associations. Usage: python prediction.py --input_csv --target_disease --output_csv """ import os import sys import argparse import numpy as np import pandas as pd import torch as th from rdkit import Chem from rdkit.Chem import DataStructs, rdMolDescriptors from rdkit import RDLogger import dgl from tqdm import tqdm import warnings warnings.filterwarnings('ignore') # Suppress RDKit warnings RDLogger.DisableLog('rdApp.*') # Import local modules from model import Model from load_data import load from utils import set_seed def calculate_tanimoto_similarity(smiles1, smiles2): """ Calculate Tanimoto similarity between two SMILES strings using Morgan fingerprints. Args: smiles1 (str): First SMILES string smiles2 (str): Second SMILES string Returns: float: Tanimoto similarity score (0-1) """ try: mol1 = Chem.MolFromSmiles(smiles1) mol2 = Chem.MolFromSmiles(smiles2) if mol1 is None or mol2 is None: return 0.0 # Generate Morgan fingerprints (radius=2, 2048 bits) fp1 = rdMolDescriptors.GetMorganFingerprintAsBitVect(mol1, radius=2, nBits=2048) fp2 = rdMolDescriptors.GetMorganFingerprintAsBitVect(mol2, radius=2, nBits=2048) # Calculate Tanimoto similarity similarity = DataStructs.TanimotoSimilarity(fp1, fp2) return similarity except: return 0.0 from concurrent.futures import ThreadPoolExecutor from functools import partial def calculate_similarity_for_drug(query_smiles, row): """ Calculate similarity for a single drug. Args: query_smiles (str): Query drug SMILES row (pd.Series): Drug row from DataFrame Returns: tuple: (drug_id, similarity_score) or None if invalid SMILES """ kg_smiles = row['SMILES'] drug_id = row['Inter_ID'] # Skip if SMILES is empty or NaN if pd.isna(kg_smiles) or kg_smiles == '': return None similarity = calculate_tanimoto_similarity(query_smiles, kg_smiles) if similarity > 0: # Only keep non-zero similarities return (drug_id, similarity) return None def find_similar_drugs(query_smiles, drug_df, top_k=20, n_jobs=8): """ Find top-k most similar drugs to the query SMILES from the knowledge graph. Uses parallel processing for faster computation. Args: query_smiles (str): Query drug SMILES drug_df (pd.DataFrame): DataFrame containing drug information with SMILES top_k (int): Number of top similar drugs to return n_jobs (int): Number of parallel threads to use Returns: list: List of tuples (drug_inter_id, similarity_score) """ # Create a partial function with fixed query_smiles calc_func = partial(calculate_similarity_for_drug, query_smiles) # Calculate similarities in parallel similarities = [] with ThreadPoolExecutor(max_workers=n_jobs) as executor: results = list(executor.map(calc_func, [row for _, row in drug_df.iterrows()])) # Filter out None results and sort similarities = [x for x in results if x is not None] similarities.sort(key=lambda x: x[1], reverse=True) return similarities[:top_k] def create_drug_id_mapping(drug_df): """ Create mapping from Inter_ID to index for drugs. Args: drug_df (pd.DataFrame): Drug dataframe Returns: dict: Mapping from Inter_ID to index """ # Create mapping and also store the reverse mapping id_to_idx = {drug_id: idx for idx, drug_id in enumerate(drug_df['Inter_ID'])} return id_to_idx def create_disease_id_mapping(disease_df): """ Create mapping from Inter_ID to index for diseases. Args: disease_df (pd.DataFrame): Disease dataframe Returns: dict: Mapping from Inter_ID to index """ return {disease_id: idx for idx, disease_id in enumerate(disease_df['Inter_ID'])} def augment_graph_with_new_drugs(input_drugs_df, drug_df, drug_id_mapping, kg_dir, output_folder): """ Augment the graph by creating temporary node and edge files with new drugs. Args: input_drugs_df (pd.DataFrame): Input drugs DataFrame drug_df (pd.DataFrame): Knowledge graph drugs DataFrame drug_id_mapping (dict): Mapping from drug ID to index kg_dir (str): Base directory of knowledge graph data output_folder (str): Path to output folder Returns: tuple: (new_drug_indices, node_path, edge_path) """ global temp_dir print("Augmenting graph with new drugs...") # Get the last drug index last_drug_id = int(drug_df['Inter_ID'].str.extract(r'DRUG_(\d+)').astype(int).max()) # Create new drug entries new_drugs = [] new_drug_indices = {} # Maps input index to numerical index in augmented graph base_idx = len(drug_df) # Start new indices after existing drugs for idx, row in input_drugs_df.iterrows(): new_drug_id = f"DRUG_{last_drug_id + idx + 1}" new_drug_indices[idx] = base_idx + idx # Store numerical index new_drug = { 'id': f"NEW_{new_drug_id}", 'name': f"New_Drug_{idx}", 'source': 'Prediction', 'type': 'drug', 'SMILES': row['SMILES'], 'Inter_ID': new_drug_id } new_drugs.append(new_drug) # Create augmented drug node file augmented_drug_df = pd.concat([ drug_df, pd.DataFrame(new_drugs) ], ignore_index=True) # Create new edges based on similarity new_edges = [] # Add similarity-based edges for new drugs for idx, row in tqdm(input_drugs_df.iterrows(), total=len(input_drugs_df), desc="Finding similar drugs"): query_smiles = row['SMILES'] new_drug_id = new_drug_indices[idx] # Find top-20 similar drugs similar_drugs = find_similar_drugs(query_smiles, drug_df, top_k=20) for similar_drug_id, similarity in similar_drugs: # Get similar drug info similar_drug_row = drug_df[drug_df['Inter_ID'] == similar_drug_id].iloc[0] similar_drug_name = similar_drug_row['name'] similar_drug_idx = drug_id_mapping[similar_drug_id] # Add bidirectional edges with similarity score new_edges.append({ 'relation': 'drug_drug', 'display_relation': 'structural_similarity', 'x_index': base_idx + idx, # New drug index 'x_id': f'DRUG_{new_drug_id}', 'x_type': 'drug', 'x_name': f"New_Drug_{idx}", 'x_source': 'Prediction', 'y_index': similar_drug_idx, # Existing drug index 'y_id': similar_drug_id, 'y_type': 'drug', 'y_name': similar_drug_name, 'y_source': 'DrugBank', 'similarity_score': similarity }) new_edges.append({ 'relation': 'drug_drug', 'display_relation': 'structural_similarity', 'x_index': similar_drug_idx, # Existing drug index 'x_id': similar_drug_id, 'x_type': 'drug', 'x_name': similar_drug_name, 'x_source': 'DrugBank', 'y_index': base_idx + idx, # New drug index 'y_id': f'DRUG_{new_drug_id}', 'y_type': 'drug', 'y_name': f"New_Drug_{idx}", 'y_source': 'Prediction', 'similarity_score': similarity }) # Create new edge DataFrame with correct format new_edge_df = pd.DataFrame(new_edges) # Create temporary directory in output folder temp_dir = os.path.join(output_folder, 'temp') os.makedirs(temp_dir, exist_ok=True) # Save augmented drug nodes and edges temp_node_path = os.path.join(temp_dir, 'drug.csv') temp_edge_path = os.path.join(temp_dir, 'drug_drug.csv') augmented_drug_df.to_csv(temp_node_path, index=False) # Read original drug_drug edges and combine with new edges orig_edges = pd.read_csv(os.path.join(kg_dir, 'edge/drug_drug.csv')) # Ensure new edges have all required columns for col in orig_edges.columns: if col not in new_edge_df.columns: new_edge_df[col] = '' # Fill missing columns with empty string # Combine original and new edges combined_edges = pd.concat([orig_edges, new_edge_df], ignore_index=True) # Save combined edges combined_edges.to_csv(temp_edge_path, index=False) # Create symlinks for all other files for file_name in os.listdir(os.path.join(kg_dir, 'node')): if file_name != 'drug.csv': src = os.path.join(kg_dir, 'node', file_name) dst = os.path.join(temp_dir, file_name) if not os.path.exists(dst): os.symlink(src, dst) for file_name in os.listdir(os.path.join(kg_dir, 'edge')): if file_name != 'drug_drug.csv': src = os.path.join(kg_dir, 'edge', file_name) dst = os.path.join(temp_dir, file_name) if not os.path.exists(dst): os.symlink(src, dst) return new_drug_indices, temp_dir, temp_dir def load_models(model_dir, g, feature_dim, device, hidden_feats=128, num_heads=5, dropout=0.0): """ Load all 5 trained model checkpoints. Args: model_dir (str): Directory containing model checkpoints g (dgl.DGLHeteroGraph): Graph for model initialization feature_dim (int): Feature dimension device (torch.device): Device to load models on hidden_feats (int): Hidden layer dimension (default: 64) num_heads (int): Number of attention heads (default: 5) dropout (float): Dropout rate (default: 0.0) Returns: list: List of loaded models """ models = [] for fold in range(1, 6): # model_1.pth to model_5.pth model_path = os.path.join(model_dir, f"model_{fold}.pth") if not os.path.exists(model_path): print(f"Warning: Model checkpoint {model_path} not found") continue # Initialize model with same architecture as training model = Model( etypes=g.etypes, ntypes=g.ntypes, in_feats=feature_dim, hidden_feats=hidden_feats, num_heads=num_heads, dropout=dropout, ).to(device) # Load checkpoint checkpoint = th.load(model_path, map_location=device) model.load_state_dict(checkpoint) model.eval() models.append(model) print(f"Loaded model from {model_path}") return models def ensemble_predict(models, g, feature, device): """ Perform ensemble prediction using multiple model checkpoints. Args: models (list): List of trained models g (dgl.DGLHeteroGraph): Input graph feature (dict): Node features device (torch.device): Device for computation Returns: tuple: (mean_predictions, std_predictions) """ predictions = [] print("Performing ensemble prediction...") with th.no_grad(): for i, model in enumerate(models): print(f"Predicting with model {i+1}/{len(models)}") pred = th.sigmoid(model(g, feature)).cpu().numpy() predictions.append(pred) # Calculate mean and standard deviation predictions = np.array(predictions) mean_pred = np.mean(predictions, axis=0) std_pred = np.std(predictions, axis=0) return mean_pred, std_pred # Global variables temp_dir = None def main(): global temp_dir parser = argparse.ArgumentParser(description='REDDA Drug Repositioning Prediction') parser.add_argument('--input_csv', type=str, required=True, help='Path to input CSV file with drug SMILES') parser.add_argument('--target_disease', type=str, required=True, help='Target disease Inter_ID (e.g., DISEASE_123)') parser.add_argument('--output_folder', type=str, required=True, help='Path to output folder for results and temporary files') parser.add_argument('--device_id', type=str, default=None, help='GPU device ID (default: CPU)') parser.add_argument('--dataset', type=str, default='KGdataset', choices=['KGdataset', 'KGdataset_tiny'], help='Dataset to use') parser.add_argument('--model_dir', type=str, default='result_kg_42', help='Directory containing model checkpoints') pred_args = parser.parse_args() # Set device if pred_args.device_id: device = th.device(f"cuda:{pred_args.device_id}") print(f"Using GPU: {pred_args.device_id}") else: device = th.device("cpu") print("Using CPU") # Set seed for reproducibility set_seed(42) print("Loading input data...") # Load input CSV input_df = pd.read_csv(pred_args.input_csv) # Validate input CSV format required_columns = ['SMILES'] if not all(col in input_df.columns for col in required_columns): raise ValueError(f"Input CSV must contain columns: {required_columns}") print(f"Input contains {len(input_df)} drugs") # Load knowledge graph data print("Loading knowledge graph data...") drug_df = pd.read_csv('/vast/yg3191/AIVS/kg/node/drug.csv') disease_df = pd.read_csv('/vast/yg3191/AIVS/kg/node/disease.csv') # Create ID mappings drug_id_mapping = create_drug_id_mapping(drug_df) disease_id_mapping = create_disease_id_mapping(disease_df) # Validate target disease if pred_args.target_disease not in disease_id_mapping: available_diseases = list(disease_id_mapping.keys())[:10] # Show first 10 raise ValueError(f"Target disease '{pred_args.target_disease}' not found in knowledge graph. " f"Available diseases (first 10): {available_diseases}") target_disease_idx = disease_id_mapping[pred_args.target_disease] print(f"Target disease: {pred_args.target_disease} (index: {target_disease_idx})") # Load original graph print("Loading original graph...") g = load(pred_args.dataset).to(device) print(f"Original graph: {g.num_nodes('drug')} drugs, {g.num_nodes('disease')} diseases") # Create output folder if not exists os.makedirs(pred_args.output_folder, exist_ok=True) # Create temporary files with augmented graph data kg_dir = '/vast/yg3191/AIVS/kg' new_drug_indices, node_path, edge_path = augment_graph_with_new_drugs( input_df, drug_df, drug_id_mapping, kg_dir, pred_args.output_folder ) # Build augmented graph using _build_heterograph from load_data import _build_heterograph # Define node and edge files (same as in load_KGdataset) node_files = { 'drug': 'drug.csv', 'disease': 'disease.csv', 'protein': 'protein.csv', # 'bioprocess': 'bioprocess.csv', # 'cellcomp': 'cellcomp.csv', # 'molfunc': 'molfunc.csv', # 'pathway': 'pathway.csv', # 'phenotype': 'phenotype.csv', # 'exposure': 'exposure.csv', # 'effect': 'effect.csv', } edge_files = { 'drug_drug': 'drug_drug.csv', # 'drug_effect': 'drug_effect.csv', 'drug_protein': 'drug_protein.csv', 'drug_disease': 'drug_disease_indication.csv', 'protein_protein': 'protein_protein.csv', # 'protein_bioprocess': 'protein_bioprocess.csv', # 'protein_cellcomp': 'protein_cellcomp.csv', # 'protein_molfunc': 'protein_molfunc.csv', # 'protein_pathway': 'protein_pathway.csv', 'protein_disease': 'protein_disease.csv', 'disease_disease': 'disease_disease.csv', # 'disease_phenotype': 'disease_phenotype_positive.csv', # 'disease_exposure': 'disease_exposure.csv', # 'bioprocess_bioprocess': 'bioprocess_bioprocess.csv', # 'cellcomp_cellcomp': 'cellcomp_cellcomp.csv', # 'molfunc_molfunc': 'molfunc_molfunc.csv', # 'pathway_pathway': 'pathway_pathway.csv', # 'phenotype_phenotype': 'phenotype_phenotype.csv', } augmented_g = _build_heterograph( node_path=node_path, edge_path=edge_path, node_files=node_files, edge_files=edge_files, feature_dim=128 ).to(device) # Extract features feature = {ntype: augmented_g.nodes[ntype].data['h'] for ntype in augmented_g.ntypes} feature_dim = feature['drug'].shape[1] # Load trained models print("Loading trained models...") # Use default model parameters that match the training configuration models = load_models(pred_args.model_dir, augmented_g, feature_dim, device, hidden_feats=128, num_heads=8, dropout=0.4) if not models: raise ValueError("No model checkpoints found!") print(f"Loaded {len(models)} model checkpoints") # Perform ensemble prediction mean_pred, std_pred = ensemble_predict(models, augmented_g, feature, device) # Extract predictions for input drugs and target disease print("Extracting predictions...") results = [] print(mean_pred.shape) print(std_pred.shape) print(new_drug_indices) for idx, row in input_df.iterrows(): # Get the numerical index for this drug in the augmented graph drug_idx = new_drug_indices[idx] # This is now a numerical index # Get prediction Score for this drug-disease pair Score_mean = float(mean_pred[drug_idx, target_disease_idx]) Score_std = float(std_pred[drug_idx, target_disease_idx]) # Create result row result_row = row.copy() result_row['Target_Disease'] = pred_args.target_disease result_row['Score_Mean'] = Score_mean result_row['Score_Std'] = Score_std results.append(result_row) # Create output DataFrame output_df = pd.DataFrame(results) # Create output folder if not exists os.makedirs(pred_args.output_folder, exist_ok=True) # Save results output_path = os.path.join(pred_args.output_folder, 'result.csv') output_df.to_csv(output_path, index=False) print(f"Predictions saved to: {output_path}") # Print summary statistics print("\nPrediction Summary:") print(f"Mean Score: {output_df['Score_Mean'].mean():.4f}") print(f"Std Score: {output_df['Score_Mean'].std():.4f}") print(f"Max Score: {output_df['Score_Mean'].max():.4f}") print(f"Min Score: {output_df['Score_Mean'].min():.4f}") # Show top predictions print(f"\nTop 5 predictions:") top_predictions = output_df.nlargest(5, 'Score_Mean') for idx, row in top_predictions.iterrows(): print(f" SMILES: {row['SMILES'][:50]}... | Score: {row['Score_Mean']:.4f} ± {row['Score_Std']:.4f}") # Clean up temporary files (optional, comment out to keep temp files for debugging) try: # Remove symlinks first for file_name in os.listdir(temp_dir): file_path = os.path.join(temp_dir, file_name) if os.path.islink(file_path): os.unlink(file_path) elif os.path.isfile(file_path): os.remove(file_path) # Remove temp directory os.rmdir(temp_dir) print("\nTemporary files cleaned up") except Exception as e: print(f"\nWarning: Could not clean up temporary files: {e}") print(f"Temporary files remain in: {temp_dir}") if __name__ == "__main__": main()