Coverage for src/lilbee/retrieval/concepts/community.py: 100%

1"""Concept community dataclass and PMI / Leiden helpers.""" 

2 

3from __future__ import annotations 

4 

5import math 

6from collections import Counter 

7from dataclasses import dataclass 

8from typing import Any, NamedTuple 

9 

10_MIN_LEIDEN_WEIGHT = 0.01 

11# Fixed seed so Leiden returns the same communities for the same edge set; 

12# the algorithm is randomized and would otherwise drift between runs. 

13_LEIDEN_SEED = 42 

14 

15 

16@dataclass 

17class Community: 

18 """A cluster of related concepts from Leiden partitioning.""" 

19 

20 cluster_id: int 

21 size: int 

22 concepts: list[str] 

23 

24 

25def _compute_pmi( 

26 cooccurrences: Counter[tuple[str, str]], 

27 concept_counts: Counter[str], 

28 total_chunks: int, 

29) -> dict[tuple[str, str], float]: 

30 """Compute PPMI (Positive PMI) weights for concept co-occurrence pairs. 

31 Based on Church & Hanks 1990, "Word Association Norms, Mutual Information, 

32 and Lexicography." Pairs with non-positive PMI (co-occurring at or below 

33 chance) are dropped entirely: a stored 0.0 would later be floored to a 

34 positive Leiden weight, turning anti-correlation into attraction. 

35 """ 

36 pmi: dict[tuple[str, str], float] = {} 

37 for (a, b), count in cooccurrences.items(): 

38 p_a = concept_counts[a] / total_chunks 

39 p_b = concept_counts[b] / total_chunks 

40 if p_a == 0 or p_b == 0: 

41 continue 

42 p_ab = count / total_chunks 

43 value = math.log2(p_ab / (p_a * p_b)) 

44 if value > 0: 

45 pmi[(a, b)] = value 

46 return pmi 

47 

48 

49class _LeidenResult(NamedTuple): 

50 """Leiden output: node -> community id, and node -> incident-edge count. 

51 

52 ``degrees`` is an unweighted count of edges incident to each node (edge 

53 weights are not summed), which is what label ranking consumes. 

54 """ 

55 

56 partition: dict[str, int] 

57 degrees: dict[str, int] 

58 

59 

60def _leiden_partition( 

61 edge_rows: list[dict[str, Any]], 

62) -> _LeidenResult: 

63 """Run Leiden clustering on edge rows. 

64 Uses graspologic-native's Rust implementation (Traag et al. 2019, 

65 "From Louvain to Leiden: guaranteeing well-connected communities"). 

66 """ 

67 from graspologic_native import leiden 

68 

69 edges: list[tuple[str, str, float]] = [ 

70 (row["source"], row["target"], max(_MIN_LEIDEN_WEIGHT, row["weight"])) for row in edge_rows 

71 ] 

72 _modularity, partition = leiden(edges=edges, seed=_LEIDEN_SEED) # type: ignore[call-arg] 

73 

74 degree_map: Counter[str] = Counter() 

75 for row in edge_rows: 

76 degree_map[row["source"]] += 1 

77 degree_map[row["target"]] += 1 

78 return _LeidenResult(partition, dict(degree_map))