torch-molecule 0.1.0.dev2__py3-none-any.whl

torch_molecule/predictor/smiles_transformer/modeling_transformer.py

@@ -0,0 +1,252 @@
+import os
+import numpy as np
+from tqdm import tqdm
+from typing import Optional, Union, Dict, Any, List, Callable, Type, Tuple
+import warnings
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, TensorDataset
+
+from .model import Transformer
+from ..lstm.modeling_lstm import LSTMMolecularPredictor
+from ...utils.search import (
+    suggest_parameter,
+    ParameterSpec,
+    ParameterType,
+)
+
+# Dictionary mapping parameter names to their types and ranges
+DEFAULT_TRANSFORMER_SEARCH_SPACES: Dict[str, ParameterSpec] = {
+    # Integer-valued parameters
+    "hidden_size": ParameterSpec(ParameterType.INTEGER, (64, 256)),
+    "n_heads": ParameterSpec(ParameterType.INTEGER, (2, 8)),
+    "num_layers": ParameterSpec(ParameterType.INTEGER, (2, 6)),
+    "dim_feedforward": ParameterSpec(ParameterType.INTEGER, (128, 512)),
+    # Float-valued parameters with log scale
+    "learning_rate": ParameterSpec(ParameterType.LOG_FLOAT, (1e-5, 1e-2)),
+    "weight_decay": ParameterSpec(ParameterType.LOG_FLOAT, (1e-8, 1e-3)),
+    "dropout": ParameterSpec(ParameterType.FLOAT, (0.0, 0.5)),
+    "scheduler_factor": ParameterSpec(ParameterType.FLOAT, (0.1, 0.5)),
+}
+
+@dataclass
+class SMILESTransformerMolecularPredictor(LSTMMolecularPredictor):
+    """This predictor implements a Transformer model for SMILES-based molecular property predictions.
+    
+    Notes
+    -----
+    This implementation uses a transformer encoder architecture to learn
+    representations of molecular structures from SMILES strings.
+
+    Parameters
+    ----------
+    num_task : int, default=1
+        Number of prediction tasks.
+    task_type : str, default="regression"
+        Type of prediction task, either "regression" or "classification".
+    input_dim : int, default=54
+        Size of vocabulary for SMILES tokenization.
+    hidden_size : int, default=128
+        Dimension of embedding vectors.
+    n_heads : int, default=4
+        Number of attention heads in transformer layers.
+    num_layers : int, default=3
+        Number of transformer encoder layers.
+    dim_feedforward : int, default=256
+        Dimension of the feedforward network in transformer layers.
+    max_input_len : int, default=200
+        Maximum length of input sequences. Shorter sequences will be padded.
+    dropout : float, default=0.1
+        Dropout rate for transformer layers.
+    batch_size : int, default=64
+        Number of samples per batch for training.
+    epochs : int, default=200
+        Maximum number of training epochs.
+    loss_criterion : callable, optional
+        Loss function for training. Defaults to MSELoss for regression.
+    evaluate_criterion : str or callable, optional
+        Metric for model evaluation.
+    evaluate_higher_better : bool, optional
+        Whether higher values of the evaluation metric are better.
+    learning_rate : float, default=0.0001
+        Learning rate for optimizer.
+    weight_decay : float, default=0.0
+        L2 regularization strength.
+    patience : int, default=20
+        Number of epochs to wait for improvement before early stopping.
+    use_lr_scheduler : bool, default=True
+        Whether to use learning rate scheduler.
+    scheduler_factor : float, default=0.5
+        Factor by which to reduce learning rate when plateau is reached.
+    scheduler_patience : int, default=5
+        Number of epochs with no improvement after which learning rate will be reduced.
+    verbose : bool, default=False
+        Whether to print progress information during training.
+    """
+    # Model parameters
+    num_task: int = 1
+    task_type: str = "regression"
+    input_dim: int = 54  # vocabulary size
+    hidden_size: int = 128
+    n_heads: int = 4
+    num_layers: int = 3
+    dim_feedforward: Optional[int] = 256
+    max_input_len: int = 200  # max token length
+    dropout: float = 0.1
+    
+    # Training parameters
+    batch_size: int = 64
+    epochs: int = 200
+    loss_criterion: Optional[Callable] = None
+    evaluate_criterion: Optional[Union[str, Callable]] = None
+    evaluate_higher_better: Optional[bool] = None
+    learning_rate: float = 0.0001
+    weight_decay: float = 0.0
+    patience: int = 20
+
+    # Scheduler parameters
+    use_lr_scheduler: bool = True
+    scheduler_factor: float = 0.5
+    scheduler_patience: int = 5
+
+    # Other parameters
+    verbose: bool = False
+    model_name: str = "SMILESTransformerMolecularPredictor"
+    
+    # Non-init fields
+    fitting_loss: List[float] = field(default_factory=list, init=False)
+    fitting_epoch: int = field(default=0, init=False)
+    model_class: Type[Transformer] = field(default=Transformer, init=False)
+
+    def __post_init__(self):
+        """Initialize after dataclass initialization."""
+        super().__post_init__()
+        # Validate n_heads is compatible with hidden_size
+        if self.hidden_size % self.n_heads != 0:
+            raise ValueError(f"hidden_size ({self.hidden_size}) must be divisible by n_heads ({self.n_heads})")
+            
+        # Setup loss criterion and evaluation
+        if self.loss_criterion is None:
+            self.loss_criterion = nn.MSELoss()
+        self._setup_evaluation(self.evaluate_criterion, self.evaluate_higher_better)
+
+    @staticmethod
+    def _get_param_names() -> List[str]:
+        """Get parameter names for the estimator.
+
+        Returns
+        -------
+        List[str]
+            List of parameter names that can be used for model configuration.
+        """
+        return [
+            # Model Hyperparameters
+            "num_task",
+            "task_type",
+            "input_dim",
+            "hidden_size",
+            "n_heads",
+            "num_layers",
+            "dim_feedforward",
+            "max_input_len",
+            "dropout",
+            # Training Parameters
+            "batch_size",
+            "epochs",
+            "learning_rate",
+            "weight_decay",
+            "patience",
+            "loss_criterion",
+            # Evaluation Parameters
+            "evaluate_name",
+            "evaluate_criterion",
+            "evaluate_higher_better",
+            # Scheduler Parameters
+            "use_lr_scheduler",
+            "scheduler_factor",
+            "scheduler_patience",
+            # Other Parameters
+            "fitting_epoch",
+            "fitting_loss",
+            "device",
+            "verbose"
+        ]
+    
+    def _get_model_params(self, checkpoint: Optional[Dict] = None) -> Dict[str, Any]:
+        """Get model parameters from checkpoint or current instance.
+        
+        Parameters
+        ----------
+        checkpoint : Optional[Dict], default=None
+            Optional dictionary containing model checkpoint data
+            
+        Returns
+        -------
+        Dict[str, Any]
+            Dictionary of model parameters
+        """
+        if checkpoint is not None:
+            if "hyperparameters" not in checkpoint:
+                raise ValueError("Checkpoint missing 'hyperparameters' key")
+                
+            hyperparameters = checkpoint["hyperparameters"]
+            
+
+            return {
+                "num_task": hyperparameters.get("num_task", self.num_task),
+                "input_dim": hyperparameters.get("input_dim", self.input_dim),
+                "hidden_size": hyperparameters.get("hidden_size", self.hidden_size),
+                "n_heads": hyperparameters.get("n_heads", self.n_heads),
+                "num_layers": hyperparameters.get("num_layers", self.num_layers),
+                "dim_feedforward": hyperparameters.get("dim_feedforward", self.dim_feedforward),
+                "max_input_len": hyperparameters.get("max_input_len", self.max_input_len),
+                "dropout": hyperparameters.get("dropout", self.dropout),
+            }
+        else:
+            return {
+                "num_task": self.num_task,
+                "input_dim": self.input_dim,
+                "hidden_size": self.hidden_size,
+                "n_heads": self.n_heads,
+                "num_layers": self.num_layers,
+                "dim_feedforward": self.dim_feedforward,
+                "max_input_len": self.max_input_len,
+                "dropout": self.dropout,
+            }
+
+    def _setup_optimizers(self) -> Tuple[torch.optim.Optimizer, Optional[Any]]:
+        """Setup optimization components including optimizer and learning rate scheduler.
+
+        Returns
+        -------
+        Tuple[optim.Optimizer, Optional[Any]]
+            A tuple containing:
+            - The configured optimizer
+            - The learning rate scheduler (if enabled, else None)
+        """
+        optimizer = torch.optim.AdamW(
+            self.model.parameters(),
+            lr=self.learning_rate,
+            weight_decay=self.weight_decay
+        )
+        
+        scheduler = None
+        if self.use_lr_scheduler:
+            scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+                optimizer,
+                mode='min' if not self.evaluate_higher_better else 'max',
+                factor=self.scheduler_factor,
+                patience=self.scheduler_patience,
+                verbose=self.verbose
+            )
+        
+        return optimizer, scheduler
+
+
+    def _get_default_search_space(self):
+        """Get the default hyperparameter search space.
+        """
+        return DEFAULT_TRANSFORMER_SEARCH_SPACES

torch_molecule/predictor/ssr/__init__.py

@@ -0,0 +1,3 @@
+from .modeling_ssr import SSRMolecularPredictor
+
+__all__ = ['SSRMolecularPredictor']

torch_molecule/predictor/ssr/model.py

@@ -0,0 +1,283 @@
+import torch
+import torch.nn as nn
+from torch_geometric.nn import global_add_pool, global_mean_pool, global_max_pool
+import copy
+# from torch_scatter import scatter
+from .utils import scatter
+
+from ...nn import GNN_node, GNN_node_Virtualnode, MLP
+from ...utils import init_weights
+
+class SSR(nn.Module):
+    def __init__(
+        self,
+        num_task,
+        num_layer,
+        coarse_pool='mean',
+        hidden_size=300,
+        gnn_type="gin-virtual",
+        drop_ratio=0.5,
+        norm_layer="batch_norm",
+        graph_pooling="mean",
+        augmented_feature=['maccs', 'morgan']
+    ):
+
+        super(SSR, self).__init__()
+        gnn_name = gnn_type.split("-")[0]
+        self.num_task = num_task
+        self.hidden_size = hidden_size
+        self.coarse_pool = coarse_pool
+
+        if "virtual" in gnn_type:
+            self.graph_encoder = GNN_node_Virtualnode(
+                num_layer,
+                hidden_size,
+                JK="last",
+                drop_ratio=drop_ratio,
+                residual=True,
+                gnn_name=gnn_name,
+                norm_layer=norm_layer,
+            )
+        else:
+            self.graph_encoder = GNN_node(
+                num_layer,
+                hidden_size,
+                JK="last",
+                drop_ratio=drop_ratio,
+                residual=True,
+                gnn_name=gnn_name,
+                norm_layer=norm_layer,
+            )
+        if graph_pooling == "sum":
+            self.pool = global_add_pool
+        elif graph_pooling == "mean":
+            self.pool = global_mean_pool
+        elif graph_pooling == "max":
+            self.pool = global_max_pool
+        else:
+            raise ValueError(f"Invalid graph pooling type {graph_pooling}.")
+
+        graph_dim = hidden_size
+        self.augmented_feature = augmented_feature
+        if augmented_feature:
+            if "morgan" in augmented_feature:
+                graph_dim += 1024
+            if "maccs" in augmented_feature:
+                graph_dim += 167
+        self.predictor = MLP(graph_dim, hidden_features=2 * hidden_size, out_features=num_task)
+    
+    def l2diff(self, x1, x2):
+        """Compute L2 difference between tensors"""
+        return torch.sqrt(torch.sum((x1 - x2) ** 2, dim=1))
+        
+    def moment_diff(self, sx1, sx2, k, og_batch, coarse_batch):
+        """Compute difference between k-th moments"""
+        ss1 = scatter(sx1**k, og_batch, dim=0, reduce='mean')
+        ss2 = scatter(sx2**k, coarse_batch, dim=0, reduce='mean')
+        return self.l2diff(ss1, ss2)
+        
+    def cmd(self, x1, x2, og_batch, coarse_batch,n_moments):
+        """Central Moment Discrepancy between x1 and x2 distributions"""
+        # Compute means
+        mx1 = scatter(x1, og_batch, dim=0, reduce='mean')
+        mx2 = scatter(x2, coarse_batch, dim=0, reduce='mean')
+        # Center features
+        sx1 = x1 - mx1[og_batch]
+        sx2 = x2 - mx2[coarse_batch]
+        
+        # First moment difference
+        dm = self.l2diff(mx1, mx2)
+        scms = dm
+        
+        # Higher order moments
+        for i in range(n_moments-1):
+            scms = scms + self.moment_diff(sx1, sx2, i+2, og_batch, coarse_batch)
+            
+        return scms.mean()
+            
+    def pool_batch_by_graph(self, clusters, original_batch):
+        """
+        Build a new batch vector for coarse nodes so that all coarse nodes derived from the same
+        original graph get the same graph label. Here, we iterate graph-by-graph.
+        
+        Args:
+            clusters (Tensor): Cluster assignments for each original node, shape [num_original_nodes].
+            original_batch (Tensor): Original batch vector (graph assignments) for each node.
+        
+        Returns:
+            Tensor: New batch vector for coarse nodes, of shape [num_coarse_nodes].
+        """
+        unique_graphs = torch.unique(original_batch)
+        new_batch_list = []
+        for g in unique_graphs:
+            # Get indices for nodes belonging to graph g.
+            mask = (original_batch == g)
+            # For these nodes, get their cluster ids.
+            clusters_in_g = clusters[mask]
+            # Determine the unique clusters in this graph.
+            unique_clusters_in_g = torch.unique(clusters_in_g)
+            # For each coarse node from graph g, assign the same graph id (g.item())
+            new_batch_list.extend([g.item()] * unique_clusters_in_g.numel())
+        return torch.tensor(new_batch_list, device=original_batch.device, dtype=torch.long)
+
+
+    def prepare_coarsened_batch(self, batch, ratio):
+        """
+        Create a coarsened version of batch for forward pass with edge attributes,
+        ensuring that the new batch assignment preserves the original graph labels.
+        """
+        new_batch = copy.deepcopy(batch)
+        coarse_ratio_postfix = str(int(ratio * 100))
+        
+        # 1. Get clusters (one per original node) and ensure long type.
+        clusters = getattr(batch, f"clusters_{coarse_ratio_postfix}").long()  # shape: [num_nodes]
+        
+        # Compute a factor to disambiguate clusters from different graphs.
+        factor = clusters.max().item() + 1  # All clusters are in [0, factor-1]
+        # Create a global key that combines the original batch assignment and clusters.
+        global_coarse_key = batch.batch * factor + clusters  # shape: [num_nodes]
+        
+        # Get the unique global keys and the inverse mapping.
+        unique_global_keys, inverse_indices = torch.unique(global_coarse_key, return_inverse=True)
+        num_coarse_nodes = unique_global_keys.numel()
+        
+        # 2. Pool node features using the global coarse key.
+        new_x = scatter(batch.x, inverse_indices, dim=0, reduce=self.coarse_pool, dim_size=num_coarse_nodes)
+        new_batch.x = new_x
+
+        # 3. Build the new batch vector for coarse nodes:
+        # The graph assignment for each coarse node is simply:
+        #   graph_id = unique_global_key // factor
+        new_batch_vec = unique_global_keys // factor
+        new_batch.batch = new_batch_vec
+        
+        # 4. Process the coarsened edge index.
+        # Assume that the stored coarsened edge index is based on original node indices.
+        # Map each node index to its global coarse key.
+        coarse_edge_index = getattr(batch, f"coarsened_edge_index_{coarse_ratio_postfix}")
+        # For each index in coarse_edge_index, replace with its global coarse key.
+        temp_edge_index = torch.stack([
+            global_coarse_key[coarse_edge_index[0]],
+            global_coarse_key[coarse_edge_index[1]]
+        ])
+        # Build a mapping from global key to new index from unique_global_keys.
+        mapping = {old.item(): new for new, old in enumerate(unique_global_keys)}
+        new_edge_index = torch.zeros_like(temp_edge_index)
+        for i in range(temp_edge_index.size(1)):
+            new_edge_index[0, i] = mapping[temp_edge_index[0, i].item()]
+            new_edge_index[1, i] = mapping[temp_edge_index[1, i].item()]
+        new_batch.edge_index = new_edge_index
+        
+        # 5. Process edge attributes, if available.
+        if hasattr(batch, f"coarsened_edge_attr_{coarse_ratio_postfix}"):
+            edge_attr = getattr(batch, f"coarsened_edge_attr_{coarse_ratio_postfix}")
+            new_batch.edge_attr = torch.round(edge_attr).long()
+        
+        new_batch.num_nodes = num_coarse_nodes
+        '''
+        # Debug prints:
+        print("\n=== Feature and Cluster Info ===")
+        print(f"Original x shape: {batch.x.shape}")
+        print(f"Original batch shape: {batch.batch.shape}")
+        print(f"Number of unique clusters (globally disambiguated): {num_coarse_nodes}")
+        print("\n=== Final Dimensions ===")
+        print(f"New x shape: {new_batch.x.shape}")      # Should be [num_coarse_nodes, feature_dim]
+        print(f"New batch shape: {new_batch.batch.shape}")  # Should be [num_coarse_nodes]
+        print(f"New edge_index shape: {new_batch.edge_index.shape}")
+        '''
+        # Validate dimensions.
+        assert new_batch.batch.numel() == num_coarse_nodes, \
+            f"Batch size {new_batch.batch.numel()} doesn't match number of coarse nodes {num_coarse_nodes}"
+        assert new_batch.x.shape[0] == num_coarse_nodes, \
+            f"Number of node features {new_batch.x.shape[0]} doesn't match number of coarse nodes {num_coarse_nodes}"
+        
+        return new_batch
+
+
+
+    def initialize_parameters(self, seed=None):
+        """
+        Randomly initialize all model parameters using the init_weights function.
+        
+        Args:
+            seed (int, optional): Random seed for reproducibility. Defaults to None.
+        """
+        if seed is not None:
+            torch.manual_seed(seed)
+        
+        # Initialize the main components
+        init_weights(self.graph_encoder)
+        init_weights(self.predictor)
+        
+        # Reset all parameters using PyTorch Geometric's reset function
+        def reset_parameters(module):
+            if hasattr(module, 'reset_parameters'):
+                module.reset_parameters()
+            elif hasattr(module, 'weight') and hasattr(module.weight, 'data'):
+                init_weights(module)
+        
+        self.apply(reset_parameters)
+
+    def _augmented_graph_features(self, batched_data, h_rep):
+        if self.augmented_feature:
+            if 'morgan' in self.augmented_feature:
+                morgan = batched_data.morgan.type_as(h_rep)
+                h_rep = torch.cat((h_rep, morgan), dim=1)
+            if 'maccs' in self.augmented_feature:
+                maccs = batched_data.maccs.type_as(h_rep)
+                h_rep = torch.cat((h_rep, maccs), dim=1)
+        return h_rep
+
+
+    def compute_loss(self, batched_data, criterion, coarse_ratios=[0.8, 0.9],cmd_coeff=0.1,fine_grained=True,n_moments=5):
+        """Compute loss with SSR regularization"""
+        # Original forward pass
+        h_node, _ = self.graph_encoder(batched_data)
+        h_rep = self.pool(h_node, batched_data.batch)
+        h_rep = self._augmented_graph_features(batched_data, h_rep)
+        prediction = self.predictor(h_rep)
+        
+        # Compute prediction loss
+        target = batched_data.y.to(torch.float32)
+        is_labeled = batched_data.y == batched_data.y  # Check for NaN
+        pred_loss = criterion(prediction.to(torch.float32)[is_labeled], target[is_labeled]).mean()
+        
+        # Initialize SSR loss
+        ssr_loss = torch.tensor(0.0, device=batched_data.x.device)
+        
+        # Process each coarsening ratio for SSR
+        for ratio in coarse_ratios:
+            # Skip if coarsened data wasn't pre-computed
+            if not hasattr(batched_data, f"coarsened_edge_index_{str(int(ratio*100))}"):
+                continue
+                
+            # Prepare coarsened batch
+            coarse_batch = self.prepare_coarsened_batch(batched_data, ratio)
+            
+            # Get node embeddings for coarsened graph
+            coarse_h_node, _ = self.graph_encoder(coarse_batch)
+            
+            # Compute CMD loss
+            if fine_grained:
+                # Node-level CMD
+                ssr_loss = ssr_loss + self.cmd(h_node, 
+                                              coarse_h_node,
+                                              batched_data.batch,
+                                              coarse_batch.batch,n_moments)
+            else:
+                # Graph-level CMD
+                coarse_h_rep = self.pool(coarse_h_node, coarse_batch.batch)
+                ssr_loss = ssr_loss + torch.norm(h_rep - coarse_h_rep, dim=1).mean()
+        
+        # Compute total loss
+        total_loss = pred_loss + cmd_coeff * ssr_loss
+        
+        return total_loss, pred_loss, ssr_loss
+        
+    def forward(self, batched_data):
+        """Standard forward pass (for inference without SSR)"""
+        h_node, _ = self.graph_encoder(batched_data)
+        h_rep = self.pool(h_node, batched_data.batch)
+        h_rep = self._augmented_graph_features(batched_data, h_rep)
+        prediction = self.predictor(h_rep)
+        return {"prediction": prediction}