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

torch_molecule/predictor/bfgnn/model.py

@@ -0,0 +1,116 @@
+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, MLP
+from ...utils import init_weights
+
+class BFGNN(nn.Module):
+    def __init__(
+        self,
+        num_task,
+        num_layer,
+        hidden_size=300,
+        gnn_type="gin-virtual",
+        drop_ratio=0.5,
+        norm_layer="batch_norm",
+        graph_pooling="max",
+        augmented_feature=['maccs', 'morgan'],
+        algorithm_aligned = 'bf'
+    ):
+        super(BFGNN, self).__init__()
+        gnn_name = gnn_type.split("-")[0]
+        self.algorithm_aligned = algorithm_aligned
+        self.num_task = num_task
+        self.hidden_size = hidden_size
+
+        if "virtual" in gnn_type:
+            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,
+                algorithm_aligned=algorithm_aligned
+            )
+        else:
+            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,
+                algorithm_aligned=algorithm_aligned
+            )
+        if graph_pooling == "sum":
+            self.pool = global_add_pool
+        elif graph_pooling == "mean":
+            self.pool = global_mean_pool
+        elif graph_pooling == "max":
+            self.pool = global_max_pool
+        else:
+            raise ValueError(f"Invalid graph pooling type {graph_pooling}.")
+
+        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.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 _augmented_graph_features(self, batched_data, h_rep):
+        if self.augmented_feature:
+            if 'morgan' in self.augmented_feature:
+                morgan = batched_data.morgan.type_as(h_rep)
+                h_rep = torch.cat((h_rep, morgan), dim=1)
+            if 'maccs' in self.augmented_feature:
+                maccs = batched_data.maccs.type_as(h_rep)
+                h_rep = torch.cat((h_rep, maccs), dim=1)
+        return h_rep
+
+    def compute_loss(self, batched_data, criterion, l1_penalty=1e-3):
+        h_node, _ = self.graph_encoder(batched_data)
+        h_rep = self.pool(h_node, batched_data.batch)
+        h_rep = self._augmented_graph_features(batched_data, h_rep)
+        prediction = self.predictor(h_rep)
+        target = batched_data.y.to(torch.float32)
+        is_labeled = batched_data.y == batched_data.y
+        loss = criterion(prediction.to(torch.float32)[is_labeled], target[is_labeled]).mean()
+        l1_loss = sum(p.abs().sum() for p in self.graph_encoder.parameters()) + sum(p.abs().sum() for p in self.predictor.parameters())
+        return loss + l1_penalty * l1_loss
+
+    def forward(self, batched_data):
+        h_node, _ = self.graph_encoder(batched_data)
+        h_rep = self.pool(h_node, batched_data.batch)
+        h_rep = self._augmented_graph_features(batched_data, h_rep)
+        prediction = self.predictor(h_rep)
+        return {"prediction": prediction}

torch_molecule/predictor/bfgnn/modeling_bfgnn.py

@@ -0,0 +1,81 @@
+import os
+import numpy as np
+import warnings
+import datetime
+from tqdm import tqdm
+from typing import Optional, Union, Dict, Any, List, Type
+from dataclasses import dataclass, field
+
+import torch
+from torch_geometric.data import Data
+
+from .model import BFGNN
+from ...utils import graph_from_smiles
+from ..gnn.modeling_gnn import GNNMolecularPredictor
+from ...utils.search import (
+    ParameterSpec,
+    ParameterType,
+)
+
+@dataclass
+class BFGNNMolecularPredictor(GNNMolecularPredictor):
+    """
+    This predictor implements algorithm alignment of Bellman-Ford algorithm with GNN.
+
+    References
+    ----------
+    - Graph neural networks extrapolate out-of-distribution for shortest paths.
+      https://arxiv.org/abs/2503.19173
+
+    :param l1_penalty: Weight for the L1 penalty
+    :type l1_penalty: float, default=1e-3
+    """
+
+
+    l1_penalty: float = 1e-3
+    # Other Non-init fields
+    model_name: str = "BFGNNMolecularPredictor"
+    model_class: Type[BFGNN] = field(default=BFGNN, init=False)
+    
+    def __post_init__(self):
+        super().__post_init__()
+    
+    @staticmethod
+    def _get_param_names() -> List[str]:
+        return GNNMolecularPredictor._get_param_names() + [
+            "l1_penalty",
+        ]
+    def _get_default_search_space(self):
+        search_space = super()._get_default_search_space().copy()
+        search_space["l1_penalty"] = ParameterSpec(ParameterType.LOG_FLOAT, (1e-6, 1))
+        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 _train_epoch(self, train_loader, optimizer, epoch):
+        self.model.train()
+        losses = []
+
+        iterator = (
+            tqdm(train_loader, desc="Training", leave=False)
+            if self.verbose
+            else train_loader
+        )
+
+        for step, batch in enumerate(iterator):
+            batch = batch.to(self.device)
+            optimizer.zero_grad()
+
+            loss = self.model.compute_loss(batch, self.loss_criterion, self.l1_penalty)
+            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()
+
+            if self.verbose:
+                iterator.set_postfix({"Epoch": epoch, "Total Loss": f"{loss.item():.4f}"})
+            losses.append(loss.item())
+
+        return losses

torch_molecule/predictor/dir/__init__.py

@@ -0,0 +1,3 @@
+from .modeling_dir import DIRMolecularPredictor
+
+__all__ = ['DIRMolecularPredictor']

torch_molecule/predictor/dir/model.py

@@ -0,0 +1,235 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from torch_geometric.nn import global_add_pool, global_mean_pool, global_max_pool
+
+from .utils import split_batch, relabel, set_masks, clear_masks
+from ...nn import GNN_node, GNN_node_Virtualnode
+from ...utils import init_weights
+
+class DIR(nn.Module):
+    def __init__(
+        self,
+        num_task,
+        num_layer,
+        causal_ratio=0.8,
+        hidden_size=300,
+        gnn_type="gin-virtual",
+        drop_ratio=0.5,
+        norm_layer="batch_norm",
+        graph_pooling="sum",
+        augmented_feature=['maccs', 'morgan']
+    ):
+        super(DIR, self).__init__()
+        gnn_name = gnn_type.split("-")[0]
+        self.num_task = num_task
+        self.hidden_size = hidden_size
+        self.causal_ratio = causal_ratio
+        if "virtual" in gnn_type:
+            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,
+                encode_atom=False
+            )
+        else:
+            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,
+                encode_atom=False
+            )
+
+        self.causal_att_net = CausalAttNet(causal_ratio=causal_ratio, hidden_size=hidden_size)
+        
+        graph_dim = hidden_size
+        self.augmented_feature = augmented_feature
+        if graph_pooling == "sum":
+            self.pool = global_add_pool
+        elif graph_pooling == "mean":
+            self.pool = global_mean_pool
+        elif graph_pooling == "max":
+            self.pool = global_max_pool
+        else:
+            raise ValueError(f"Invalid graph pooling type {graph_pooling}.")
+        
+        if augmented_feature:
+            if "morgan" in augmented_feature:
+                graph_dim += 1024
+            if "maccs" in augmented_feature:
+                graph_dim += 167
+        
+        self.causal_lin = torch.nn.Linear(graph_dim, num_task)
+        self.conf_lin = torch.nn.Linear(hidden_size, num_task)
+    
+    def get_graph_rep(self, x, edge_index, edge_attr, batch):
+        h_node, _ = self.graph_encoder(x, edge_index, edge_attr, batch)
+        h_graph = self.pool(h_node, batch)
+        return h_graph
+    
+    def get_causal_pred(self, batched_data, h_graph):
+        h_graph = self._augment_graph_features(batched_data, h_graph)
+        return self.causal_lin(h_graph)
+    
+    def get_conf_pred(self, conf_graph_x):
+        return self.conf_lin(conf_graph_x)
+    
+    def get_comb_pred(self, batched_data, causal_graph_x, conf_graph_x):
+        causal_graph_x = self._augment_graph_features(batched_data, causal_graph_x)
+        causal_pred = self.causal_lin(causal_graph_x)
+        conf_pred = self.conf_lin(conf_graph_x).detach()
+        return torch.sigmoid(conf_pred) * causal_pred
+
+    def initialize_parameters(self, seed=None):
+        if seed is not None:
+            torch.manual_seed(seed)
+        
+        # Initialize the main components
+        init_weights(self.graph_encoder)
+        init_weights(self.causal_att_net)
+        init_weights(self.causal_lin)
+        init_weights(self.conf_lin)
+        
+        # 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):
+        if self.augmented_feature:
+            if 'morgan' in self.augmented_feature:
+                morgan = batched_data.morgan.type_as(h)
+                h = torch.cat((h, morgan), dim=1)
+            if 'maccs' in self.augmented_feature:
+                maccs = batched_data.maccs.type_as(h)
+                h = torch.cat((h, maccs), dim=1)
+        return h
+    
+    def compute_loss(self, batched_data, criterion, alpha_prime):
+        (causal_x, causal_edge_index, causal_edge_attr, causal_edge_weight, causal_batch),\
+        (conf_x, conf_edge_index, conf_edge_attr, conf_edge_weight, conf_batch), pred_edge_weight = self.causal_att_net(batched_data)
+            
+        set_masks(causal_edge_weight, self.graph_encoder)
+        set_masks(causal_edge_weight, self.causal_lin)
+        set_masks(causal_edge_weight, self.conf_lin)
+        causal_rep = self.get_graph_rep(x=causal_x, edge_index=causal_edge_index, edge_attr=causal_edge_attr, batch=causal_batch)
+        causal_out = self.get_causal_pred(batched_data, causal_rep)
+        clear_masks(self.graph_encoder)
+        clear_masks(self.causal_lin)
+        clear_masks(self.conf_lin)
+
+        set_masks(conf_edge_weight, self.graph_encoder)
+        set_masks(conf_edge_weight, self.causal_lin)
+        set_masks(conf_edge_weight, self.conf_lin)
+        conf_rep = self.get_graph_rep(x=conf_x, edge_index=conf_edge_index, edge_attr=conf_edge_attr, batch=conf_batch).detach()
+        conf_out = self.get_conf_pred(conf_rep)
+        clear_masks(self.graph_encoder)
+        clear_masks(self.causal_lin)
+        clear_masks(self.conf_lin)
+
+        batched_data.y = batched_data.y[0:causal_out.size(0),:]
+        is_labeled = batched_data.y == batched_data.y
+        
+        causal_loss = criterion(
+            causal_out.to(torch.float32)[is_labeled], 
+            batched_data.y.to(torch.float32)[is_labeled]
+            ).mean()
+        conf_loss = criterion(
+            conf_out.to(torch.float32)[is_labeled], 
+            batched_data.y.to(torch.float32)[is_labeled]
+            ).mean()
+
+        env_loss = torch.tensor([]).to(batched_data.y.device)
+        for conf in conf_rep:
+            rep_out = self.get_comb_pred(batched_data, causal_rep, conf)
+            tmp = criterion(rep_out.to(torch.float32)[is_labeled], batched_data.y.to(torch.float32)[is_labeled]).mean()
+            env_loss = torch.cat([env_loss, tmp.unsqueeze(0)])
+        causal_loss += alpha_prime * env_loss.mean()
+        env_loss = alpha_prime * torch.var(env_loss * conf_rep.size(0))
+        
+        return causal_loss, conf_loss, env_loss
+
+    def forward(self, batched_data):
+        (causal_x, causal_edge_index, causal_edge_attr, causal_edge_weight, causal_batch), _, _ = self.causal_att_net(batched_data)
+        set_masks(causal_edge_weight, self.graph_encoder)
+        set_masks(causal_edge_weight, self.causal_lin)
+        set_masks(causal_edge_weight, self.conf_lin)
+
+        causal_rep = self.get_graph_rep(x=causal_x, edge_index=causal_edge_index, edge_attr=causal_edge_attr, batch=causal_batch)
+        causal_out = self.get_causal_pred(batched_data, causal_rep)
+        clear_masks(self.graph_encoder)
+        clear_masks(self.causal_lin)
+        clear_masks(self.conf_lin)
+
+        return {"prediction": causal_out, "representation": causal_rep}
+
+    @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 CausalAttNet(nn.Module):
+    def __init__(self, causal_ratio, hidden_size):
+        super(CausalAttNet, self).__init__()
+        self.gnn_node = GNN_node_Virtualnode(2, hidden_size, JK="last", drop_ratio=0., residual=True, gnn_name="gin", norm_layer="batch_norm")
+        self.linear = nn.Linear(hidden_size*2, 1)
+        self.ratio = causal_ratio
+        
+    def forward(self, data):
+        x, _ = self.gnn_node(data.x, data.edge_index, data.edge_attr, data.batch)
+
+        row, col = data.edge_index
+        edge_rep = torch.cat([x[row], x[col]], dim=-1)
+        pred_edge_weight = self.linear(edge_rep).view(-1)
+        causal_edge_index = torch.LongTensor([[],[]]).to(x.device)
+        causal_edge_weight = torch.tensor([]).to(x.device)
+        causal_edge_attr = torch.LongTensor([]).to(x.device)
+        conf_edge_index = torch.LongTensor([[],[]]).to(x.device)
+        conf_edge_weight = torch.tensor([]).to(x.device)
+        conf_edge_attr = torch.LongTensor([]).to(x.device)
+
+        edge_indices, _, _, num_edges, cum_edges = split_batch(data)
+        for edge_index, N, C in zip(edge_indices, num_edges, cum_edges):
+            n_reserve =  int(self.ratio * N)
+            edge_attr = data.edge_attr[C:C+N]
+            single_mask = pred_edge_weight[C:C+N]
+            single_mask_detach = pred_edge_weight[C:C+N].detach().cpu().numpy()
+            rank = np.argpartition(-single_mask_detach, n_reserve)
+            idx_reserve, idx_drop = rank[:n_reserve], rank[n_reserve:]
+
+            causal_edge_index = torch.cat([causal_edge_index, edge_index[:, idx_reserve]], dim=1)
+            conf_edge_index = torch.cat([conf_edge_index, edge_index[:, idx_drop]], dim=1)
+            causal_edge_weight = torch.cat([causal_edge_weight, single_mask[idx_reserve]])
+            conf_edge_weight = torch.cat([conf_edge_weight, -1 * single_mask[idx_drop]])
+            causal_edge_attr = torch.cat([causal_edge_attr, edge_attr[idx_reserve]])
+            conf_edge_attr = torch.cat([conf_edge_attr, edge_attr[idx_drop]])
+
+        causal_x, causal_edge_index, causal_batch, _ = relabel(x, causal_edge_index, data.batch)
+        conf_x, conf_edge_index, conf_batch, _ = relabel(x, conf_edge_index, data.batch)
+
+        return (causal_x, causal_edge_index, causal_edge_attr, causal_edge_weight, causal_batch),\
+                (conf_x, conf_edge_index, conf_edge_attr, conf_edge_weight, conf_batch),\
+                pred_edge_weight

torch_molecule/predictor/dir/modeling_dir.py

@@ -0,0 +1,161 @@
+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 .model import DIR
+from ..gnn.modeling_gnn import GNNMolecularPredictor
+from ...utils.search import (
+    ParameterSpec,
+    ParameterType,
+)
+
+@dataclass
+class DIRMolecularPredictor(GNNMolecularPredictor):
+    """This predictor implements the DIR for molecular property prediction tasks.
+
+    The full name of DIR is Discovering Invariant Rationales.
+
+    References
+    ----------
+    - Discovering Invariant Rationales for Graph Neural Networks.
+      https://openreview.net/forum?id=hGXij5rfiHw
+    - Code: https://github.com/Wuyxin/DIR-GNN
+        
+    Parameters
+    ----------
+    causal_ratio : float, default=0.8
+        The ratio of causal edges to keep during training. A higher ratio means more edges
+        are considered causal/important for the prediction. This controls the sparsity of
+        the learned rationales.
+        
+    lw_invariant : float, default=1e-4
+        The weight of the invariance loss term. This loss encourages the model to learn
+        rationales that are invariant across different environments/perturbations. A higher
+        value puts more emphasis on learning invariant features.
+    """
+    
+    # DIR-specific parameters
+    causal_ratio: float = 0.8
+    lw_invariant: float = 1e-4
+    # Override parent defaults
+    model_name: str = "DIRMolecularPredictor"
+    model_class: Type[DIR] = field(default=DIR, init=False)
+    
+    def __post_init__(self):
+        super().__post_init__()
+
+    @staticmethod
+    def _get_param_names() -> List[str]:
+        return ["causal_ratio", "lw_invariant"] + GNNMolecularPredictor._get_param_names()
+
+    def _get_default_search_space(self):
+        search_space = super()._get_default_search_space().copy()
+        search_space["causal_ratio"] = ParameterSpec(ParameterType.FLOAT, (0.1, 0.9))
+        search_space["lw_invariant"] = ParameterSpec(ParameterType.FLOAT, (1e-5, 1e-2))
+        return search_space
+
+    def _setup_optimizers(self) -> Tuple[Dict[str, torch.optim.Optimizer], Optional[Any]]:
+        model_optimizer = torch.optim.Adam(self.model.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay)
+        conf_optimizer = torch.optim.Adam(self.model.conf_lin.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay)
+        if self.grad_clip_value is not None:
+            for group in model_optimizer.param_groups:
+                group.setdefault("max_norm", self.grad_clip_value)
+                group.setdefault("norm_type", 2.0)
+
+        scheduler = None
+        if self.use_lr_scheduler:
+            scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+                model_optimizer,
+                mode="min",
+                factor=self.scheduler_factor,
+                patience=self.scheduler_patience,
+                min_lr=1e-6,
+                cooldown=0,
+                eps=1e-8,
+            )
+        optimizer = {"model": model_optimizer, "conf": conf_optimizer}
+
+        return optimizer, scheduler
+    
+    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["causal_ratio"] = checkpoint["hyperparameters"]["causal_ratio"]
+        else:
+            base_params["causal_ratio"] = self.causal_ratio
+        return base_params
+
+    def _train_epoch(self, train_loader, optimizer, epoch):
+        self.model.train()
+        losses = []
+
+        iterator = (
+            tqdm(train_loader, desc="Training", leave=False)
+            if self.verbose
+            else train_loader
+        )
+
+        alpha_prime = self.lw_invariant * (epoch ** 1.6)
+        conf_opt = optimizer["conf"]
+        model_optimizer = optimizer["model"]
+
+        for batch in iterator:
+
+            batch = batch.to(self.device)
+
+            # Forward pass and loss computation
+            causal_loss, conf_loss, env_loss = self.model.compute_loss(batch, self.loss_criterion, alpha_prime)
+
+            conf_opt.zero_grad()
+            conf_loss.backward()
+            conf_opt.step()
+
+            model_optimizer.zero_grad()
+            (causal_loss + env_loss).backward()
+            model_optimizer.step()
+
+            loss = causal_loss + env_loss + conf_loss
+            losses.append(loss.item())
+
+            # Update progress bar if using tqdm
+            if self.verbose:
+                iterator.set_postfix({"Epoch": epoch, "Causal Loss": f"{causal_loss.item():.4f}", "Conf Loss": f"{conf_loss.item():.4f}", "Env Loss": f"{env_loss.item():.4f}", "Total Loss": f"{loss.item():.4f}"})
+
+        return losses
+    
+    def predict(self, X: List[str]) -> Dict[str, Union[np.ndarray, List[List]]]:
+        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 = []
+        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())
+
+        if predictions:
+            return {
+                "prediction": np.concatenate(predictions, axis=0),
+            }
+        else:
+            warnings.warn(
+                "No valid predictions could be made from the input data. Returning empty results."
+            )
+            return {"prediction": np.array([])}
+
+

torch_molecule/predictor/dir/utils.py

@@ -0,0 +1,40 @@
+import torch
+from torch_geometric.utils import degree
+from torch_geometric.nn.conv import MessagePassing
+
+def split_batch(g):
+    split = degree(g.batch[g.edge_index[0]], dtype=torch.long).tolist()
+    edge_indices = torch.split(g.edge_index, split, dim=1)
+    num_nodes = degree(g.batch, dtype=torch.long)
+    cum_nodes = torch.cat([g.batch.new_zeros(1), num_nodes.cumsum(dim=0)[:-1]])
+    num_edges = torch.tensor([e.size(1) for e in edge_indices], dtype=torch.long).to(g.x.device)
+    cum_edges = torch.cat([g.batch.new_zeros(1), num_edges.cumsum(dim=0)[:-1]])
+
+    return edge_indices, num_nodes, cum_nodes, num_edges, cum_edges
+
+def relabel(x, edge_index, batch, pos=None):
+    num_nodes = x.size(0)
+    sub_nodes = torch.unique(edge_index)
+    x = x[sub_nodes]
+    batch = batch[sub_nodes]
+    row, col = edge_index
+    # remapping the nodes in the explanatory subgraph to new ids.
+    node_idx = row.new_full((num_nodes,), -1)
+    node_idx[sub_nodes] = torch.arange(sub_nodes.size(0), device=row.device)
+    edge_index = node_idx[edge_index]
+    if pos is not None:
+        pos = pos[sub_nodes]
+    return x, edge_index, batch, pos
+
+
+def set_masks(mask: torch.Tensor, model: torch.nn.Module):
+    for module in model.modules():
+        if isinstance(module, MessagePassing):
+            module.explain = True
+            module._edge_mask = mask
+        
+def clear_masks(model: torch.nn.Module):
+    for module in model.modules():
+        if isinstance(module, MessagePassing):
+            module.explain = None
+            module._edge_mask = None

torch_molecule/predictor/gnn/__init__.py

@@ -0,0 +1,3 @@
+from .modeling_gnn import GNNMolecularPredictor
+
+__all__ = ['GNNMolecularPredictor']