Implementing Neural Fingerprint from Duvenaud et al.

CHANGELOG.md

@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added the `NeuralFingerprint` model for learning fingerprints of molecules ([#7919](https://github.com/pyg-team/pytorch_geometric/pull/7919))
 - Added `SparseTensor` support to `WLConvContinuous`, `GeneralConv`, `PDNConv` and `ARMAConv` ([#8013](https://github.com/pyg-team/pytorch_geometric/pull/8013))
 - Added `LCMAggregation`, an implementation of Learnable Communitive Monoids ([#7976](https://github.com/pyg-team/pytorch_geometric/pull/7976))
 - Added a warning for isolated/non-existing node types in `HeteroData.validate()` ([#7995](https://github.com/pyg-team/pytorch_geometric/pull/7995))

test/nn/models/test_neural_fingerprint.py

@@ -0,0 +1,28 @@
+import pytest
+import torch
+
+import torch_geometric.typing
+from torch_geometric.nn import NeuralFingerprint
+from torch_geometric.testing import is_full_test
+from torch_geometric.typing import SparseTensor
+
+
+@pytest.mark.parametrize('batch', [None, torch.tensor([0, 1, 1])])
+def test_neural_fingerprint(batch):
+    x = torch.randn(3, 7)
+    edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
+
+    model = NeuralFingerprint(7, 16, out_channels=5, num_layers=4)
+    assert str(model) == 'NeuralFingerprint(7, 5, num_layers=4)'
+    model.reset_parameters()
+
+    out = model(x, edge_index, batch)
+    assert out.size() == (1, 5) if batch is None else (2, 5)
+
+    if torch_geometric.typing.WITH_TORCH_SPARSE:
+        adj = SparseTensor.from_edge_index(edge_index, sparse_sizes=(3, 3))
+        assert torch.allclose(model(x, adj.t(), batch), out)
+
+    if is_full_test():
+        jit = torch.jit.export(model)
+        assert torch.allclose(jit(x, edge_index, batch), out)

torch_geometric/nn/models/__init__.py

@@ -24,6 +24,7 @@
 from .rev_gnn import GroupAddRev
 from .gnnff import GNNFF
 from .pmlp import PMLP
+from .neural_fingerprint import NeuralFingerprint
 
 __all__ = classes = [
     'MLP',
@@ -64,4 +65,5 @@
     'GroupAddRev',
     'GNNFF',
     'PMLP',
+    'NeuralFingerprint',
 ]

torch_geometric/nn/models/neural_fingerprint.py

@@ -0,0 +1,72 @@
+from typing import Optional
+
+import torch
+from torch import Tensor
+
+from torch_geometric.nn import Linear, MFConv, global_add_pool
+from torch_geometric.typing import Adj
+
+
+class NeuralFingerprint(torch.nn.Module):
+    r"""The Neural Fingerprint model from the
+    `"Convolutional Networks on Graphs for Learning Molecular Fingerprints"
+    <https://arxiv.org/abs/1509.09292>`__ paper to generate fingerprints
+    of molecules.
+
+    Args:
+        in_channels (int): Size of each input sample.
+        hidden_channels (int): Size of each hidden sample.
+        out_channels (int): Size of each output fingerprint.
+        num_layers (int): Number of layers.
+        **kwargs (optional): Additional arguments of
+            :class:`torch_geometric.nn.conv.MFConv`.
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: int,
+        out_channels: int,
+        num_layers: int,
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.out_channels = out_channels
+        self.num_layers = num_layers
+
+        self.convs = torch.nn.ModuleList()
+        for i in range(self.num_layers):
+            in_channels = self.in_channels if i == 0 else self.hidden_channels
+            self.convs.append(MFConv(in_channels, hidden_channels, **kwargs))
+
+        self.lins = torch.nn.ModuleList()
+        for _ in range(self.num_layers):
+            self.lins.append(Linear(hidden_channels, out_channels, bias=False))
+
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        for conv in self.convs:
+            conv.reset_parameters()
+        for lin in self.lins:
+            lin.reset_parameters()
+
+    def forward(
+        self,
+        x: Tensor,
+        edge_index: Adj,
+        batch: Optional[Tensor] = None,
+        batch_size: Optional[int] = None,
+    ) -> Tensor:
+        """"""
+        outs = []
+        for conv, lin in zip(self.convs, self.lins):
+            x = conv(x, edge_index).sigmoid()
+            y = lin(x).softmax(dim=-1)
+            outs.append(global_add_pool(y, batch, batch_size))
+        return sum(outs)
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}({self.in_channels}, '
+                f'{self.out_channels}, num_layers={self.num_layers})')