Request for Enhancement: Support for Weighted Graphs in ClusterData

[Amandayoula]

🚀 The feature, motivation and pitch

Firstly, I want to express my appreciation for the valuable contributions made to this excellent library.

I am currently utilizing the ClusterData functionality within the library, and I am interested in applying it to weighted graphs, where the clustering of nodes is based on the weights of the edges connecting them.

Upon reviewing the source code of ClusterData, I noticed that it only utilizes the edge index of the graph. Here is the relevant snippet from the main code:

cluster = self._metis(data.edge_index, data.num_nodes)
self.partition = self._partition(data.edge_index, cluster)

The METIS method used by ClusterData provides parameters for node weights (vweights) and edge weights (eweights), which suggests that it should be capable of handling weighted graphs.

Here is the METIS function for reference:

pymetis.part_graph(nparts, adjacency=None, xadj=None, adjncy=None, vweights=None, eweights=None, recursive=None, contiguous=None)

I apologize if I have overlooked any important code or if this functionality is already available. I would greatly appreciate it if you could confirm whether weighted graph support exists in ClusterData or if it could be added in a future release.

Thank you for considering my request.

Alternatives

No response

Additional context

No response

[Amandayoula]

After revisiting the source code of ClusterData, I discovered an additional code snippet:

if cluster is None and torch_geometric.typing.WITH_METIS:
            cluster = pyg_lib.partition.metis(
                indptr.cpu(),
                index.cpu(),
                self.num_parts,
                recursive=self.recursive,
            ).to(edge_index.device)

Upon further investigation into the pyg_lib.partition.metis() function, I found that it indeed supports edge weights. However, it seems that this functionality is not currently utilized in the ClusterData module.

The metis() function signature suggests that it could easily accommodate edge weights with a slight modification to the source code:

metis(rowptr: Tensor, col: Tensor, num_partitions: int, node_weight: Optional[Tensor] = None, edge_weight: Optional[Tensor] = None, recursive: bool = False)

[rusty1s]

Yeah, this is well supported within pyg-lib, just not really exposed in PyG. This is not super high on my TODO, so I would appreciate some help on this one to expose this additional arguments.