Merge branch 'martin-carrasco-feat_nc_rg'

Author: gbg141

configs/transforms/feature_liftings/element_wise_mean.yaml

@@ -0,0 +1,2 @@
+transform_name: "ElementwiseMean"
+transform_type: null

configs/transforms/liftings/graph2simplicial/neighborhood_complex_lifting.yaml

@@ -0,0 +1,6 @@
+transform_type: 'lifting'
+transform_name: "NeighborhoodComplexLifting"
+preserve_edge_attr: False
+signed: True
+feature_lifting: ElementwiseMean
+complex_dim: 5

modules/transforms/data_transform.py

@@ -7,7 +7,10 @@
     NodeFeaturesToFloat,
     OneHotDegreeFeatures,
 )
-from modules.transforms.feature_liftings.feature_liftings import ProjectionSum
+from modules.transforms.feature_liftings.feature_liftings import (
+    ElementwiseMean,
+    ProjectionSum,
+)
 from modules.transforms.liftings.graph2cell.cycle_lifting import (
     CellCycleLifting,
 )
@@ -59,6 +62,9 @@
 from modules.transforms.liftings.graph2simplicial.line_lifting import (
     SimplicialLineLifting,
 )
+from modules.transforms.liftings.graph2simplicial.neighborhood_complex_lifting import (
+    NeighborhoodComplexLifting,
+)
 from modules.transforms.liftings.graph2simplicial.vietoris_rips_lifting import (
     SimplicialVietorisRipsLifting,
 )
@@ -106,6 +112,7 @@
     "SimplicialVietorisRipsLifting": SimplicialVietorisRipsLifting,
     "LatentCliqueLifting": LatentCliqueLifting,
     "SimplicialDnDLifting": SimplicialDnDLifting,
+    "NeighborhoodComplexLifting": NeighborhoodComplexLifting,
     # Graph -> Cell Complex
     "CellCycleLifting": CellCycleLifting,
     "DiscreteConfigurationComplexLifting": DiscreteConfigurationComplexLifting,
@@ -130,6 +137,7 @@
     "CofaceCCLifting": CofaceCCLifting,
     # Feature Liftings
     "ProjectionSum": ProjectionSum,
+    "ElementwiseMean": ElementwiseMean,
     # Data Manipulations
     "Identity": IdentityTransform,
     "NodeDegrees": NodeDegrees,

modules/transforms/feature_liftings/feature_liftings.py

@@ -1,4 +1,5 @@
 import torch
+import torch.nn.functional as F
 import torch_geometric
 
 
@@ -56,3 +57,104 @@ def forward(
             The lifted data.
         """
         return self.lift_features(data)
+
+
+class ElementwiseMean(torch_geometric.transforms.BaseTransform):
+    r"""Lifts r-cell features to r+1-cells by taking the mean of the lower
+    dimensional features.
+
+    Parameters
+    ----------
+    **kwargs : optional
+        Additional arguments for the class.
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__()
+
+    def lift_features(
+        self, data: torch_geometric.data.Data | dict
+    ) -> torch_geometric.data.Data | dict:
+        r"""Projects r-cell features of a graph to r+1-cell structures using the incidence matrix.
+
+        Parameters
+        ----------
+        data : torch_geometric.data.Data | dict
+            The input data to be lifted.
+
+        Returns
+        -------
+        torch_geometric.data.Data | dict
+            The lifted data."""
+
+        # Find the maximum dimension of the input data
+        max_dim = max(
+            [int(key.split("_")[-1]) for key in data if "x_idx" in key]
+        )
+
+        # Create a list of all x_idx tensors
+        x_idx_tensors = [data[f"x_idx_{i}"] for i in range(max_dim + 1)]
+
+        # Find the maximum sizes
+        max_simplices = max(tensor.size(0) for tensor in x_idx_tensors)
+        max_nodes = max(tensor.size(1) for tensor in x_idx_tensors)
+
+        # Pad tensors to have the same size
+        padded_tensors = [
+            F.pad(
+                tensor,
+                (
+                    0,
+                    max_nodes - tensor.size(1),
+                    0,
+                    max_simplices - tensor.size(0),
+                ),
+            )
+            for tensor in x_idx_tensors
+        ]
+
+        # Stack all x_idx tensors
+        all_indices = torch.stack(padded_tensors)
+
+        # Create a mask for valid indices
+        mask = all_indices != 0
+
+        # Replace 0s with a valid index (e.g., 0) to avoid indexing errors
+        all_indices = all_indices.clamp(min=0)
+
+        # Get all embeddings at once
+        all_embeddings = data["x_0"][all_indices]
+
+        # Apply mask to set padded embeddings to 0
+        all_embeddings = all_embeddings * mask.unsqueeze(-1).float()
+
+        # Compute sum and count of non-zero elements
+        embedding_sum = all_embeddings.sum(dim=2)
+        count = mask.sum(dim=2).clamp(min=1)  # Avoid division by zero
+
+        # Compute mean
+        mean_embeddings = embedding_sum / count.unsqueeze(-1)
+
+        # Assign results back to data dictionary
+        for i in range(1, max_dim + 1):
+            original_size = x_idx_tensors[i].size(0)
+            data[f"x_{i}"] = mean_embeddings[i, :original_size]
+
+        return data
+
+    def forward(
+        self, data: torch_geometric.data.Data | dict
+    ) -> torch_geometric.data.Data | dict:
+        r"""Applies the lifting to the input data.
+
+        Parameters
+        ----------
+        data : torch_geometric.data.Data | dict
+            The input data to be lifted.
+
+        Returns
+        -------
+        torch_geometric.data.Data | dict
+            The lifted data.
+        """
+        return self.lift_features(data)

modules/transforms/liftings/graph2simplicial/base.py

@@ -47,7 +47,7 @@ def _get_lifted_topology(
             list(simplicial_complex.get_simplex_attributes("features", 0).values())
         )
         # If new edges have been added during the lifting process, we discard the edge attributes
-        if self.contains_edge_attr and simplicial_complex.shape[1] == (
+        if self.preserve_edge_attr and simplicial_complex.shape[1] == (
             graph.number_of_edges()
         ):
             lifted_topology["x_1"] = torch.stack(

modules/transforms/liftings/graph2simplicial/neighborhood_complex_lifting.py

@@ -0,0 +1,67 @@
+import networkx as nx
+import torch
+from toponetx.classes import SimplicialComplex
+from torch_geometric.data import Data
+from torch_geometric.utils.convert import to_networkx
+
+from modules.transforms.liftings.graph2simplicial.base import (
+    Graph2SimplicialLifting,
+)
+
+
+class NeighborhoodComplexLifting(Graph2SimplicialLifting):
+    """Lifts graphs to a simplicial complex domain by identifying the neighborhood complex as k-simplices.
+    The neighborhood complex of a node u is the set of nodes that share a neighbor with u.
+
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def lift_topology(self, data: Data) -> dict:
+        graph: nx.Graph = to_networkx(data, to_undirected=True)
+        simplicial_complex = SimplicialComplex(simplices=graph)
+
+        # For every node u
+        for u in graph.nodes:
+            neighbourhood_complex = set()
+            neighbourhood_complex.add(u)
+            # Check it's neighbours
+            for v in graph.neighbors(u):
+                # For every other node w != u ^ w != v
+                for w in graph.nodes:
+                    # w == u
+                    if w == u:
+                        continue
+                    # w == v
+                    if w == v:
+                        continue
+
+                    # w and u share v as it's neighbour
+                    if v in graph.neighbors(w):
+                        neighbourhood_complex.add(w)
+            # Do not add 0-simplices
+            if len(neighbourhood_complex) < 2:
+                continue
+            # Do not add i-simplices if the maximum dimension is lower
+            if len(neighbourhood_complex) > self.complex_dim + 1:
+                continue
+            simplicial_complex.add_simplex(neighbourhood_complex)
+
+        feature_dict = {i: f for i, f in enumerate(data["x"])}
+
+        simplicial_complex.set_simplex_attributes(
+            feature_dict, name="features"
+        )
+
+        return self._get_lifted_topology(simplicial_complex, graph)
+
+    def _get_lifted_topology(
+        self, simplicial_complex: SimplicialComplex, graph: nx.Graph
+    ) -> dict:
+        data = super()._get_lifted_topology(simplicial_complex, graph)
+
+        for r in range(simplicial_complex.dim + 1):
+            data[f"x_idx_{r}"] = torch.tensor(simplicial_complex.skeleton(r))
+
+        return data

modules/transforms/liftings/lifting.py

@@ -7,11 +7,15 @@
 from modules.transforms.data_manipulations.manipulations import (
     IdentityTransform,
 )
-from modules.transforms.feature_liftings.feature_liftings import ProjectionSum
+from modules.transforms.feature_liftings.feature_liftings import (
+    ElementwiseMean,
+    ProjectionSum,
+)
 
 # Implemented Feature Liftings
 FEATURE_LIFTINGS = {
     "ProjectionSum": ProjectionSum,
+    "ElementwiseMean": ElementwiseMean,
     None: IdentityTransform,
 }