Repository: pyg-team/pytorch_geometric
Pull request #9776 (2024-11-12 10:03:21)

Added functionality of `FaceToEdge` transform to work for 3D tetrahedral mesh elements
Now transforms 2D triangular elements/faces with shape [3,n] to  as well as 3D tetrahedral elements with shape [4,n] to edges of shape [2,n]. 

Including pytest case for face input with shape [4,n]


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index d491af50d..73c595492 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 3D tetrahedral mesh elements of shape `[4, num_faces]` in the `FaceToEdge` transformation ([#9776](https://github.com/pyg-team/pytorch_geometric/pull/9776))
 - Added the `use_pcst` option to `WebQSPDataset` ([#9722](https://github.com/pyg-team/pytorch_geometric/pull/9722))
 - Allowed users to pass `edge_weight` to `GraphUNet` models ([#9737](https://github.com/pyg-team/pytorch_geometric/pull/9737))
 - Consolidated `examples/ogbn_{papers_100m,products_gat,products_sage}.py` into `examples/ogbn_train.py` ([#9467](https://github.com/pyg-team/pytorch_geometric/pull/9467))
@@ -19,6 +20,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Fixed `utils.group_cat` concatenating dimension ([#9766](https://github.com/pyg-team/pytorch_geometric/pull/9766))
 - Fixed `WebQSDataset.process` raising exceptions ([#9665](https://github.com/pyg-team/pytorch_geometric/pull/9665))
 
 ### Removed
diff --git a/examples/multi_gpu/distributed_sampling.py b/examples/multi_gpu/distributed_sampling.py
index d72a8051f..805e6b305 100644
--- a/examples/multi_gpu/distributed_sampling.py
+++ b/examples/multi_gpu/distributed_sampling.py
@@ -1,4 +1,5 @@
 import os
+from math import ceil
 
 import torch
 import torch.distributed as dist
@@ -56,11 +57,11 @@ def run(rank, world_size, dataset):
 
     # Split indices into `world_size` many chunks:
     train_idx = data.train_mask.nonzero(as_tuple=False).view(-1)
-    train_idx = train_idx.split(train_idx.size(0) // world_size)[rank]
+    train_idx = train_idx.split(ceil(train_idx.size(0) / world_size))[rank]
     val_idx = data.val_mask.nonzero(as_tuple=False).view(-1)
-    val_idx = val_idx.split(val_idx.size(0) // world_size)[rank]
+    val_idx = val_idx.split(ceil(val_idx.size(0) / world_size))[rank]
     test_idx = data.val_mask.nonzero(as_tuple=False).view(-1)
-    test_idx = test_idx.split(test_idx.size(0) // world_size)[rank]
+    test_idx = test_idx.split(ceil(test_idx.size(0) / world_size))[rank]
 
     kwargs = dict(
         data=data,
diff --git a/examples/multi_gpu/distributed_sampling_multinode.py b/examples/multi_gpu/distributed_sampling_multinode.py
index 3de476d6c..b83131082 100644
--- a/examples/multi_gpu/distributed_sampling_multinode.py
+++ b/examples/multi_gpu/distributed_sampling_multinode.py
@@ -1,5 +1,6 @@
 import copy
 import os
+from math import ceil
 
 import torch
 import torch.distributed as dist
@@ -86,7 +87,7 @@ def run(world_size: int, rank: int, local_rank: int):
 
     # Split training indices into `world_size` many chunks:
     train_idx = data.train_mask.nonzero(as_tuple=False).view(-1)
-    train_idx = train_idx.split(train_idx.size(0) // world_size)[rank]
+    train_idx = train_idx.split(ceil(train_idx.size(0) / world_size))[rank]
 
     kwargs = dict(batch_size=1024, num_workers=4, persistent_workers=True)
     train_loader = NeighborLoader(
diff --git a/examples/multi_gpu/pcqm4m_ogb.py b/examples/multi_gpu/pcqm4m_ogb.py
index caf30e713..03566e84b 100644
--- a/examples/multi_gpu/pcqm4m_ogb.py
+++ b/examples/multi_gpu/pcqm4m_ogb.py
@@ -1,6 +1,7 @@
 # Code adapted from OGB.
 # https://github.com/snap-stanford/ogb/tree/master/examples/lsc/pcqm4m-v2
 import argparse
+import math
 import os
 
 import torch
@@ -433,7 +434,8 @@ def run(rank, dataset, args):
         train_idx = split_idx["train"]
 
     if num_devices > 1:
-        train_idx = train_idx.split(train_idx.size(0) // num_devices)[rank]
+        num_splits = math.ceil(train_idx.size(0) / num_devices)
+        train_idx = train_idx.split(num_splits)[rank]
 
     if args.train_subset:
         subset_ratio = 0.1
diff --git a/examples/quiver/multi_gpu_quiver.py b/examples/quiver/multi_gpu_quiver.py
index a87fadaff..2700c2114 100644
--- a/examples/quiver/multi_gpu_quiver.py
+++ b/examples/quiver/multi_gpu_quiver.py
@@ -1,6 +1,7 @@
 # This script shows how to use Quiver in an existing PyG example:
 # https://github.com/pyg-team/pytorch_geometric/blob/master/examples/multi_gpu/distributed_sampling.py
 import os
+from math import ceil
 
 import quiver
 import torch
@@ -69,7 +70,7 @@ def run(rank, world_size, dataset, quiver_feature, quiver_sampler):
 
     data = dataset[0]
     train_idx = data.train_mask.nonzero(as_tuple=False).view(-1)
-    train_idx = train_idx.split(train_idx.size(0) // world_size)[rank]
+    train_idx = train_idx.split(ceil(train_idx.size(0) / world_size))[rank]
 
     train_loader = torch.utils.data.DataLoader(train_idx, batch_size=1024,
                                                shuffle=True, num_workers=0)
diff --git a/torch_geometric/data/data.py b/torch_geometric/data/data.py
index 3a1abb5e7..690b35698 100644
--- a/torch_geometric/data/data.py
+++ b/torch_geometric/data/data.py
@@ -659,7 +659,13 @@ class Data(BaseData, FeatureStore, GraphStore):
                 return value.get_dim_size()
             return int(value.max()) + 1
         elif 'index' in key or key == 'face':
-            return self.num_nodes
+            num_nodes = self.num_nodes
+            if num_nodes is None:
+                raise RuntimeError(f"Unable to infer 'num_nodes' from the "
+                                   f"attribute '{key}'. Please explicitly set "
+                                   f"'num_nodes' as an attribute of 'data' to "
+                                   f"prevent this error")
+            return num_nodes
         else:
             return 0
 
diff --git a/torch_geometric/transforms/face_to_edge.py b/torch_geometric/transforms/face_to_edge.py
index cd0aceb2f..f2234328a 100644
--- a/torch_geometric/transforms/face_to_edge.py
+++ b/torch_geometric/transforms/face_to_edge.py
@@ -8,8 +8,15 @@ from torch_geometric.utils import to_undirected
 
 @functional_transform('face_to_edge')
 class FaceToEdge(BaseTransform):
-    r"""Converts mesh faces :obj:`[3, num_faces]` to edge indices
-    :obj:`[2, num_edges]` (functional name: :obj:`face_to_edge`).
+    r"""Converts mesh faces of shape :obj:`[3, num_faces]` or
+    :obj:`[4, num_faces]` to edge indices of shape :obj:`[2, num_edges]`
+    (functional name: :obj:`face_to_edge`).
+
+    This transform supports both 2D triangular faces, represented by a
+    tensor of shape :obj:`[3, num_faces]`, and 3D tetrahedral mesh faces,
+    represented by a tensor of shape :obj:`[4, num_faces]`. It will convert
+    these faces into edge indices, where each edge is defined by the indices
+    of its two endpoints.
 
     Args:
         remove_faces (bool, optional): If set to :obj:`False`, the face tensor
@@ -22,7 +29,29 @@ class FaceToEdge(BaseTransform):
         if hasattr(data, 'face'):
             assert data.face is not None
             face = data.face
-            edge_index = torch.cat([face[:2], face[1:], face[::2]], dim=1)
+
+            if face.size(0) not in [3, 4]:
+                raise RuntimeError(f"Expected 'face' tensor with shape "
+                                   f"[3, num_faces] or [4, num_faces] "
+                                   f"(got {list(face.size())})")
+
+            if face.size()[0] == 3:
+                edge_index = torch.cat([
+                    face[:2],
+                    face[1:],
+                    face[::2],
+                ], dim=1)
+            else:
+                assert face.size()[0] == 4
+                edge_index = torch.cat([
+                    face[:2],
+                    face[1:3],
+                    face[2:4],
+                    face[::2],
+                    face[1::2],
+                    face[::3],
+                ], dim=1)
+
             edge_index = to_undirected(edge_index, num_nodes=data.num_nodes)
 
             data.edge_index = edge_index
diff --git a/torch_geometric/utils/_scatter.py b/torch_geometric/utils/_scatter.py
index cc081e922..2c867fd1a 100644
--- a/torch_geometric/utils/_scatter.py
+++ b/torch_geometric/utils/_scatter.py
@@ -351,5 +351,5 @@ def group_cat(
     """
     assert len(tensors) == len(indices)
     index, perm = torch.cat(indices).sort(stable=True)
-    out = torch.cat(tensors, dim=0)[perm]
+    out = torch.cat(tensors, dim=dim).index_select(dim, perm)
     return (out, index) if return_index else out
diff --git a/torch_geometric/visualization/graph.py b/torch_geometric/visualization/graph.py
index c48d54fbb..beeceb1a2 100644
--- a/torch_geometric/visualization/graph.py
+++ b/torch_geometric/visualization/graph.py
@@ -132,7 +132,7 @@ def _visualize_graph_via_networkx(
             xy=pos[src],
             xytext=pos[dst],
             arrowprops=dict(
-                arrowstyle="->",
+                arrowstyle="<-",
                 alpha=data['alpha'],
                 shrinkA=sqrt(node_size) / 2.0,
                 shrinkB=sqrt(node_size) / 2.0,


Test patch:
diff --git a/test/data/test_data.py b/test/data/test_data.py
index 3b94d4169..9c95fcb58 100644
--- a/test/data/test_data.py
+++ b/test/data/test_data.py
@@ -582,3 +582,14 @@ def test_data_time_handling(num_nodes, num_edges):
 
     out = out.sort_by_time()
     assert torch.equal(out.time, data.time.repeat_interleave(2))
+
+
+def test_data_inc():
+    data = Data(edge_index=torch.tensor([[0, 1], [1, 0]]))
+    with pytest.warns(UserWarning, match="Unable to accurately infer"):
+        assert data.__inc__('edge_index', data.edge_index) == 2
+
+        data = Data(index=torch.empty(2, 0, dtype=torch.long))
+    with pytest.raises(RuntimeError, match="Unable to infer"):
+        with pytest.warns(UserWarning, match="Unable to accurately infer"):
+            data.__inc__('index', data.edge_index)
diff --git a/test/transforms/test_face_to_edge.py b/test/transforms/test_face_to_edge.py
index d8add6776..d0338510c 100644
--- a/test/transforms/test_face_to_edge.py
+++ b/test/transforms/test_face_to_edge.py
@@ -4,7 +4,7 @@ from torch_geometric.data import Data
 from torch_geometric.transforms import FaceToEdge
 
 
-def test_face_to_edge():
+def test_2d_face_to_edge() -> None:
     transform = FaceToEdge()
     assert str(transform) == 'FaceToEdge()'
 
@@ -13,6 +13,23 @@ def test_face_to_edge():
 
     data = transform(data)
     assert len(data) == 2
-    assert data.edge_index.tolist() == [[0, 0, 0, 1, 1, 1, 2, 2, 3, 3],
-                                        [1, 2, 3, 0, 2, 3, 0, 1, 0, 1]]
+    assert data.edge_index.tolist() == [
+        [0, 0, 0, 1, 1, 1, 2, 2, 3, 3],
+        [1, 2, 3, 0, 2, 3, 0, 1, 0, 1],
+    ]
     assert data.num_nodes == 4
+
+
+def test_3d_face_to_edge() -> None:
+    transform = FaceToEdge()
+    assert str(transform) == 'FaceToEdge()'
+
+    face = torch.tensor([[0, 1, 2, 3], [1, 2, 3, 4]]).t()
+    data = Data(face=face, num_nodes=5)
+
+    data = transform(data)
+    assert data.edge_index.tolist() == [
+        [0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4],
+        [1, 2, 3, 0, 2, 3, 4, 0, 1, 3, 4, 0, 1, 2, 4, 1, 2, 3],
+    ]
+    assert data.num_nodes == 5