Repository: pyg-team/pytorch_geometric
Pull request #6790 (2023-02-24 21:57:37)

`get_mesh_laplacian` fix and `normalization_type` customization
Fixed get_mesh_laplacian, and user can select its normalization type ([#6037](https://github.com/pyg-team/pytorch_geometric/issues/6037))

codecov[bot]: # [Codecov](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6790?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#6790](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6790?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (78e0751) into [master](https://codecov.io/gh/pyg-team/pytorch_geometric/commit/eac6cf1f487ccf5eb11fe52b0c7cfcd0c15aad38?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (eac6cf1) will **decrease** coverage by `0.01%`.
> The diff coverage is `90.24%`.

```diff
@@            Coverage Diff             @@
##           master    #6790      +/-   ##
==========================================
- Coverage   91.37%   91.36%   -0.01%     
==========================================
  Files         425      425              
  Lines       22950    22971      +21     
==========================================
+ Hits        20970    20987      +17     
- Misses       1980     1984       +4     
```


| [Impacted Files](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6790?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/utils/get\_mesh\_laplacian.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6790?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL2dldF9tZXNoX2xhcGxhY2lhbi5weQ==) | `92.30% <90.24%> (-7.70%)` | :arrow_down: |

: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)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 49e99894d..5428b2ee3 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
 
+- Support `normalization` customization in `get_mesh_laplacian` ([#6037](https://github.com/pyg-team/pytorch_geometric/issues/6037))
 - Added the `TemporalEncoding` module ([#6785](https://github.com/pyg-team/pytorch_geometric/pull/6785))
 - Added CPU-optimized `spmm_reduce` functionality via CSR format ([#6699](https://github.com/pyg-team/pytorch_geometric/pull/6699), [#6759](https://github.com/pyg-team/pytorch_geometric/pull/6759))
 - Added support for the revised version of the `MD17` dataset ([#6734](https://github.com/pyg-team/pytorch_geometric/pull/6734))
@@ -91,6 +92,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Fixed a bug in `get_mesh_laplacian` ([#6037](https://github.com/pyg-team/pytorch_geometric/issues/6037))
 - Fixed a bug in which masks were not properly masked in `GNNExplainer` on link prediction tasks ([#6787](https://github.com/pyg-team/pytorch_geometric/pull/6787))
 - Allow the usage of `ChebConv` within `GNNExplainer` ([#6778](https://github.com/pyg-team/pytorch_geometric/pull/6778))
 - Allow setting the `EdgeStorage.num_edges` property ([#6710](https://github.com/pyg-team/pytorch_geometric/pull/6710))
diff --git a/torch_geometric/utils/get_mesh_laplacian.py b/torch_geometric/utils/get_mesh_laplacian.py
index fef86a714..c739d715f 100644
--- a/torch_geometric/utils/get_mesh_laplacian.py
+++ b/torch_geometric/utils/get_mesh_laplacian.py
@@ -1,74 +1,109 @@
-from typing import Tuple
+from typing import Optional, Tuple
 
 import torch
 from torch import Tensor
 
-from torch_geometric.utils import add_self_loops, coalesce, scatter
+from torch_geometric.utils import add_self_loops, scatter, to_undirected
 
 
-def get_mesh_laplacian(pos: Tensor, face: Tensor) -> Tuple[Tensor, Tensor]:
+def get_mesh_laplacian(
+    pos: Tensor,
+    face: Tensor,
+    normalization: Optional[str] = None,
+) -> Tuple[Tensor, Tensor]:
     r"""Computes the mesh Laplacian of a mesh given by :obj:`pos` and
-    :obj:`face`.
-    It is computed as
+    :obj:`face`. Computation is based on the cotangent matrix defined as
 
     .. math::
-        \mathbf{L}_{ij} = \begin{cases}
-            \frac{\cot \angle_{ikj} + \cot \angle_{ilj}}{2 a_{ij}} &
-            \mbox{if } i, j \mbox{ is an edge} \\
-            \sum_{j \in N(i)}{L_{ij}} &
-            \mbox{if } i \mbox{ is in the diagonal} \\
-            0 \mbox{ otherwise}
+        \mathbf{C}_{ij} = \begin{cases}
+            \frac{\cot \angle_{ikj}~+\cot \angle_{ilj}}{2} &
+            \text{if } i, j \text{ is an edge} \\
+            -\sum_{j \in N(i)}{C_{ij}} &
+            \text{if } i \text{ is in the diagonal} \\
+            0 & \text{otherwise}
       \end{cases}
 
-    where :math:`a_{ij}` is the local area element, *i.e.* one-third of the
-    neighboring triangle's area.
+    Normalization depends on the mass matrix defined as
+
+    .. math::
+        \mathbf{M}_{ij} = \begin{cases}
+            a(i) & \text{if } i \text{ is in the diagonal} \\
+            0 & \text{otherwise}
+      \end{cases}
+
+    where :math:`a(i)` is obtained by joining the barycenters of the
+    triangles around vertex :math:`i`.
 
     Args:
         pos (Tensor): The node positions.
         face (LongTensor): The face indices.
+        normalization (str, optional): The normalization scheme for the mesh
+            Laplacian (default: :obj:`None`):
+
+            1. :obj:`None`: No normalization
+            :math:`\mathbf{L} = \mathbf{C}`
+
+            2. :obj:`"sym"`: Symmetric normalization
+            :math:`\mathbf{L} = \mathbf{M}^{-1/2} \mathbf{C}\mathbf{M}^{-1/2}`
+
+            3. :obj:`"rw"`: Row-wise normalization
+            :math:`\mathbf{L} = \mathbf{M}^{-1} \mathbf{C}`
     """
     assert pos.size(1) == 3 and face.size(0) == 3
 
     num_nodes = pos.shape[0]
 
-    def add_angles(left, centre, right):
+    def get_cots(left, centre, right):
         left_pos, central_pos, right_pos = pos[left], pos[centre], pos[right]
         left_vec = left_pos - central_pos
         right_vec = right_pos - central_pos
         dot = torch.einsum('ij, ij -> i', left_vec, right_vec)
         cross = torch.norm(torch.cross(left_vec, right_vec, dim=1), dim=1)
         cot = dot / cross  # cot = cos / sin
-        area = cross / 6.0  # one-third of a triangle's area is cross / 6.0
-        return cot / 2.0, area / 2.0
-
-    # For each triangle face, add all three angles:
-    cot_201, area_201 = add_angles(face[2], face[0], face[1])
-    cot_012, area_012 = add_angles(face[0], face[1], face[2])
-    cot_120, area_120 = add_angles(face[1], face[2], face[0])
-
-    cot_weight = torch.cat(
-        [cot_201, cot_201, cot_012, cot_012, cot_120, cot_120])
-    area_weight = torch.cat(
-        [area_201, area_201, area_012, area_012, area_120, area_120])
-
-    edge_index = torch.cat([
-        torch.stack([face[2], face[1]], dim=0),
-        torch.stack([face[1], face[2]], dim=0),
-        torch.stack([face[0], face[2]], dim=0),
-        torch.stack([face[2], face[0]], dim=0),
-        torch.stack([face[1], face[0]], dim=0),
-        torch.stack([face[0], face[1]], dim=0)
-    ], dim=1)
-
-    edge_index, weight = coalesce(edge_index, [cot_weight, area_weight])
-    cot_weight, area_weight = weight
+        return cot / 2.0  # by definition
+
+    # For each triangle face, get all three cotangents:
+    cot_021 = get_cots(face[0], face[2], face[1])
+    cot_102 = get_cots(face[1], face[0], face[2])
+    cot_012 = get_cots(face[0], face[1], face[2])
+    cot_weight = torch.cat([cot_021, cot_102, cot_012])
+
+    # Face to edge:
+    cot_index = torch.cat([face[:2], face[1:], face[::2]], dim=1)
+    cot_index, cot_weight = to_undirected(cot_index, cot_weight)
 
     # Compute the diagonal part:
-    row, col = edge_index
-    cot_deg = scatter(cot_weight, row, 0, num_nodes, reduce='sum')
-    area_deg = scatter(area_weight, row, 0, num_nodes, reduce='sum')
-    deg = cot_deg / area_deg
-    edge_index, _ = add_self_loops(edge_index, num_nodes=num_nodes)
-    edge_weight = torch.cat([-cot_weight, deg], dim=0)
+    cot_deg = scatter(cot_weight, cot_index[0], 0, num_nodes, reduce='sum')
+    edge_index, _ = add_self_loops(cot_index, num_nodes=num_nodes)
+    edge_weight = torch.cat([cot_weight, -cot_deg], dim=0)
+
+    if normalization is not None:
+
+        def get_areas(left, centre, right):
+            central_pos = pos[centre]
+            left_vec = pos[left] - central_pos
+            right_vec = pos[right] - central_pos
+            cross = torch.norm(torch.cross(left_vec, right_vec, dim=1), dim=1)
+            area = cross / 6.0  # one-third of a triangle's area is cross / 6.0
+            return area / 2.0  # since each corresponding area is counted twice
+
+        # Like before, but here we only need the diagonal (the mass matrix):
+        area_021 = get_areas(face[0], face[2], face[1])
+        area_102 = get_areas(face[1], face[0], face[2])
+        area_012 = get_areas(face[0], face[1], face[2])
+        area_weight = torch.cat([area_021, area_102, area_012])
+        area_index = torch.cat([face[:2], face[1:], face[::2]], dim=1)
+        area_index, area_weight = to_undirected(area_index, area_weight)
+        area_deg = scatter(area_weight, area_index[0], 0, num_nodes, 'sum')
+
+        if normalization == 'sym':
+            area_deg_inv_sqrt = area_deg.pow_(-0.5)
+            area_deg_inv_sqrt[area_deg_inv_sqrt == float('inf')] = 0.0,
+            edge_weight = (area_deg_inv_sqrt[edge_index[0]] * edge_weight *
+                           area_deg_inv_sqrt[edge_index[1]])
+        elif normalization == 'rw':
+            area_deg_inv = 1.0 / area_deg
+            area_deg_inv[area_deg_inv == float('inf')] = 0.0
+            edge_weight = area_deg_inv[edge_index[0]] * edge_weight
 
     return edge_index, edge_weight


Test patch:
diff --git a/test/utils/test_get_mesh_laplacian.py b/test/utils/test_get_mesh_laplacian.py
index f438a26c3..04a6bdb83 100644
--- a/test/utils/test_get_mesh_laplacian.py
+++ b/test/utils/test_get_mesh_laplacian.py
@@ -30,7 +30,8 @@ def test_get_mesh_laplacian_of_cube():
         [1, 2, 6],
     ])
 
-    edge_index, edge_weight = get_mesh_laplacian(pos, face.t())
+    edge_index, edge_weight = get_mesh_laplacian(pos, face.t(),
+                                                 normalization='rw')
 
     assert edge_index.tolist() == [
         [
@@ -46,10 +47,10 @@ def test_get_mesh_laplacian_of_cube():
     assert torch.allclose(
         edge_weight,
         torch.tensor([
-            -1.0, -0.0, -1.0, -1.0, -1.0, -1.0, -0.0, -1.0, -0.0, -0.0, -1.0,
-            -1.0, -1.0, -1.0, -1.0, -0.0, -0.0, -1.0, -1.0, -0.0, -0.0, -1.0,
-            -1.0, -1.0, -1.0, -1.0, -0.0, -0.0, -1.0, -0.0, -1.0, -1.0, -1.0,
-            -1.0, -0.0, -1.0, 1.125, 0.9, 1.125, 0.9, 0.9, 1.125, 0.9, 1.125
+            0.375, 0.0, 0.375, 0.375, 0.3, 0.3, 0.0, 0.3, 0.0, 0.0, 0.375,
+            0.375, 0.375, 0.3, 0.3, 0.0, 0.0, 0.3, 0.3, 0.0, 0.0, 0.3, 0.3,
+            0.375, 0.375, 0.375, 0.0, 0.0, 0.3, 0.0, 0.3, 0.3, 0.375, 0.375,
+            0.0, 0.375, -1.125, -0.9, -1.125, -0.9, -0.9, -1.125, -0.9, -1.125
         ]))
 
 
@@ -74,7 +75,8 @@ def test_get_mesh_laplacian_of_irregular_triangular_prism():
         [2, 5, 3],
     ])
 
-    edge_index, edge_weight = get_mesh_laplacian(pos, face.t())
+    edge_index, edge_weight = get_mesh_laplacian(pos, face.t(),
+                                                 normalization='rw')
 
     assert edge_index.tolist() == [
         [
@@ -90,9 +92,10 @@ def test_get_mesh_laplacian_of_irregular_triangular_prism():
     assert torch.allclose(
         edge_weight,
         torch.tensor([
-            -0.938834, -1.451131, -0.490290, -0.000000, -0.938834, -0.378790,
-            -0.577017, -0.077878, -1.451131, -0.378790, -0.163153, -0.344203,
-            -0.490290, -0.163153, -1.421842, -2.387739, -0.000000, -0.577017,
-            -1.421842, -2.550610, -0.077878, -0.344203, -2.387739, -2.550610,
-            0.298518, 0.183356, 0.233502, 0.761257, 0.688181, 0.768849
+            0.09730332, 0.15039921, 0.05081503, 0.00000000, 0.08726977,
+            0.03521059, 0.05363689, 0.00723919, 0.14497279, 0.03784235,
+            0.01629947, 0.03438699, 0.08362866, 0.02782887, 0.24252312,
+            0.40727590, 0.00000000, 0.08728313, 0.21507657, 0.38582093,
+            0.01117009, 0.04936920, 0.34247482, 0.36583540, -0.29851755,
+            -0.18335645, -0.23350160, -0.76125660, -0.68818060, -0.76884955
         ]))