Repository: pyg-team/pytorch_geometric
Pull request #7919 (2023-08-23 09:25:02)

Implementing Neural Fingerprint from Duvenaud et al.
Implemented the Neural Fingerprint Model as suggested in the paper "[Convolutional Networks on Graphs for Learning Molecular Fingerprints](https://arxiv.org/pdf/1509.09292.pdf)" and suggested in the issue https://github.com/pyg-team/pytorch_geometric/issues/6077.

Unit test has been added and all tests are passing.



codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7919?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7919](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7919?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (76a12d1) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/1dd13de89a5168b950ca5bf37717fea4d2a95f7f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (1dd13de) will **decrease** coverage by `0.53%`.
> The diff coverage is `100.00%`.

> :exclamation: Current head 76a12d1 differs from pull request most recent head 733bdbc. Consider uploading reports for the commit 733bdbc to get more accurate results

```diff
@@            Coverage Diff             @@
##           master    #7919      +/-   ##
==========================================
- Coverage   89.90%   89.38%   -0.53%     
==========================================
  Files         457      458       +1     
  Lines       26835    26857      +22     
==========================================
- Hits        24125    24005     -120     
- Misses       2710     2852     +142     
```


| [Files Changed](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7919?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/nn/models/\_\_init\_\_.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7919?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL21vZGVscy9fX2luaXRfXy5weQ==) | `100.00% <100.00%> (ø)` | |
| [torch\_geometric/nn/models/neural\_fingerprint.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7919?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL21vZGVscy9uZXVyYWxfZmluZ2VycHJpbnQucHk=) | `100.00% <100.00%> (ø)` | |

... and [24 files with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7919/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

harshit5674: @wsad1 I request you to please review the PR.
harshit5674: @wsad1 @EdisonLeeeee I request you to please review the PR.

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index fb96d7d1d..5e939d3c3 100644
--- a/CHANGELOG.md
+++ b/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))
diff --git a/torch_geometric/nn/models/__init__.py b/torch_geometric/nn/models/__init__.py
index 9b716e21e..4fbbffb4b 100644
--- a/torch_geometric/nn/models/__init__.py
+++ b/torch_geometric/nn/models/__init__.py
@@ -24,6 +24,7 @@ from .mask_label import MaskLabel
 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 @@ __all__ = classes = [
     'GroupAddRev',
     'GNNFF',
     'PMLP',
+    'NeuralFingerprint',
 ]
diff --git a/torch_geometric/nn/models/neural_fingerprint.py b/torch_geometric/nn/models/neural_fingerprint.py
new file mode 100644
index 000000000..62fa898aa
--- /dev/null
+++ b/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})')


Test patch:
diff --git a/test/nn/models/test_neural_fingerprint.py b/test/nn/models/test_neural_fingerprint.py
new file mode 100644
index 000000000..f32d7e72c
--- /dev/null
+++ b/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)