Upload files to 'Scenario2/new'

Scenario2/new/main_new.py

@@ -0,0 +1,111 @@
+import argparse
+import torch
+import numpy as np
+from sklearn.metrics import roc_auc_score
+from data_loader import load_data
+from model import DeepTraCDR
+
+
+def parse_arguments() -> argparse.Namespace:
+    """
+    Parses command-line arguments for running the DeepTraCDR model.
+
+    Returns:
+        argparse.Namespace: Parsed arguments containing configuration parameters.
+    """
+    parser = argparse.ArgumentParser(description="Run DeepTraCDR model for cell-drug interaction prediction.")
+    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu",
+                        help="Device to run the model on (e.g., 'cuda:0' or 'cpu').")
+    parser.add_argument("--data", type=str, default="gdsc",
+                        help="Dataset to use (e.g., 'gdsc').")
+    parser.add_argument("--wd", type=float, default=1e-5,
+                        help="Weight decay for the optimizer.")
+    parser.add_argument("--layer_size", nargs="+", type=int, default=[512],
+                        help="List of layer sizes for the model.")
+    parser.add_argument("--gamma", type=float, default=15.0,
+                        help="Gamma parameter for decoder scaling.")
+    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("--lr", type=float, default=0.0001,
+                        help="Learning rate for the optimizer.")
+    return parser.parse_args()
+
+def main():
+    """
+    Main function to run the DeepTraCDR model, including data loading, model training,
+    and evaluation across specified targets.
+    """
+    # Parse command-line arguments
+    args = parse_arguments()
+
+    # Load dataset
+    adj_matrix, drug_fingerprints, cell_expressions, null_mask, pos_num, args = load_data(args)
+
+    # Compute interaction sums for filtering
+    cell_interaction_sums = np.sum(adj_matrix, axis=1)  # Sum of interactions per cell line
+    drug_interaction_sums = np.sum(adj_matrix, axis=0)  # Sum of interactions per drug
+
+    # Log dataset and configuration details
+    print(f"\n{'='*40}")
+    print(f"Dataset: {args.data} | Adjacency Matrix Shape: {adj_matrix.shape}")
+    print(f"Device: {args.device} | Layer Sizes: {args.layer_size} | Learning Rate: {args.lr}")
+    print(f"{'='*40}\n")
+
+    # Define target dimension (0 for cell lines)
+    target_dimensions = [0]  # Only process cell lines
+    adj_matrix = torch.from_numpy(adj_matrix).float() if isinstance(adj_matrix, np.ndarray) else adj_matrix
+
+    num_folds = 1  # Number of cross-validation folds
+    total_targets_processed = 0
+    auc_scores = []
+    auprc_scores = []
+
+    # Process each target dimension
+    for dim in target_dimensions:
+        dim_name = "Cell Lines" if dim == 0 else "Drugs"
+        # Filter valid targets with at least 10 interactions
+        valid_indices = [
+            idx for idx in range(adj_matrix.shape[dim])
+            if (drug_interaction_sums[idx] >= 10 if dim == 1 else cell_interaction_sums[idx] >= 10)
+        ]
+
+        print(f"Processing {dim_name} ({len(valid_indices)} valid targets):")
+
+        # Process each valid target
+        for target_index in valid_indices:
+            total_targets_processed += 1
+            print(f"  Target {target_index} - ", end="", flush=True)
+
+            # Run model for each fold
+            for fold in range(num_folds):
+                best_auc, best_auprc = DeepTraCDR(
+                    cell_exprs=cell_expressions,
+                    drug_finger=drug_fingerprints,
+                    res_mat=adj_matrix,
+                    null_mask=null_mask,
+                    target_dim=dim,
+                    target_index=target_index,
+                    evaluate_fun=roc_auc_score,
+                    args=args
+                )
+                auc_scores.append(best_auc)
+                auprc_scores.append(best_auprc)
+
+            print(f"AUC: {best_auc:.4f}, AUPRC: {best_auprc:.4f}")
+
+    # Compute and display average metrics
+    mean_auc = np.mean(auc_scores)
+    std_auc = np.std(auc_scores)
+    mean_auprc = np.mean(auprc_scores)
+    std_auprc = np.std(auprc_scores)
+
+    print(f"\n{'='*40}")
+    print(f"Average AUC: {mean_auc:.4f} ± {std_auc:.4f}")
+    print(f"Average AUPRC: {mean_auprc:.4f} ± {std_auprc:.4f}")
+    print(f"{'='*40}\n")
+
+
+if __name__ == "__main__":
+    main()

Scenario2/new/model.py

@@ -0,0 +1,383 @@
+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 sklearn.metrics import roc_auc_score, average_precision_score
+from utils import torch_corr_x_y, cross_entropy_loss, prototypical_loss
+
+
+class AdjacencyMatrixConstructor(nn.Module):
+    """
+    Constructs normalized adjacency matrices for graph-based computations.
+    These matrices are used for aggregating cell and drug features in the GCN model.
+    """
+    def __init__(self, original_adj_mat, device="cpu"):
+        super().__init__()
+        # Convert numpy array to torch tensor if necessary and move to specified device
+        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):
+        """
+        Computes normalized adjacency matrices for cell and drug aggregations.
+        Returns four matrices: aggregated cell, aggregated drug, self-cell, and self-drug.
+        """
+        with torch.no_grad():
+            # Compute degree normalization matrices
+            degree_x = torch.pow(torch.sum(self.adj, dim=1) + 1, -0.5)
+            degree_y = torch.pow(torch.sum(self.adj, dim=0) + 1, -0.5)
+            d_x = torch.diag(degree_x)
+            d_y = torch.diag(degree_y)
+
+            # Compute aggregated Laplacian matrices
+            agg_cell_lp = torch.mm(torch.mm(d_x, self.adj), d_y)  # [num_cells x num_drugs]
+            agg_drug_lp = torch.mm(torch.mm(d_y, self.adj.T), d_x)  # [num_drugs x num_cells]
+
+            # Compute self-loop Laplacian matrices
+            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 FeatureLoader(nn.Module):
+    """
+    Loads and preprocesses cell expression and drug fingerprint features.
+    Applies transformations to project features into a common embedding space.
+    """
+    def __init__(self, cell_exprs, drug_fingerprints, device="cpu"):
+        super().__init__()
+        self.device = device
+
+        # Convert input features to torch tensors and move to 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.5)
+            ).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=1
+        ).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.5),
+            nn.Linear(1024, 512)
+        ).to(device)
+
+    def forward(self):
+        """
+        Processes cell and drug features to produce encoded representations.
+        Returns encoded cell and drug features in a common 512-dimensional space.
+        """
+        # Normalize and encode cell features
+        cell_feat = self.cell_norm(self.cell_exprs)  # [num_cells x num_cell_features]
+        cell_encoded = self.cell_encoder(cell_feat)  # [num_cells x 512]
+
+        # Project and transform drug fingerprints
+        projected = [proj(fp) for proj, fp in zip(self.drug_proj, self.drug_fingerprints)]  # List of [num_samples x 512]
+        stacked = torch.stack(projected, dim=1)  # [num_samples x num_drugs x 512]
+        drug_feat = self.transformer(stacked)  # [num_samples x num_drugs x 512]
+        drug_feat = self.drug_norm(drug_feat.mean(dim=1))  # [num_samples x 512]
+
+        return cell_encoded, drug_feat
+
+
+class GraphEncoder(nn.Module):
+    """
+    Encodes cell and drug features using graph-based aggregation and attention mechanisms.
+    Produces final embeddings for cells and drugs.
+    """
+    def __init__(self, agg_cell_lp, agg_drug_lp, self_cell_lp, self_drug_lp, device="cpu"):
+        super().__init__()
+        self.agg_cell_lp = agg_cell_lp
+        self.agg_drug_lp = agg_drug_lp
+        self.self_cell_lp = self_cell_lp
+        self.self_drug_lp = self_drug_lp
+        self.device = device
+
+        # Encoder for aggregated cell features
+        self.cell_encoder = nn.Sequential(
+            nn.Linear(512, 1024),
+            nn.BatchNorm1d(1024),
+            nn.GELU(),
+            nn.Dropout(0.5),
+            nn.Linear(1024, 512)
+        ).to(device)
+
+        # Encoder for aggregated drug features
+        self.drug_encoder = nn.Sequential(
+            nn.Linear(512, 1024),
+            nn.BatchNorm1d(1024),
+            nn.GELU(),
+            nn.Dropout(0.5),
+            nn.Linear(1024, 512)
+        ).to(device)
+
+        # Attention mechanism for cell-drug interactions
+        self.attention = MultiheadAttention(embed_dim=512, num_heads=8, batch_first=True).to(device)
+        self.residual = nn.Linear(512, 512).to(device)
+
+        # Final fully connected layer
+        self.fc = nn.Sequential(
+            nn.Linear(1024, 512),
+            nn.BatchNorm1d(512),
+            nn.GELU(),
+            nn.Dropout(0.5)
+        ).to(device)
+
+    def forward(self, cell_features, drug_features):
+        """
+        Encodes cell and drug features using graph aggregation and attention.
+        Returns final cell and drug embeddings.
+        """
+        # Aggregate features using Laplacian matrices
+        cell_agg = torch.mm(self.agg_cell_lp, drug_features) + torch.mm(self.self_cell_lp, cell_features)  # [num_cells x 512]
+        drug_agg = torch.mm(self.agg_drug_lp, cell_features) + torch.mm(self.self_drug_lp, drug_features)  # [num_drugs x 512]
+
+        # Encode aggregated features
+        cell_fc = self.cell_encoder(cell_agg)  # [num_cells x 512]
+        drug_fc = self.drug_encoder(drug_agg)  # [num_drugs x 512]
+
+        # Apply attention mechanism
+        attn_output, _ = self.attention(
+            query=cell_fc.unsqueeze(0),  # [1 x num_cells x 512]
+            key=drug_fc.unsqueeze(0),    # [1 x num_drugs x 512]
+            value=drug_fc.unsqueeze(0)   # [1 x num_drugs x 512]
+        )
+        attn_output = attn_output.squeeze(0)  # [num_cells x 512]
+
+        # Combine attention output with residual connection
+        cell_emb = cell_fc + self.residual(attn_output)  # [num_cells x 512]
+
+        # Apply final activation
+        cell_emb = F.gelu(cell_emb)  # [num_cells x 512]
+        drug_emb = F.gelu(drug_fc)  # [num_drugs x 512]
+
+        return cell_emb, drug_emb
+
+
+class GraphDecoder(nn.Module):
+    """
+    Decodes cell and drug embeddings to predict interaction scores.
+    Combines embeddings and applies a correlation-based adjustment.
+    """
+    def __init__(self, emb_dim, gamma):
+        super().__init__()
+        self.gamma = gamma
+
+        # Decoder network for combined embeddings
+        self.decoder = nn.Sequential(
+            nn.Linear(2 * emb_dim, 1024),
+            nn.BatchNorm1d(1024),
+            nn.GELU(),
+            nn.Dropout(0.2),
+            nn.Linear(1024, 1)
+        )
+
+        # Learnable weight for balancing scores and correlation
+        self.corr_weight = nn.Parameter(torch.tensor(0.5))
+
+    def forward(self, cell_emb, drug_emb):
+        """
+        Decodes cell and drug embeddings to produce interaction scores.
+        Returns a matrix of interaction probabilities.
+        """
+        # Expand embeddings for pairwise combinations
+        cell_exp = cell_emb.unsqueeze(1).repeat(1, drug_emb.size(0), 1)  # [num_cells x num_drugs x emb_dim]
+        drug_exp = drug_emb.unsqueeze(0).repeat(cell_emb.size(0), 1, 1)  # [num_cells x num_drugs x emb_dim]
+
+        # Combine cell and drug embeddings
+        combined = torch.cat([cell_exp, drug_exp], dim=-1)  # [num_cells x num_drugs x 2*emb_dim]
+
+        # Compute interaction scores
+        scores = self.decoder(combined.view(-1, 2 * cell_emb.size(1))).view(cell_emb.size(0), drug_emb.size(0))  # [num_cells x num_drugs]
+
+        # Compute correlation between embeddings
+        corr = torch_corr_x_y(cell_emb, drug_emb)  # [num_cells x num_drugs]
+
+        # Combine scores and correlation with learnable weight
+        return torch.sigmoid(self.gamma * (self.corr_weight * scores + (1 - self.corr_weight) * corr))
+
+
+class DeepTraCDR(nn.Module):
+    """
+    Main Graph Convolutional Network model for predicting cell-drug interactions.
+    Combines feature loading, graph encoding, and decoding.
+    """
+    def __init__(self, adj_mat, cell_exprs, drug_fingerprints, layer_size, gamma, device="cpu"):
+        super().__init__()
+        self.device = device
+
+        # Convert adjacency matrix to tensor if necessary
+        if isinstance(adj_mat, np.ndarray):
+            adj_mat = torch.from_numpy(adj_mat).float()
+        self.adj_mat = adj_mat.to(device)
+
+        # Initialize components
+        self.construct_adj = AdjacencyMatrixConstructor(self.adj_mat, device=device)
+        self.load_feat = FeatureLoader(cell_exprs, drug_fingerprints, device=device)
+
+        # Compute fixed adjacency matrices
+        agg_cell, agg_drug, self_cell, self_drug = self.construct_adj()
+
+        # Initialize encoder and decoder
+        self.encoder = GraphEncoder(agg_cell, agg_drug, self_cell, self_drug, device=device).to(device)
+        self.decoder = GraphDecoder(512, gamma).to(device)  # emb_dim fixed to 512
+
+    def forward(self):
+        """
+        Performs a full forward pass through the DeepTraCDR model.
+        Returns predicted interaction scores and final embeddings.
+        """
+        # Load and encode features
+        cell_features, drug_features = self.load_feat()
+
+        # Encode features using graph structure
+        cell_emb, drug_emb = self.encoder(cell_features, drug_features)
+
+        # Decode to predict interaction scores
+        return self.decoder(cell_emb, drug_emb), cell_emb, drug_emb
+
+
+class ModelOptimizer:
+    """
+    Handles training and evaluation of the DeepTraCDR model.
+    Implements early stopping and tracks best performance metrics.
+    """
+    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=100, device="gpu"):
+        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.device = device
+
+        # Convert adjacency matrix to tensor if necessary
+        if isinstance(adj_matrix, np.ndarray):
+            adj_matrix = torch.from_numpy(adj_matrix).float()
+        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  # Early stopping counter
+        self.best_epoch_auc = None
+        self.best_epoch_auprc = None
+
+    def train(self):
+        """
+        Trains the model with early stopping and evaluates performance.
+        Returns the best AUC and AUPRC achieved during training.
+        """
+        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
+
+            # Backward pass and optimization
+            self.optimizer.zero_grad()
+            total_loss.backward()
+            self.optimizer.step()
+
+            # Evaluate model
+            self.model.eval()
+            with torch.no_grad():
+                # Compute metrics for training data
+                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
+
+                # Compute metrics for test data
+                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}, 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
+
+        # Load 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)
+
+        # Print 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 self.best_auc, self.best_auprc