Add LaplacianPE transform

spektral/transforms/__init__.py

@@ -4,6 +4,7 @@
 from .degree import Degree
 from .delaunay import Delaunay
 from .gcn_filter import GCNFilter
+from .laplacian_pe import LaplacianPE
 from .layer_preprocess import LayerPreprocess
 from .normalize_adj import NormalizeAdj
 from .normalize_one import NormalizeOne

spektral/transforms/laplacian_pe.py

@@ -0,0 +1,34 @@
+from scipy.sparse.linalg import eigsh
+import numpy as np
+
+from spektral.utils import normalized_laplacian
+
+
+class LaplacianPE:
+    r"""
+    Adds Laplacian positional encodings to the nodes.
+
+    The first `k` eigenvectors are computed and concatenated to the node features.
+    Each node will be extended with its corresponding entries in the first k
+    eigenvectors.
+    """
+    def __init__(self, k):
+        assert k > 0, "k must be greater than 0"
+        self.k = k
+
+    def __call__(self, graph):
+        if "a" not in graph:
+            raise ValueError("The graph must have an adjacency matrix")
+        assert (
+            self.k < graph.n_nodes
+        ), f"k = {self.k} must be smaller than graph.n_nodes = {graph.n_nodes}"
+
+        l = normalized_laplacian(graph.a)
+        _, eigvec = eigsh(l, k=self.k, which="SM")
+
+        if "x" not in graph:
+            graph.x = eigvec
+        else:
+            graph.x = np.concatenate((graph.x, eigvec), axis=-1)
+
+        return graph

tests/test_transforms/test_transforms.py

@@ -75,6 +75,17 @@ def test_gcn_filter():
     t(g)
 
 
+def test_laplacian_pe():
+    k = 3
+    t = tr.LaplacianPE(k)
+    assert callable(t)
+    g = Graph(x=x, a=a, e=e, y=y_gl)
+    t(g)
+    g = Graph(x=None, a=a, e=e, y=y_gl)
+    g_ = t(g)
+    assert g_.x.shape[-1] == k
+
+
 def test_layer_preprocess():
     from spektral.layers import GCNConv