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

torch_molecule/encoder/contextpred/modeling_contextpred.py

@@ -0,0 +1,299 @@
+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 ContextPredMolecularEncoder(BaseMolecularEncoder):
+    """This encoder implements a GNN-based model for molecular representation learning
+    using the context prediction pretraining strategy.
+
+    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
+    ----------
+    mode : str, default="cbow"
+        Type of context prediction task. One of ["cbow", "skipgram"].
+    context_size : int, default=2
+        Size of the context window used for predicting node-level features.
+    neg_samples : int, default=1
+        Number of negative samples used in the training objective.
+    num_layer : int, default=3
+        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. One of ["gin-virtual", "gcn-virtual", "gin", "gcn"].
+    readout : str, default="sum"
+        Method for aggregating node features to obtain graph-level representations. One of ["sum", "mean", "max"].
+    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 the 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 a learning rate scheduler during training.
+    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.
+    """
+    # Task related parameters
+    mode: str = "cbow" # cbow or skipgram
+    context_size: int = 2
+    neg_samples: int = 1
+    
+    # Model parameters
+    num_layer: int = 3
+    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 = "ContextPredMolecularEncoder"
+    
+    # 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 ["mode", "context_size", "neg_samples"] + 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,
+            "drop_ratio": self.drop_ratio,
+            "norm_layer": self.norm_layer,
+            "readout": self.readout,
+            "encoder_type": self.encoder_type,
+            "mode": self.mode,
+            "context_size": self.context_size,
+            "neg_samples": self.neg_samples
+        }
+        if checkpoint is not None:
+            if "hyperparameters" not in checkpoint:
+                raise ValueError("Checkpoint missing 'hyperparameters' key")                
+            hyperparameters = checkpoint["hyperparameters"]
+            params = {k: 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]
+    ) -> "ContextPredMolecularEncoder":
+        """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 : ContextPredMolecularEncoder
+            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/contextpred/utils.py

@@ -0,0 +1,206 @@
+import torch
+import copy
+import random
+import networkx as nx
+import numpy as np
+
+from torch_geometric.data import Data
+
+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", "edge_index_substruct", "edge_index_context", "overlap_context_substruct_idx", "center_substruct_idx"]
+
+class ExtractSubstructureContextPair:
+    def __init__(self, k, l1, l2):
+        """
+        Randomly selects a node from the data object, and adds attributes
+        that contain the substructure that corresponds to k hop neighbours
+        rooted at the node, and the context substructures that corresponds to
+        the subgraph that is between l1 and l2 hops away from the
+        root node.
+        :param k:
+        :param l1:
+        :param l2:
+        """
+        self.k = k
+        self.l1 = l1
+        self.l2 = l2
+
+        # for the special case of 0, addresses the quirk with
+        # single_source_shortest_path_length
+        if self.k == 0:
+            self.k = -1
+        if self.l1 == 0:
+            self.l1 = -1
+        if self.l2 == 0:
+            self.l2 = -1
+
+    def __call__(self, data, root_idx=None):
+        """
+
+        :param data: pytorch geometric data object
+        :param root_idx: If None, then randomly samples an atom idx.
+        Otherwise sets atom idx of root (for debugging only)
+        :return: None. Creates new attributes in original data object:
+        data.center_substruct_idx
+        data.x_substruct
+        data.edge_attr_substruct
+        data.edge_index_substruct
+        data.x_context
+        data.edge_attr_context
+        data.edge_index_context
+        data.overlap_context_substruct_idx
+        """
+        num_atoms = data.x.size()[0]
+        if root_idx == None:
+            root_idx = random.sample(range(num_atoms), 1)[0]
+
+        G = graph_data_obj_to_nx_simple(data)  # same ordering as input data obj
+
+        # Get k-hop subgraph rooted at specified atom idx
+        substruct_node_idxes = nx.single_source_shortest_path_length(G,
+                                                                     root_idx,
+                                                                     self.k).keys()
+        if len(substruct_node_idxes) > 0:
+            substruct_G = G.subgraph(substruct_node_idxes)
+            substruct_G, substruct_node_map = reset_idxes(substruct_G)  # need
+            # to reset node idx to 0 -> num_nodes - 1, otherwise data obj does not
+            # make sense, since the node indices in data obj must start at 0
+            substruct_data = nx_to_graph_data_obj_simple(substruct_G)
+            data.substruct_node_idxes = torch.tensor(list(substruct_node_idxes))
+            data.x_substruct = substruct_data.x
+            data.edge_attr_substruct = substruct_data.edge_attr
+            data.edge_index_substruct = substruct_data.edge_index
+            data.center_substruct_idx = torch.tensor([substruct_node_map[
+                                                          root_idx]])  # need
+            # to convert center idx from original graph node ordering to the
+            # new substruct node ordering
+
+        # Get subgraphs that is between l1 and l2 hops away from the root node
+        l1_node_idxes = nx.single_source_shortest_path_length(G, root_idx,
+                                                              self.l1).keys()
+        l2_node_idxes = nx.single_source_shortest_path_length(G, root_idx,
+                                                              self.l2).keys()
+        context_node_idxes = set(l1_node_idxes).symmetric_difference(
+            set(l2_node_idxes))
+        if len(context_node_idxes) > 0:
+            context_G = G.subgraph(context_node_idxes)
+            context_G, context_node_map = reset_idxes(context_G)  # need to
+            # reset node idx to 0 -> num_nodes - 1, otherwise data obj does not
+            # make sense, since the node indices in data obj must start at 0
+            context_data = nx_to_graph_data_obj_simple(context_G)
+            data.context_node_idxes = torch.tensor(list(context_node_idxes))
+            data.x_context = context_data.x
+            data.edge_attr_context = context_data.edge_attr
+            data.edge_index_context = context_data.edge_index
+
+        # Get indices of overlapping nodes between substruct and context,
+        # WRT context ordering
+        context_substruct_overlap_idxes = list(set(
+            context_node_idxes).intersection(set(substruct_node_idxes)))
+        if len(context_substruct_overlap_idxes) > 0:
+            context_substruct_overlap_idxes_reorder = [context_node_map[old_idx]
+                                                       for
+                                                       old_idx in
+                                                       context_substruct_overlap_idxes]
+            # need to convert the overlap node idxes, which is from the
+            # original graph node ordering to the new context node ordering
+            data.overlap_context_substruct_idx = \
+                torch.tensor(context_substruct_overlap_idxes_reorder)
+
+        return data
+    
+def reset_idxes(G):
+    """
+    Resets node indices such that they are numbered from 0 to num_nodes - 1
+    :param G:
+    :return: copy of G with relabelled node indices, mapping
+    """
+    mapping = {}
+    for new_idx, old_idx in enumerate(G.nodes()):
+        mapping[old_idx] = new_idx
+    new_G = nx.relabel_nodes(G, mapping, copy=True)
+    return new_G, mapping
+
+def graph_data_obj_to_nx_simple(data):
+    """
+    Converts graph Data object required by the pytorch geometric package to
+    network x data object. NB: Uses simplified atom and bond features,
+    and represent as indices. NB: possible issues with recapitulating relative
+    stereochemistry since the edges in the nx object are unordered.
+    :param data: pytorch geometric Data object
+    :return: network x object
+    """
+    G = nx.Graph()
+
+    # atoms
+    atom_features = data.x.cpu().numpy()
+    num_atoms = atom_features.shape[0]
+    for i in range(num_atoms):
+        atomic_num_idx, chirality_tag_idx, _, _, _, _, _, _, _ = atom_features[i]
+        G.add_node(i, atom_num_idx=atomic_num_idx, chirality_tag_idx=chirality_tag_idx)
+        pass
+
+    # bonds
+    edge_index = data.edge_index.cpu().numpy()
+    edge_attr = data.edge_attr.cpu().numpy()
+    num_bonds = edge_index.shape[1]
+    for j in range(0, num_bonds, 2):
+        begin_idx = int(edge_index[0, j])
+        end_idx = int(edge_index[1, j])
+        bond_type_idx, bond_dir_idx, _ = edge_attr[j]
+        if not G.has_edge(begin_idx, end_idx):
+            G.add_edge(begin_idx, end_idx, bond_type_idx=bond_type_idx,
+                       bond_dir_idx=bond_dir_idx)
+
+    return G
+
+def nx_to_graph_data_obj_simple(G):
+    """
+    Converts nx graph to pytorch geometric Data object. Assume node indices
+    are numbered from 0 to num_nodes - 1. NB: Uses simplified atom and bond
+    features, and represent as indices. NB: possible issues with
+    recapitulating relative stereochemistry since the edges in the nx
+    object are unordered.
+    :param G: nx graph obj
+    :return: pytorch geometric Data object
+    """
+    # atoms
+    num_atom_features = 2  # atom type,  chirality tag
+    atom_features_list = []
+    for _, node in G.nodes(data=True):
+        atom_feature = [node['atom_num_idx'], node['chirality_tag_idx']]
+        atom_features_list.append(atom_feature)
+    x = torch.tensor(np.array(atom_features_list), dtype=torch.long)
+
+    # bonds
+    num_bond_features = 2  # bond type, bond direction
+    if len(G.edges()) > 0:  # mol has bonds
+        edges_list = []
+        edge_features_list = []
+        for i, j, edge in G.edges(data=True):
+            edge_feature = [edge['bond_type_idx'], edge['bond_dir_idx']]
+            edges_list.append((i, j))
+            edge_features_list.append(edge_feature)
+            edges_list.append((j, i))
+            edge_features_list.append(edge_feature)
+
+        # data.edge_index: Graph connectivity in COO format with shape [2, num_edges]
+        edge_index = torch.tensor(np.array(edges_list).T, dtype=torch.long)
+
+        # data.edge_attr: Edge feature matrix with shape [num_edges, num_edge_features]
+        edge_attr = torch.tensor(np.array(edge_features_list),
+                                 dtype=torch.long)
+    else:   # mol has no bonds
+        edge_index = torch.empty((2, 0), dtype=torch.long)
+        edge_attr = torch.empty((0, num_bond_features), dtype=torch.long)
+
+    data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr)
+
+    return data

torch_molecule/encoder/edgepred/__init__.py

@@ -0,0 +1,3 @@
+from .modeling_edgepred import EdgePredMolecularEncoder
+
+__all__ = ['EdgePredMolecularEncoder']

torch_molecule/encoder/edgepred/model.py

@@ -0,0 +1,104 @@
+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
+from ...utils import init_weights
+
+criterion = torch.nn.BCEWithLogitsLoss()
+
+class GNN(nn.Module):
+    def __init__(
+        self,
+        num_layer,
+        hidden_size,
+        drop_ratio=0.5,
+        norm_layer="batch_norm",
+        encoder_type="gin-virtual",
+        readout="max"
+    ):
+        super(GNN, self).__init__()
+        gnn_name = encoder_type.split("-")[0]
+        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}.")
+
+    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)
+        
+        # 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):
+        # sample negative edges
+        num_nodes = batched_data.num_nodes
+        num_edges = batched_data.num_edges
+
+        edge_set = set([str(batched_data.edge_index[0, i].cpu().item()) + "," + str(
+            batched_data.edge_index[1, i].cpu().item()) for i in range(batched_data.edge_index.shape[1])])
+
+        redandunt_sample = torch.randint(0, num_nodes, (2, 5 * num_edges))
+        sampled_ind = []
+        sampled_edge_set = set([])
+        for i in range(5 * num_edges):
+            node1 = redandunt_sample[0, i].cpu().item()
+            node2 = redandunt_sample[1, i].cpu().item()
+            edge_str = str(node1) + "," + str(node2)
+            if not edge_str in edge_set and not edge_str in sampled_edge_set and not node1 == node2:
+                sampled_edge_set.add(edge_str)
+                sampled_ind.append(i)
+            if len(sampled_ind) == num_edges / 2:
+                break
+
+        batched_data.negative_edge_index = redandunt_sample[:, sampled_ind]
+    
+        # generate predictions
+        h_node, _ = self.graph_encoder(batched_data)
+        h_rep = self.pool(h_node, batched_data.batch)
+        
+        positive_score = torch.sum(h_node[batched_data.edge_index[0, ::2]] * h_node[batched_data.edge_index[1, ::2]], dim = 1)
+        negative_score = torch.sum(h_node[batched_data.negative_edge_index[0]] * h_node[batched_data.negative_edge_index[1]], dim = 1)
+        
+        loss_class = criterion(positive_score, torch.ones_like(positive_score)) + criterion(negative_score, torch.zeros_like(negative_score))
+        
+        return loss_class
+
+    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}