Repository: deepchem/deepchem

Issue:
Scaffold splitting of BBBP dataset
Scaffolding splitting of BBBP dataset throws a ValueError.

Steps to reproduct:
```py
import deepchem as dc

tasks, dataset, transformer = dc.molnet.load_bbbp(featurizer=dc.feat.DummyFeaturizer(), 
                                                  splitter=dc.splits.ScaffoldSplitter())
```

Error traceback:
```py
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[6], line 3
      1 import deepchem as dc
----> 3 tasks, dataset, transformer = dc.molnet.load_bbbp(featurizer=dc.feat.DummyFeaturizer(), 
      4                                                   splitter=dc.splits.ScaffoldSplitter())

...

File ~/dc/deepchem/splits/splitters.py:1590, in ScaffoldSplitter.generate_scaffolds(self, dataset, 
log_every_n)
   1588 if ind % log_every_n == 0:
   1589     logger.info("Generating scaffold %d/%d" % (ind, data_len))
-> 1590 scaffold = _generate_scaffold(smiles)
   1591 if scaffold not in scaffolds:
   1592     scaffolds[scaffold] = [ind]

File ~//dc/deepchem/splits/splitters.py:1312, in _generate_scaffold(smiles, include_chirality)
   1309     raise ImportError("This function requires RDKit to be installed.")
   1311 mol = Chem.MolFromSmiles(smiles)
-> 1312 scaffold = MurckoScaffoldSmiles(mol=mol, includeChirality=include_chirality)
   1313 return scaffold

File ~/Applications/miniconda3/envs/dc/lib/python3.9/site-packages/rdkit/Chem/Scaffolds/MurckoScaffold.py:114, in 
MurckoScaffoldSmiles(smiles, mol, includeChirality)
    112   mol = Chem.MolFromSmiles(smiles)
    113 if mol is None:
--> 114   raise ValueError('No molecule provided')
    115 scaffold = GetScaffoldForMol(mol)
    116 if not scaffold:

ValueError: No molecule provided
```

Expected result: scaffold splitted BBBP dataset.


Discussion:

Patch:
diff --git a/deepchem/models/torch_models/layers.py b/deepchem/models/torch_models/layers.py
index 74aa2320a..f8524e990 100644
--- a/deepchem/models/torch_models/layers.py
+++ b/deepchem/models/torch_models/layers.py
@@ -1,4 +1,5 @@
 import math
+from math import pi as PI
 import numpy as np
 from typing import Any, Tuple, Optional, Sequence, List, Union, Callable, Dict, TypedDict
 from collections.abc import Sequence as SequenceCollection
@@ -18,6 +19,8 @@ except ModuleNotFoundError:
 from deepchem.utils.typing import OneOrMany, ActivationFn, ArrayLike
 from deepchem.utils.pytorch_utils import get_activation, segment_sum
 from torch.nn import init as initializers
+from torch_geometric.nn import MessagePassing
+from torch_geometric.utils import add_self_loops
 
 
 class MultilayerPerceptron(nn.Module):
@@ -4536,3 +4539,256 @@ class _MXMNetEnvelope(torch.nn.Module):
         zero: torch.Tensor = torch.zeros_like(x)
         output: torch.Tensor = torch.where(x < 1, env_val, zero)
         return output
+
+
+class MXMNetGlobalMessagePassing(MessagePassing):
+    """This class implements the Global Message Passing Layer from the Molecular Mechanics-Driven Graph Neural Network
+    with Multiplex Graph for Molecular Structures(MXMNet) paper [1]_.
+
+    This layer consists of two message passing steps and an update step between them.
+
+    Let:
+        - **x_i** : ``The node to be updated``
+        - **h_i** : ``The hidden state of x_i``
+        - **x_j** : ``The neighbour node connected to x_i by edge e_ij``
+        - **h_j** : ``The hidden state of x_j``
+        - **W** : ``The edge weights``
+        - **m_ij** : ``The message between x_i and x_j``
+        - **h_j (self_loop)** : ``The set of hidden states of atom features``
+        - **mlp** : ``MultilayerPerceptron``
+        - **res** : ``ResidualBlock``
+
+    **In each message passing step**
+
+        .. code-block:: python
+            
+            m_ij = mlp1([h_i || h_j || e_ij])*(e_ij W)
+
+        **To handle self loops**
+
+            .. code-block:: python
+            
+                m_ij = m_ij + h_j(self_loop)
+
+    **In each update step**
+
+        .. code-block:: python
+
+            hm_j = res1(sum(m_ij)) 
+            h_j_new = mlp2(hm_j) + h_j 
+            h_j_new = res2(h_j_new) 
+            h_j_new = res3(h_j_new) 
+
+    .. note::
+    Message passing and message aggregation(sum) is handled by ``self.propagate()``.
+
+    References
+    ----------
+    .. [1] Molecular Mechanics-Driven Graph Neural Network with Multiplex Graph for Molecular Structures. https://arxiv.org/pdf/2011.07457.pdf
+
+
+    Examples
+    --------
+    The provided example demonstrates how to use the GlobalMessagePassing layer by creating an instance, passing input tensors (node_features, edge_attributes, edge_indices) through it, and checking the shape of the output.
+
+    Initializes variables and creates a configuration dictionary with specific values.
+
+    >>> dim = 1
+    >>> node_features = torch.tensor([[0.8343], [1.2713], [1.2713], [1.2713], [1.2713]])
+    >>> edge_attributes = torch.tensor([[1.0004], [1.0004], [1.0005], [1.0004], [1.0004],[-0.2644], [-0.2644], [-0.2644], [1.0004],[-0.2644], [-0.2644], [-0.2644], [1.0005],[-0.2644], [-0.2644], [-0.2644], [1.0004],[-0.2644], [-0.2644], [-0.2644]])
+    >>> edge_indices = torch.tensor([[0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4],[1, 2, 3, 4, 0, 2, 3, 4, 0, 1, 3, 4, 0, 1, 2, 4, 0, 1, 2, 3]])
+    >>> out = MXMNetGlobalMessagePassing(dim)
+    >>> output = out(node_features, edge_attributes, edge_indices)
+    >>> output.shape
+    torch.Size([5, 1])
+
+    """
+
+    def __init__(self, dim: int, activation_fn: Union[Callable, str] = 'silu'):
+        """Initializes the MXMNETGlobalMessagePassing layer.
+
+        Parameters
+        -----------
+        dim: int
+            The dimension of the input and output features.
+        """
+
+        super(MXMNetGlobalMessagePassing, self).__init__()
+        activation_fn = get_activation(activation_fn)
+
+        self.h_mlp: MultilayerPerceptron = MultilayerPerceptron(
+            d_input=dim, d_output=dim, activation_fn=activation_fn)
+
+        self.res1: MultilayerPerceptron = MultilayerPerceptron(
+            d_input=dim,
+            d_hidden=(dim,),
+            d_output=dim,
+            activation_fn=activation_fn,
+            skip_connection=True,
+            weighted_skip=False)
+        self.res2: MultilayerPerceptron = MultilayerPerceptron(
+            d_input=dim,
+            d_hidden=(dim,),
+            d_output=dim,
+            activation_fn=activation_fn,
+            skip_connection=True,
+            weighted_skip=False)
+        self.res3: MultilayerPerceptron = MultilayerPerceptron(
+            d_input=dim,
+            d_hidden=(dim,),
+            d_output=dim,
+            activation_fn=activation_fn,
+            skip_connection=True,
+            weighted_skip=False)
+
+        self.mlp: MultilayerPerceptron = MultilayerPerceptron(
+            d_input=dim, d_output=dim, activation_fn=activation_fn)
+
+        self.x_edge_mlp: MultilayerPerceptron = MultilayerPerceptron(
+            d_input=dim * 3, d_output=dim, activation_fn=activation_fn)
+        self.linear: nn.Linear = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, node_features: torch.Tensor,
+                edge_attributes: torch.Tensor,
+                edge_indices: torch.Tensor) -> torch.Tensor:
+        """
+        Performs the forward pass of the GlobalMessagePassing layer.
+
+        Parameters
+        -----------
+        node_features: torch.Tensor
+            The input node features tensor of shape (num_nodes, feature_dim).
+        edge_attributes: torch.Tensor
+            The input edge attribute tensor of shape (num_edges, attribute_dim).
+        edge_indices: torch.Tensor
+            The input edge index tensor of shape (2, num_edges).
+
+        Returns
+        --------
+        torch.Tensor
+            The updated node features tensor after message passing of shape (num_nodes, feature_dim).
+        """
+        edge_indices, _ = add_self_loops(edge_indices,
+                                         num_nodes=node_features.size(0))
+
+        residual_node_features: torch.Tensor = node_features
+
+        # Integrate the Cross Layer Mapping inside the Global Message Passing
+        node_features = self.h_mlp(node_features)
+
+        # Message Passing operation
+        node_features = self.propagate(edge_indices,
+                                       x=node_features,
+                                       num_nodes=node_features.size(0),
+                                       edge_attr=edge_attributes)
+
+        # Update function f_u
+        node_features = self.res1(node_features)
+        node_features = self.mlp(node_features) + residual_node_features
+        node_features = self.res2(node_features)
+        node_features = self.res3(node_features)
+
+        # Message Passing operation
+        node_features = self.propagate(edge_indices,
+                                       x=node_features,
+                                       num_nodes=node_features.size(0),
+                                       edge_attr=edge_attributes)
+
+        return node_features
+
+    def message(self, x_i: torch.Tensor, x_j: torch.Tensor,
+                edge_attr: torch.Tensor) -> torch.Tensor:
+        """Constructs messages to be passed along the edges in the graph.
+
+        Parameters
+        -----------
+        x_i: torch.Tensor
+            The source node features tensor of shape (num_edges+num_nodes, feature_dim).
+        x_j: torch.Tensor
+            The target node features tensor of shape (num_edges+num_nodes, feature_dim).
+        edge_attributes: torch.Tensor
+            The edge attribute tensor of shape (num_edges, attribute_dim).
+
+        Returns
+        --------
+        torch.Tensor
+            The constructed messages tensor.
+        """
+        num_edge: int = edge_attr.size()[0]
+
+        x_edge: torch.Tensor = torch.cat(
+            (x_i[:num_edge], x_j[:num_edge], edge_attr), -1)
+        x_edge = self.x_edge_mlp(x_edge)
+
+        x_j = torch.cat((self.linear(edge_attr) * x_edge, x_j[num_edge:]),
+                        dim=0)
+
+        return x_j
+
+
+class MXMNetBesselBasisLayer(torch.nn.Module):
+    """This layer implements a basis layer for the MXMNet model using Bessel functions.
+    The basis layer is used to model radial symmetry in molecular systems.
+
+    The output of the layer is given by:
+    output = envelope(dist / cutoff) * (freq * dist / cutoff).sin()
+
+    Examples
+    --------
+    >>> radial_layer = MXMNetBesselBasisLayer(num_radial=2, cutoff=2.0, envelope_exponent=2)
+    >>> distances = torch.tensor([0.5, 1.0, 2.0, 3.0])
+    >>> output = radial_layer(distances)
+    >>> output.shape
+    torch.Size([4, 2])
+    """
+
+    def __init__(self,
+                 num_radial: int,
+                 cutoff: float = 5.0,
+                 envelope_exponent: int = 5):
+        """Initialize the MXMNet Bessel Basis Layer.
+
+        Parameters
+        ----------
+        num_radial: int
+            The number of radial basis functions to use.
+        cutoff: float, optional (default 5.0)
+            The radial cutoff distance used to scale the distances.
+        envelope_exponent: int, optional (default 5)
+            The exponent of the envelope function.
+        """
+
+        super(MXMNetBesselBasisLayer, self).__init__()
+        self.cutoff = cutoff
+        self.envelope: _MXMNetEnvelope = _MXMNetEnvelope(envelope_exponent)
+        self.freq: torch.Tensor = torch.nn.Parameter(torch.empty(num_radial))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        """Reset and initialize the learnable parameters of the MXMNet Bessel Basis Layer.
+
+        The 'freq' tensor, representing the frequencies of the Bessel functions, is set up with initial values proportional to π (PI) and becomes a learnable parameter.
+
+        The 'freq' tensor will be updated during the training process to optimize the performance of the MXMNet model for the specific task it is being trained on.
+        """
+
+        with torch.no_grad():
+            torch.arange(1, self.freq.numel() + 1, out=self.freq).mul_(PI)
+        self.freq.requires_grad_()
+
+    def forward(self, dist: torch.Tensor) -> torch.Tensor:
+        """Compute the output of the MXMNet Bessel Basis Layer.
+
+        Parameters
+        ----------
+        dist: torch.Tensor
+            The input tensor representing the pairwise distances between atoms.
+
+        Returns
+        -------
+        output: torch.Tensor
+            The output tensor representing the radial basis functions applied to the input distances.
+        """
+        dist = dist.unsqueeze(-1) / self.cutoff
+        output: torch.Tensor = self.envelope(dist) * (self.freq * dist).sin()
+        return output
diff --git a/deepchem/molnet/load_function/tox21_datasets.py b/deepchem/molnet/load_function/tox21_datasets.py
index 4b58e984f..83e71903b 100644
--- a/deepchem/molnet/load_function/tox21_datasets.py
+++ b/deepchem/molnet/load_function/tox21_datasets.py
@@ -34,6 +34,7 @@ def load_tox21(
     reload: bool = True,
     data_dir: Optional[str] = None,
     save_dir: Optional[str] = None,
+    tasks: List[str] = TOX21_TASKS,
     **kwargs
 ) -> Tuple[List[str], Tuple[Dataset, ...], List[dc.trans.Transformer]]:
     """Load Tox21 dataset
@@ -75,11 +76,15 @@ def load_tox21(
         a directory to save the raw data in
     save_dir: str
         a directory to save the dataset in
+    tasks: List[str], (optional)
+        Specify the set of tasks to load. If no task is specified, then it loads
+    the default set of tasks which are NR-AR, NR-AR-LBD, NR-AhR, NR-Aromatase, NR-ER,
+    NR-ER-LBD, NR-PPAR-gamma, SR-ARE, SR-ATAD5, SR-HSE, SR-MMP, SR-p53.
 
     References
     ----------
     .. [1] Tox21 Challenge. https://tripod.nih.gov/tox21/challenge/
     """
-    loader = _Tox21Loader(featurizer, splitter, transformers, TOX21_TASKS,
-                          data_dir, save_dir, **kwargs)
+    loader = _Tox21Loader(featurizer, splitter, transformers, tasks, data_dir,
+                          save_dir, **kwargs)
     return loader.load_dataset('tox21', reload)
diff --git a/deepchem/splits/splitters.py b/deepchem/splits/splitters.py
index 294b5967a..cffd25a45 100644
--- a/deepchem/splits/splitters.py
+++ b/deepchem/splits/splitters.py
@@ -2,17 +2,16 @@
 Contains an abstract base class that supports chemically aware data splits.
 """
 import inspect
+import itertools
+import logging
 import os
 import random
 import tempfile
-import itertools
-import logging
-from typing import Any, Dict, List, Iterator, Optional, Sequence, Tuple
+from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple, Union
 
+import deepchem as dc
 import numpy as np
 import pandas as pd
-
-import deepchem as dc
 from deepchem.data import Dataset, DiskDataset
 from deepchem.utils import get_print_threshold
 
@@ -1274,7 +1273,8 @@ class ButinaSplitter(Splitter):
         return train_inds, valid_inds, test_inds
 
 
-def _generate_scaffold(smiles: str, include_chirality: bool = False) -> str:
+def _generate_scaffold(smiles: str,
+                       include_chirality: bool = False) -> Union[str, None]:
     """Compute the Bemis-Murcko scaffold for a SMILES string.
 
     Bemis-Murcko scaffolds are described in DOI: 10.1021/jm9602928.
@@ -1309,6 +1309,12 @@ def _generate_scaffold(smiles: str, include_chirality: bool = False) -> str:
         raise ImportError("This function requires RDKit to be installed.")
 
     mol = Chem.MolFromSmiles(smiles)
+    if mol is None:
+        logger.info(
+            'Not generating scaffold for smiles %s - invalid smiles string' %
+            smiles)
+        return None
+
     scaffold = MurckoScaffoldSmiles(mol=mol, includeChirality=include_chirality)
     return scaffold
 
@@ -1502,9 +1508,13 @@ class ScaffoldSplitter(Splitter):
     .. [1] Bemis, Guy W., and Mark A. Murcko. "The properties of known drugs.
         1. Molecular frameworks." Journal of medicinal chemistry 39.15 (1996): 2887-2893.
 
-    Note
-    ----
-    This class requires RDKit to be installed.
+    Notes
+    -----
+    - This class requires RDKit to be installed.
+
+    - When a SMILES representation of a molecule is invalid, the splitter skips processing
+    the datapoint i.e it will not include the molecule in any splits.
+
     """
 
     def split(
@@ -1588,10 +1598,11 @@ class ScaffoldSplitter(Splitter):
             if ind % log_every_n == 0:
                 logger.info("Generating scaffold %d/%d" % (ind, data_len))
             scaffold = _generate_scaffold(smiles)
-            if scaffold not in scaffolds:
-                scaffolds[scaffold] = [ind]
-            else:
-                scaffolds[scaffold].append(ind)
+            if scaffold is not None:
+                if scaffold not in scaffolds:
+                    scaffolds[scaffold] = [ind]
+                else:
+                    scaffolds[scaffold].append(ind)
 
         # Sort from largest to smallest scaffold sets
         scaffolds = {key: sorted(value) for key, value in scaffolds.items()}
diff --git a/docs/source/api_reference/layers.rst b/docs/source/api_reference/layers.rst
index e6a673162..b1d13da50 100644
--- a/docs/source/api_reference/layers.rst
+++ b/docs/source/api_reference/layers.rst
@@ -233,6 +233,14 @@ Torch Layers
 
 .. autoclass:: deepchem.models.torch_models.layers.WeaveGather
   :members:
+  
+.. autoclass:: deepchem.models.torch_models.layers.MXMNetGlobalMessagePassing
+  :members:
+
+
+.. autoclass:: deepchem.models.torch_models.layers.MXMNetBesselBasisLayer
+  :members:
+
 Grover Layers
 ^^^^^^^^^^^^^
 
diff --git a/docs/source/api_reference/torch_layers.csv b/docs/source/api_reference/torch_layers.csv
index f796aad2c..46ed26ebb 100644
--- a/docs/source/api_reference/torch_layers.csv
+++ b/docs/source/api_reference/torch_layers.csv
@@ -42,3 +42,5 @@ MolGANEncoderLayer, `ref <https://arxiv.org/abs/1805.11973>`_, MolGan
 DTNNEmbedding, ref`<https://arxiv.org/abs/1609.08259>`_, DTNNModel
 DTNNStep, ref`<https://arxiv.org/abs/1609.08259>`_, DTNNModel
 DTNNGather, ref`<https://arxiv.org/abs/1609.08259>`_, DTNNModel
+MXMNetGlobalMessagePassing, `ref <https://arxiv.org/pdf/2011.07457>`_, MXMNetModel
+MXMNetBesselBasisLayer, `ref <https://arxiv.org/pdf/2011.07457>`_, MXMNetModel
\ No newline at end of file


Test patch:
diff --git a/deepchem/models/tests/test_layers.py b/deepchem/models/tests/test_layers.py
index e809ec28d..c0855e350 100644
--- a/deepchem/models/tests/test_layers.py
+++ b/deepchem/models/tests/test_layers.py
@@ -1215,3 +1215,50 @@ def test_mxmnet_envelope():
     output = output.detach().numpy()
     result = np.array([1.3125, 0.0000, 0.0000, 0.0000])
     assert np.allclose(result, output, atol=1e-04)
+
+
+@pytest.mark.torch
+def test_mxmnet_global_message_passing():
+    """ Test for MXMNetGlobalMessagePassing Layer."""
+    seed = 123
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+    dim = 1
+    node_features = torch.tensor([[0.8343], [1.2713], [1.2713], [1.2713],
+                                  [1.2713]])
+
+    edge_attr = torch.tensor([[1.0004], [1.0004], [1.0005], [1.0004], [1.0004],
+                              [-0.2644], [-0.2644], [-0.2644], [1.0004],
+                              [-0.2644], [-0.2644], [-0.2644], [1.0005],
+                              [-0.2644], [-0.2644], [-0.2644], [1.0004],
+                              [-0.2644], [-0.2644], [-0.2644]])
+
+    edge_indices = torch.tensor(
+        [[0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4],
+         [1, 2, 3, 4, 0, 2, 3, 4, 0, 1, 3, 4, 0, 1, 2, 4, 0, 1, 2, 3]])
+
+    out = dc.models.torch_models.layers.MXMNetGlobalMessagePassing(dim)
+    output = out(node_features, edge_attr, edge_indices)
+    output = output.detach().numpy()
+    result = np.array([[-0.27947044], [2.417905], [2.417905], [2.4178727],
+                       [2.417905]])
+    print(output)
+    assert np.allclose(output, result, atol=1e-04)
+    assert output.shape == (5, 1)
+
+
+@pytest.mark.torch
+def test_mxmnet_besselbasis():
+    """Test for MXMNetBesselBasisLayer"""
+
+    radial_layer = dc.models.torch_models.layers.MXMNetBesselBasisLayer(
+        num_radial=2, cutoff=2.0, envelope_exponent=2)
+    distances = torch.tensor([0.5, 1.0, 2.0, 3.0])
+    output = radial_layer(distances)
+    output = output.detach().numpy()
+    result = np.array([[2.6434e+00, 3.7383e+00], [1.3125e+00, -1.1474e-07],
+                       [-0.0000e+00, 0.0000e+00], [-0.0000e+00, -0.0000e+00]])
+    assert np.allclose(result, output, atol=1e-04)
diff --git a/deepchem/splits/tests/test_scaffold_splitter.py b/deepchem/splits/tests/test_scaffold_splitter.py
index 35586bea4..e92c6e110 100644
--- a/deepchem/splits/tests/test_scaffold_splitter.py
+++ b/deepchem/splits/tests/test_scaffold_splitter.py
@@ -23,3 +23,14 @@ class TestScaffoldSplitter(unittest.TestCase):
         # has to be smaller or equal than number of total molecules
         scaffolds_separate_cnt = len(scaffolds_separate)
         self.assertTrue(scaffolds_separate_cnt <= train_dataset.X.shape[0])
+
+    def test_generate_scaffold(self):
+        from deepchem.splits.splitters import _generate_scaffold
+        valid_smiles = r's1cc(nc1\[N]=C(\N)N)C'
+        scaffold = _generate_scaffold(valid_smiles)
+        self.assertTrue(scaffold == 'c1cscn1')
+
+        # Invalid because valence for atom 5 N is greater than permitted (4)
+        invalid_smiles = r's1cc(nc1\[NH]=C(\N)N)C'
+        scaffold = _generate_scaffold(invalid_smiles)
+        self.assertIsNone(scaffold)