105 lines
4.1 KiB
Python
105 lines
4.1 KiB
Python
"""Leiden community detection on NetworkX graphs. Splits oversized communities. Returns cohesion scores."""
|
|
from __future__ import annotations
|
|
import networkx as nx
|
|
|
|
|
|
def build_graph(nodes: list[dict], edges: list[dict]) -> nx.Graph:
|
|
"""Build a NetworkX graph from graphify node/edge dicts.
|
|
|
|
Preserves original edge direction as _src/_tgt attributes so that
|
|
display functions can show relationships in the correct direction,
|
|
even though the graph is undirected for structural analysis.
|
|
"""
|
|
G = nx.Graph()
|
|
for n in nodes:
|
|
G.add_node(n["id"], **{k: v for k, v in n.items() if k != "id"})
|
|
for e in edges:
|
|
attrs = {k: v for k, v in e.items() if k not in ("source", "target")}
|
|
attrs["_src"] = e["source"]
|
|
attrs["_tgt"] = e["target"]
|
|
G.add_edge(e["source"], e["target"], **attrs)
|
|
return G
|
|
|
|
_MAX_COMMUNITY_FRACTION = 0.25 # communities larger than 25% of graph get split
|
|
_MIN_SPLIT_SIZE = 10 # only split if community has at least this many nodes
|
|
|
|
|
|
def cluster(G: nx.Graph) -> dict[int, list[str]]:
|
|
"""Run Leiden community detection. Returns {community_id: [node_ids]}.
|
|
|
|
Community IDs are stable across runs: 0 = largest community after splitting.
|
|
Oversized communities (> 25% of graph nodes, min 10) are split by running
|
|
a second Leiden pass on the subgraph.
|
|
"""
|
|
if G.number_of_nodes() == 0:
|
|
return {}
|
|
if G.number_of_edges() == 0:
|
|
return {i: [n] for i, n in enumerate(sorted(G.nodes))}
|
|
|
|
from graspologic.partition import leiden # lazy - avoids 15s numba JIT on import
|
|
|
|
# Leiden warns and drops isolates - handle them separately
|
|
isolates = [n for n in G.nodes() if G.degree(n) == 0]
|
|
connected_nodes = [n for n in G.nodes() if G.degree(n) > 0]
|
|
connected = G.subgraph(connected_nodes)
|
|
|
|
raw: dict[int, list[str]] = {}
|
|
if connected.number_of_nodes() > 0:
|
|
partition: dict[str, int] = leiden(connected)
|
|
for node, cid in partition.items():
|
|
raw.setdefault(cid, []).append(node)
|
|
|
|
# Each isolate becomes its own single-node community
|
|
next_cid = max(raw.keys(), default=-1) + 1
|
|
for node in isolates:
|
|
raw[next_cid] = [node]
|
|
next_cid += 1
|
|
|
|
# Split oversized communities
|
|
max_size = max(_MIN_SPLIT_SIZE, int(G.number_of_nodes() * _MAX_COMMUNITY_FRACTION))
|
|
final_communities: list[list[str]] = []
|
|
for nodes in raw.values():
|
|
if len(nodes) > max_size:
|
|
final_communities.extend(_split_community(G, nodes))
|
|
else:
|
|
final_communities.append(nodes)
|
|
|
|
# Re-index by size descending for deterministic ordering
|
|
final_communities.sort(key=len, reverse=True)
|
|
return {i: sorted(nodes) for i, nodes in enumerate(final_communities)}
|
|
|
|
|
|
def _split_community(G: nx.Graph, nodes: list[str]) -> list[list[str]]:
|
|
"""Run a second Leiden pass on a community subgraph to split it further."""
|
|
subgraph = G.subgraph(nodes)
|
|
if subgraph.number_of_edges() == 0:
|
|
# No edges - split into individual nodes
|
|
return [[n] for n in sorted(nodes)]
|
|
try:
|
|
from graspologic.partition import leiden
|
|
sub_partition: dict[str, int] = leiden(subgraph)
|
|
sub_communities: dict[int, list[str]] = {}
|
|
for node, cid in sub_partition.items():
|
|
sub_communities.setdefault(cid, []).append(node)
|
|
if len(sub_communities) <= 1:
|
|
# Leiden couldn't split it - return as-is
|
|
return [sorted(nodes)]
|
|
return [sorted(v) for v in sub_communities.values()]
|
|
except Exception:
|
|
return [sorted(nodes)]
|
|
|
|
|
|
def cohesion_score(G: nx.Graph, community_nodes: list[str]) -> float:
|
|
"""Ratio of actual intra-community edges to maximum possible."""
|
|
n = len(community_nodes)
|
|
if n <= 1:
|
|
return 1.0
|
|
subgraph = G.subgraph(community_nodes)
|
|
actual = subgraph.number_of_edges()
|
|
possible = n * (n - 1) / 2
|
|
return round(actual / possible, 2) if possible > 0 else 0.0
|
|
|
|
|
|
def score_all(G: nx.Graph, communities: dict[int, list[str]]) -> dict[int, float]:
|
|
return {cid: cohesion_score(G, nodes) for cid, nodes in communities.items()}
|