feat: split community with max_size

Author: ChenZiHong-Gavin

graphgen/models/community/community_detector.py

@@ -1,5 +1,6 @@
+from collections import defaultdict
 from dataclasses import dataclass
-from typing import Any, Dict
+from typing import Any, Dict, List
 
 from graphgen.models.storage.networkx_storage import NetworkXStorage
 
@@ -13,45 +14,40 @@ class CommunityDetector:
     method_params: Dict[str, Any] = None
 
     async def detect_communities(self) -> Dict[str, int]:
-        """
-        Detect communities based on the chosen method.
-        """
         if self.method == "leiden":
             return await self._leiden_communities(**self.method_params or {})
         raise ValueError(f"Unknown community detection method: {self.method}")
 
     async def get_graph(self):
-        """
-        Asynchronously get the graph from the storage.
-        """
         return await self.graph_storage.get_graph()
 
-    async def _leiden_communities(self, **kwargs) -> Dict[str, int]:
+    async def _leiden_communities(
+        self, max_size: int = None, **kwargs
+    ) -> Dict[str, int]:
         """
         Detect communities using the Leiden algorithm.
+        If max_size is given, any community larger than max_size will be split
+        into smaller sub-communities each having at most max_size nodes.
         """
         import igraph as ig
         import networkx as nx
         from leidenalg import ModularityVertexPartition, find_partition
 
         graph = await self.get_graph()
-        # Filter out isolated nodes
         graph.remove_nodes_from(list(nx.isolates(graph)))
 
-        # Convert NetworkX graph to igraph graph
         ig_graph = ig.Graph.TupleList(graph.edges(), directed=False)
 
         random_seed = kwargs.get("random_seed", 42)
         use_lcc = kwargs.get("use_lcc", False)
 
-        communities = {}
+        communities: Dict[str, int] = {}
         if use_lcc:
-            # Use the largest connected component
             lcc = ig_graph.components().giant()
             partition = find_partition(lcc, ModularityVertexPartition, seed=random_seed)
             for part, cluster in enumerate(partition):
                 for v in cluster:
-                    communities[v] = part
+                    communities[lcc.vs[v]["name"]] = part
         else:
             offset = 0
             for component in ig_graph.components():
@@ -64,4 +60,36 @@ async def _leiden_communities(self, **kwargs) -> Dict[str, int]:
                         original_node = subgraph.vs[v]["name"]
                         communities[original_node] = part + offset
                 offset += len(partition)
-        return communities
+
+        # split large communities if max_size is specified
+        if max_size is None or max_size <= 0:
+            return communities
+
+        return await self._split_communities(communities, max_size)
+
+    @staticmethod
+    async def _split_communities(
+        communities: Dict[str, int], max_size: int
+    ) -> Dict[str, int]:
+        """
+        Split communities larger than max_size into smaller sub-communities.
+        """
+        cid2nodes: Dict[int, List[str]] = defaultdict(list)
+        for node, cid in communities.items():
+            cid2nodes[cid].append(node)
+
+        new_communities: Dict[str, int] = {}
+        new_cid = 0
+        for cid, nodes in cid2nodes.items():
+            if len(nodes) <= max_size:
+                for n in nodes:
+                    new_communities[n] = new_cid
+                new_cid += 1
+            else:
+                for start in range(0, len(nodes), max_size):
+                    sub = nodes[start : start + max_size]
+                    for n in sub:
+                        new_communities[n] = new_cid
+                    new_cid += 1
+
+        return new_communities

graphgen/models/vis/community_visualizer.py

@@ -31,7 +31,6 @@ def visualize(self, save_path: str = None):
         plt.figure(figsize=(10, 10))
 
         node_colors = [self.communities.get(node, 0) for node in self.graph.nodes()]
-        print(node_colors)
 
         nx.draw_networkx_nodes(
             self.graph,