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

torch_molecule/encoder/moama/modeling_moama.py

@@ -0,0 +1,307 @@
+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 MoamaMolecularEncoder(BaseMolecularEncoder):
+    """This encoder implements Motif-aware Attribute Masking for Molecular Graph Pre-training.
+
+    References
+    ----------
+    - Paper: https://openreview.net/forum?id=uqPnesiGGi
+    - Code: https://github.com/einae-nd/MoAMa-dev
+
+    Parameters
+    ----------
+    mask_rate : float, default=0.15 
+        Fraction of nodes to mask in each graph
+    lw_rec : float, default=0.5
+        Weight balancing between reconstruction loss and fingerprint loss.
+        Higher values emphasize reconstruction, lower values emphasize fingerprint matching.
+    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. 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 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.
+    scheduler_factor : float, default=0.5
+        Factor by which to reduce the learning rate when plateau is detected.
+    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_name : str, default="MoamaMolecularEncoder"
+        Name of the encoder model.
+    """
+    # Pretraining task parameters
+    mask_rate: float = 0.15
+    lw_rec: float = 0.5
+    
+    # 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 = "MoamaMolecularEncoder"
+    
+    # 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]:
+        """Get parameter names for the estimator.
+
+        Returns
+        -------
+        List[str]
+            List of parameter names that can be used for model configuration.
+        """
+        return ["mask_rate", "lw_rec"] + 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,
+            "mask_rate": self.mask_rate,
+            "lw_rec": self.lw_rec
+        }
+        
+        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],
+    ) -> "MoamaMolecularEncoder":
+        """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 : MoamaMolecularEncoder
+            Fitted estimator
+        """
+
+        self._initialize_model(self.model_class)
+        self.model.initialize_parameters()
+        optimizer, scheduler = self._setup_optimizers()
+        
+        # Prepare datasets and loaders
+        smiles = X_train
+        X_train, _ = self._validate_inputs(X_train, return_rdkit_mol=True)
+        train_dataset = self._convert_to_pytorch_data(X_train)
+        for i in range(len(X_train)):
+            train_dataset[i].smiles = smiles[i]
+        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/moama/utils.py

@@ -0,0 +1,143 @@
+import torch
+import random
+import torch.nn as nn
+from rdkit import Chem, DataStructs
+from rdkit.Chem import BRICS
+
+def get_fingerprint_loss(smiles, h_rep):
+    cos = nn.CosineSimilarity(dim=0, eps=1e-6)
+    fingerprint_list = [Chem.RDKFingerprint(Chem.MolFromSmiles(x)) for x in smiles]
+    fingerprint_loss = 0
+    
+    for i in range(len(h_rep)):
+        for j in range(len(h_rep)-i-1):
+            finger_sim = DataStructs.FingerprintSimilarity(fingerprint_list[i], fingerprint_list[i+j+1]) #* 10
+            emb_sim = (cos(h_rep[i], h_rep[i+j+1]) + 1)/2
+            fingerprint_loss += (finger_sim - emb_sim)**2
+
+    return fingerprint_loss
+
+def get_mask_indices(batched_data, mask_rate=0.15):
+    masked_node_indices_full = list()
+    offset = 0
+    
+    for i in range(len(batched_data)):
+        masked_node_indices = list()
+        data = batched_data[i]
+        mol = Chem.MolFromSmiles(data.smiles)
+        motifs_list, edges = get_motifs_edges(mol)
+        random.shuffle(motifs_list)
+        num_atoms = mol.GetNumAtoms()
+        sample_size = int(num_atoms * mask_rate + 1)
+        
+        for motif in motifs_list:
+            edge_index_mask = sum(data.edge_index[0]==x for x in motif).bool()
+            neigh = torch.masked_select(data.edge_index[1], edge_index_mask)
+
+            if any(x in neigh for x in masked_node_indices):
+                continue
+
+            if len(masked_node_indices) + len(motif) > sample_size:
+                break
+            if len(motif) > 9:
+                del motif[5::6]
+            for atom in motif:
+                masked_node_indices.append(atom + offset)
+            
+        atom_candidates = [x for x in range(0, num_atoms) if x not in masked_node_indices]
+        random.shuffle(atom_candidates)
+        if len(masked_node_indices) < sample_size:
+            for atom in atom_candidates[:sample_size-len(masked_node_indices)]:
+                masked_node_indices.append(atom + offset)
+                
+        offset += num_atoms
+            
+        masked_node_indices_full += masked_node_indices
+
+    return masked_node_indices_full
+
+def get_motifs_edges(data):
+    Chem.SanitizeMol(data)
+    motifs, edges = brics_decomp(data)
+    return motifs, edges
+
+def brics_decomp(mol):
+    n_atoms = mol.GetNumAtoms()
+    if n_atoms == 1:
+        return [[0]], []
+
+    cliques = []
+    breaks = []
+    for bond in mol.GetBonds():
+        a1 = bond.GetBeginAtom().GetIdx()
+        a2 = bond.GetEndAtom().GetIdx()
+        cliques.append([a1, a2])
+
+    res = list(BRICS.FindBRICSBonds(mol))
+    if len(res) == 0:
+        return [list(range(n_atoms))], []
+    else:
+        for bond in res:
+            if [bond[0][0], bond[0][1]] in cliques:
+                cliques.remove([bond[0][0], bond[0][1]])
+            else:
+                cliques.remove([bond[0][1], bond[0][0]])
+            cliques.append([bond[0][0]])
+            cliques.append([bond[0][1]])
+
+    # break bonds between rings and non-ring atoms
+    for c in cliques:
+        if len(c) > 1:
+            if mol.GetAtomWithIdx(c[0]).IsInRing() and not mol.GetAtomWithIdx(c[1]).IsInRing():
+                cliques.remove(c)
+                cliques.append([c[1]])
+                breaks.append(c)
+            if mol.GetAtomWithIdx(c[1]).IsInRing() and not mol.GetAtomWithIdx(c[0]).IsInRing():
+                cliques.remove(c)
+                cliques.append([c[0]])
+                breaks.append(c)
+
+    # select atoms at intersections as motif
+    for atom in mol.GetAtoms():
+        if len(atom.GetNeighbors()) > 2 and not atom.IsInRing():
+            cliques.append([atom.GetIdx()])
+            for nei in atom.GetNeighbors():
+                if [nei.GetIdx(), atom.GetIdx()] in cliques:
+                    cliques.remove([nei.GetIdx(), atom.GetIdx()])
+                    breaks.append([nei.GetIdx(), atom.GetIdx()])
+                elif [atom.GetIdx(), nei.GetIdx()] in cliques:
+                    cliques.remove([atom.GetIdx(), nei.GetIdx()])
+                    breaks.append([atom.GetIdx(), nei.GetIdx()])
+                cliques.append([nei.GetIdx()])
+
+    # merge cliques
+    for c in range(len(cliques) - 1):
+        if c >= len(cliques):
+            break
+        for k in range(c + 1, len(cliques)):
+            if k >= len(cliques):
+                break
+            if len(set(cliques[c]) & set(cliques[k])) > 0:
+                cliques[c] = list(set(cliques[c]) | set(cliques[k]))
+                cliques[k] = []
+        cliques = [c for c in cliques if len(c) > 0]
+    cliques = [c for c in cliques if len(c) > 0]
+
+    # edges
+    edges = []
+    for bond in res:
+        for c in range(len(cliques)):
+            if bond[0][0] in cliques[c]:
+                c1 = c
+            if bond[0][1] in cliques[c]:
+                c2 = c
+        edges.append((c1, c2))
+    for bond in breaks:
+        for c in range(len(cliques)):
+            if bond[0] in cliques[c]:
+                c1 = c
+            if bond[1] in cliques[c]:
+                c2 = c
+        edges.append((c1, c2))
+
+    return cliques, edges

torch_molecule/encoder/pretrained/__init__.py

@@ -0,0 +1,3 @@
+from .modeling_pretrained import HFPretrainedMolecularEncoder
+
+__all__ = ['HFPretrainedMolecularEncoder']