AI-Powered Drug Discovery Pipelines: Accelerating the Path from Molecule to Medicine

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import MessagePassing, global_mean_pool

class MPNNLayer(MessagePassing):
    """Message Passing Neural Network layer for molecular graphs."""
    
    def __init__(self, node_dim, edge_dim, hidden_dim):
        super().__init__(aggr='add')  # Aggregate messages by summation
        
        # Edge network: transforms edge features
        self.edge_network = nn.Sequential(
            nn.Linear(2 * node_dim + edge_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, node_dim)
        )
        
        # Node update network
        self.node_network = nn.Sequential(
            nn.Linear(2 * node_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, node_dim)
        )
        
    def forward(self, x, edge_index, edge_attr):
        """
        x: node features [num_nodes, node_dim]
        edge_index: graph connectivity [2, num_edges]
        edge_attr: edge features [num_edges, edge_dim]
        """
        # Propagate messages from neighbors
        aggregated = self.propagate(edge_index, x=x, edge_attr=edge_attr)
        
        # Update node representations
        x_updated = self.node_network(torch.cat([x, aggregated], dim=-1))
        return x_updated + x  # Residual connection
    
    def message(self, x_i, x_j, edge_attr):
        """Compute messages from node j to node i."""
        # Concatenate source node, target node, and edge features
        edge_input = torch.cat([x_i, x_j, edge_attr], dim=-1)
        return self.edge_network(edge_input)


class MolecularPropertyPredictor(nn.Module):
    """Graph neural network for predicting molecular properties."""
    
    def __init__(self, node_features=9, edge_features=3, hidden_dim=128, 
                 num_layers=6, num_properties=15, dropout=0.2):
        super().__init__()
        
        # Initial node embedding
        self.node_embedding = nn.Linear(node_features, hidden_dim)
        
        # Message passing layers
        self.mp_layers = nn.ModuleList([
            MPNNLayer(hidden_dim, edge_features, hidden_dim) 
            for _ in range(num_layers)
        ])
        
        # Readout function: graph-level representation
        self.dropout = nn.Dropout(dropout)
        
        # Property prediction heads (multi-task learning)
        self.property_heads = nn.ModuleDict({
            'binding_affinity': nn.Linear(hidden_dim, 1),
            'solubility': nn.Linear(hidden_dim, 1),
            'permeability': nn.Linear(hidden_dim, 1),
            'metabolic_stability': nn.Linear(hidden_dim, 1),
            'toxicity_herg': nn.Linear(hidden_dim, 1),
            'toxicity_ames': nn.Linear(hidden_dim, 1),
            # ... additional property heads
        })
        
    def forward(self, data, return_uncertainty=False):
        """
        data: PyTorch Geometric data object containing:
            - x: node features
            - edge_index: graph connectivity
            - edge_attr: edge features
            - batch: batch assignment for each node
        """
        x, edge_index, edge_attr, batch = data.x, data.edge_index, data.edge_attr, data.batch
        
        # Initial embedding
        x = self.node_embedding(x)
        
        # Message passing
        for mp_layer in self.mp_layers:
            x = mp_layer(x, edge_index, edge_attr)
            x = F.relu(x)
        
        # Graph-level pooling (aggregate node features)
        graph_embedding = global_mean_pool(x, batch)
        graph_embedding = self.dropout(graph_embedding)
        
        # Predict multiple properties
        predictions = {}
        for property_name, head in self.property_heads.items():
            predictions[property_name] = head(graph_embedding)
        
        if return_uncertainty:
            # Monte Carlo dropout for uncertainty estimation
            uncertainties = self.estimate_uncertainty(data, num_samples=20)
            return predictions, uncertainties
        
        return predictions
    
    def estimate_uncertainty(self, data, num_samples=20):
        """Estimate epistemic uncertainty via MC dropout."""
        self.train()  # Enable dropout
        samples = []
        
        with torch.no_grad():
            for _ in range(num_samples):
                preds = self.forward(data, return_uncertainty=False)
                samples.append(preds)
        
        # Compute variance across samples
        uncertainties = {}
        for property_name in self.property_heads.keys():
            property_samples = torch.stack([s[property_name] for s in samples])
            uncertainties[property_name] = property_samples.var(dim=0)
        
        self.eval()
        return uncertainties


# Example: Bayesian optimization acquisition function
def expected_improvement(predictions, uncertainties, best_value, xi=0.01):
    """
    Expected Improvement acquisition function for Bayesian optimization.
    
    Args:
        predictions: predicted property values
        uncertainties: epistemic uncertainty estimates
        best_value: current best observed value
        xi: exploration parameter
    """
    mean = predictions
    std = torch.sqrt(uncertainties)
    
    # Compute improvement over current best
    improvement = mean - best_value - xi
    Z = improvement / (std + 1e-9)
    
    # Expected improvement = E[max(0, improvement)]
    ei = improvement * torch.distributions.Normal(0, 1).cdf(Z) + \
         std * torch.distributions.Normal(0, 1).log_prob(Z).exp()
    
    return ei
                    

class MolecularVAE(nn.Module):
    """Variational Autoencoder for molecular SMILES strings."""
    
    def __init__(self, vocab_size=42, max_length=120, latent_dim=128, hidden_dim=256):
        super().__init__()
        self.latent_dim = latent_dim
        
        # Encoder: SMILES string ' latent distribution
        self.encoder_lstm = nn.LSTM(vocab_size, hidden_dim, num_layers=3, batch_first=True)
        self.fc_mu = nn.Linear(hidden_dim, latent_dim)
        self.fc_logvar = nn.Linear(hidden_dim, latent_dim)
        
        # Decoder: latent vector ' SMILES string
        self.decoder_fc = nn.Linear(latent_dim, hidden_dim)
        self.decoder_lstm = nn.LSTM(hidden_dim, hidden_dim, num_layers=3, batch_first=True)
        self.output_fc = nn.Linear(hidden_dim, vocab_size)
        
    def encode(self, x):
        """Encode SMILES to latent distribution parameters."""
        _, (h_n, _) = self.encoder_lstm(x)
        h = h_n[-1]  # Use final hidden state
        mu = self.fc_mu(h)
        logvar = self.fc_logvar(h)
        return mu, logvar
    
    def reparameterize(self, mu, logvar):
        """Sample from latent distribution using reparameterization trick."""
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std
    
    def decode(self, z, max_length=120):
        """Decode latent vector to SMILES string."""
        batch_size = z.size(0)
        h = self.decoder_fc(z).unsqueeze(1).repeat(1, max_length, 1)
        output, _ = self.decoder_lstm(h)
        logits = self.output_fc(output)
        return logits
    
    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        reconstruction = self.decode(z, x.size(1))
        return reconstruction, mu, logvar
    
    def generate(self, num_samples=10, temperature=1.0):
        """Generate novel molecules by sampling from prior."""
        z = torch.randn(num_samples, self.latent_dim) * temperature
        logits = self.decode(z)
        probs = F.softmax(logits, dim=-1)
        # Sample SMILES characters from probability distribution
        samples = torch.multinomial(probs.view(-1, probs.size(-1)), 1).view(num_samples, -1)
        return samples
    
    def optimize_molecule(self, initial_smiles, target_property_fn, num_steps=100, lr=0.01):
        """Optimize molecule in latent space toward target property."""
        # Encode initial molecule
        x = smiles_to_tensor(initial_smiles)
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        z.requires_grad = True
        
        optimizer = torch.optim.Adam([z], lr=lr)
        
        for step in range(num_steps):
            optimizer.zero_grad()
            
            # Decode to SMILES
            logits = self.decode(z)
            smiles = tensor_to_smiles(logits)
            
            # Evaluate target property (e.g., binding affinity)
            property_value = target_property_fn(smiles)
            loss = -property_value  # Maximize property
            
            loss.backward()
            optimizer.step()
        
        # Return optimized molecule
        final_logits = self.decode(z)
        return tensor_to_smiles(final_logits)