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

torch_molecule/predictor/rpgnn/model.py

@@ -0,0 +1,185 @@
+import torch
+import torch
+import torch.nn as nn
+from torch_geometric.nn import global_add_pool, global_mean_pool, global_max_pool
+from ...nn import GNN_node, GNN_node_Virtualnode, MLP
+from ...utils import init_weights
+
+
+class RPGNN(nn.Module):
+    def __init__(
+        self,
+        num_task,
+        num_layer,
+        num_node_feature,
+        fixed_size=32,
+        num_perm=3,
+        hidden_size=300,
+        gnn_type="gin-virtual",
+        drop_ratio=0.5,
+        norm_layer="batch_norm",
+        graph_pooling="mean",
+        augmented_feature=['maccs', 'morgan']
+    ):
+        """
+        Random Permutation Graph Neural Network (RPGNN)
+        
+        Args:
+            num_task: Number of prediction tasks
+            num_layer: Number of GNN layers
+            num_node_feature: Number of node features
+            fixed_size: Size of position encoding vectors
+            num_perm: Number of random permutations to average over
+            hidden_size: Hidden dimension size
+            gnn_type: Type of GNN to use (gin, gin-virtual, gcn, gcn-virtual)
+            drop_ratio: Dropout ratio
+            norm_layer: Type of normalization layer
+            graph_pooling: Type of graph pooling (mean, sum, max)
+            augmented_feature: List of additional molecular features to use
+        """
+        super(RPGNN, self).__init__()
+        
+        gnn_name = gnn_type.split("-")[0]
+        self.num_task = num_task
+        self.hidden_size = hidden_size
+        self.fixed_size = fixed_size
+        self.num_perm = num_perm
+        self.augmented_feature = augmented_feature
+        self.num_node_feature = num_node_feature        
+        # Register fixed-size node IDs (position encoding)
+        self.register_buffer("node_ids", torch.eye(self.fixed_size))
+        
+        # GNN selection
+        if "virtual" in gnn_type:
+            self.gnn = GNN_node_Virtualnode(
+                num_layer=num_layer,
+                hidden_size=hidden_size,
+                JK="last",
+                drop_ratio=drop_ratio,
+                residual=True,
+                gnn_name=gnn_name,
+                norm_layer=norm_layer
+            )
+        else:
+            self.gnn = GNN_node(
+                num_layer=num_layer,
+                hidden_size=hidden_size,
+                JK="last",
+                drop_ratio=drop_ratio,
+                residual=True,
+                gnn_name=gnn_name,
+                norm_layer=norm_layer
+            )
+        
+        # Graph pooling selection
+        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}")
+            
+        # Calculate final prediction dimension
+        graph_dim = hidden_size+self.fixed_size
+        if augmented_feature:
+            if "morgan" in augmented_feature:
+                graph_dim += 1024
+            if "maccs" in augmented_feature:
+                graph_dim += 167
+                
+        # Final predictor MLP
+        self.predictor = MLP(
+            in_features=graph_dim,
+            hidden_features=2 * hidden_size,
+            out_features=num_task
+        )
+
+    def initialize_parameters(self, seed=None):
+        """Initialize all model parameters."""
+        if seed is not None:
+            torch.manual_seed(seed)
+            
+        # Initialize GNN and predictor
+        init_weights(self.gnn)
+        init_weights(self.predictor)
+        
+        # Apply reset_parameters to all modules that have it
+        def reset_parameters(module):
+            if hasattr(module, 'reset_parameters'):
+                module.reset_parameters()
+        
+        self.apply(reset_parameters)
+
+    def _augmented_graph_features(self, batched_data, h_rep):
+        """Add augmented molecular features to the graph representation."""
+        if not self.augmented_feature:
+            return h_rep
+            
+        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 forward(self, batched_data):
+        """Forward pass of RPGNN."""
+        # Extract original data
+        x, edge_index, batch = batched_data.x, batched_data.edge_index, batched_data.batch
+        
+        # Initialize final output
+        out = None
+        
+        # Average over multiple permutations
+        for _ in range(self.num_perm):
+            temp_data = batched_data.clone()
+            # Create new_x with same dtype as x
+            new_x = torch.empty(x.size(0), x.size(1) + self.fixed_size, dtype=x.dtype, device=x.device)
+            
+            for graph_idx in range(batch.max().item() + 1):
+                node_indices = (batch == graph_idx).nonzero().squeeze(1)
+                graph_size = node_indices.size(0)
+                perm = torch.randperm(graph_size)
+
+                node_ids = self.node_ids.repeat(
+                    graph_size // self.fixed_size + 1, 1
+                )[:graph_size]
+                permuted_node_ids = node_ids[perm, :]
+                
+                # Convert both tensors to the same type as x before concatenation
+                new_x[node_indices] = torch.cat([x[node_indices].to(x.dtype), permuted_node_ids.to(x.dtype)], dim=1)
+            temp_data.x = new_x[:,:self.num_node_feature]
+            # Process through GNN
+            h_v, _ = self.gnn(temp_data)
+
+            fused_h_v = torch.cat([h_v, new_x[:,self.num_node_feature:]], dim=1)
+
+            # Pool node representations to graph representations
+            h_rep = self.pool(fused_h_v, temp_data.batch)
+            
+            # Add augmented features
+            h_rep = self._augmented_graph_features(temp_data, h_rep)
+            
+            # Accumulate output
+            if out is None:
+                out = h_rep / self.num_perm
+            else:
+                out += h_rep / self.num_perm
+        
+        # Final prediction
+        prediction = self.predictor(out)
+        
+        return {"prediction": prediction}
+
+    def compute_loss(self, batched_data, criterion):
+        out = self.forward(batched_data)
+        prediction = out['prediction']
+        target = batched_data.y.to(torch.float32)
+        is_labeled = batched_data.y == batched_data.y
+        loss = criterion(prediction.to(torch.float32)[is_labeled], target[is_labeled]).mean()
+        
+        return loss

torch_molecule/predictor/rpgnn/modeling_rpgnn.py

@@ -0,0 +1,76 @@
+import os
+import numpy as np
+import warnings
+import datetime
+from tqdm import tqdm
+from typing import Optional, Union, Dict, Any, Tuple, List, Callable, Literal, Type
+from dataclasses import dataclass, field
+
+import torch
+from torch_geometric.loader import DataLoader
+
+from .model import RPGNN
+from ..gnn.modeling_gnn import GNNMolecularPredictor
+from ...utils.search import (
+    ParameterSpec,
+    ParameterType,
+)
+
+@dataclass
+class RPGNNMolecularPredictor(GNNMolecularPredictor):
+    """This predictor implements a GNN model based on Relational pooling.
+
+    The full name of RPGNN is Relational Pooling for Graph Representations.
+
+    References
+    ----------
+    - Relational Pooling for Graph Representations.
+      https://arxiv.org/abs/1903.02541
+    - Reference Code: https://github.com/PurdueMINDS/RelationalPooling/tree/master?tab=readme-ov-file
+
+    Parameters
+    ----------
+    num_perm : int, default=3
+        Number of random permutations to use for relational pooling. 
+    fixed_size : int, default=10 
+        Maximum number of nodes to consider in the graph. 
+    num_node_feature : int, default=9
+        Dimension of the input node features. This should match the number of atomic features used to represent
+        each node in the molecular graph (e.g., atomic number, degree, hybridization, etc.).
+    """
+    
+    # RPGNN-specific parameter
+    num_perm: int = 3
+    fixed_size: int = 10
+    num_node_feature: int = 9
+    # Override parent defaults
+    model_name: str = "RPGNNMolecularPredictor"
+    model_class: Type[RPGNN] = field(default=RPGNN, init=False)
+    
+    def __post_init__(self):
+        super().__post_init__()
+
+    @staticmethod
+    def _get_param_names() -> List[str]:
+        return ["num_perm", "fixed_size", "num_node_feature"] + GNNMolecularPredictor._get_param_names()
+
+    def _get_default_search_space(self):
+        search_space = super()._get_default_search_space().copy()
+        search_space["num_perm"] = ParameterSpec(ParameterType.INTEGER, (1, 10))
+        search_space["fixed_size"] = ParameterSpec(ParameterType.INTEGER, (1, 10))
+        return search_space
+
+    def _get_model_params(self, checkpoint: Optional[Dict] = None) -> Dict[str, Any]:
+        base_params = super()._get_model_params(checkpoint)
+        if checkpoint and "hyperparameters" in checkpoint:
+            base_params["num_perm"] = checkpoint["hyperparameters"]["num_perm"]
+            base_params["fixed_size"] = checkpoint["hyperparameters"]["fixed_size"]
+            base_params["num_node_feature"] = checkpoint["hyperparameters"]["num_node_feature"]
+        else:
+            base_params["num_perm"] = self.num_perm
+            base_params["fixed_size"] = self.fixed_size
+            base_params["num_node_feature"] = self.num_node_feature
+        base_params.pop("graph_pooling", None)
+        return base_params
+
+    

torch_molecule/predictor/sgir/__init__.py

@@ -0,0 +1,3 @@
+from .modeling_sgir import SGIRMolecularPredictor
+
+__all__ = ['SGIRMolecularPredictor']

torch_molecule/predictor/sgir/modeling_sgir.py

@@ -0,0 +1,260 @@
+import os
+import numpy as np
+import warnings
+import datetime
+from tqdm import tqdm
+from typing import Optional, Union, Dict, Any, Tuple, List, Callable, Literal
+from dataclasses import dataclass
+
+import torch
+from torch_geometric.loader import DataLoader
+
+from .strategy import build_selection_dataset, build_augmentation_dataset
+from ..grea.modeling_grea import GREAMolecularPredictor
+from ..grea.model import GREA
+
+from ...utils.search import (
+    ParameterSpec,
+    ParameterType,
+)
+
+@dataclass
+class SGIRMolecularPredictor(GREAMolecularPredictor):
+    """
+    This predictor implements SGIR for semi-supervised graph imbalanced regression.
+
+    It trains the GREA model based on pseudo-labeling and data augmentation.
+
+    References
+    ----------
+    - Semi-Supervised Graph Imbalanced Regression.
+      https://dl.acm.org/doi/10.1145/3580305.3599497
+    - Code: https://github.com/liugangcode/SGIR
+
+    :param num_anchor: Number of anchor points used to split the label space during pseudo-labeling
+    :type num_anchor: int, default=10
+    :param warmup_epoch: Number of epochs to train before starting pseudo-labeling and data augmentation
+    :type warmup_epoch: int, default=20
+    :param labeling_interval: Interval (in epochs) between pseudo-labeling steps
+    :type labeling_interval: int, default=5
+    :param augmentation_interval: Interval (in epochs) between data augmentation steps
+    :type augmentation_interval: int, default=5
+    :param top_quantile: Quantile threshold for selecting high confidence predictions during pseudo-labeling
+    :type top_quantile: float, default=0.1
+    :param label_logscale: Whether to use log scale for the label space during pseudo-labeling and data augmentation
+    :type label_logscale: bool, default=False
+    :param lw_aug: Weight for the data augmentation loss
+    :type lw_aug: float, default=1
+    """
+    # SGIR-specific parameters
+    num_anchor: int = 10
+    warmup_epoch: int = 20
+    labeling_interval: int = 5
+    augmentation_interval: int = 5
+    top_quantile: float = 0.1
+    label_logscale: bool = False
+    lw_aug: float = 1
+    # Override parent defaults
+    task_type: str = "regression"
+    model_name: str = "SGIRMolecularPredictor"
+    
+    def __post_init__(self):
+        super().__post_init__()
+        
+        if self.task_type != "regression" or self.num_task != 1:
+            raise ValueError("SGIR only supports regression tasks with 1 task")
+        
+    @staticmethod
+    def _get_param_names():
+        grea_params = [
+            "num_anchor", "warmup_epoch", "labeling_interval",
+            "augmentation_interval", "top_quantile", "label_logscale", "lw_aug"
+        ]
+        return grea_params + GREAMolecularPredictor._get_param_names()
+
+    def _get_default_search_space(self):
+        search_space = super()._get_default_search_space().copy()
+        search_space["num_anchor"] = ParameterSpec(ParameterType.INTEGER, (10, 100))
+        search_space["labeling_interval"] = ParameterSpec(ParameterType.INTEGER, (10, 20))
+        search_space["augmentation_interval"] = ParameterSpec(ParameterType.INTEGER, (10, 20))
+        search_space["top_quantile"] = ParameterSpec(ParameterType.LOG_FLOAT, (0.01, 0.5))
+        search_space["lw_aug"] = ParameterSpec(ParameterType.FLOAT, (0.1, 1))
+        return search_space
+
+    def fit(
+        self,
+        X_train: List[str],
+        y_train: Optional[Union[List, np.ndarray]],
+        X_val: Optional[List[str]] = None,
+        y_val: Optional[Union[List, np.ndarray]] = None,
+        X_unlbl: Optional[List[str]] = None,
+    ) -> "SGIRMolecularPredictor":
+        """Fit the model to training data with optional validation set.
+        """
+        if (X_val is None) != (y_val is None):
+            raise ValueError("X_val and y_val must both be provided for validation")
+        if X_unlbl is None:
+            raise ValueError("X_unlbl (unlabeled SMILES strings) must be provided in SGIR")
+        if len(X_unlbl) == 0:
+            raise ValueError("X_unlbl (unlabeled SMILES strings) must not be empty")
+
+        # Initialize model and optimization
+        self._initialize_model(self.model_class)
+        self.model.initialize_parameters()
+        optimizer, scheduler = self._setup_optimizers()
+        
+        # Prepare datasets
+        X_train, y_train = self._validate_inputs(X_train, y_train)
+        train_dataset = self._convert_to_pytorch_data(X_train, y_train)
+        train_loader = DataLoader(
+            train_dataset,
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=0
+        )
+
+        X_unlbl, _ = self._validate_inputs(X_unlbl, None)
+        unlbl_dataset = self._convert_to_pytorch_data(X_unlbl)
+
+        if X_val is None:
+            val_loader = train_loader
+            warnings.warn(
+                "No validation set provided. Using training set for validation.",
+                UserWarning
+            )
+        else:
+            X_val, y_val = self._validate_inputs(X_val, y_val)
+            val_dataset = self._convert_to_pytorch_data(X_val, y_val)
+            val_loader = DataLoader(
+                val_dataset,
+                batch_size=self.batch_size,
+                shuffle=False,
+                num_workers=0
+            )
+
+        # Training loop
+        augmented_dataset = None
+        self.fitting_loss = []
+        self.fitting_epoch = 0
+        best_state_dict = None
+        best_eval = float('-inf') if self.evaluate_higher_better else float('inf')
+        cnt_wait = 0
+
+        self.model.train()
+        for epoch in range(self.epochs):
+            # Training phase
+            train_losses = self._train_epoch(train_loader, augmented_dataset, optimizer, epoch)
+            
+            # Update datasets after warmup
+            if epoch > self.warmup_epoch:
+                if epoch % self.labeling_interval == 0:
+                    train_loader = build_selection_dataset(
+                        self.model, train_dataset, unlbl_dataset,
+                        self.batch_size, self.num_anchor, self.top_quantile,
+                        self.device, self.label_logscale
+                    )
+
+                if epoch % self.augmentation_interval == 0:
+                    augmented_dataset = build_augmentation_dataset(
+                        self.model, train_dataset, unlbl_dataset,
+                        self.batch_size, self.num_anchor, self.device, 
+                        self.label_logscale
+                    )
+
+            self.fitting_loss.append(np.mean(train_losses))
+
+            # Validation and model selection
+            current_eval = self._evaluation_epoch(val_loader)
+            if scheduler:
+                scheduler.step(current_eval)
+            
+            is_better = (
+                current_eval > best_eval if self.evaluate_higher_better
+                else current_eval < best_eval
+            )
+            
+            if is_better:
+                self.fitting_epoch = epoch
+                best_eval = current_eval
+                best_state_dict = self.model.state_dict()
+                cnt_wait = 0
+            else:
+                cnt_wait += 1
+                if cnt_wait > self.patience:
+                    if self.verbose:
+                        print(f"Early stopping at epoch {epoch}")
+                    break
+            
+            if self.verbose and epoch % 10 == 0:
+                print(
+                    f"Epoch {epoch}: Loss = {np.mean(train_losses):.4f}, "
+                    f"{self.evaluate_name} = {current_eval:.4f}, "
+                    f"Best {self.evaluate_name} = {best_eval:.4f}"
+                )
+
+        # Restore best model
+        if best_state_dict is not None:
+            self.model.load_state_dict(best_state_dict)
+        else:
+            warnings.warn(
+                "No improvement achieved during training.",
+                UserWarning
+            )
+
+        self.is_fitted_ = True
+        return self
+    
+    def _train_epoch(self, train_loader, augmented_dataset, optimizer, epoch):
+        losses = []
+
+        if augmented_dataset is not None and self.lw_aug != 0:
+            aug_reps = augmented_dataset['representations']
+            aug_targets = augmented_dataset['labels']
+            random_inds = torch.randperm(aug_reps.size(0))
+            aug_reps = aug_reps[random_inds]
+            aug_targets = aug_targets[random_inds]
+            num_step = len(train_loader)
+            aug_batch_size = aug_reps.size(0) // max(1, num_step)
+            aug_inputs = list(torch.split(aug_reps, aug_batch_size))
+            aug_outputs = list(torch.split(aug_targets, aug_batch_size))
+        else:
+            aug_inputs = None
+            aug_outputs = None
+
+        iterator = (
+            tqdm(train_loader, desc="Training", leave=False)
+            if self.verbose
+            else train_loader
+        )
+
+        for batch_idx, batch in enumerate(iterator):
+            batch = batch.to(self.device)
+            optimizer.zero_grad()
+
+            # augmentation loss
+            if aug_inputs is not None and aug_outputs is not None and aug_inputs[batch_idx].size(0) != 1:
+                self.model._disable_batchnorm_tracking(self.model)
+                pred_aug = self.model.predictor(aug_inputs[batch_idx])
+                self.model._enable_batchnorm_tracking(self.model)
+                targets_aug = aug_outputs[batch_idx]
+                Laug = self.loss_criterion(pred_aug.view(targets_aug.size()).to(torch.float32), targets_aug).mean()
+            else:
+                Laug = torch.tensor(0.)      
+            Lx = self.model.compute_loss(batch, self.loss_criterion)
+            loss = Lx + Laug * self.lw_aug
+
+            loss.backward()
+
+            # Compute gradient norm if gradient clipping is enabled
+            if self.grad_clip_value is not None:
+                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip_value)
+
+            optimizer.step()
+
+            losses.append(loss.item())
+
+            # Update progress bar if using tqdm
+            if self.verbose:
+                iterator.set_postfix({"Epoch": epoch, "Total Loss": f"{loss.item():.4f}", "Lbls Loss": f"{Lx.item():.4f}", "Aug Loss": f"{Laug.item():.4f}",})
+
+        return losses