Infomap in Python
Finds communities that minimise the description length of random-walk trajectories.
- File
infomap.py, 150 lines- Level
- I0 not checked yet
- Toolchain
- Python 3.14, pytest 9
- Test framework
- pytest
- Run the tests
python3 -m pytest
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/infomap/. The authors are listed in prov/R7CYRH.jsonl.
from fractions import Fraction
from math import log2
def _pw(a: int) -> int:
"""2^(a log2 a) = a^a, with 0^0 = 1."""
return a**a
def map_equation(n: int, edges: list[tuple[int, int, int]], modules: list[list[int]]) -> float:
"""The two-level map equation L(M) in bits, for an undirected graph without self-loops.
The walker is at u with probability k_u / 2W, and module i is left with
probability q_i = cut(i) / 2W. Then
L = q log2 q - 2 sum q_i log2 q_i - sum p_u log2 p_u + sum (q_i + p_i) log2 (q_i + p_i),
where q is the sum of the q_i and p_i the sum of p_u over i.
"""
deg = [0] * n
where = {v: i for i, m in enumerate(modules) for v in m}
cut = [0] * len(modules)
tot = [0] * len(modules)
for u, v, w in edges:
deg[u] += w
deg[v] += w
if where[u] != where[v]:
cut[where[u]] += w
cut[where[v]] += w
two_w = sum(deg)
if two_w == 0:
return 0.0
for u in range(n):
tot[where[u]] += deg[u]
def plogp(x: float) -> float:
return x * log2(x) if x > 0 else 0.0
q = [c / two_w for c in cut]
return (
plogp(sum(q))
- 2 * sum(plogp(x) for x in q)
- sum(plogp(d / two_w) for d in deg)
+ sum(plogp(q[i] + tot[i] / two_w) for i in range(len(modules)))
)
def _one_level(
k: int, adj: list[dict[int, int]], loops: list[int]
) -> tuple[list[int], bool]:
"""Local moving on a graph with k vertices; returns (module of each vertex, moved)."""
degree = [2 * loops[v] + sum(adj[v].values()) for v in range(k)]
own_cut = [degree[v] - 2 * loops[v] for v in range(k)]
comm = list(range(k)) # @line L3
tot = degree[:]
cut = own_cut[:]
exit_sum = sum(cut)
moved_any = False
improved = True
while improved: # @line L4
improved = False # @line L5
for v in range(k): # @line L6
old = comm[v] # @line L7
links: dict[int, int] = {}
for w, x in adj[v].items():
links[comm[w]] = links.get(comm[w], 0) + x
tot[old] -= degree[v]
new_cut_old = cut[old] - own_cut[v] + 2 * links.get(old, 0)
exit_sum += new_cut_old - cut[old]
cut[old] = new_cut_old
def score(c: int) -> Fraction:
"""2 to the power 2W times the change of L when v joins c."""
c_in = cut[c] + own_cut[v] - 2 * links.get(c, 0)
t_in = tot[c] + degree[v]
big = exit_sum + c_in - cut[c]
return Fraction(
_pw(big) * _pw(cut[c]) ** 2 * _pw(c_in + t_in),
_pw(c_in) ** 2 * _pw(cut[c] + tot[c]),
)
best = old # @line L8
best_score = score(old)
for c in sorted(links):
if c != old:
s = score(c)
if s < best_score:
best, best_score = c, s
c_in = cut[best] + own_cut[v] - 2 * links.get(best, 0) # @line L9
exit_sum += c_in - cut[best]
cut[best] = c_in
tot[best] += degree[v]
comm[v] = best
if best != old:
improved = True
moved_any = True
return comm, moved_any # @line L10
def infomap(n: int, edges: list[tuple[int, int, int]]) -> tuple[list[list[int]], list[list[list[int]]]]:
"""Return (modules, levels) for an undirected graph on 0 .. n-1 by the two-level map equation.
edges holds pairs (u, v, w) with u != v and integer weight w > 0; a pair
given twice adds its weights. Each level moves vertices in index order to
the neighbouring module that lowers the map equation most (ties to the
smallest module id, and a vertex stays when its own module is among the
best), then contracts each module to one vertex. levels lists the
partition of the original vertices after each level that changed
something; modules is the last. Comparisons are exact.
"""
adj: list[dict[int, int]] = [{} for _ in range(n)]
total = 0
for u, v, w in edges:
if u == v:
raise ValueError("self-loops are not allowed")
total += w
adj[u][v] = adj[u].get(v, 0) + w
adj[v][u] = adj[v].get(u, 0) + w
loops = [0] * n
parts = [[v] for v in range(n)] # @line L1
levels: list[list[list[int]]] = []
if total == 0:
return parts, levels
while True: # @line L2
k = len(parts)
comm, moved = _one_level(k, adj, loops)
if not moved: # @line L11
return sorted(sorted(p) for p in parts), levels
index: dict[int, int] = {} # @line L12
for v in range(k):
if comm[v] not in index:
index[comm[v]] = len(index)
new_parts: list[list[int]] = [[] for _ in index]
for v in range(k):
new_parts[index[comm[v]]] += parts[v]
new_adj: list[dict[int, int]] = [{} for _ in index] # @line L13
new_loops = [0] * len(index)
for v in range(k):
a = index[comm[v]]
new_loops[a] += loops[v]
for w, x in adj[v].items():
b = index[comm[w]]
if a == b:
if v < w:
new_loops[a] += x
else:
new_adj[a][b] = new_adj[a].get(b, 0) + x
parts, adj, loops = new_parts, new_adj, new_loops
levels.append(sorted(sorted(p) for p in parts))
test_infomap.py 172 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/infomap/. The authors are listed in prov/R7CYRH.jsonl.
import random
from fractions import Fraction
from itertools import combinations
from math import isclose, log2
from infomap import infomap, map_equation
def exact_power(n, edges, groups):
"""2^(2W L(M)) from the definition, as an exact fraction of integers."""
deg = [0] * n
for u, v, w in edges:
deg[u] += w
deg[v] += w
where = {v: i for i, g in enumerate(groups) for v in g}
cut = [0] * len(groups)
for u, v, w in edges:
if where[u] != where[v]:
cut[where[u]] += w
cut[where[v]] += w
tot = [sum(deg[v] for v in g) for g in groups]
big = sum(cut)
num = big**big
den = 1
for c, t in zip(cut, tot):
num *= (c + t) ** (c + t)
den *= c ** (2 * c)
for d in deg:
den *= d**d
return Fraction(num, den)
def reference(n, edges):
"""The same rules on blocks of original vertices, each score computed from the definition."""
if not edges:
return [[v] for v in range(n)], []
linked = set()
for u, v, _ in edges:
linked.add((u, v))
linked.add((v, u))
blocks = [[v] for v in range(n)]
levels = []
while True:
k = len(blocks)
comm = list(range(k))
moved = improved = False
improved = True
while improved:
improved = False
for b in range(k):
old = comm[b]
near = {
comm[c]
for c in range(k)
if c != b and any((x, y) in linked for x in blocks[b] for y in blocks[c])
}
def score(target):
trial = comm[:]
trial[b] = target
groups = {}
for c in range(k):
groups.setdefault(trial[c], []).extend(blocks[c])
return exact_power(n, edges, list(groups.values()))
best, best_s = old, score(old)
for c in sorted(near):
if c != old and score(c) < best_s:
best, best_s = c, score(c)
comm[b] = best
if best != old:
improved = moved = True
if not moved:
return sorted(sorted(p) for p in blocks), levels
order = []
for b in range(k):
if comm[b] not in order:
order.append(comm[b])
blocks = [sum((blocks[b] for b in range(k) if comm[b] == c), []) for c in order]
levels.append(sorted(sorted(p) for p in blocks))
def all_partitions(items):
if not items:
yield []
return
first, rest = items[0], items[1:]
for p in all_partitions(rest):
for i in range(len(p)):
yield p[:i] + [[first] + p[i]] + p[i + 1 :]
yield [[first]] + p
def random_weighted(rng, n, p):
return [(u, v, rng.choice([1, 1, 1, 2, 3])) for u, v in combinations(range(n), 2) if rng.random() < p]
EXAMPLE = [(0, 1), (0, 2), (1, 2), (1, 3), (2, 3), (3, 4), (4, 5), (4, 6), (5, 6)]
def test_example():
edges = [(u, v, 1) for u, v in EXAMPLE]
modules, levels = infomap(7, edges)
assert modules == [[0, 1, 2, 3], [4, 5, 6]]
assert levels == [[[0, 1, 2, 3], [4, 5, 6]]]
assert isclose(map_equation(7, edges, modules), 2.44444, abs_tol=5e-5)
assert isclose(map_equation(7, edges, [list(range(7))]), 2.77995, abs_tol=5e-5)
assert isclose(map_equation(7, edges, [[v] for v in range(7)]), 4.77995, abs_tol=5e-5)
def test_edge_cases():
assert infomap(0, []) == ([], [])
assert infomap(3, []) == ([[0], [1], [2]], [])
assert infomap(2, [(0, 1, 1)])[0] == [[0, 1]]
assert infomap(4, [(0, 1, 1), (2, 3, 1)])[0] == [[0, 1], [2, 3]]
assert infomap(4, [(u, v, 1) for u, v in combinations(range(4), 2)])[0] == [[0, 1, 2, 3]]
try:
infomap(2, [(1, 1, 1)])
except ValueError:
pass
else:
raise AssertionError("a self-loop must be refused")
def test_formula_matches_exact_form():
rng = random.Random(1)
for _ in range(200):
n = rng.randint(1, 8)
edges = random_weighted(rng, n, 0.5)
if not edges:
continue
two_w = 2 * sum(w for _, _, w in edges)
labels = [rng.randrange(3) for _ in range(n)]
groups = [[v for v in range(n) if labels[v] == c] for c in range(3)]
groups = [g for g in groups if g]
x = exact_power(n, edges, groups)
exact = (log2(x.numerator) - log2(x.denominator)) / two_w
assert isclose(map_equation(n, edges, groups), exact, abs_tol=1e-9)
def test_random_against_reference():
rng = random.Random(20080129)
for _ in range(300):
n = rng.randint(0, 8)
edges = random_weighted(rng, n, rng.choice([0.2, 0.4, 0.6]))
assert infomap(n, edges) == reference(n, edges)
def test_levels_lower_the_codelength():
rng = random.Random(118)
for _ in range(150):
n = rng.randint(1, 7)
edges = random_weighted(rng, n, 0.45)
if not edges:
continue
modules, levels = infomap(n, edges)
xs = [exact_power(n, edges, [[v] for v in range(n)])] + [exact_power(n, edges, p) for p in levels]
assert all(a > b for a, b in zip(xs, xs[1:]))
best = min(exact_power(n, edges, p) for p in all_partitions(list(range(n))))
assert exact_power(n, edges, modules) >= best
if __name__ == "__main__":
test_example()
test_edge_cases()
test_formula_matches_exact_form()
test_random_against_reference()
test_levels_lower_the_codelength()
print("ok")
What the next level needs
- the tests did not run in this build