Browse Source

Upload files to 'case_study'

master
Zahra Asgari 6 days ago
parent
commit
96bf72c8b4
2 changed files with 598 additions and 0 deletions
  1. 172
    0
      case_study/main_case_study.py
  2. 426
    0
      case_study/model.py

+ 172
- 0
case_study/main_case_study.py View File

@@ -0,0 +1,172 @@
import argparse
import numpy as np
import pandas as pd
import torch
import scipy.sparse as sp
from sklearn.model_selection import KFold
from sklearn.metrics import roc_auc_score, average_precision_score
from model import DeepTraCDR, Optimizer
from utils import evaluate_auc, common_data_index
from data_sampler import TargetSampler
from data_loader import load_data
# Clear CUDA cache to optimize memory usage
torch.cuda.empty_cache()
def main():
# Parse command-line arguments for model configuration
parser = argparse.ArgumentParser(description="DeepTraCDR Case Study for Drug Response Prediction")
parser.add_argument('-device', type=str, default="cuda:0" if torch.cuda.is_available() else "cpu",
help="Device to run the model (cuda:0 or cpu)")
parser.add_argument('-data', type=str, default='gdsc',
help="Dataset to use (e.g., gdsc or ccle)")
parser.add_argument('--wd', type=float, default=1e-4, help="Weight decay for optimizer")
parser.add_argument('--layer_size', nargs='+', type=int, default=[512],
help="List of layer sizes for the GCN model")
parser.add_argument('--gamma', type=float, default=15, help="Gamma parameter for loss function")
parser.add_argument('--epochs', type=int, default=1000, help="Number of training epochs")
parser.add_argument('--test_freq', type=int, default=50, help="Frequency of evaluation during training")
parser.add_argument('--patience', type=int, default=100, help="Patience for early stopping")
parser.add_argument('--lr', type=float, default=0.0005, help="Learning rate for optimizer")
parser.add_argument('--k_fold', type=int, default=5, help="Number of folds for cross-validation")
args = parser.parse_args()
# Load dataset-specific drug response data
if args.data == "gdsc":
# Define target drug CIDs (e.g., Dasatinib=5330286, GSK690693=11338033)
target_drug_cids = np.array([5330286, 11338033, 24825971])
# Load cell-drug binary response matrix
cell_drug = pd.read_csv(
"/media/external_16TB_1/ali_kianfar/Data/GDSC/cell_drug_binary.csv",
index_col=0, header=0
)
cell_drug.columns = cell_drug.columns.astype(np.int32)
drug_cids = cell_drug.columns.values
# Extract target drug responses and compute positive sample count
cell_target_drug = np.array(cell_drug.loc[:, target_drug_cids], dtype=np.float32)
target_pos_num = sp.coo_matrix(cell_target_drug).data.shape[0]
target_indexes = common_data_index(drug_cids, target_drug_cids)
elif args.data == "ccle":
# Define target drug CIDs for CCLE dataset
target_drug_cids = np.array([5330286])
# Load cell-drug binary response matrix
cell_drug = pd.read_csv(
"/media/external_16TB_1/ali_kianfar/Data/CCLE/cell_drug_binary.csv",
index_col=0, header=0
)
cell_drug.columns = cell_drug.columns.astype(np.int32)
drug_cids = cell_drug.columns.values
# Extract target drug responses and compute positive sample count
cell_target_drug = np.array(cell_drug.loc[:, target_drug_cids], dtype=np.float32)
target_pos_num = sp.coo_matrix(cell_target_drug).data.shape[0]
target_indexes = common_data_index(drug_cids, target_drug_cids)
# Load additional data (adjacency matrix, fingerprints, expression, etc.)
full_adj, drug_fingerprints, exprs, null_mask, pos_num, args = load_data(args)
full_adj_np = full_adj.copy() # Copy for sampler usage
# Print data shapes for verification
print(f"Adjacency matrix shape: {full_adj.shape}")
print(f"Expression data shape: {exprs.shape}")
print(f"Null mask shape: {null_mask.shape}")
# Convert adjacency matrix to PyTorch tensor
if isinstance(full_adj, np.ndarray):
full_adj = torch.from_numpy(full_adj).float().to(args.device)
# Initialize k-fold cross-validation
k = args.k_fold
n_kfolds = 5 # Number of k-fold iterations
all_metrics = {'auc': [], 'auprc': []}
# Perform k-fold cross-validation
for n_kfold in range(n_kfolds):
kfold = KFold(n_splits=k, shuffle=True, random_state=n_kfold)
idx_all = np.arange(target_pos_num)
for fold, (train_idx, test_idx) in enumerate(kfold.split(idx_all)):
print(f"\n--- Fold {fold+1}/{k} (Iteration {n_kfold+1}/{n_kfolds}) ---")
# Initialize data sampler for training and testing
sampler = TargetSampler(
response_mat=full_adj_np,
null_mask=null_mask,
target_indexes=target_indexes,
pos_train_index=train_idx,
pos_test_index=test_idx
)
# Initialize DeepTraCDR model
model = DeepTraCDR(
adj_mat=full_adj,
cell_exprs=exprs,
drug_finger=drug_fingerprints,
layer_size=args.layer_size,
gamma=args.gamma,
device=args.device
)
# Initialize optimizer for training
opt = Optimizer(
model=model,
train_data=sampler.train_data,
test_data=sampler.test_data,
test_mask=sampler.test_mask,
train_mask=sampler.train_mask,
adj_matrix=full_adj,
evaluate_fun=evaluate_auc,
lr=args.lr,
wd=args.wd,
epochs=args.epochs,
test_freq=args.test_freq,
patience=args.patience,
device=args.device
)
# Train the model and retrieve best metrics
true, pred, best_auc, best_auprc = opt.train()
all_metrics['auc'].append(best_auc)
all_metrics['auprc'].append(best_auprc)
print(f"Fold {fold+1}: AUC={best_auc:.4f}, AUPRC={best_auprc:.4f}")
# Compute and display average metrics across all folds
print(f"\nFinal Average Metrics (Across {n_kfolds*k} Folds):")
for metric, values in all_metrics.items():
mean = np.mean(values)
std = np.std(values)
print(f"{metric.upper()}: {mean:.4f} ± {std:.4f}")
# Perform case study: Predict missing responses for target drugs
print("\n--- Case Study: Predicting Missing Responses for Target Drugs ---")
model.eval()
with torch.no_grad():
final_pred, cell_emb, drug_emb = model() # Shape: [num_cells, num_drugs]
# Create a DataFrame to sort cell lines by predicted sensitivity
num_cells, num_drugs = final_pred.size()
cell_names = cell_drug.index.values # Cell line names
cid_list = cell_drug.columns.values # Drug CIDs
# Identify top 10 sensitive cell lines for each target drug
for d in range(num_drugs):
cid = cid_list[d]
if cid in [5330286, 11338033]: # Focus on Dasatinib or GSK690693
drug_preds = final_pred[:, d].cpu().numpy()
sorted_idx = np.argsort(-drug_preds) # Sort in descending order
top_10_cells = [(cell_names[i], drug_preds[i]) for i in sorted_idx[:10]]
drug_name = "Dasatinib" if cid == 5330286 else "GSK690693"
print(f"\nTop 10 Sensitive Cell Lines for {drug_name} (CID={cid}):")
for rank, (cell, score) in enumerate(top_10_cells, start=1):
print(f"{rank}. Cell: {cell}, Score: {score:.4f}")
if __name__ == "__main__":
# Set high precision for matrix multiplication
torch.set_float32_matmul_precision('high')
main()

+ 426
- 0
case_study/model.py View File

@@ -0,0 +1,426 @@
# model.py
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import MultiheadAttention, TransformerEncoder, TransformerEncoderLayer
from utils import torch_corr_x_y, cross_entropy_loss, prototypical_loss
from sklearn.metrics import roc_auc_score, average_precision_score
class ConstructAdjMatrix(nn.Module):
"""Module to construct normalized adjacency matrices for graph-based operations."""
def __init__(self, original_adj_mat, device="cpu"):
"""
Initialize the adjacency matrix constructor.
Args:
original_adj_mat (np.ndarray or torch.Tensor): Original adjacency matrix.
device (str): Device to perform computations on (e.g., 'cpu' or 'cuda:0').
"""
super().__init__()
# Convert NumPy array to PyTorch tensor if necessary
if isinstance(original_adj_mat, np.ndarray):
original_adj_mat = torch.from_numpy(original_adj_mat).float()
self.adj = original_adj_mat.to(device)
self.device = device
def forward(self):
"""
Compute normalized adjacency matrices for cells and drugs.
Returns:
tuple: (agg_cell_lp, agg_drug_lp, self_cell_lp, self_drug_lp) normalized matrices.
"""
with torch.no_grad():
# Compute degree-normalized matrices
d_x = torch.diag(torch.pow(torch.sum(self.adj, dim=1) + 1, -0.5))
d_y = torch.diag(torch.pow(torch.sum(self.adj, dim=0) + 1, -0.5))
agg_cell_lp = torch.mm(torch.mm(d_x, self.adj), d_y)
agg_drug_lp = torch.mm(torch.mm(d_y, self.adj.T), d_x)
self_cell_lp = torch.diag(torch.add(torch.pow(torch.sum(self.adj, dim=1) + 1, -1), 1))
self_drug_lp = torch.diag(torch.add(torch.pow(torch.sum(self.adj, dim=0) + 1, -1), 1))
return (
agg_cell_lp.to(self.device),
agg_drug_lp.to(self.device),
self_cell_lp.to(self.device),
self_drug_lp.to(self.device)
)
class LoadFeature(nn.Module):
"""Module to load and preprocess cell and drug features."""
def __init__(self, cell_exprs, drug_fingerprints, device="cpu"):
"""
Initialize feature loading and preprocessing layers.
Args:
cell_exprs (np.ndarray): Cell expression data.
drug_fingerprints (list): List of drug fingerprint arrays.
device (str): Device to perform computations on.
"""
super().__init__()
self.device = device
self.cell_exprs = torch.from_numpy(cell_exprs).float().to(device)
self.drug_fingerprints = [torch.from_numpy(fp).float().to(device) for fp in drug_fingerprints]
# Projection layers for drug fingerprints
self.drug_proj = nn.ModuleList([
nn.Sequential(
nn.Linear(fp.shape[1], 512),
nn.BatchNorm1d(512),
nn.GELU(),
nn.Dropout(0.3)
).to(device) for fp in drug_fingerprints
])
# Transformer encoder for drug features
self.transformer = TransformerEncoder(
TransformerEncoderLayer(
d_model=512,
nhead=8,
dim_feedforward=2048,
batch_first=True
),
num_layers=3
).to(device)
# Normalization layers
self.cell_norm = nn.LayerNorm(cell_exprs.shape[1]).to(device)
self.drug_norm = nn.LayerNorm(512).to(device)
# Encoder for cell features
self.cell_encoder = nn.Sequential(
nn.Linear(cell_exprs.shape[1], 1024),
nn.BatchNorm1d(1024),
nn.GELU(),
nn.Dropout(0.3),
nn.Linear(1024, 512)
).to(device)
def forward(self):
"""
Process cell and drug features through encoding and transformation.
Returns:
tuple: (cell_encoded, drug_feat) encoded cell and drug features.
"""
# Normalize and encode cell features
cell_feat = self.cell_norm(self.cell_exprs)
cell_encoded = self.cell_encoder(cell_feat)
# Project and transform drug fingerprints
projected = [proj(fp) for proj, fp in zip(self.drug_proj, self.drug_fingerprints)]
stacked = torch.stack(projected, dim=1)
drug_feat = self.transformer(stacked)
drug_feat = self.drug_norm(drug_feat.mean(dim=1))
return cell_encoded, drug_feat
class GEncoder(nn.Module):
"""Graph encoder module for processing cell and drug embeddings."""
def __init__(self, agg_c_lp, agg_d_lp, self_c_lp, self_d_lp, device="cpu"):
"""
Initialize the graph encoder.
Args:
agg_c_lp (torch.Tensor): Aggregated cell Laplacian matrix.
agg_d_lp (torch.Tensor): Aggregated drug Laplacian matrix.
self_c_lp (torch.Tensor): Self-loop cell Laplacian matrix.
self_d_lp (torch.Tensor): Self-loop drug Laplacian matrix.
device (str): Device to perform computations on.
"""
super().__init__()
self.agg_c_lp = agg_c_lp
self.agg_d_lp = agg_d_lp
self.self_c_lp = self_c_lp
self.self_d_lp = self_d_lp
self.device = device
# Cell feature encoder
self.cell_encoder = nn.Sequential(
nn.Linear(512, 1024),
nn.BatchNorm1d(1024),
nn.GELU(),
nn.Dropout(0.3),
nn.Linear(1024, 512)
).to(device)
# Drug feature encoder
self.drug_encoder = nn.Sequential(
nn.Linear(512, 1024),
nn.BatchNorm1d(1024),
nn.GELU(),
nn.Dropout(0.3),
nn.Linear(1024, 512)
).to(device)
# Attention mechanism
self.attention = MultiheadAttention(embed_dim=512, num_heads=8, batch_first=True).to(device)
self.residual = nn.Linear(512, 512).to(device)
self.fc = nn.Sequential(
nn.Linear(1024, 512),
nn.BatchNorm1d(512),
nn.GELU(),
nn.Dropout(0.2)
).to(device)
def forward(self, cell_f, drug_f):
"""
Encode cell and drug features using graph-based aggregation and attention.
Args:
cell_f (torch.Tensor): Cell features.
drug_f (torch.Tensor): Drug features.
Returns:
tuple: (cell_emb, drug_emb) encoded embeddings.
"""
# Aggregate features using Laplacian matrices
cell_agg = torch.mm(self.agg_c_lp, drug_f) + torch.mm(self.self_c_lp, cell_f)
drug_agg = torch.mm(self.agg_d_lp, cell_f) + torch.mm(self.self_d_lp, drug_f)
# Encode aggregated features
cell_fc = self.cell_encoder(cell_agg)
drug_fc = self.drug_encoder(drug_agg)
# Apply attention mechanism
attn_output, _ = self.attention(
query=cell_fc.unsqueeze(0),
key=drug_fc.unsqueeze(0),
value=drug_fc.unsqueeze(0)
)
attn_output = attn_output.squeeze(0)
cell_emb = cell_fc + self.residual(attn_output)
# Apply final activation
cell_emb = F.gelu(cell_emb)
drug_emb = F.gelu(drug_fc)
return cell_emb, drug_emb
class GDecoder(nn.Module):
"""Decoder module to predict interaction scores from embeddings."""
def __init__(self, emb_dim, gamma):
"""
Initialize the decoder.
Args:
emb_dim (int): Embedding dimension.
gamma (float): Scaling factor for output scores.
"""
super().__init__()
self.gamma = gamma
self.decoder = nn.Sequential(
nn.Linear(2 * emb_dim, 1024),
nn.BatchNorm1d(1024),
nn.GELU(),
nn.Dropout(0.2),
nn.Linear(1024, 1)
)
self.corr_weight = nn.Parameter(torch.tensor(0.5))
def forward(self, cell_emb, drug_emb):
"""
Decode embeddings to predict interaction scores.
Args:
cell_emb (torch.Tensor): Cell embeddings.
drug_emb (torch.Tensor): Drug embeddings.
Returns:
torch.Tensor: Predicted interaction scores.
"""
# Expand embeddings for pairwise combinations
cell_exp = cell_emb.unsqueeze(1).repeat(1, drug_emb.size(0), 1)
drug_exp = drug_emb.unsqueeze(0).repeat(cell_emb.size(0), 1, 1)
combined = torch.cat([cell_exp, drug_exp], dim=-1)
# Decode combined embeddings
scores = self.decoder(combined.view(-1, 2 * cell_emb.size(1))).view(cell_emb.size(0), drug_emb.size(0))
corr = torch_corr_x_y(cell_emb, drug_emb)
# Combine scores with correlation and apply sigmoid
return torch.sigmoid(self.gamma * (self.corr_weight * scores + (1 - self.corr_weight) * corr))
class DeepTraCDR(nn.Module):
"""DeepTraCDR model for cell-drug interaction prediction."""
def __init__(self, adj_mat, cell_exprs, drug_finger, layer_size, gamma, device="cpu"):
"""
Initialize the DeepTraCDR model.
Args:
adj_mat (np.ndarray or torch.Tensor): Adjacency matrix.
cell_exprs (np.ndarray): Cell expression data.
drug_finger (list): Drug fingerprints.
layer_size (list): Layer sizes for the model.
gamma (float): Scaling factor for decoder.
device (str): Device to perform computations on.
"""
super().__init__()
self.device = device
if isinstance(adj_mat, np.ndarray):
adj_mat = torch.from_numpy(adj_mat).float()
self.adj_mat = adj_mat.to(device)
# Initialize submodules
self.construct_adj = ConstructAdjMatrix(self.adj_mat, device=device)
self.load_feat = LoadFeature(cell_exprs, drug_finger, device=device)
# Compute fixed adjacency matrices
agg_c, agg_d, self_c, self_d = self.construct_adj()
self.encoder = GEncoder(agg_c, agg_d, self_c, self_d, device=device).to(device)
self.decoder = GDecoder(512, gamma).to(device) # Fixed emb_dim to 512
def forward(self):
"""
Forward pass through the DeepTraCDR model.
Returns:
tuple: (predicted_scores, cell_emb, drug_emb) predicted scores and embeddings.
"""
cell_f, drug_f = self.load_feat()
cell_emb, drug_emb = self.encoder(cell_f, drug_f)
return self.decoder(cell_emb, drug_emb), cell_emb, drug_emb
def get_cell_embeddings(self):
"""
Retrieve cell embeddings from the model.
Returns:
torch.Tensor: Cell embeddings.
"""
cell_f, drug_f = self.load_feat()
cell_emb, _ = self.encoder(cell_f, drug_f)
return cell_emb
class Optimizer:
"""Optimizer class for training the DeepTraCDR model with early stopping."""
def __init__(self, model, train_data, test_data, test_mask, train_mask, adj_matrix, evaluate_fun,
lr=0.001, wd=1e-05, epochs=200, test_freq=20, patience=50, device="cpu"):
"""
Initialize the optimizer.
Args:
model (nn.Module): DeepTraCDR model to optimize.
train_data (torch.Tensor): Training data.
test_data (torch.Tensor): Test data.
test_mask (torch.Tensor): Test mask.
train_mask (torch.Tensor): Training mask.
adj_matrix (torch.Tensor): Adjacency matrix.
evaluate_fun (callable): Evaluation function.
lr (float): Learning rate.
wd (float): Weight decay.
epochs (int): Number of training epochs.
test_freq (int): Frequency of evaluation.
patience (int): Patience for early stopping.
device (str): Device to perform computations on.
"""
self.model = model.to(device)
self.train_data = train_data.float().to(device)
self.test_data = test_data.float().to(device)
self.train_mask = train_mask.to(device)
self.test_mask_bool = test_mask.to(device).bool()
self.adj_matrix = adj_matrix.to(device)
self.evaluate_fun = evaluate_fun
self.optimizer = torch.optim.Adam(self.model.parameters(), lr=lr, weight_decay=wd)
self.epochs = epochs
self.test_freq = test_freq
self.patience = patience
self.best_auc = 0.0
self.best_auprc = 0.0
self.best_weights = None
self.counter = 0
self.device = device
self.best_epoch_auc = None
self.best_epoch_auprc = None
def train(self):
"""
Train the model with early stopping and evaluate performance.
Returns:
tuple: (true_data, final_pred_masked, best_auc, best_auprc) evaluation results.
"""
true_data = torch.masked_select(self.test_data, self.test_mask_bool).cpu().numpy()
for epoch in range(self.epochs):
self.model.train()
# Forward pass and compute loss
pred_train, cell_emb, drug_emb = self.model()
ce_loss = cross_entropy_loss(self.train_data, pred_train, self.train_mask)
proto_loss = prototypical_loss(cell_emb, drug_emb, self.adj_matrix)
total_loss = 0.7 * ce_loss + 0.3 * proto_loss
# Backpropagation
self.optimizer.zero_grad()
total_loss.backward()
self.optimizer.step()
self.model.eval()
with torch.no_grad():
# Evaluate training performance
train_pred, _, _ = self.model()
train_pred_masked = torch.masked_select(train_pred, self.train_mask).cpu().numpy()
train_true_data = torch.masked_select(self.train_data, self.train_mask).cpu().numpy()
try:
train_auc = roc_auc_score(train_true_data, train_pred_masked)
train_auprc = average_precision_score(train_true_data, train_pred_masked)
except ValueError:
train_auc, train_auprc = 0.0, 0.0
# Evaluate test performance
pred_eval, _, _ = self.model()
pred_masked = torch.masked_select(pred_eval, self.test_mask_bool).cpu().numpy()
try:
auc = roc_auc_score(true_data, pred_masked)
auprc = average_precision_score(true_data, pred_masked)
except ValueError:
auc, auprc = 0.0, 0.0
# Update best metrics and weights
if auc > self.best_auc:
self.best_auc = auc
self.best_auprc = auprc
self.best_weights = self.model.state_dict().copy()
self.counter = 0
self.best_epoch_auc = auc
self.best_epoch_auprc = auprc
else:
self.counter += 1
# Log progress
if epoch % self.test_freq == 0 or epoch == self.epochs - 1:
print(f"Epoch {epoch}: Loss={total_loss.item():.4f}, Train AUC={train_auc:.4f}, "
f"Train AUPRC={train_auprc:.4f}, Test AUC={auc:.4f}, Test AUPRC={auprc:.4f}")
# Check early stopping
if self.counter >= self.patience:
print(f"\nEarly stopping triggered at epoch {epoch}!")
print(f"No improvement in AUC for {self.patience} consecutive epochs.")
break
# Restore best weights
if self.best_weights is not None:
self.model.load_state_dict(self.best_weights)
# Final evaluation
self.model.eval()
with torch.no_grad():
final_pred, _, _ = self.model()
final_pred_masked = torch.masked_select(final_pred, self.test_mask_bool).cpu().numpy()
best_auc = roc_auc_score(true_data, final_pred_masked)
best_auprc = average_precision_score(true_data, final_pred_masked)
# Log final results
print("\nBest Metrics After Training (on Test Data):")
print(f"AUC: {self.best_auc:.4f}")
print(f"AUPRC: {self.best_auprc:.4f}")
return true_data, final_pred_masked, self.best_auc, self.best_auprc

Loading…
Cancel
Save