[Refactor] Refactor Laplacian positional encoding transform (#4628)

* refactor the function and module of laplacian_pe

* add blank lines between design doc

* 1st approving review

* update the unittest

* fix lint issues

* fix trailing space

* update by mufei's comments and fix backend bugs

* del test file
Co-authored-by: rudongyu <ru_dongyu@outlook.com>
Co-authored-by: Rhett Ying <85214957+Rhett-Ying@users.noreply.github.com>

python/dgl/transforms/functional.py

@@ -3650,44 +3650,63 @@ def random_walk_pe(g, k, eweight_name=None):
 
     return PE
 
-def laplacian_pe(g, k):
+def laplacian_pe(g, k, padding=False, return_eigval=False):
     r"""Laplacian Positional Encoding, as introduced in
     `Benchmarking Graph Neural Networks
     <https://arxiv.org/abs/2003.00982>`__
 
     This function computes the laplacian positional encodings as the
-    k smallest non-trivial eigenvectors (k << n). k and n are the positional
-    encoding dimensions and the number of nodes in the given graph.
+    k smallest non-trivial eigenvectors.
 
     Parameters
     ----------
     g : DGLGraph
-        The input graph. Must be homogeneous.
+        The input graph. Must be homogeneous and bidirected.
     k : int
-        Number of smallest non-trivial eigenvectors to use for positional encoding
-        (smaller than the number of nodes).
+        Number of smallest non-trivial eigenvectors to use for positional encoding.
+    padding : bool, optional
+        If False, raise an exception when k>=n.
+        Otherwise, add zero paddings in the end of eigenvectors and 'nan' paddings
+        in the end of eigenvalues when k>=n.
+        Default: False.
+        n is the number of nodes in the given graph.
+    return_eigval : bool, optional
+        If True, return laplacian eigenvalues together with eigenvectors.
+        Otherwise, return laplacian eigenvectors only.
+        Default: False.
 
     Returns
     -------
-    Tensor
-        The laplacian positional encodings of shape :math:`(N, k)`, where :math:`N` is the
-        number of nodes in the input graph.
+    Tensor or (Tensor, Tensor)
+        Return the laplacian positional encodings of shape :math:`(N, k)`, where :math:`N` is the
+        number of nodes in the input graph, when :attr:`return_eigval` is False. The eigenvalues
+        of shape :math:`N` is additionally returned as the second element when :attr:`return_eigval`
+        is True.
 
     Example
     -------
     >>> import dgl
-    >>> g = dgl.rand_graph(6, 12)
+    >>> g = dgl.graph(([0,1,2,3,1,2,3,0], [1,2,3,0,0,1,2,3]))
     >>> dgl.laplacian_pe(g, 2)
-    tensor([[-0.8931, -0.7713],
-            [-0.0000,  0.6198],
-            [ 0.2704, -0.0138],
-            [-0.0000,  0.0554],
-            [ 0.3595, -0.0477],
-            [-0.0000,  0.1240]])
+    tensor([[ 7.0711e-01, -6.4921e-17],
+            [ 3.0483e-16, -7.0711e-01],
+            [-7.0711e-01, -2.4910e-16],
+            [ 9.9288e-17,  7.0711e-01]])
+    >>> dgl.laplacian_pe(g, 5, padding=True)
+    tensor([[ 7.0711e-01, -6.4921e-17,  5.0000e-01,  0.0000e+00,  0.0000e+00],
+            [ 3.0483e-16, -7.0711e-01, -5.0000e-01,  0.0000e+00,  0.0000e+00],
+            [-7.0711e-01, -2.4910e-16,  5.0000e-01,  0.0000e+00,  0.0000e+00],
+            [ 9.9288e-17,  7.0711e-01, -5.0000e-01,  0.0000e+00,  0.0000e+00]])
+    >>> dgl.laplacian_pe(g, 5, padding=True, return_eigval=True)
+    (tensor([[-7.0711e-01,  6.4921e-17, -5.0000e-01,  0.0000e+00,  0.0000e+00],
+             [-3.0483e-16,  7.0711e-01,  5.0000e-01,  0.0000e+00,  0.0000e+00],
+             [ 7.0711e-01,  2.4910e-16, -5.0000e-01,  0.0000e+00,  0.0000e+00],
+             [-9.9288e-17, -7.0711e-01,  5.0000e-01,  0.0000e+00,  0.0000e+00]]),
+     tensor([1., 1., 2., nan, nan]))
     """
     # check for the "k < n" constraint
     n = g.num_nodes()
-    if n <= k:
+    if not padding and n <= k:
         assert "the number of eigenvectors k must be smaller than the number of nodes n, " + \
             f"{k} and {n} detected."
 
@@ -3698,15 +3717,26 @@ def laplacian_pe(g, k):
 
     # select eigenvectors with smaller eigenvalues O(n + klogk)
     EigVal, EigVec = np.linalg.eig(L.toarray())
-    kpartition_indices = np.argpartition(EigVal, k+1)[:k+1]
+    max_freqs = min(n-1, k)
+    kpartition_indices = np.argpartition(EigVal, max_freqs)[:max_freqs+1]
     topk_eigvals = EigVal[kpartition_indices]
     topk_indices = kpartition_indices[topk_eigvals.argsort()][1:]
-    topk_EigVec = np.real(EigVec[:, topk_indices])
+    topk_EigVec = EigVec[:, topk_indices]
+    eigvals = F.tensor(EigVal[topk_indices], dtype=F.float32)
 
     # get random flip signs
-    rand_sign = 2 * (np.random.rand(k) > 0.5) - 1.
+    rand_sign = 2 * (np.random.rand(max_freqs) > 0.5) - 1.
     PE = F.astype(F.tensor(rand_sign * topk_EigVec), F.float32)
 
+    # add paddings
+    if n <= k:
+        temp_EigVec = F.zeros([n, k-n+1], dtype=F.float32, ctx=F.context(PE))
+        PE = F.cat([PE, temp_EigVec], dim=1)
+        temp_EigVal = F.tensor(np.full(k-n+1, np.nan), F.float32)
+        eigvals = F.cat([eigvals, temp_EigVal], dim=0)
+
+    if return_eigval:
+        return PE, eigvals
     return PE
 
 def to_half(g):

python/dgl/transforms/module.py

@@ -372,33 +372,70 @@ class LaplacianPE(BaseTransform):
     Parameters
     ----------
     k : int
-        Number of smallest non-trivial eigenvectors to use for positional encoding
-        (smaller than the number of nodes).
+        Number of smallest non-trivial eigenvectors to use for positional encoding.
     feat_name : str, optional
         Name to store the computed positional encodings in ndata.
+    eigval_name : str, optional
+        If None, store laplacian eigenvectors only.
+        Otherwise, it's the name to store corresponding laplacian eigenvalues in ndata.
+        Default: None.
+    padding : bool, optional
+        If False, raise an exception when k>=n.
+        Otherwise, add zero paddings in the end of eigenvectors and 'nan' paddings
+        in the end of eigenvalues when k>=n.
+        Default: False.
+        n is the number of nodes in the given graph.
 
     Example
     -------
-
     >>> import dgl
     >>> from dgl import LaplacianPE
-
-    >>> transform = LaplacianPE(k=3)
-    >>> g = dgl.rand_graph(5, 10)
-    >>> g = transform(g)
-    >>> print(g.ndata['PE'])
-    tensor([[ 0.0000, -0.3646,  0.3646],
-            [ 0.0000,  0.2825, -0.2825],
-            [ 1.0000, -0.6315,  0.6315],
-            [ 0.0000,  0.3739, -0.3739],
-            [ 0.0000, -0.1663,  0.1663]])
+    >>> transform1 = LaplacianPE(k=3)
+    >>> transform2 = LaplacianPE(k=5, padding=True)
+    >>> transform3 = LaplacianPE(k=5, feat_name='eigvec', eigval_name='eigval', padding=True)
+    >>> g = dgl.graph(([0,1,2,3,4,2,3,1,4,0], [2,3,1,4,0,0,1,2,3,4]))
+    >>> g1 = transform1(g)
+    >>> print(g1.ndata['PE'])
+    tensor([[ 0.6325,  0.1039,  0.3489],
+            [-0.5117,  0.2826,  0.6095],
+            [ 0.1954,  0.6254, -0.5923],
+            [-0.5117, -0.4508, -0.3938],
+            [ 0.1954, -0.5612,  0.0278]])
+    >>> g2 = transform2(g)
+    >>> print(g2.ndata['PE'])
+    tensor([[-0.6325, -0.1039,  0.3489, -0.2530,  0.0000],
+            [ 0.5117, -0.2826,  0.6095,  0.4731,  0.0000],
+            [-0.1954, -0.6254, -0.5923, -0.1361,  0.0000],
+            [ 0.5117,  0.4508, -0.3938, -0.6295,  0.0000],
+            [-0.1954,  0.5612,  0.0278,  0.5454,  0.0000]])
+    >>> g3 = transform3(g)
+    >>> print(g3.ndata['eigval'])
+    tensor([[0.6910, 0.6910, 1.8090, 1.8090,    nan],
+            [0.6910, 0.6910, 1.8090, 1.8090,    nan],
+            [0.6910, 0.6910, 1.8090, 1.8090,    nan],
+            [0.6910, 0.6910, 1.8090, 1.8090,    nan],
+            [0.6910, 0.6910, 1.8090, 1.8090,    nan]])
+    >>> print(g3.ndata['eigvec'])
+    tensor([[ 0.6325, -0.1039,  0.3489,  0.2530,  0.0000],
+            [-0.5117, -0.2826,  0.6095, -0.4731,  0.0000],
+            [ 0.1954, -0.6254, -0.5923,  0.1361,  0.0000],
+            [-0.5117,  0.4508, -0.3938,  0.6295,  0.0000],
+            [ 0.1954,  0.5612,  0.0278, -0.5454,  0.0000]])
     """
-    def __init__(self, k, feat_name='PE'):
+    def __init__(self, k, feat_name='PE', eigval_name=None, padding=False):
         self.k = k
         self.feat_name = feat_name
+        self.eigval_name = eigval_name
+        self.padding = padding
 
     def __call__(self, g):
-        PE = functional.laplacian_pe(g, k=self.k)
+        if self.eigval_name:
+            PE, eigval = functional.laplacian_pe(g, k=self.k, padding=self.padding,
+                                                 return_eigval=True)
+            eigval = F.repeat(F.reshape(eigval, [1,-1]), g.num_nodes(), dim=0)
+            g.ndata[self.eigval_name] = F.copy_to(eigval, g.device)
+        else:
+            PE = functional.laplacian_pe(g, k=self.k, padding=self.padding)
         g.ndata[self.feat_name] = F.copy_to(PE, g.device)
 
         return g

tests/compute/test_transform.py

@@ -2439,19 +2439,45 @@ def test_module_random_walk_pe(idtype):
 
 @parametrize_idtype
 def test_module_laplacian_pe(idtype):
-    transform = dgl.LaplacianPE(2, 'lappe')
-    g = dgl.graph(([2, 1, 0, 3, 1, 1],[3, 0, 1, 3, 3, 1]), idtype=idtype, device=F.ctx())
+    g = dgl.graph(([2, 1, 0, 3, 1, 1],[3, 1, 1, 2, 1, 0]), idtype=idtype, device=F.ctx())
+    tgt_eigval = F.copy_to(F.repeat(F.tensor([[1.1534e-17, 1.3333e+00, 2., np.nan, np.nan]]),
+        g.num_nodes(), dim=0), g.device)
+    tgt_pe = F.copy_to(F.tensor([[0.5, 0.86602539, 0., 0., 0.],
+        [0.86602539, 0.5, 0., 0., 0.],
+        [0., 0., 0.70710677, 0., 0.],
+        [0., 0., 0.70710677, 0., 0.]]), g.device)
+
+    # without padding (k<n)
+    transform = dgl.LaplacianPE(2, feat_name='lappe')
     new_g = transform(g)
-    tgt = F.copy_to(F.tensor([[ 0.24971116, 0.],
-        [ 0.11771496, 0.],
-        [ 0.83237050, 1.],
-        [ 0.48056933, 0.]]), g.device)
     # tensorflow has no abs() api
     if dgl.backend.backend_name == 'tensorflow':
-        assert F.allclose(new_g.ndata['lappe'].__abs__(), tgt)
+        assert F.allclose(new_g.ndata['lappe'].__abs__(), tgt_pe[:,:2])
     # pytorch & mxnet
     else:
-        assert F.allclose(new_g.ndata['lappe'].abs(), tgt)
+        assert F.allclose(new_g.ndata['lappe'].abs(), tgt_pe[:,:2])
+
+    # with padding (k>=n)
+    transform = dgl.LaplacianPE(5, feat_name='lappe', padding=True)
+    new_g = transform(g)
+    # tensorflow has no abs() api
+    if dgl.backend.backend_name == 'tensorflow':
+        assert F.allclose(new_g.ndata['lappe'].__abs__(), tgt_pe)
+    # pytorch & mxnet
+    else:
+        assert F.allclose(new_g.ndata['lappe'].abs(), tgt_pe)
+
+    # with eigenvalues
+    transform = dgl.LaplacianPE(5, feat_name='lappe', eigval_name='eigval', padding=True)
+    new_g = transform(g)
+    # tensorflow has no abs() api
+    if dgl.backend.backend_name == 'tensorflow':
+        assert F.allclose(new_g.ndata['eigval'][:,:3], tgt_eigval[:,:3])
+        assert F.allclose(new_g.ndata['lappe'].__abs__(), tgt_pe)
+    # pytorch & mxnet
+    else:
+        assert F.allclose(new_g.ndata['eigval'][:,:3], tgt_eigval[:,:3])
+        assert F.allclose(new_g.ndata['lappe'].abs(), tgt_pe)
 
 @unittest.skipIf(dgl.backend.backend_name != 'pytorch', reason='Only support PyTorch for now')
 @pytest.mark.parametrize('g', get_cases(['has_scalar_e_feature']))