Update ece_partitioner.py (#123)

mdjhacker · web-flow · commit 658397d928c4 · 2025-12-18T18:54:56.000+08:00
diff --git a/graphgen/models/partitioner/ece_partitioner.py b/graphgen/models/partitioner/ece_partitioner.py
@@ -142,7 +142,7 @@ def _add_unit(u):
             return Community(
                 id=seed_unit[1],
                 nodes=list(community_nodes.keys()),
-                edges=[(u, v) for (u, v), _ in community_edges.items()],
+                edges=[tuple(edge) for edge in community_edges if isinstance(edge, frozenset) and len(edge)==2],
             )
 
         for unit in tqdm(all_units, desc="ECE partition"):

Original file line number	Diff line number	Diff line change
`@@ -142,7 +142,7 @@ def _add_unit(u):`
`142`	`142`	`return Community(`
`143`	`143`	`id=seed_unit[1],`
`144`	`144`	`nodes=list(community_nodes.keys()),`
`145`		`- edges=[(u, v) for (u, v), _ in community_edges.items()],`
	`145`	`+ edges=[tuple(edge) for edge in community_edges if isinstance(edge, frozenset) and len(edge)==2],`
`146`	`146`	`)`
`147`	`147`
`148`	`148`	`for unit in tqdm(all_units, desc="ECE partition"):`