|
|
@@ -0,0 +1,461 @@ |
|
|
|
# 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 typing import Tuple, List
|
|
|
|
from sklearn.metrics import roc_auc_score, average_precision_score
|
|
|
|
from utils import torch_corr_x_y, cross_entropy_loss, prototypical_loss
|
|
|
|
|
|
|
|
class ConstructAdjMatrix(nn.Module):
|
|
|
|
"""
|
|
|
|
Constructs normalized adjacency matrices for graph-based operations.
|
|
|
|
"""
|
|
|
|
def __init__(self, original_adj_mat: torch.Tensor | np.ndarray, device: str = "cpu"):
|
|
|
|
"""
|
|
|
|
Initialize the adjacency matrix construction module.
|
|
|
|
|
|
|
|
Args:
|
|
|
|
original_adj_mat (torch.Tensor | np.ndarray): Input adjacency matrix.
|
|
|
|
device (str): Device to run computations on (e.g., 'cuda:0' or 'cpu').
|
|
|
|
"""
|
|
|
|
super().__init__()
|
|
|
|
self.device = device
|
|
|
|
# Convert to tensor if input is NumPy array
|
|
|
|
if isinstance(original_adj_mat, np.ndarray):
|
|
|
|
original_adj_mat = torch.from_numpy(original_adj_mat).float()
|
|
|
|
self.adj = original_adj_mat.to(device)
|
|
|
|
|
|
|
|
def forward(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
|
|
"""
|
|
|
|
Compute normalized adjacency matrices for cells and drugs.
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
Tuple of aggregated and self-loop adjacency matrices for cells and drugs.
|
|
|
|
"""
|
|
|
|
with torch.no_grad():
|
|
|
|
# Degree normalization for cells and drugs
|
|
|
|
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))
|
|
|
|
|
|
|
|
# Aggregated Laplacian matrices
|
|
|
|
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-loop 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 LoadFeature(nn.Module):
|
|
|
|
"""
|
|
|
|
Loads and processes cell expression and drug fingerprint features.
|
|
|
|
"""
|
|
|
|
def __init__(self, cell_exprs: np.ndarray, drug_fingerprints: List[np.ndarray], device: str = "cpu"):
|
|
|
|
"""
|
|
|
|
Initialize the feature loading module.
|
|
|
|
|
|
|
|
Args:
|
|
|
|
cell_exprs (np.ndarray): Cell expression data.
|
|
|
|
drug_fingerprints (List[np.ndarray]): List of drug fingerprint matrices.
|
|
|
|
device (str): Device to run computations on (e.g., 'cuda:0' or 'cpu').
|
|
|
|
"""
|
|
|
|
super().__init__()
|
|
|
|
self.device = device
|
|
|
|
|
|
|
|
# Convert inputs to tensors
|
|
|
|
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]
|
|
|
|
|
|
|
|
# Drug projection layers
|
|
|
|
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)
|
|
|
|
|
|
|
|
# Cell encoder
|
|
|
|
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) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
|
|
"""
|
|
|
|
Process cell and drug features.
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
Tuple of encoded cell features and drug features.
|
|
|
|
"""
|
|
|
|
# Process cell features
|
|
|
|
cell_feat = self.cell_norm(self.cell_exprs) # [num_cells x num_features]
|
|
|
|
cell_encoded = self.cell_encoder(cell_feat) # [num_cells x 512]
|
|
|
|
|
|
|
|
# Process drug features
|
|
|
|
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 GEncoder(nn.Module):
|
|
|
|
"""
|
|
|
|
Graph encoder for combining cell and drug features with graph structure.
|
|
|
|
"""
|
|
|
|
def __init__(
|
|
|
|
self,
|
|
|
|
agg_c_lp: torch.Tensor,
|
|
|
|
agg_d_lp: torch.Tensor,
|
|
|
|
self_c_lp: torch.Tensor,
|
|
|
|
self_d_lp: torch.Tensor,
|
|
|
|
device: str = "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 run computations on (e.g., 'cuda:0' or 'cpu').
|
|
|
|
"""
|
|
|
|
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 and drug encoders
|
|
|
|
self.cell_encoder = nn.Sequential(
|
|
|
|
nn.Linear(512, 1024),
|
|
|
|
nn.BatchNorm1d(1024),
|
|
|
|
nn.GELU(),
|
|
|
|
nn.Dropout(0.5),
|
|
|
|
nn.Linear(1024, 512)
|
|
|
|
).to(device)
|
|
|
|
|
|
|
|
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
|
|
|
|
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_f: torch.Tensor, drug_f: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
|
|
"""
|
|
|
|
Encode cell and drug features using graph structure and attention.
|
|
|
|
|
|
|
|
Args:
|
|
|
|
cell_f (torch.Tensor): Cell features.
|
|
|
|
drug_f (torch.Tensor): Drug features.
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
Tuple of encoded cell and drug embeddings.
|
|
|
|
"""
|
|
|
|
# Aggregate features using Laplacian matrices
|
|
|
|
cell_agg = torch.mm(self.agg_c_lp, drug_f) + torch.mm(self.self_c_lp, cell_f) # [num_cells x 512]
|
|
|
|
drug_agg = torch.mm(self.agg_d_lp, cell_f) + torch.mm(self.self_d_lp, drug_f) # [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]
|
|
|
|
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 GDecoder(nn.Module):
|
|
|
|
"""
|
|
|
|
Decoder to predict interaction scores between cells and drugs.
|
|
|
|
"""
|
|
|
|
def __init__(self, emb_dim: int, gamma: float):
|
|
|
|
"""
|
|
|
|
Initialize the decoder.
|
|
|
|
|
|
|
|
Args:
|
|
|
|
emb_dim (int): Embedding dimension (default: 512).
|
|
|
|
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: torch.Tensor, drug_emb: torch.Tensor) -> torch.Tensor:
|
|
|
|
"""
|
|
|
|
Predict interaction scores between cells and drugs.
|
|
|
|
|
|
|
|
Args:
|
|
|
|
cell_emb (torch.Tensor): Cell embeddings.
|
|
|
|
drug_emb (torch.Tensor): Drug embeddings.
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
torch.Tensor: Predicted interaction scores.
|
|
|
|
"""
|
|
|
|
# Expand dimensions for pairwise combination
|
|
|
|
cell_exp = cell_emb.unsqueeze(1).repeat(1, drug_emb.size(0), 1) # [num_cells x num_drugs x 512]
|
|
|
|
drug_exp = drug_emb.unsqueeze(0).repeat(cell_emb.size(0), 1, 1) # [num_cells x num_drugs x 512]
|
|
|
|
combined = torch.cat([cell_exp, drug_exp], dim=-1) # [num_cells x num_drugs x 1024]
|
|
|
|
|
|
|
|
# Decode combined features
|
|
|
|
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]
|
|
|
|
corr = torch_corr_x_y(cell_emb, drug_emb) # [num_cells x num_drugs]
|
|
|
|
|
|
|
|
# Combine scores and correlation with learned weighting
|
|
|
|
return torch.sigmoid(self.gamma * (self.corr_weight * scores + (1 - self.corr_weight) * corr))
|
|
|
|
|
|
|
|
class DeepTraCDR(nn.Module):
|
|
|
|
"""
|
|
|
|
DeepTraCDR model for predicting cell-drug interactions.
|
|
|
|
"""
|
|
|
|
def __init__(
|
|
|
|
self,
|
|
|
|
adj_mat: torch.Tensor | np.ndarray,
|
|
|
|
cell_exprs: np.ndarray,
|
|
|
|
drug_finger: List[np.ndarray],
|
|
|
|
layer_size: List[int],
|
|
|
|
gamma: float,
|
|
|
|
device: str = "cpu"
|
|
|
|
):
|
|
|
|
"""
|
|
|
|
Initialize the DeepTraCDR model.
|
|
|
|
|
|
|
|
Args:
|
|
|
|
adj_mat (torch.Tensor | np.ndarray): Adjacency matrix.
|
|
|
|
cell_exprs (np.ndarray): Cell expression data.
|
|
|
|
drug_finger (List[np.ndarray]): List of drug fingerprint matrices.
|
|
|
|
layer_size (List[int]): Sizes of hidden layers.
|
|
|
|
gamma (float): Scaling factor for decoder.
|
|
|
|
device (str): Device to run computations on (e.g., 'cuda:0' or 'cpu').
|
|
|
|
"""
|
|
|
|
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()
|
|
|
|
|
|
|
|
# Initialize encoder and decoder
|
|
|
|
self.encoder = GEncoder(agg_c, agg_d, self_c, self_d, device=device).to(device)
|
|
|
|
self.decoder = GDecoder(emb_dim=512, gamma=gamma).to(device)
|
|
|
|
|
|
|
|
def forward(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
|
|
"""
|
|
|
|
Forward pass of the DeepTraCDR model.
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
Tuple of predicted scores, cell embeddings, and drug embeddings.
|
|
|
|
"""
|
|
|
|
cell_f, drug_f = self.load_feat()
|
|
|
|
cell_emb, drug_emb = self.encoder(cell_f, drug_f)
|
|
|
|
scores = self.decoder(cell_emb, drug_emb)
|
|
|
|
return scores, cell_emb, drug_emb
|
|
|
|
|
|
|
|
class Optimizer:
|
|
|
|
"""
|
|
|
|
Optimizer for training the DeepTraCDR model with early stopping and evaluation.
|
|
|
|
"""
|
|
|
|
def __init__(
|
|
|
|
self,
|
|
|
|
model: DeepTraCDR,
|
|
|
|
train_data: torch.Tensor,
|
|
|
|
test_data: torch.Tensor,
|
|
|
|
test_mask: torch.Tensor,
|
|
|
|
train_mask: torch.Tensor,
|
|
|
|
adj_matrix: torch.Tensor | np.ndarray,
|
|
|
|
evaluate_fun,
|
|
|
|
lr: float = 0.001,
|
|
|
|
wd: float = 1e-05,
|
|
|
|
epochs: int = 200,
|
|
|
|
test_freq: int = 20,
|
|
|
|
patience: int = 200,
|
|
|
|
device: str = "cpu"
|
|
|
|
):
|
|
|
|
"""
|
|
|
|
Initialize the optimizer.
|
|
|
|
|
|
|
|
Args:
|
|
|
|
model (DeepTraCDR): The DeepTraCDR model to train.
|
|
|
|
train_data (torch.Tensor): Training data.
|
|
|
|
test_data (torch.Tensor): Test data.
|
|
|
|
test_mask (torch.Tensor): Mask for test data.
|
|
|
|
train_mask (torch.Tensor): Mask for training data.
|
|
|
|
adj_matrix (torch.Tensor | np.ndarray): Adjacency matrix.
|
|
|
|
evaluate_fun: Function to evaluate model performance.
|
|
|
|
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 run computations on (e.g., 'cuda:0' or 'cpu').
|
|
|
|
"""
|
|
|
|
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 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
|
|
|
|
|
|
|
|
def train(self) -> Tuple[np.ndarray, np.ndarray, float, float]:
|
|
|
|
"""
|
|
|
|
Train the DeepTraCDR model and evaluate performance.
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
Tuple of true labels, predicted scores, best AUC, and best AUPRC.
|
|
|
|
"""
|
|
|
|
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
|
|
|
|
self.optimizer.zero_grad()
|
|
|
|
total_loss.backward()
|
|
|
|
self.optimizer.step()
|
|
|
|
|
|
|
|
# Evaluate model
|
|
|
|
self.model.eval()
|
|
|
|
with torch.no_grad():
|
|
|
|
# Training metrics
|
|
|
|
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
|
|
|
|
|
|
|
|
# Test metrics
|
|
|
|
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
|
|
|
|
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}, "
|
|
|
|
f"Train AUC={train_auc:.4f}, Train AUPRC={train_auprc:.4f}, "
|
|
|
|
f"Test AUC={auc:.4f}, Test AUPRC={auprc:.4f}")
|
|
|
|
|
|
|
|
# Early stopping
|
|
|
|
if self.counter >= self.patience:
|
|
|
|
print(f"\nEarly stopping at epoch {epoch}: No AUC improvement for {self.patience} 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("\nBest Metrics (Test Data):")
|
|
|
|
print(f"AUC: {best_auc:.4f}")
|
|
|
|
print(f"AUPRC: {best_auprc:.4f}")
|
|
|
|
|
|
|
|
return true_data, final_pred_masked, best_auc, best_auprc |