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

torch_molecule/predictor/ssr/modeling_ssr.py

@@ -0,0 +1,267 @@
+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 torch_geometric.data import Data
+
+from .model import SSR
+from ...utils import graph_from_smiles
+from ..gnn.modeling_gnn import GNNMolecularPredictor
+from ...utils.search import (
+    ParameterSpec,
+    ParameterType,
+)
+
+@dataclass
+class SSRMolecularPredictor(GNNMolecularPredictor):
+    """This predictor implements a SizeShiftReg model with the GNN.
+    
+    References
+    ----------
+    - Paper: SizeShiftReg: a Regularization Method for Improving Size-Generalization in Graph Neural Networks.
+      https://arxiv.org/abs/2207.07888
+    - Reference Code: https://github.com/DavideBuffelli/SizeShiftReg/tree/main
+    
+    Parameters
+    ----------
+    coarse_ratios : List[float], default=[0.8, 0.9]
+        List of ratios for graph coarsening. Each ratio determines the percentage of 
+        nodes to keep in the coarsened graph.
+    cmd_coeff : float, default=0.1
+        Weight for CMD (Central Moment Discrepancy) loss. Controls the strength of 
+        the size-shift regularization.
+    fine_grained : bool, default=True
+        Whether to use fine-grained CMD. When True, matches distributions at a more 
+        detailed level.
+    n_moments : int, default=5
+        Number of moments to match in the CMD calculation. Higher values capture 
+        more complex distribution characteristics.
+    """
+    # SSR-specific parameters
+    coarse_ratios: List[float] = field(default_factory=lambda: [0.8, 0.9])
+    cmd_coeff: float = field(default=0.1)
+    fine_grained: bool = field(default=True)
+    n_moments: int = field(default=5)
+    coarse_pool: str = field(default='mean')
+
+    # Other Non-init fields
+    model_name: str = "SSRMolecularPredictor"
+    model_class: Type[SSR] = field(default=SSR, init=False)
+
+    def __post_init__(self):
+        super().__post_init__()
+
+    @staticmethod
+    def _get_param_names() -> List[str]:
+        return GNNMolecularPredictor._get_param_names() + [
+            "coarse_ratios",
+            "cmd_coeff",
+            "fine_grained",
+            "n_moments",
+            "coarse_pool",
+        ]
+
+    def _get_default_search_space(self):
+        search_space = super()._get_default_search_space().copy()
+        search_space["cmd_coeff"] = ParameterSpec(ParameterType.FLOAT, (0.01, 1.0))
+        search_space["n_moments"] = 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)
+        return base_params
+
+
+    def _convert_to_pytorch_data(self, X, y=None):
+        """Convert SMILES to PyTorch Geometric data with coarsened versions, preserving edge attributes."""
+        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:
+            if y is not None:
+                properties = y[idx]
+            else:
+                properties = None
+                
+            # Convert SMILES to graph
+            graph = graph_from_smiles(smiles_or_mol, properties, self.augmented_feature)
+            g = Data()
+            g.num_nodes = graph["num_nodes"]
+            g.edge_index = torch.from_numpy(graph["edge_index"])
+
+            # Standard attributes
+            if graph["edge_feat"] is not None:
+                g.edge_attr = torch.from_numpy(graph["edge_feat"])
+            
+            if graph["node_feat"] is not None:
+                g.x = torch.from_numpy(graph["node_feat"])
+
+            if graph["y"] is not None:
+                g.y = torch.from_numpy(graph["y"])
+                
+            if graph.get("morgan") is not None:
+                g.morgan = torch.tensor(graph["morgan"], dtype=torch.int8).view(1, -1)
+                
+            if graph.get("maccs") is not None:
+                g.maccs = torch.tensor(graph["maccs"], dtype=torch.int8).view(1, -1)
+
+            # Add coarsened versions
+            for ratio in self.coarse_ratios:
+                num_clusters = max(1, int(g.num_nodes * ratio))
+                coarse_ratio_postfix = str(int(ratio*100))
+                
+                # Get coarsened graph with edge attributes
+                coarse_edge_index, coarse_edge_attr, clusters = self.spectral_graph_coarsening(g, num_clusters)
+                
+                # Add attributes to graph
+                setattr(g, f"coarsened_edge_index_{coarse_ratio_postfix}", coarse_edge_index)
+                if hasattr(g, 'edge_attr'):
+                    setattr(g, f"coarsened_edge_attr_{coarse_ratio_postfix}", coarse_edge_attr)
+                setattr(g, f"num_coarse_nodes_{coarse_ratio_postfix}", torch.tensor(num_clusters))
+                setattr(g, f"clusters_{coarse_ratio_postfix}", clusters)
+            
+            if self.verbose:
+                print(f"Processed molecule {idx}: {g.num_nodes} nodes, coarsened versions added")
+                
+            pyg_graph_list.append(g)
+            
+        return pyg_graph_list
+
+    def spectral_graph_coarsening(self, graph, num_clusters):
+        """Coarsen graph based on spectral clustering while preserving edge attributes"""
+        # Extract node features and edge information
+        edge_index = graph.edge_index
+        edge_attr = graph.edge_attr if hasattr(graph, 'edge_attr') else None
+        num_nodes = graph.num_nodes
+        edge_attr = edge_attr.float()
+
+        # Convert to adjacency matrix
+        adj = torch.zeros((num_nodes, num_nodes), device=edge_index.device)
+        adj[edge_index[0], edge_index[1]] = 1
+        
+        # Calculate Laplacian
+        degree = adj.sum(dim=1)
+        degree_mat = torch.diag(degree)
+        laplacian = degree_mat - adj
+        
+        # Compute eigenvectors
+        if num_clusters < num_nodes - 1:
+            eigvals, eigvecs = torch.linalg.eigh(laplacian)
+            # Use smallest non-zero eigenvalues
+            indices = torch.argsort(eigvals)[1:num_clusters+1]  # Skip first eigenvector
+            fiedler_vectors = eigvecs[:, indices]
+            
+            # Use k-means for clustering
+            from sklearn.cluster import KMeans
+            kmeans = KMeans(n_clusters=num_clusters, random_state=42)
+            clusters = kmeans.fit_predict(fiedler_vectors.cpu().numpy())
+            clusters = torch.tensor(clusters, device=edge_index.device)
+        else:
+            clusters = torch.arange(num_nodes, device=edge_index.device)
+        
+        # Create coarsened edge index and attributes
+        coarse_edge_index = []
+        coarse_edge_attr = []
+        
+        # Store mapping from (cluster_i, cluster_j) to list of edge indices
+        cluster_edges = {}
+        
+        # Group edges by their clusters
+        for e_idx in range(edge_index.shape[1]):
+            src, dst = edge_index[0, e_idx], edge_index[1, e_idx]
+            c_src, c_dst = clusters[src], clusters[dst]
+            
+            # Skip self-loops in coarsened graph
+            if c_src == c_dst:
+                continue
+                
+            # Sort clusters to avoid duplicates
+            if c_src > c_dst:
+                c_src, c_dst = c_dst, c_src
+                
+            key = (c_src.item(), c_dst.item())
+            if key not in cluster_edges:
+                cluster_edges[key] = []
+            
+            if edge_attr is not None:
+                cluster_edges[key].append(edge_attr[e_idx])
+        
+        # Create coarsened edges
+        for (c_src, c_dst), edge_attrs in cluster_edges.items():
+            coarse_edge_index.append([c_src, c_dst])
+            coarse_edge_index.append([c_dst, c_src])  # Add both directions
+            
+            if edge_attr is not None:
+                # Aggregate edge attributes using mean
+                mean_attr = torch.stack(edge_attrs).mean(dim=0)
+                coarse_edge_attr.append(mean_attr)
+                coarse_edge_attr.append(mean_attr)  # Same for both directions
+        
+        # Convert to tensors
+        if coarse_edge_index:
+            coarse_edge_index = torch.tensor(coarse_edge_index, dtype=torch.long, 
+                                            device=edge_index.device).t()
+            if edge_attr is not None:
+                coarse_edge_attr = torch.stack(coarse_edge_attr)
+        else:
+            coarse_edge_index = torch.zeros((2, 0), dtype=torch.long, 
+                                        device=edge_index.device)
+            if edge_attr is not None:
+                coarse_edge_attr = torch.zeros((0, edge_attr.size(1)), 
+                                            device=edge_attr.device)
+        
+        return coarse_edge_index, coarse_edge_attr, clusters
+
+    def _train_epoch(self, train_loader, optimizer, epoch):
+        """Training logic for one epoch with SSR."""
+        self.model.train()
+        losses = []
+        pred_losses = []
+        ssr_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()
+
+            # Forward pass and loss computation
+            total_loss, pred_loss, ssr_loss = self.model.compute_loss(batch, self.loss_criterion, self.coarse_ratios, self.cmd_coeff, self.fine_grained, self.n_moments)
+
+            # Backward pass
+            total_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(total_loss.item())
+            pred_losses.append(pred_loss.item())
+            ssr_losses.append(ssr_loss.item())
+
+            # Update progress bar if using tqdm
+            if self.verbose:
+                iterator.set_postfix({
+                    "Epoch": epoch,
+                    "Total loss": f"{total_loss.item():.4f}",
+                    "Pred loss": f"{pred_loss.item():.4f}",
+                    "ssr_loss": f"{ssr_loss.item():.4f}"
+                })
+
+        # Return all loss components for logging
+        return losses

torch_molecule/predictor/ssr/utils.py

@@ -0,0 +1,185 @@
+from typing import Optional, Tuple
+
+import torch
+
+def scatter_sum(src: torch.Tensor,
+                index: torch.Tensor,
+                dim: int = -1,
+                out: Optional[torch.Tensor] = None,
+                dim_size: Optional[int] = None) -> torch.Tensor:
+    index = broadcast(index, src, dim)
+    if out is None:
+        size = list(src.size())
+        if dim_size is not None:
+            size[dim] = dim_size
+        elif index.numel() == 0:
+            size[dim] = 0
+        else:
+            size[dim] = int(index.max()) + 1
+        out = torch.zeros(size, dtype=src.dtype, device=src.device)
+        return out.scatter_add_(dim, index, src)
+    else:
+        return out.scatter_add_(dim, index, src)
+
+
+def scatter_add(src: torch.Tensor,
+                index: torch.Tensor,
+                dim: int = -1,
+                out: Optional[torch.Tensor] = None,
+                dim_size: Optional[int] = None) -> torch.Tensor:
+    return scatter_sum(src, index, dim, out, dim_size)
+
+
+def scatter_mul(src: torch.Tensor,
+                index: torch.Tensor,
+                dim: int = -1,
+                out: Optional[torch.Tensor] = None,
+                dim_size: Optional[int] = None) -> torch.Tensor:
+    return torch.ops.torch_scatter.scatter_mul(src, index, dim, out, dim_size)
+
+
+def scatter_mean(src: torch.Tensor,
+                 index: torch.Tensor,
+                 dim: int = -1,
+                 out: Optional[torch.Tensor] = None,
+                 dim_size: Optional[int] = None) -> torch.Tensor:
+    out = scatter_sum(src, index, dim, out, dim_size)
+    dim_size = out.size(dim)
+
+    index_dim = dim
+    if index_dim < 0:
+        index_dim = index_dim + src.dim()
+    if index.dim() <= index_dim:
+        index_dim = index.dim() - 1
+
+    ones = torch.ones(index.size(), dtype=src.dtype, device=src.device)
+    count = scatter_sum(ones, index, index_dim, None, dim_size)
+    count[count < 1] = 1
+    count = broadcast(count, out, dim)
+    if out.is_floating_point():
+        out.true_divide_(count)
+    else:
+        out.div_(count, rounding_mode='floor')
+    return out
+
+
+def scatter_min(
+        src: torch.Tensor,
+        index: torch.Tensor,
+        dim: int = -1,
+        out: Optional[torch.Tensor] = None,
+        dim_size: Optional[int] = None) -> Tuple[torch.Tensor, torch.Tensor]:
+    return torch.ops.torch_scatter.scatter_min(src, index, dim, out, dim_size)
+
+
+def scatter_max(
+        src: torch.Tensor,
+        index: torch.Tensor,
+        dim: int = -1,
+        out: Optional[torch.Tensor] = None,
+        dim_size: Optional[int] = None) -> Tuple[torch.Tensor, torch.Tensor]:
+    return torch.ops.torch_scatter.scatter_max(src, index, dim, out, dim_size)
+
+
+def scatter(src: torch.Tensor,
+            index: torch.Tensor,
+            dim: int = -1,
+            out: Optional[torch.Tensor] = None,
+            dim_size: Optional[int] = None,
+            reduce: str = "sum") -> torch.Tensor:
+    r"""
+    |
+
+    .. image:: https://raw.githubusercontent.com/rusty1s/pytorch_scatter/
+            master/docs/source/_figures/add.svg?sanitize=true
+        :align: center
+        :width: 400px
+
+    |
+
+    Reduces all values from the :attr:`src` tensor into :attr:`out` at the
+    indices specified in the :attr:`index` tensor along a given axis
+    :attr:`dim`.
+    For each value in :attr:`src`, its output index is specified by its index
+    in :attr:`src` for dimensions outside of :attr:`dim` and by the
+    corresponding value in :attr:`index` for dimension :attr:`dim`.
+    The applied reduction is defined via the :attr:`reduce` argument.
+
+    Formally, if :attr:`src` and :attr:`index` are :math:`n`-dimensional
+    tensors with size :math:`(x_0, ..., x_{i-1}, x_i, x_{i+1}, ..., x_{n-1})`
+    and :attr:`dim` = `i`, then :attr:`out` must be an :math:`n`-dimensional
+    tensor with size :math:`(x_0, ..., x_{i-1}, y, x_{i+1}, ..., x_{n-1})`.
+    Moreover, the values of :attr:`index` must be between :math:`0` and
+    :math:`y - 1`, although no specific ordering of indices is required.
+    The :attr:`index` tensor supports broadcasting in case its dimensions do
+    not match with :attr:`src`.
+
+    For one-dimensional tensors with :obj:`reduce="sum"`, the operation
+    computes
+
+    .. math::
+        \mathrm{out}_i = \mathrm{out}_i + \sum_j~\mathrm{src}_j
+
+    where :math:`\sum_j` is over :math:`j` such that
+    :math:`\mathrm{index}_j = i`.
+
+    .. note::
+
+        This operation is implemented via atomic operations on the GPU and is
+        therefore **non-deterministic** since the order of parallel operations
+        to the same value is undetermined.
+        For floating-point variables, this results in a source of variance in
+        the result.
+
+    :param src: The source tensor.
+    :param index: The indices of elements to scatter.
+    :param dim: The axis along which to index. (default: :obj:`-1`)
+    :param out: The destination tensor.
+    :param dim_size: If :attr:`out` is not given, automatically create output
+        with size :attr:`dim_size` at dimension :attr:`dim`.
+        If :attr:`dim_size` is not given, a minimal sized output tensor
+        according to :obj:`index.max() + 1` is returned.
+    :param reduce: The reduce operation (:obj:`"sum"`, :obj:`"mul"`,
+        :obj:`"mean"`, :obj:`"min"` or :obj:`"max"`). (default: :obj:`"sum"`)
+
+    :rtype: :class:`Tensor`
+
+    .. code-block:: python
+
+        from torch_scatter import scatter
+
+        src = torch.randn(10, 6, 64)
+        index = torch.tensor([0, 1, 0, 1, 2, 1])
+
+        # Broadcasting in the first and last dim.
+        out = scatter(src, index, dim=1, reduce="sum")
+
+        print(out.size())
+
+    .. code-block::
+
+        torch.Size([10, 3, 64])
+    """
+    if reduce == 'sum' or reduce == 'add':
+        return scatter_sum(src, index, dim, out, dim_size)
+    if reduce == 'mul':
+        return scatter_mul(src, index, dim, out, dim_size)
+    elif reduce == 'mean':
+        return scatter_mean(src, index, dim, out, dim_size)
+    elif reduce == 'min':
+        return scatter_min(src, index, dim, out, dim_size)[0]
+    elif reduce == 'max':
+        return scatter_max(src, index, dim, out, dim_size)[0]
+    else:
+        raise ValueError
+
+def broadcast(src: torch.Tensor, other: torch.Tensor, dim: int):
+    if dim < 0:
+        dim = other.dim() + dim
+    if src.dim() == 1:
+        for _ in range(0, dim):
+            src = src.unsqueeze(0)
+    for _ in range(src.dim(), other.dim()):
+        src = src.unsqueeze(-1)
+    src = src.expand(other.size())
+    return src

torch_molecule/utils/__init__.py

@@ -0,0 +1,31 @@
+from .generic.weights import init_weights
+from .generic.metrics import (
+    roc_auc_score,
+    accuracy_score,
+    mean_squared_error,
+    mean_absolute_error,
+    root_mean_squared_error,
+    r2_score,
+)
+from .generic.pseudo_tasks import PSEUDOTASK
+from .graph.graph_from_smiles import graph_from_smiles
+from .graph.graph_to_smiles import graph_to_smiles
+from .graph.features import get_atom_feature_dims, get_bond_feature_dims
+
+__all__ = [
+    "init_weights",
+    # metric
+    "roc_auc_score",
+    "accuracy_score",
+    "mean_squared_error", 
+    "mean_absolute_error",
+    "root_mean_squared_error", 
+    "r2_score",
+    # graph
+    "graph_from_smiles",
+    "graph_to_smiles",
+    "get_atom_feature_dims",
+    "get_bond_feature_dims",
+    # pseudo_tasks
+    "PSEUDOTASK",
+]

torch_molecule/utils/checker.py

@@ -0,0 +1,158 @@
+import warnings
+import numpy as np
+from rdkit import Chem
+from typing import Optional, Union, List, Tuple
+
+class MolecularInputChecker:
+    """
+    Class for validating input data used in molecular models.
+    """
+
+    @staticmethod
+    def validate_smiles(
+        smiles: str, 
+        idx: int
+    ) -> Tuple[bool, Optional[str], Optional[Chem.Mol]]:
+        """Validate a single SMILES string at a given index.
+
+        Parameters
+        ----------
+        smiles : str
+            The SMILES string to validate
+        idx : int
+            The index of the SMILES string in the original list
+
+        Returns
+        -------
+        Tuple[bool, Optional[str], Optional[Chem.Mol]]
+            A tuple containing:
+            
+            - A boolean indicating whether the SMILES string is valid
+            - A string describing the error if the SMILES is invalid, or None if valid
+            - The RDKit Mol object if valid, or None if invalid
+        """
+        if not smiles or not smiles.strip():
+            return False, f"Empty SMILES at index {idx}", None
+
+        try:
+            mol = Chem.MolFromSmiles(smiles)
+            if mol is None:
+                return False, f"Invalid SMILES structure at index {idx}: {smiles}", None
+            return True, None, mol
+        except Exception as e:
+            return False, f"RDKit error at index {idx}: {str(e)}", None
+
+    @staticmethod
+    def validate_inputs(
+        X: List[str],
+        y: Optional[Union[List, np.ndarray]] = None,
+        num_task: int = 0,
+        num_pretask: int = 0,
+        return_rdkit_mol: bool = True
+    ) -> Tuple[Union[List[str], List["Chem.Mol"]], Optional[np.ndarray]]:
+        """Validate a list of SMILES strings, and optionally validate a target array.
+
+        Parameters
+        ----------
+        X : List[str]
+            List of SMILES strings
+        y : Optional[Union[List, np.ndarray]], optional
+            Optional target values, by default None
+        num_task : int, optional
+            Total number of tasks; used to check dimensions of y, by default 0
+        num_pretask : int, optional
+            Number of (pseudo)-tasks that are predefined in the modeling; 
+            used to check dimensions of y. Preliminarily used in supervised pretraining,
+            by default 0
+        return_rdkit_mol : bool, optional
+            If True, convert SMILES to RDKit Mol objects, by default True
+
+        Returns
+        -------
+        Tuple[Union[List[str], List["Chem.Mol"]], Optional[np.ndarray]]
+            A tuple containing:
+            
+            - The original or converted SMILES (RDKit Mol objects if return_rdkit_mol=True)
+            - The target array as a numpy array, or None if y was not provided
+
+        Raises
+        ------
+        ValueError
+            If SMILES or target dimensions are invalid
+        """
+        if not isinstance(X, list):
+            raise ValueError("X must be a list of SMILES strings.")
+
+        if not all(isinstance(s, str) for s in X):
+            raise ValueError("All elements in X must be strings.")
+
+        invalid_smiles = []
+        rdkit_mols = []
+        for i, smiles in enumerate(X):
+            is_valid, error_msg, mol = MolecularInputChecker.validate_smiles(smiles, i)
+            if not is_valid:
+                invalid_smiles.append(error_msg)
+            else:
+                rdkit_mols.append(mol)
+
+        if invalid_smiles:
+            raise ValueError("Invalid SMILES found:\n" + "\n".join(invalid_smiles))
+
+        if y is not None:
+            try:
+                y = np.asarray(y, dtype=np.float32)
+            except Exception as e:
+                raise ValueError(f"Could not convert y to numpy array: {str(e)}")
+
+            if len(y.shape) == 1:
+                if num_task - num_pretask != 1:
+                    raise ValueError(
+                        f"1D target array provided but num_task is {num_task - num_pretask}. "
+                        "For multiple tasks, y must be 2D."
+                    )
+                y = y.reshape(-1, 1)
+
+            if len(y.shape) != 2:
+                raise ValueError(
+                    "y must be 1D (single task) or 2D (multiple tasks). "
+                    f"Got shape {y.shape}."
+                )
+
+            if y.shape[0] != len(X):
+                raise ValueError(
+                    f"Number of samples in y ({y.shape[0]}) must match length of X ({len(X)})."
+                )
+
+            if y.shape[1] != num_task - num_pretask:
+                raise ValueError(
+                    f"Second dimension of y ({y.shape[1]}) must match num_task ({num_task - num_pretask})."
+                )
+
+            inf_mask = np.isinf(y)
+            if np.any(inf_mask):
+                inf_indices = np.where(inf_mask)
+                warnings.warn(
+                    f"Infinite values found in y at indices: {list(zip(*inf_indices))}. "
+                    "Converting to NaN.",
+                    RuntimeWarning,
+                )
+                y = y.astype(float)
+                y[inf_mask] = np.nan
+
+            # nan_mask = np.isnan(y)
+            # if np.any(nan_mask):
+            #     nan_counts = np.sum(nan_mask, axis=0)
+            #     nan_percentages = (nan_counts / len(X)) * 100
+            #     task_warnings = []
+            #     for task_idx, (count, percentage) in enumerate(zip(nan_counts, nan_percentages)):
+            #         if count > 0:
+            #             task_warnings.append(f"Task {task_idx}: {count} NaNs ({percentage:.1f}%)")
+
+            #     warnings.warn(
+            #         "NaN values present in y:\n"
+            #         + "\n".join(task_warnings)
+            #         + "\nSamples with NaN will be ignored or cause issues unless handled.",
+            #         RuntimeWarning,
+            #     )
+
+        return rdkit_mols if return_rdkit_mol else X, y