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

torch_molecule/predictor/grea/__init__.py

@@ -0,0 +1,3 @@
+from .modeling_grea import GREAMolecularPredictor
+
+__all__ = ['GREAMolecularPredictor']

torch_molecule/predictor/grea/model.py

@@ -0,0 +1,190 @@
+import os
+import torch
+import torch.nn as nn
+# from torch_scatter import scatter_add
+from .utils import scatter_add
+
+from ...nn import GNN_node, GNN_node_Virtualnode, MLP
+from ...utils import init_weights
+
+class GREA(nn.Module):
+    def __init__(
+        self,
+        num_task,
+        num_layer,
+        gamma=0.4,
+        hidden_size=300,
+        gnn_type="gin-virtual",
+        drop_ratio=0.5,
+        norm_layer="batch_norm",
+        augmented_feature=['maccs', 'morgan']
+    ):
+        super(GREA, self).__init__()
+        gnn_name = gnn_type.split("-")[0]
+        self.num_task = num_task
+        self.gamma = gamma
+        self.hidden_size = hidden_size
+
+        if "virtual" in gnn_type:
+            rationale_encoder = GNN_node_Virtualnode(
+                2, hidden_size, JK="last", drop_ratio=drop_ratio, residual=True, gnn_name=gnn_name
+            )
+            self.graph_encoder = GNN_node_Virtualnode(
+                num_layer,
+                hidden_size,
+                JK="last",
+                drop_ratio=drop_ratio,
+                residual=True,
+                gnn_name=gnn_name,
+                norm_layer=norm_layer,
+            )
+        else:
+            rationale_encoder = GNN_node(
+                2, hidden_size, JK="last", drop_ratio=drop_ratio, residual=True, gnn_name=gnn_name
+            )
+            self.graph_encoder = GNN_node(
+                num_layer,
+                hidden_size,
+                JK="last",
+                drop_ratio=drop_ratio,
+                residual=True,
+                gnn_name=gnn_name,
+                norm_layer=norm_layer,
+            )
+
+        self.separator = Separator(
+            rationale_encoder=rationale_encoder,
+            gate_nn=torch.nn.Sequential(
+                torch.nn.Linear(hidden_size, 2 * hidden_size),
+                torch.nn.BatchNorm1d(2 * hidden_size),
+                torch.nn.ReLU(),
+                torch.nn.Dropout(),
+                torch.nn.Linear(2 * hidden_size, 1),
+            ),
+        )
+        
+        graph_dim = hidden_size
+        self.augmented_feature = augmented_feature
+        if augmented_feature:
+            if "morgan" in augmented_feature:
+                graph_dim += 1024
+            if "maccs" in augmented_feature:
+                graph_dim += 167
+        self.predictor = MLP(graph_dim, hidden_features=2 * hidden_size, out_features=num_task)
+
+    
+    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.separator)
+        init_weights(self.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 _augment_graph_features(self, batched_data, h_r, h_rep):
+        if self.augmented_feature:
+            if 'morgan' in self.augmented_feature:
+                morgan = batched_data.morgan.type_as(h_r)
+                h_r = torch.cat((h_r, morgan), dim=1)
+                morgans = morgan.repeat_interleave(batched_data.batch[-1]+1, dim=0)
+                h_rep = torch.cat((h_rep, morgans), dim=1)
+            if 'maccs' in self.augmented_feature:
+                maccs = batched_data.maccs.type_as(h_r)
+                h_r = torch.cat((h_r, maccs), dim=1)
+                maccses = maccs.repeat_interleave(batched_data.batch[-1]+1, dim=0)
+                h_rep = torch.cat((h_rep, maccses), dim=1)
+        return h_r, h_rep
+    
+    def compute_loss(self, batched_data, criterion):
+        h_node, _ = self.graph_encoder(batched_data)
+        h_r, h_env, rationale_size, envir_size, _ = self.separator(batched_data, h_node)
+        h_rep = (h_r.unsqueeze(1) + h_env.unsqueeze(0)).view(-1, self.hidden_size)
+        h_r, h_rep = self._augment_graph_features(batched_data, h_r, h_rep)
+        pred_rem = self.predictor(h_r)
+        pred_rep = self.predictor(h_rep)
+        loss = torch.abs(
+            rationale_size / (rationale_size + envir_size)
+            - self.gamma * torch.ones_like(rationale_size)
+        ).mean()
+        target = batched_data.y.to(torch.float32)
+        is_labeled = batched_data.y == batched_data.y
+        loss += criterion(pred_rem.to(torch.float32)[is_labeled], target[is_labeled]).mean()
+        target_rep = batched_data.y.to(torch.float32).repeat_interleave(
+            batched_data.batch[-1] + 1, dim=0
+        )
+        is_labeled_rep = target_rep == target_rep
+        loss += criterion(pred_rep.to(torch.float32)[is_labeled_rep], target_rep[is_labeled_rep]).mean()
+        return loss
+
+    def forward(self, batched_data):
+        h_node, _ = self.graph_encoder(batched_data)
+        h_r, h_env, _, _, node_score = self.separator(batched_data, h_node)
+        h_rep = (h_r.unsqueeze(1) + h_env.unsqueeze(0)).view(-1, self.hidden_size)
+        h_r, h_rep = self._augment_graph_features(batched_data, h_r, h_rep)
+        prediction = self.predictor(h_r)
+        pred_rep = self.predictor(h_rep).view(h_r.size(0), -1)
+        if pred_rep.size(1) > 1:
+            variance = pred_rep.var(dim=-1, keepdim=True)
+        else:
+            variance = torch.zeros_like(pred_rep)
+        num_graphs = batched_data.batch.max().item() + 1
+        score_by_graph = [node_score[batched_data.batch == i].view(-1).tolist() for i in range(num_graphs)]
+        return {"prediction": prediction, "variance": variance, "score": score_by_graph, "representation": h_r}
+
+    @staticmethod
+    def _disable_batchnorm_tracking(model):
+        def fn(module):
+            if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
+                module.track_running_stats = False
+
+        model.apply(fn)
+    
+    @staticmethod
+    def _enable_batchnorm_tracking(model):
+        def fn(module):
+            if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
+                module.track_running_stats = True
+
+        model.apply(fn)
+
+class Separator(torch.nn.Module):
+    def __init__(self, rationale_encoder, gate_nn, nn=None):
+        super(Separator, self).__init__()
+        self.rationale_encoder = rationale_encoder
+        self.gate_nn = gate_nn
+        self.nn = nn
+
+    def forward(self, batched_data, h_node, size=None):
+        x, _ = self.rationale_encoder(batched_data)
+        batch = batched_data.batch
+        x = x.unsqueeze(-1) if x.dim() == 1 else x
+        size = batch[-1].item() + 1 if size is None else size
+
+        gate = self.gate_nn(x).view(-1, 1)
+        h_node = self.nn(h_node) if self.nn is not None else h_node
+        assert gate.dim() == h_node.dim() and gate.size(0) == h_node.size(0)
+        gate = torch.sigmoid(gate)
+
+        h_out = scatter_add(gate * h_node, batch, dim=0, dim_size=size)
+        c_out = scatter_add((1 - gate) * h_node, batch, dim=0, dim_size=size)
+
+        rationale_size = scatter_add(gate, batch, dim=0, dim_size=size)
+        envir_size = scatter_add((1 - gate), batch, dim=0, dim_size=size)
+
+        return h_out, c_out, rationale_size + 1e-8, envir_size + 1e-8, gate

torch_molecule/predictor/grea/modeling_grea.py

@@ -0,0 +1,107 @@
+import numpy as np
+import warnings
+from tqdm import tqdm
+from typing import Optional, Union, Dict, Any, List, Type
+from dataclasses import dataclass, field
+
+import torch
+from torch_geometric.loader import DataLoader
+
+from .model import GREA
+from ..gnn.modeling_gnn import GNNMolecularPredictor
+from ...utils.search import (
+    ParameterSpec,
+    ParameterType,
+)
+
+@dataclass
+class GREAMolecularPredictor(GNNMolecularPredictor):
+    """This predictor implements a Graph Rationalization model called GREA.
+
+    The full name of GREA is Graph Rationalization with Environment-based Augmentations. During model training, it learns the rationales (explainable subgraphs) and use them for molecular property prediction tasks.
+
+    References
+    ----------
+    - Graph Rationalization with Environment-based Augmentations.
+      https://dl.acm.org/doi/10.1145/3534678.3539347
+    - Code: https://github.com/liugangcode/GREA
+
+    :param gamma: GREA-specific parameter. Default is 0.4.
+    :type gamma: float
+    """
+    
+    # GREA-specific parameter
+    gamma: float = 0.4
+    # Override parent defaults
+    model_name: str = "GREAMolecularPredictor"
+    model_class: Type[GREA] = field(default=GREA, init=False)
+    
+    def __post_init__(self):
+        super().__post_init__()
+
+    @staticmethod
+    def _get_param_names() -> List[str]:
+        return ["gamma"] + GNNMolecularPredictor._get_param_names()
+
+    def _get_default_search_space(self):
+        search_space = super()._get_default_search_space().copy()
+        search_space["gamma"] = ParameterSpec(ParameterType.FLOAT, (0.1, 0.9))
+        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["gamma"] = checkpoint["hyperparameters"].get("gamma", self.gamma)
+        else:
+            base_params["gamma"] = self.gamma
+        base_params.pop("graph_pooling", None)
+        return base_params
+
+    def predict(self, X: List[str]) -> Dict[str, Union[np.ndarray, List[List]]]:
+        """Make predictions using the fitted model.
+
+        Parameters
+        ----------
+        X : List[str]
+            List of SMILES strings to make predictions for
+
+        Returns
+        -------
+        Dict[str, np.ndarray]
+            Dictionary containing:
+                - 'prediction': Model predictions (shape: [n_samples, n_tasks])
+                - 'variance': Prediction variances (shape: [n_samples, n_tasks])
+                - 'node_importance': A nested list where the outer list has length n_samples and each inner list has length n_nodes for that molecule
+        """
+        self._check_is_fitted()
+
+        # Convert to PyTorch Geometric format and create loader
+        X, _ = self._validate_inputs(X)
+        dataset = self._convert_to_pytorch_data(X)
+        loader = DataLoader(dataset, batch_size=self.batch_size, shuffle=False)
+
+        # Make predictions
+        self.model = self.model.to(self.device)
+        self.model.eval()
+        predictions = []
+        variances = []
+        node_scores = []
+        with torch.no_grad():
+            for batch in tqdm(loader, disable=not self.verbose):
+                batch = batch.to(self.device)
+                out = self.model(batch)
+                predictions.append(out["prediction"].cpu().numpy())
+                variances.append(out["variance"].cpu().numpy())
+                node_scores.extend(out["score"])
+
+        if predictions and variances:
+            return {
+                "prediction": np.concatenate(predictions, axis=0),
+                "variance": np.concatenate(variances, axis=0),
+                "node_importance": node_scores,
+            }
+        else:
+            warnings.warn(
+                "No valid predictions could be made from the input data. Returning empty results."
+            )
+            return {"prediction": np.array([]), "variance": np.array([]), "node_importance": np.array([])}

torch_molecule/predictor/grea/old1_features.py

@@ -0,0 +1,229 @@
+# allowable multiple choice node and edge features 
+allowable_features = {
+    'possible_atomic_num_list' : list(range(1, 119)) + ['misc'],
+    'possible_chirality_list' : [
+        'CHI_UNSPECIFIED',
+        'CHI_TETRAHEDRAL_CW',
+        'CHI_TETRAHEDRAL_CCW',
+        'CHI_OTHER',
+        'misc'
+    ],
+    'possible_degree_list' : [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 'misc'],
+    'possible_formal_charge_list' : [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 'misc'],
+    'possible_numH_list' : [0, 1, 2, 3, 4, 5, 6, 7, 8, 'misc'],
+    'possible_number_radical_e_list': [0, 1, 2, 3, 4, 'misc'],
+    'possible_hybridization_list' : [
+        'SP', 'SP2', 'SP3', 'SP3D', 'SP3D2', 'misc'
+        ],
+    'possible_is_aromatic_list': [False, True],
+    'possible_is_in_ring_list': [False, True],
+    'possible_bond_type_list' : [
+        'SINGLE',
+        'DOUBLE',
+        'TRIPLE',
+        'AROMATIC',
+        'misc'
+    ],
+    'possible_bond_stereo_list': [
+        'STEREONONE',
+        'STEREOZ',
+        'STEREOE',
+        'STEREOCIS',
+        'STEREOTRANS',
+        'STEREOANY',
+    ], 
+    'possible_is_conjugated_list': [False, True],
+}
+
+def safe_index(l, e):
+    """
+    Return index of element e in list l. If e is not present, return the last index
+    """
+    try:
+        return l.index(e)
+    except:
+        return len(l) - 1
+# # miscellaneous case
+# i = safe_index(allowable_features['possible_atomic_num_list'], 'asdf')
+# assert allowable_features['possible_atomic_num_list'][i] == 'misc'
+# # normal case
+# i = safe_index(allowable_features['possible_atomic_num_list'], 2)
+# assert allowable_features['possible_atomic_num_list'][i] == 2
+
+def atom_to_feature_vector(atom):
+    """
+    Converts rdkit atom object to feature list of indices
+    :param mol: rdkit atom object
+    :return: list
+    """
+    atom_feature = [
+            safe_index(allowable_features['possible_atomic_num_list'], atom.GetAtomicNum()),
+            safe_index(allowable_features['possible_chirality_list'], str(atom.GetChiralTag())),
+            safe_index(allowable_features['possible_degree_list'], atom.GetTotalDegree()),
+            safe_index(allowable_features['possible_formal_charge_list'], atom.GetFormalCharge()),
+            safe_index(allowable_features['possible_numH_list'], atom.GetTotalNumHs()),
+            safe_index(allowable_features['possible_number_radical_e_list'], atom.GetNumRadicalElectrons()),
+            safe_index(allowable_features['possible_hybridization_list'], str(atom.GetHybridization())),
+            allowable_features['possible_is_aromatic_list'].index(atom.GetIsAromatic()),
+            allowable_features['possible_is_in_ring_list'].index(atom.IsInRing()),
+            ]
+    return atom_feature
+# from rdkit import Chem
+# mol = Chem.MolFromSmiles('Cl[C@H](/C=C/C)Br')
+# atom = mol.GetAtomWithIdx(1)  # chiral carbon
+# atom_feature = atom_to_feature_vector(atom)
+# assert atom_feature == [5, 2, 4, 5, 1, 0, 2, 0, 0]
+
+
+def get_atom_feature_dims():
+    return list(map(len, [
+        allowable_features['possible_atomic_num_list'],
+        allowable_features['possible_chirality_list'],
+        allowable_features['possible_degree_list'],
+        allowable_features['possible_formal_charge_list'],
+        allowable_features['possible_numH_list'],
+        allowable_features['possible_number_radical_e_list'],
+        allowable_features['possible_hybridization_list'],
+        allowable_features['possible_is_aromatic_list'],
+        allowable_features['possible_is_in_ring_list']
+        ]))
+
+def bond_to_feature_vector(bond):
+    """
+    Converts rdkit bond object to feature list of indices
+    :param mol: rdkit bond object
+    :return: list
+    """
+    bond_feature = [
+                safe_index(allowable_features['possible_bond_type_list'], str(bond.GetBondType())),
+                allowable_features['possible_bond_stereo_list'].index(str(bond.GetStereo())),
+                allowable_features['possible_is_conjugated_list'].index(bond.GetIsConjugated()),
+            ]
+    return bond_feature
+# uses same molecule as atom_to_feature_vector test
+# bond = mol.GetBondWithIdx(2)  # double bond with stereochem
+# bond_feature = bond_to_feature_vector(bond)
+# assert bond_feature == [1, 2, 0]
+
+def get_bond_feature_dims():
+    return list(map(len, [
+        allowable_features['possible_bond_type_list'],
+        allowable_features['possible_bond_stereo_list'],
+        allowable_features['possible_is_conjugated_list']
+        ]))
+
+def atom_feature_vector_to_dict(atom_feature):
+    [atomic_num_idx, 
+    chirality_idx,
+    degree_idx,
+    formal_charge_idx,
+    num_h_idx,
+    number_radical_e_idx,
+    hybridization_idx,
+    is_aromatic_idx,
+    is_in_ring_idx] = atom_feature
+
+    feature_dict = {
+        'atomic_num': allowable_features['possible_atomic_num_list'][atomic_num_idx],
+        'chirality': allowable_features['possible_chirality_list'][chirality_idx],
+        'degree': allowable_features['possible_degree_list'][degree_idx],
+        'formal_charge': allowable_features['possible_formal_charge_list'][formal_charge_idx],
+        'num_h': allowable_features['possible_numH_list'][num_h_idx],
+        'num_rad_e': allowable_features['possible_number_radical_e_list'][number_radical_e_idx],
+        'hybridization': allowable_features['possible_hybridization_list'][hybridization_idx],
+        'is_aromatic': allowable_features['possible_is_aromatic_list'][is_aromatic_idx],
+        'is_in_ring': allowable_features['possible_is_in_ring_list'][is_in_ring_idx]
+    }
+
+    return feature_dict
+# # uses same atom_feature as atom_to_feature_vector test
+# atom_feature_dict = atom_feature_vector_to_dict(atom_feature)
+# assert atom_feature_dict['atomic_num'] == 6
+# assert atom_feature_dict['chirality'] == 'CHI_TETRAHEDRAL_CCW'
+# assert atom_feature_dict['degree'] == 4
+# assert atom_feature_dict['formal_charge'] == 0
+# assert atom_feature_dict['num_h'] == 1
+# assert atom_feature_dict['num_rad_e'] == 0
+# assert atom_feature_dict['hybridization'] == 'SP3'
+# assert atom_feature_dict['is_aromatic'] == False
+# assert atom_feature_dict['is_in_ring'] == False
+
+def bond_feature_vector_to_dict(bond_feature):
+    [bond_type_idx, 
+    bond_stereo_idx,
+    is_conjugated_idx] = bond_feature
+
+    feature_dict = {
+        'bond_type': allowable_features['possible_bond_type_list'][bond_type_idx],
+        'bond_stereo': allowable_features['possible_bond_stereo_list'][bond_stereo_idx],
+        'is_conjugated': allowable_features['possible_is_conjugated_list'][is_conjugated_idx]
+    }
+
+    return feature_dict
+# # uses same bond as bond_to_feature_vector test
+# bond_feature_dict = bond_feature_vector_to_dict(bond_feature)
+# assert bond_feature_dict['bond_type'] == 'DOUBLE'
+# assert bond_feature_dict['bond_stereo'] == 'STEREOE'
+# assert bond_feature_dict['is_conjugated'] == False
+
+
+import numpy as np
+from rdkit import Chem
+
+def smiles_to_graph(smiles_string, properties):
+    """
+    Converts SMILES string to graph Data object
+    :input: SMILES string (str)
+    :return: graph object
+    """
+    try:
+        mol = Chem.MolFromSmiles(smiles_string)
+        # atoms
+        atom_features_list = []
+        # atom_label = []
+        for atom in mol.GetAtoms():
+            atom_features_list.append(atom_to_feature_vector(atom))
+            # atom_label.append(atom.GetSymbol())
+
+        x = np.array(atom_features_list, dtype=np.int64)
+
+        # bonds
+        num_bond_features = 3  # bond type, bond stereo, is_conjugated
+        if len(mol.GetBonds()) > 0:  # mol has bonds
+            edges_list = []
+            edge_features_list = []
+            for bond in mol.GetBonds():
+                i = bond.GetBeginAtomIdx()
+                j = bond.GetEndAtomIdx()
+
+                edge_feature = bond_to_feature_vector(bond)
+                # add edges in both directions
+                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 = np.array(edges_list, dtype=np.int64).T
+
+            # data.edge_attr: Edge feature matrix with shape [num_edges, num_edge_features]
+            edge_attr = np.array(edge_features_list, dtype=np.int64)
+
+        else:  # mol has no bonds
+            edge_index = np.empty((2, 0), dtype=np.int64)
+            edge_attr = np.empty((0, num_bond_features), dtype=np.int64)
+
+        graph = dict()
+        graph["edge_index"] = edge_index
+        graph["edge_feat"] = edge_attr
+        graph["node_feat"] = x
+        graph["num_nodes"] = len(x)
+        if properties is None:
+            y = np.full(1, np.nan, dtype=np.float32)
+            graph['y'] = y.reshape(1, -1)
+        else:
+            graph['y'] = np.array(properties, dtype=np.float32).reshape(1, -1)
+
+        return graph
+    except:
+        return None