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

torch_molecule/encoder/edgepred/modeling_edgepred.py

@@ -0,0 +1,285 @@
+import os
+import numpy as np
+from tqdm import tqdm
+from typing import Optional, Union, Dict, Any, Tuple, List, Callable, Literal, Type
+import warnings
+from dataclasses import dataclass, field
+
+import torch
+from torch_geometric.loader import DataLoader
+from torch_geometric.data import Data
+
+from .model import GNN
+from ..constant import GNN_ENCODER_MODELS, GNN_ENCODER_READOUTS, GNN_ENCODER_PARAMS
+from ...base import BaseMolecularEncoder
+from ...utils import graph_from_smiles
+
+ALLOWABLE_ENCODER_MODELS = GNN_ENCODER_MODELS
+ALLOWABLE_ENCODER_READOUTS = GNN_ENCODER_READOUTS
+
+@dataclass
+class EdgePredMolecularEncoder(BaseMolecularEncoder):
+    """This encoder implements a GNN-based model for molecular representation learning
+    using the edge prediction.
+    
+    References
+    ----------
+    - Paper: Strategies for Pre-training Graph Neural Networks (ICLR 2020) https://arxiv.org/abs/1905.12265
+    - Code: https://github.com/snap-stanford/pretrain-gnns/tree/master/chem
+
+    Parameters
+    ----------
+    num_layer : int, default=5
+        Number of GNN layers.
+    hidden_size : int, default=300
+        Dimension of hidden node features.
+    drop_ratio : float, default=0.5
+        Dropout probability.
+    norm_layer : str, default="batch_norm"
+        Type of normalization layer to use. One of ["batch_norm", "layer_norm", "instance_norm", "graph_norm", "size_norm", "pair_norm"].
+    encoder_type : str, default="gin-virtual"
+        Type of GNN architecture to use.
+    readout : str, default="sum"
+        Method for aggregating node features to obtain graph-level representations.
+    batch_size : int, default=128
+        Number of samples per batch for training.
+    epochs : int, default=500
+        Maximum number of training epochs.
+    learning_rate : float, default=0.001
+        Learning rate for optimizer.
+    grad_clip_value : float, optional
+        Maximum norm of gradients for gradient clipping.
+    weight_decay : float, default=0.0
+        L2 regularization strength.
+    use_lr_scheduler : bool, default=False
+        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_layer: int = 5
+    hidden_size: int = 300
+    drop_ratio: float = 0.5
+    norm_layer: str = "batch_norm"
+    
+    encoder_type: str = "gin-virtual"
+    readout: str = "sum"
+    
+    # Training parameters
+    batch_size: int = 128
+    epochs: int = 500
+    learning_rate: float = 0.001
+    grad_clip_value: Optional[float] = None
+    weight_decay: float = 0.0
+    
+    # Scheduler parameters
+    use_lr_scheduler: bool = False
+    scheduler_factor: float = 0.5
+    scheduler_patience: int = 5
+    
+    # Other parameters
+    verbose: bool = False
+    model_name: str = "EdgePredMolecularEncoder"
+    
+    # Non-init fields
+    fitting_loss: List[float] = field(default_factory=list, init=False)
+    fitting_epoch: int = field(default=0, init=False)
+    model_class: Type[GNN] = field(default=GNN, init=False)
+    
+    def __post_init__(self):
+        """Initialize the model after dataclass initialization."""
+        super().__post_init__()
+        if self.encoder_type not in ALLOWABLE_ENCODER_MODELS:
+            raise ValueError(f"Invalid encoder_model: {self.encoder_type}. Currently only {ALLOWABLE_ENCODER_MODELS} are supported.")
+        if self.readout not in ALLOWABLE_ENCODER_READOUTS:
+            raise ValueError(f"Invalid encoder_readout: {self.readout}. Currently only {ALLOWABLE_ENCODER_READOUTS} are supported.")
+        
+    @staticmethod
+    def _get_param_names() -> List[str]:
+        return GNN_ENCODER_PARAMS.copy()
+    
+    def _get_model_params(self, checkpoint: Optional[Dict] = None) -> Dict[str, Any]:
+        params = {
+            "num_layer": self.num_layer,
+            "hidden_size": self.hidden_size,
+            "encoder_type": self.encoder_type,
+            "drop_ratio": self.drop_ratio,
+            "norm_layer": self.norm_layer,
+            "readout": self.readout,
+        }
+        
+        if checkpoint is not None:
+            if "hyperparameters" not in checkpoint:
+                raise ValueError("Checkpoint missing 'hyperparameters' key")
+            params = {k: checkpoint["hyperparameters"].get(k, v) for k, v in params.items()}
+            
+        return params
+    
+    def _convert_to_pytorch_data(self, X):
+        """Convert numpy arrays to PyTorch Geometric data format.
+        """
+        if self.verbose:
+            iterator = tqdm(enumerate(X), desc="Converting molecules to graphs", total=len(X))
+        else:
+            iterator = enumerate(X)
+
+        pyg_graph_list = []
+        for idx, smiles_or_mol in iterator:
+            graph = graph_from_smiles(smiles_or_mol, None)
+            g = Data()
+            g.num_nodes = graph["num_nodes"]
+            g.edge_index = torch.from_numpy(graph["edge_index"])
+
+            del graph["num_nodes"]
+            del graph["edge_index"]
+
+            if graph["edge_feat"] is not None:
+                g.edge_attr = torch.from_numpy(graph["edge_feat"])
+                del graph["edge_feat"]
+
+            if graph["node_feat"] is not None:
+                g.x = torch.from_numpy(graph["node_feat"])
+                del graph["node_feat"]
+
+            pyg_graph_list.append(g)
+
+        return pyg_graph_list
+    
+    def _setup_optimizers(self) -> Tuple[torch.optim.Optimizer, Optional[Any]]:
+        """Setup optimization components including optimizer and learning rate scheduler.
+        """
+
+        optimizer = torch.optim.Adam(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",
+                factor=self.scheduler_factor,
+                patience=self.scheduler_patience,
+                min_lr=1e-6,
+                cooldown=0,
+                eps=1e-8,
+            )
+
+        return optimizer, scheduler
+    
+    def fit(
+        self,
+        X_train: List[str],
+    ) -> "EdgePredMolecularEncoder":
+        """Fit the model to the training data with optional validation set.
+
+        Parameters
+        ----------
+        X_train : List[str]
+            Training set input molecular structures as SMILES strings
+        Returns
+        -------
+        self : EdgePredMolecularEncoder
+            Fitted estimator
+        """
+
+        self._initialize_model(self.model_class)
+        self.model.initialize_parameters()
+        optimizer, scheduler = self._setup_optimizers()
+        
+        # Prepare datasets and loaders
+        X_train, _ = self._validate_inputs(X_train, return_rdkit_mol=True)
+        train_dataset = self._convert_to_pytorch_data(X_train)
+        train_loader = DataLoader(
+            train_dataset,
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=0
+        )
+        self.fitting_loss = []
+
+        for epoch in range(self.epochs):
+            # Training phase
+            train_losses = self._train_epoch(train_loader, optimizer, epoch)
+            self.fitting_loss.append(np.mean(train_losses))
+            if scheduler:
+                scheduler.step(np.mean(train_losses))
+
+        self.fitting_epoch = epoch
+        self.is_fitted_ = True
+        return self
+
+    def _train_epoch(self, train_loader, optimizer, epoch):
+        """Training logic for one epoch.
+
+        Args:
+            train_loader: DataLoader containing training data
+            optimizer: Optimizer instance for model parameter updates
+
+        Returns:
+            list: List of loss values for each training step
+        """
+        self.model.train()
+        losses = []
+
+        iterator = (
+            tqdm(train_loader, desc="Training", leave=False)
+            if self.verbose
+            else train_loader
+        )
+
+        for batch in iterator:
+            batch = batch.to(self.device)
+            optimizer.zero_grad()
+            loss = self.model.compute_loss(batch)
+            loss.backward()
+            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": f"{epoch}", "Loss": f"{loss.item():.4f}"})
+
+        return losses
+
+    def encode(self, X: List[str], return_type: Literal["np", "pt"] = "pt") -> Union[np.ndarray, torch.Tensor]:
+        """Encode molecules into vector representations.
+
+        Parameters
+        ----------
+        X : List[str]
+            List of SMILES strings
+        return_type : Literal["np", "pt"], default="pt"
+            Return type of the representations
+
+        Returns
+        -------
+        representations : ndarray or torch.Tensor
+            Molecular representations
+        """
+        self._check_is_fitted()
+
+        # Convert to PyTorch Geometric format and create loader
+        X, _ = self._validate_inputs(X, return_rdkit_mol=True)
+        dataset = self._convert_to_pytorch_data(X)
+        loader = DataLoader(dataset, batch_size=self.batch_size, shuffle=False)
+
+        if self.model is None:
+            raise RuntimeError("Model not initialized")
+
+        # Generate encodings
+        self.model = self.model.to(self.device)
+        self.model.eval()
+        encodings = []
+        with torch.no_grad():
+            for batch in tqdm(loader, disable=not self.verbose):
+                batch = batch.to(self.device)
+                out = self.model(batch)
+                encodings.append(out["graph"].cpu())
+
+        # Concatenate and convert to requested format
+        encodings = torch.cat(encodings, dim=0)
+        return encodings if return_type == "pt" else encodings.numpy()

torch_molecule/encoder/graphmae/__init__.py

@@ -0,0 +1,3 @@
+from .modeling_graphmae import GraphMAEMolecularEncoder
+
+__all__ = ['GraphMAEMolecularEncoder']

torch_molecule/encoder/graphmae/dataloader.py

@@ -0,0 +1,206 @@
+import random
+import copy
+import torch
+import torch.nn.functional as F
+from torch_geometric.data import Data
+
+class DataLoaderMaskingPred(torch.utils.data.DataLoader):
+    r"""Data loader which merges data objects from a
+    :class:`torch_geometric.data.dataset` to a mini-batch.
+    Args:
+        dataset (Dataset): The dataset from which to load the data.
+        batch_size (int, optional): How may samples per batch to load.
+            (default: :obj:`1`)
+        shuffle (bool, optional): If set to :obj:`True`, the data will be
+            reshuffled at every epoch (default: :obj:`True`)
+    """
+
+    def __init__(self, dataset, batch_size=1, shuffle=True, mask_rate=0.0, mask_edge=False, **kwargs):
+        self._transform = MaskAtom(num_atom_type = 119, num_edge_type = 5, mask_rate = mask_rate, mask_edge=mask_edge)
+        super(DataLoaderMaskingPred, self).__init__(
+            dataset,
+            batch_size,
+            shuffle,
+            collate_fn=self.collate_fn,
+            **kwargs)
+    
+    def collate_fn(self, batches):
+        batchs = [self._transform(x) for x in batches]
+        return BatchMasking.from_data_list(batchs)
+
+class MaskAtom:
+    def __init__(self, num_atom_type, num_edge_type, mask_rate, mask_edge=True):
+        """
+        Randomly masks an atom, and optionally masks edges connecting to it.
+        The mask atom type index is num_possible_atom_type
+        The mask edge type index in num_possible_edge_type
+        :param num_atom_type:
+        :param num_edge_type:
+        :param mask_rate: % of atoms to be masked
+        :param mask_edge: If True, also mask the edges that connect to the
+        masked atoms
+        """
+        self.num_atom_type = num_atom_type
+        self.num_edge_type = num_edge_type
+        self.mask_rate = mask_rate
+        self.mask_edge = mask_edge
+        # self.num_chirality_tag = 3
+        # self.num_bond_direction = 3 
+
+    def __call__(self, input_data, masked_atom_indices=None):
+        """
+
+        :param data: pytorch geometric data object. Assume that the edge
+        ordering is the default pytorch geometric ordering, where the two
+        directions of a single edge occur in pairs.
+        Eg. data.edge_index = tensor([[0, 1, 1, 2, 2, 3],
+                                     [1, 0, 2, 1, 3, 2]])
+        :param masked_atom_indices: If None, then randomly samples num_atoms
+        * mask rate number of atom indices
+        Otherwise a list of atom idx that sets the atoms to be masked (for
+        debugging only)
+        :return: None, Creates new attributes in original data object:
+        data.mask_node_idx
+        data.mask_node_label
+        data.mask_edge_idx
+        data.mask_edge_label
+        """
+        data = copy.deepcopy(input_data)
+        if masked_atom_indices == None:
+            # sample x distinct atoms to be masked, based on mask rate. But
+            # will sample at least 1 atom
+            num_atoms = data.x.size()[0]
+            sample_size = int(num_atoms * self.mask_rate + 1)
+            masked_atom_indices = random.sample(range(num_atoms), sample_size)
+
+        # create mask node label by copying atom feature of mask atom
+        mask_node_labels_list = []
+        for atom_idx in masked_atom_indices:
+            mask_node_labels_list.append(data.x[atom_idx].view(1, -1))
+        data.mask_node_label = torch.cat(mask_node_labels_list, dim=0)
+        data.masked_atom_indices = torch.tensor(masked_atom_indices)
+
+        # ----------- graphMAE -----------
+        atom_type = F.one_hot(data.mask_node_label[:, 0], num_classes=self.num_atom_type).float()
+        # atom_chirality = F.one_hot(data.mask_node_label[:, 1], num_classes=self.num_chirality_tag).float()
+        # data.node_attr_label = torch.cat((atom_type,atom_chirality), dim=1)
+        data.node_attr_label = atom_type
+
+        # modify the original node feature of the masked node
+        for atom_idx in masked_atom_indices:
+            feat_dim = data.x[atom_idx].shape[0]
+            data.x[atom_idx] = torch.zeros(feat_dim)
+            data.x[atom_idx][0] = self.num_atom_type - 1
+
+        if self.mask_edge:
+            # create mask edge labels by copying edge features of edges that are bonded to
+            # mask atoms
+            connected_edge_indices = []
+            for bond_idx, (u, v) in enumerate(data.edge_index.cpu().numpy().T):
+                for atom_idx in masked_atom_indices:
+                    if atom_idx in set((u, v)) and \
+                        bond_idx not in connected_edge_indices:
+                        connected_edge_indices.append(bond_idx)
+
+            if len(connected_edge_indices) > 0:
+                # create mask edge labels by copying bond features of the bonds connected to
+                # the mask atoms
+                mask_edge_labels_list = []
+                for bond_idx in connected_edge_indices[::2]: # because the
+                    # edge ordering is such that two directions of a single
+                    # edge occur in pairs, so to get the unique undirected
+                    # edge indices, we take every 2nd edge index from list
+                    mask_edge_labels_list.append(
+                        data.edge_attr[bond_idx].view(1, -1))
+
+                data.mask_edge_label = torch.cat(mask_edge_labels_list, dim=0)
+                # modify the original bond features of the bonds connected to the mask atoms
+                for bond_idx in connected_edge_indices:
+                    feat_dim = data.edge_attr[bond_idx].shape[0]
+                    data.edge_attr[bond_idx] = torch.zeros(feat_dim)
+                    data.edge_attr[bond_idx][0] = self.num_edge_type - 1
+                    # print('feat_dim', feat_dim, 'data.edge_attr[bond_idx].shape', data.edge_attr[bond_idx].shape)
+
+                data.connected_edge_indices = torch.tensor(
+                    connected_edge_indices[::2])
+            else:
+                data.mask_edge_label = torch.empty((0, 2)).to(torch.int64)
+                data.connected_edge_indices = torch.tensor(
+                    connected_edge_indices).to(torch.int64)
+
+            edge_type = F.one_hot(data.mask_edge_label[:, 0], num_classes=self.num_edge_type).float()
+            # bond_direction = F.one_hot(data.mask_edge_label[:, 1], num_classes=self.num_bond_direction).float()
+            data.edge_attr_label = edge_type
+            # data.edge_attr_label = edge_type
+
+        return data
+
+    def __repr__(self):
+        return '{}(num_atom_type={}, num_edge_type={}, mask_rate={}, mask_edge={})'.format(
+            self.__class__.__name__, self.num_atom_type, self.num_edge_type,
+            self.mask_rate, self.mask_edge)
+
+class BatchMasking(Data):
+    r"""A plain old python object modeling a batch of graphs as one big
+    (dicconnected) graph. With :class:`torch_geometric.data.Data` being the
+    base class, all its methods can also be used here.
+    In addition, single graphs can be reconstructed via the assignment vector
+    :obj:`batch`, which maps each node to its respective graph identifier.
+    """
+
+    def __init__(self, batch=None, **kwargs):
+        super(BatchMasking, self).__init__(**kwargs)
+        self.batch = batch
+
+    @staticmethod
+    def from_data_list(data_list):
+        r"""Constructs a batch object from a python list holding
+        :class:`torch_geometric.data.Data` objects.
+        The assignment vector :obj:`batch` is created on the fly."""
+        keys = [set(data.keys()) for data in data_list]
+        keys = list(set.union(*keys))
+        assert 'batch' not in keys
+
+        batch = BatchMasking()
+
+        for key in keys:
+            batch[key] = []
+        batch.batch = []
+
+        cumsum_node = 0
+        cumsum_edge = 0
+
+        for i, data in enumerate(data_list):
+            num_nodes = data.num_nodes
+            batch.batch.append(torch.full((num_nodes, ), i, dtype=torch.long))
+            for key in data.keys():
+                item = data[key]
+                if key in ['edge_index', 'masked_atom_indices']:
+                    item = item + cumsum_node
+                elif key  == 'connected_edge_indices':
+                    item = item + cumsum_edge
+                batch[key].append(item)
+
+            cumsum_node += num_nodes
+            cumsum_edge += data.edge_index.shape[1]
+
+        for key in keys:
+            batch[key] = torch.cat(
+                batch[key], dim=data_list[0].__cat_dim__(key, batch[key][0]))
+        batch.batch = torch.cat(batch.batch, dim=-1)
+        return batch.contiguous()
+
+    def cumsum(self, key, item):
+        r"""If :obj:`True`, the attribute :obj:`key` with content :obj:`item`
+        should be added up cumulatively before concatenated together.
+        .. note::
+            This method is for internal use only, and should only be overridden
+            if the batch concatenation process is corrupted for a specific data
+            attribute.
+        """
+        return key in ['edge_index', 'face', 'masked_atom_indices', 'connected_edge_indices']
+
+    @property
+    def num_graphs(self):
+        """Returns the number of graphs in the batch."""
+        return self.batch[-1].item() + 1

torch_molecule/encoder/graphmae/model.py

@@ -0,0 +1,137 @@
+from functools import partial
+
+import torch
+import torch.nn as nn
+from torch_geometric.nn import global_add_pool, global_mean_pool, global_max_pool
+import torch.nn.functional as F
+
+from ...nn import GNN_node, GNN_node_Virtualnode
+from ...nn.gnn import GINConv, GCNConv
+from ...utils import init_weights
+
+def sce_loss(x, y, alpha=1):
+    x = F.normalize(x, p=2, dim=-1)
+    y = F.normalize(y, p=2, dim=-1)
+    loss = (1 - (x * y).sum(dim=-1)).pow_(alpha)
+    loss = loss.mean()
+    return loss
+
+class GNN(nn.Module):
+    def __init__(
+        self,
+        num_layer,
+        hidden_size,
+        drop_ratio=0.5,
+        norm_layer="batch_norm",
+        encoder_type="gin-virtual",
+        readout="sum",
+        mask_edge=False,
+        predictor_type="gin",
+    ):
+        super(GNN, self).__init__()
+        gnn_name = encoder_type.split("-")[0]
+        self.num_atom_type = 119
+        self.num_edge_type = 5
+        self.hidden_size = hidden_size
+
+        encoder_params = {
+            "num_layer": num_layer,
+            "hidden_size": hidden_size,
+            "JK": "last", 
+            "drop_ratio": drop_ratio,
+            "residual": True,
+            "gnn_name": gnn_name,
+            "norm_layer": norm_layer
+        }
+        
+        # Choose encoder type based on encoder_model
+        encoder_class = GNN_node_Virtualnode if "virtual" in encoder_type else GNN_node
+        self.graph_encoder = encoder_class(**encoder_params)
+        pooling_funcs = {
+            "sum": global_add_pool,
+            "mean": global_mean_pool,
+            "max": global_max_pool
+        }
+        self.pool = pooling_funcs.get(readout)
+        if self.pool is None:
+            raise ValueError(f"Invalid graph pooling type {readout}.")
+
+        self.atom_predictor = Predictor(hidden_size, self.num_atom_type, nn_type=predictor_type)
+        self.mask_edge = mask_edge
+        if mask_edge:
+            self.bond_predictor = Predictor(hidden_size, self.num_edge_type, nn_type="linear")
+    
+    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.atom_predictor)
+        if self.mask_edge:
+            init_weights(self.bond_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 compute_loss(self, batched_data, alpha_l=1):
+        criterion = partial(sce_loss, alpha=alpha_l)
+
+        h_node, _ = self.graph_encoder(batched_data)
+
+        node_attr_label = batched_data.node_attr_label
+        masked_node_indices = batched_data.masked_atom_indices
+        pred_node = self.atom_predictor(h_node, batched_data.edge_index, batched_data.edge_attr, masked_node_indices)
+        loss_atom = criterion(pred_node[masked_node_indices], node_attr_label.to(pred_node.dtype))
+        if self.mask_edge:
+            masked_edge_index = batched_data.edge_index[:, batched_data.connected_edge_indices]
+            edge_rep = h_node[masked_edge_index[0]] + h_node[masked_edge_index[1]]
+            pred_edge = self.bond_predictor(edge_rep)
+            loss_edge = criterion(pred_edge, batched_data.edge_attr_label.to(pred_edge.dtype))
+        else:
+            loss_edge = torch.tensor(0.0)
+        return loss_atom, loss_edge
+
+    def forward(self, batched_data):
+        h_node, _ = self.graph_encoder(batched_data)
+        h_rep = self.pool(h_node, batched_data.batch)
+        return {"graph": h_rep, "node": h_node}
+
+class Predictor(torch.nn.Module):
+    def __init__(self, hidden_size, output_size, nn_type = "gin"):
+        super().__init__()
+        self._dec_type = nn_type 
+        if "gin" in nn_type:
+            self.conv = GINConv(hidden_size, output_size)
+        elif "gcn" in nn_type:
+            self.conv = GCNConv(hidden_size, output_size)
+        elif "linear" in nn_type:
+            self.dec = torch.nn.Linear(hidden_size, output_size)
+        else:
+            raise NotImplementedError(f"{nn_type}")
+        if "linear" not in nn_type: 
+            self.activation = torch.nn.PReLU() 
+            self.enc_to_dec = torch.nn.Linear(hidden_size, hidden_size, bias=False)
+
+    def forward(self, x, edge_index=None, edge_attr=None, mask_node_indices=None):
+        if self._dec_type == "linear":
+            out = self.dec(x)
+        else:
+            assert edge_index is not None and edge_attr is not None and mask_node_indices is not None
+            x = self.activation(x)
+            x = self.enc_to_dec(x)
+            x[mask_node_indices] = 0
+            out = self.conv(x, edge_index, edge_attr)
+        return out