diff --git a/graphify/cluster.py b/graphify/cluster.py index 68221070..db476b5b 100644 --- a/graphify/cluster.py +++ b/graphify/cluster.py @@ -19,11 +19,86 @@ def _suppress_output(): return contextlib.redirect_stdout(io.StringIO()) +def _native_leiden(stable: nx.Graph, resolution: float) -> dict[str, int] | None: + """Call graspologic_native.leiden() directly, bypassing graspologic's own + package import. + + graspologic.partition.leiden() is a thin wrapper around exactly this + native (Rust) call. Importing the *package* — as opposed to the native + extension module it depends on — pulls in graspologic.layouts, which + imports umap, which imports pynndescent, which numba-JIT-compiles at + import time for a layout algorithm this function never calls: measured + at 7-19s of one-time import cost against a ~1s native call and a ~1.4s + full round trip (conversion + call + map-back) — see the "third update" + in GRAPHIFY_BUILD_PERF.md for the measurements this is based on. + + Returns None (the caller falls through to the graspologic.partition.leiden + path, then to the networkx Louvain fallback) if graspologic_native isn't + installed, or if `stable` isn't the plain undirected, non-multigraph + input leiden actually supports — the same shape check + graspologic.partition.leiden itself makes before calling the same native + function. + """ + try: + import graspologic_native as gn + except ImportError: + return None + + if stable.is_directed() or stable.is_multigraph(): + return None + + # graspologic_native identifies nodes by their string form; two DISTINCT + # node objects that happen to stringify the same way would silently merge + # under it (this is exactly what graspologic.partition.leiden's own + # _IdentityMapper guards against). Graphify's own node IDs are already + # unique strings by construction — extractors/resolution.py's + # _disambiguate_colliding_node_ids salts any two distinct nodes that would + # otherwise share a string id before the graph is ever built — so this is + # a defensive check on an assumption that should never actually trip, not + # an expected path. One pass over the nodes, cheaper than an + # _IdentityMapper-style dict-store-per-edge-endpoint. + id_to_node: dict[str, object] = {} + for node in stable.nodes(): + key = str(node) + existing = id_to_node.get(key) + if existing is not None and existing != node: + return None # let graspologic.partition.leiden's own check handle/raise on this + id_to_node[key] = node + + edges = [ + (str(u), str(v), float(attrs.get("weight", 1.0))) + for u, v, attrs in stable.edges(data=True) + ] + + try: + old_stderr = sys.stderr + try: + sys.stderr = io.StringIO() + with _suppress_output(): + _quality, native_partitions = gn.leiden( + edges=edges, + starting_communities=None, + resolution=resolution, + randomness=0.001, + iterations=1, + use_modularity=True, + seed=42, + trials=1, + ) + finally: + sys.stderr = old_stderr + except Exception: + return None + + return {id_to_node[node_id]: community for node_id, community in native_partitions.items()} + + def _partition(G: nx.Graph, resolution: float = 1.0) -> dict[str, int]: """Run community detection. Returns {node_id: community_id}. - Tries Leiden (graspologic) first — best quality. - Falls back to Louvain (built into networkx) if graspologic is not installed. + Tries Leiden (graspologic_native directly, then graspologic) first — best + quality. Falls back to Louvain (built into networkx) if neither is + installed. resolution > 1.0 → more, smaller communities. resolution < 1.0 → fewer, larger communities. @@ -44,6 +119,10 @@ def _partition(G: nx.Graph, resolution: float = 1.0) -> dict[str, int]: for src, tgt, attrs in edge_rows: stable.add_edge(src, tgt, **attrs) + native_result = _native_leiden(stable, resolution) + if native_result is not None: + return native_result + try: from graspologic.partition import leiden lsig = inspect.signature(leiden).parameters