Merge branch 'PatRyg99-pr/voronoi-lifting'

Author: gbg141

configs/datasets/random_pointcloud.yaml

@@ -0,0 +1,15 @@
+data_domain: pointcloud
+data_type: toy_dataset
+data_name: random_pointcloud
+data_dir: datasets/${data_domain}/${data_type}/${data_name}
+
+# Dataset parameters
+pos_to_x: True
+num_features: 3
+num_points: 20
+dim: 3
+num_classes: 2
+task: classification
+loss_type: cross_entropy
+monitor_metric: accuracy
+task_level: node

configs/datasets/stanford_bunny.yaml

@@ -0,0 +1,13 @@
+data_domain: pointcloud
+data_type: toy_dataset
+data_name: stanford_bunny
+data_dir: datasets/${data_domain}/${data_type}/${data_name}
+
+# Dataset parameters
+pos_to_x: True
+num_features: 3
+num_classes: 1
+task: regression
+loss_type: mse
+monitor_metric: mae
+task_level: graph

configs/transforms/liftings/pointcloud2hypergraph/voronoi_lifting.yaml

@@ -0,0 +1,4 @@
+transform_type: 'lifting'
+transform_name: "VoronoiLifting"
+support_ratio: 0.005
+feature_lifting: ProjectionSum

modules/data/load/loaders.py

@@ -24,6 +24,7 @@
     load_manual_mol,
     load_manual_points,
     load_point_cloud,
+    load_pointcloud_dataset,
     load_random_points,
     load_simplicial_dataset,
 )
@@ -278,7 +279,6 @@ def load(self) -> torch_geometric.data.Dataset:
         ----------
         None
 
-
         Returns
         -------
         torch_geometric.data.Dataset
@@ -316,3 +316,8 @@ def load(self) -> torch_geometric.data.Dataset:
             )
 
         return CustomDataset([data], self.data_dir)
+        self.data_dir = os.path.join(root_folder, self.parameters["data_dir"])
+
+        return CustomDataset(
+            [load_pointcloud_dataset(self.parameters)], self.data_dir
+        )

modules/data/utils/utils.py

@@ -405,7 +405,6 @@ def load_manual_graph():
     )
 
 
-
 def load_k4_graph() -> torch_geometric.data.Data:
     """K_4 is a complete graph with 4 vertices."""
     vertices = [i for i in range(4)]
@@ -475,9 +474,13 @@ def load_8_vertex_cubic_graphs() -> list[torch_geometric.data.Data]:
         n = 8 if i < 5 else 10
         vertices = [i for i in range(n)]
         x = (
-            torch.tensor([1, 5, 10, 50, 100, 500, 1000, 5000]).unsqueeze(1).float()
+            torch.tensor([1, 5, 10, 50, 100, 500, 1000, 5000])
+            .unsqueeze(1)
+            .float()
             if i < 5
-            else torch.tensor([1, 5, 10, 50, 100, 500, 1000, 5000, 10000, 50000])
+            else torch.tensor(
+                [1, 5, 10, 50, 100, 500, 1000, 5000, 10000, 50000]
+            )
             .unsqueeze(1)
             .float()
         )
@@ -494,11 +497,14 @@ def load_8_vertex_cubic_graphs() -> list[torch_geometric.data.Data]:
         G.to_undirected()
         edge_list = torch.Tensor(list(G.edges())).T.long()
 
-        data = torch_geometric.data.Data(x=x, edge_index=edge_list, num_nodes=n, y=y)
+        data = torch_geometric.data.Data(
+            x=x, edge_index=edge_list, num_nodes=n, y=y
+        )
 
         list_data.append(data)
     return list_data
 
+
 def load_manual_mol():
     """Create a manual graph for testing the ring implementation.
     Actually is the 471 molecule of QM9 dataset."""
@@ -620,7 +626,6 @@ def load_manual_mol():
     )
 
 
-
 def get_Planetoid_pyg(cfg):
     r"""Loads Planetoid graph datasets from torch_geometric.
 
@@ -822,3 +827,71 @@ def load_manual_points():
     x = torch.ones_like(pos, dtype=torch.float)
     y = torch.randint(0, 2, (pos.shape[0],), dtype=torch.float)
     return torch_geometric.data.Data(x=x, y=y, pos=pos, complex_dim=0)
+
+
+def load_pointcloud_dataset(cfg):
+    r"""Loads point cloud datasets.
+
+    Parameters
+    ----------
+    cfg : DictConfig
+        Configuration parameters.
+
+    Returns
+    -------
+    torch_geometric.data.Data
+        Point cloud dataset.
+    """
+    # Define the path to the data directory
+    root_folder = rootutils.find_root()
+    data_dir = osp.join(root_folder, cfg["data_dir"])
+
+    if cfg["data_name"] == "random_pointcloud":
+        num_points, dim = cfg["num_points"], cfg["dim"]
+        pos = torch.rand((num_points, dim))
+    elif cfg["data_name"] == "stanford_bunny":
+        pos = fetch_bunny(
+            file_path=osp.join(data_dir, "stanford_bunny.npy"),
+            accept_license=False,
+        )
+        num_points = len(pos)
+        pos = torch.tensor(pos)
+
+    if cfg.pos_to_x:
+        return torch_geometric.data.Data(
+            x=pos, pos=pos, num_nodes=num_points, num_features=pos.size(1)
+        )
+
+    return torch_geometric.data.Data(
+        pos=pos, num_nodes=num_points, num_features=0
+    )
+
+
+def load_manual_pointcloud(pos_to_x: bool = False):
+    """Create a manual pointcloud for testing purposes."""
+    # Define the positions
+    pos = torch.tensor(
+        [
+            [0, 0, 0],
+            [0, 0, 1],
+            [0, 1, 0],
+            [10, 0, 0],
+            [10, 0, 1],
+            [10, 1, 0],
+            [10, 1, 1],
+            [20, 0, 0],
+            [20, 0, 1],
+            [20, 1, 0],
+            [20, 1, 1],
+            [30, 0, 0],
+        ]
+    ).float()
+
+    if pos_to_x:
+        return torch_geometric.data.Data(
+            x=pos, pos=pos, num_nodes=pos.size(0), num_features=pos.size(1)
+        )
+
+    return torch_geometric.data.Data(
+        pos=pos, num_nodes=pos.size(0), num_features=0
+    )

modules/transforms/data_transform.py

@@ -44,6 +44,9 @@
 from modules.transforms.liftings.graph2simplicial.vietoris_rips_lifting import (
     SimplicialVietorisRipsLifting,
 )
+from modules.transforms.liftings.pointcloud2hypergraph.voronoi_lifting import (
+    VoronoiLifting,
+)
 from modules.transforms.liftings.pointcloud2simplicial.alpha_complex_lifting import (
     AlphaComplexLifting,
 )
@@ -72,6 +75,8 @@
     "AlphaComplexLifting": AlphaComplexLifting,
     # Point-cloud -> Simplicial Complex
     "DelaunayLifting": DelaunayLifting,
+    # Pointcloud -> Hypergraph
+    "VoronoiLifting": VoronoiLifting,
     # Feature Liftings
     "ProjectionSum": ProjectionSum,
     # Data Manipulations

modules/transforms/liftings/pointcloud2hypergraph/voronoi_lifting.py

@@ -0,0 +1,54 @@
+import torch
+import torch_geometric
+from torch_cluster import fps, knn
+
+from modules.transforms.liftings.pointcloud2hypergraph.base import (
+    PointCloud2HypergraphLifting,
+)
+
+
+class VoronoiLifting(PointCloud2HypergraphLifting):
+    r"""Lifts pointcloud to Farthest-point Voronoi graph.
+
+    Parameters
+    ----------
+    support_ratio : float
+        Ratio of points to sample with FPS to form voronoi support set.
+    **kwargs : optional
+        Additional arguments for the class.
+    """
+
+    def __init__(self, support_ratio: float, **kwargs):
+        super().__init__(**kwargs)
+        self.support_ratio = support_ratio
+
+    def lift_topology(self, data: torch_geometric.data.Data) -> dict:
+        r"""Lifts pointcloud to voronoi graph induced by Farthest Point Sampling (FPS) support set.
+
+        Parameters
+        ----------
+        data : torch_geometric.data.Data
+            The input data to be lifted.
+
+        Returns
+        -------
+        dict
+            The lifted topology.
+        """
+
+        # Sample FPS induced Voronoi graph
+        support_idcs = fps(data.pos, ratio=self.support_ratio)
+        target_idcs, source_idcs = knn(data.pos[support_idcs], data.pos, k=1)
+
+        # Construct incidence matrix
+        incidence_matrix = torch.sparse_coo_tensor(
+            torch.stack((target_idcs, source_idcs)),
+            torch.ones(source_idcs.numel()),
+            size=(data.num_nodes, support_idcs.numel()),
+        )
+
+        return {
+            "incidence_hyperedges": incidence_matrix,
+            "num_hyperedges": incidence_matrix.size(1),
+            "x_0": data.x,
+        }

modules/utils/utils.py

@@ -10,6 +10,7 @@
 import rootutils
 import torch
 import torch_geometric
+from matplotlib import cm
 from matplotlib.patches import Polygon
 
 plt.rcParams["text.usetex"] = bool(shutil.which("latex"))
@@ -148,19 +149,22 @@ def describe_data(dataset: torch_geometric.data.Dataset, idx_sample: int = 0):
         if hasattr(data, "num_nodes"):
             complex_dim.append(data.num_nodes)
             features_dim.append(data.num_node_features)
-        elif hasattr(data, "x"):
+        elif hasattr(data, "x") and data.x is not None:
             complex_dim.append(data.x.shape[0])
             features_dim.append(data.x.shape[1])
         else:
             raise ValueError(
                 "Data object does not contain any vertices/points."
             )
-        if hasattr(data, "num_edges") and hasattr(data, "num_edge_features"):
+
+        if hasattr(data, "num_edges") and data.num_edges > 0:
             complex_dim.append(data.num_edges)
             features_dim.append(data.num_edge_features)
-        elif hasattr(data, "edge_index") and (data.edge_index is not None):
+
+        elif hasattr(data, "edge_index") and data.edge_index is not None:
             complex_dim.append(data.edge_index.shape[1])
             features_dim.append(data.edge_attr.shape[1])
+
     # Check if the data object contains hyperedges
     hyperedges = False
     if hasattr(data, "x_hyperedges"):
@@ -201,6 +205,7 @@ def describe_data(dataset: torch_geometric.data.Dataset, idx_sample: int = 0):
                 hasattr(data, "edge_index")
                 and hasattr(data, "x")
                 and (data.edge_index is not None)
+                and data.x is not None
             ):
                 connected_nodes = torch.unique(data.edge_index)
                 isolated_nodes = []
@@ -587,3 +592,22 @@ def describe_hypergraph(data: torch_geometric.data.Data):
         if he_idx >= 10:
             print("...")
             break
+
+
+def plot_pointcloud_voronoi(
+    dataset, idx_sample: int = 0, azim: float = 180, roll: float = -90
+):
+    fig = plt.figure(figsize=(8, 8))
+    ax = fig.add_subplot(111, projection="3d")
+
+    data = dataset.get(idx_sample % len(dataset))
+    points = data.pos
+
+    if hasattr(data, "incidence_hyperedges"):
+        color = np.array(data.incidence_hyperedges.coalesce().indices()[1])
+    else:
+        color = np.ones(len(points))
+
+    ax.scatter(*points.T, s=1.0, c=color, cmap=cm.flag)
+    ax.view_init(elev=10.0, azim=azim, roll=roll)
+    plt.show()

test/transforms/liftings/pointcloud2hypergraph/test_voronoi_lifting.py

@@ -0,0 +1,37 @@
+"""Test the message passing module."""
+
+import torch
+
+from modules.data.utils.utils import load_manual_pointcloud
+from modules.transforms.liftings.pointcloud2hypergraph.voronoi_lifting import (
+    VoronoiLifting,
+)
+
+
+class TestVoronoiLifting:
+    """Test the SimplicialCliqueLifting class."""
+
+    def setup_method(self):
+        # Load the graph
+        self.data = load_manual_pointcloud(pos_to_x=True)
+
+        # Initialise the VoronoiLifting class
+        self.lifting = VoronoiLifting(support_ratio=0.26)
+
+    def test_lift_topology(self):
+        """Test the lift_topology method."""
+
+        # Test the lift_topology method
+        lifted_data = self.lifting.forward(self.data.clone())
+
+        expected_cluster_sizes = torch.tensor([1, 3, 4, 4])
+        cluster_sizes = torch.sort(
+            torch.unique(
+                lifted_data.incidence_hyperedges.coalesce().indices()[1],
+                return_counts=True,
+            )[1]
+        )[0]
+
+        assert (
+            expected_cluster_sizes == cluster_sizes
+        ).all(), "Something is wrong with edge_index (mismatched cluster sizes)."