k-dimensional Weisfeiler-Leman algorithm in Python
Refines colorings of k-tuples of vertices; higher k distinguishes more graphs but costs more time.
- File
k_dimensional_weisfeiler_leman.py, 110 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/k-dimensional-weisfeiler-leman/. The authors are listed in prov/G42WBZ.jsonl.
"""The k-dimensional Weisfeiler-Leman algorithm (k-WL) for k >= 2.
Graphs are simple and undirected, given as (n, edges) with vertices
0..n-1 and optional integer vertex colors. k-WL colors all k-tuples of
vertices. A tuple starts with its atomic type (colors, equalities and
edges among its entries) and is refined by the multiset, over all
vertices w, of the colors of the k tuples obtained by putting w in each
position. This is the version for which k-WL has the power of the counting
logic with k + 1 variables; k = 2 is the classical algorithm of Weisfeiler
and Leman (1968). Color refinement (1-WL) is on its own page.
"""
import itertools
Edges = list[tuple[int, int]]
Graph = tuple[int, Edges]
def _atomic(t: tuple[int, ...], adj: list[set[int]], col: list[int]) -> tuple:
"""Atomic type of a tuple: colors, equality pattern, adjacency pattern."""
k = len(t)
eq = tuple(t[i] == t[j] for i in range(k) for j in range(i + 1, k))
ed = tuple(t[j] in adj[t[i]] for i in range(k) for j in range(i + 1, k))
return (tuple(col[v] for v in t), eq, ed)
def _setup(g: Graph, colors: list[int] | None, k: int
) -> tuple[list[tuple[int, ...]], dict[tuple[int, ...], int], list[tuple]]:
n, edges = g
adj = [set() for _ in range(n)]
for u, v in edges:
adj[u].add(v)
adj[v].add(u)
col = colors if colors is not None else [0] * n
tuples = list(itertools.product(range(n), repeat=k))
index = {t: i for i, t in enumerate(tuples)}
return tuples, index, [_atomic(t, adj, col) for t in tuples]
def _rename(keys: list[list]) -> list[list[int]]:
"""Map keys (shared by several graphs) to 0, 1, ... in sorted order."""
table = {key: i for i, key in enumerate(sorted({x for ks in keys for x in ks}))}
return [[table[x] for x in ks] for ks in keys]
def _step(n: int, k: int, tuples: list[tuple[int, ...]],
index: dict[tuple[int, ...], int], chi: list[int]) -> list[tuple]:
"""Signatures of one refinement round. Cost O(k n^(k+1) log n)."""
sigs = []
for t in tuples:
rows = []
for w in range(n):
rows.append(tuple(chi[index[t[:i] + (w,) + t[i + 1:]]] for i in range(k)))
rows.sort()
sigs.append((chi[index[t]], tuple(rows)))
return sigs
def kwl_colors(g: Graph, k: int, colors: list[int] | None = None,
) -> tuple[dict[tuple[int, ...], int], int]:
"""Stable k-WL coloring of the k-tuples of g and the number of rounds.
Cost: at most n^k rounds of O(k n^(k+1) log n) each.
"""
if k < 2:
raise ValueError("k must be at least 2")
n = g[0]
# @line L1
tuples, index, atoms = _setup(g, colors, k)
chi = _rename([atoms])[0]
rounds = 0
# @line L2
while True:
# @line L3
sigs = _step(n, k, tuples, index, chi)
# @line L4
new = _rename([sigs])[0]
# @line L5
if len(set(new)) == len(set(chi)):
break
chi = new
rounds += 1
# @line L6
return {t: chi[i] for i, t in enumerate(tuples)}, rounds
def kwl_distinguishes(g: Graph, h: Graph, k: int,
cg: list[int] | None = None, ch: list[int] | None = None) -> bool:
"""True if k-WL tells g and h apart (then they are not isomorphic).
Both graphs are refined in parallel with one shared table of colors;
they are told apart when the color histograms differ in some round.
"""
if k < 2:
raise ValueError("k must be at least 2")
if g[0] != h[0]:
return True
n = g[0]
tg, ig, ag = _setup(g, cg, k)
th, ih, ah = _setup(h, ch, k)
chi_g, chi_h = _rename([ag, ah])
while True:
if sorted(chi_g) != sorted(chi_h):
return True
new_g, new_h = _rename([_step(n, k, tg, ig, chi_g), _step(n, k, th, ih, chi_h)])
if len(set(new_g) | set(new_h)) == len(set(chi_g) | set(chi_h)):
return sorted(new_g) != sorted(new_h)
chi_g, chi_h = new_g, new_h
test_k_dimensional_weisfeiler_leman.py 116 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/k-dimensional-weisfeiler-leman/. The authors are listed in prov/G42WBZ.jsonl.
import itertools
import random
from collections import Counter
import pytest
from k_dimensional_weisfeiler_leman import kwl_colors, kwl_distinguishes
Edges = list[tuple[int, int]]
def brute_iso(n: int, a: Edges, b: Edges, ca: list[int] | None = None,
cb: list[int] | None = None) -> bool:
if len(a) != len(b):
return False
ca = ca or [0] * n
cb = cb or [0] * n
s = {frozenset(e) for e in b}
return any(all(cb[p[v]] == ca[v] for v in range(n))
and all(frozenset((p[u], p[v])) in s for u, v in a)
for p in itertools.permutations(range(n)))
def rand_graph(rng: random.Random, n: int, p: float) -> Edges:
return [(u, v) for u in range(n) for v in range(u + 1, n) if rng.random() < p]
def rook4() -> tuple[int, Edges]:
c = [(i, j) for i in range(4) for j in range(4)]
return 16, [(a, b) for a in range(16) for b in range(a + 1, 16) if c[a][0] == c[b][0] or c[a][1] == c[b][1]]
def shrikhande() -> tuple[int, Edges]:
c = [(i, j) for i in range(4) for j in range(4)]
d = {(0, 1), (0, 3), (1, 0), (3, 0), (1, 1), (3, 3)}
return 16, [(a, b) for a in range(16) for b in range(a + 1, 16)
if ((c[b][0] - c[a][0]) % 4, (c[b][1] - c[a][1]) % 4) in d]
def test_page_example() -> None:
c6 = (6, [(i, (i + 1) % 6) for i in range(6)])
col, rounds = kwl_colors(c6, 2)
assert rounds == 1
assert Counter(col.values()) == Counter({0: 12, 1: 6, 2: 12, 3: 6})
assert [col[t] for t in [(0, 0), (0, 1), (0, 2), (0, 3)]] == [3, 2, 0, 1]
for i in range(6):
for j in range(6):
d = min((i - j) % 6, (j - i) % 6)
assert col[(i, j)] == col[(0, d)]
two_k3 = (6, [(0, 1), (1, 2), (0, 2), (3, 4), (4, 5), (3, 5)])
assert kwl_distinguishes(c6, two_k3, 2)
def test_edge_cases() -> None:
with pytest.raises(ValueError):
kwl_colors((2, []), 1)
assert not kwl_distinguishes((0, []), (0, []), 2)
assert not kwl_distinguishes((1, []), (1, []), 2)
assert kwl_distinguishes((1, []), (2, []), 2)
assert kwl_distinguishes((2, []), (2, [(0, 1)]), 2)
assert kwl_distinguishes((2, []), (2, []), 2, [0, 1], [0, 0])
assert not kwl_distinguishes((3, [(0, 1)]), (3, [(1, 2)]), 3, [1, 1, 0], [0, 1, 1])
def test_strongly_regular_pair() -> None:
r, s = rook4(), shrikhande()
for g in (r, s):
col, _ = kwl_colors(g, 2)
assert sorted(Counter(col.values()).values()) == [16, 96, 144]
assert not kwl_distinguishes(r, s, 2)
assert kwl_distinguishes(r, s, 3)
assert len(set(kwl_colors(r, 3)[0].values())) == 15
assert len(set(kwl_colors(s, 3)[0].values())) == 31
def test_random_against_brute_force() -> None:
rng = random.Random(20261105)
checked = 0
for trial in range(800):
n = rng.randint(1, 6)
p = rng.choice([0.3, 0.5, 0.7])
a = rand_graph(rng, n, p)
ca = [rng.randint(0, 1) for _ in range(n)] if trial % 4 == 0 else None
if trial % 2 == 0:
perm = list(range(n))
rng.shuffle(perm)
b = [(perm[u], perm[v]) for u, v in a]
cb = None
if ca is not None:
cb = [0] * n
for v in range(n):
cb[perm[v]] = ca[v]
else:
b = rand_graph(rng, n, p)
cb = [rng.randint(0, 1) for _ in range(n)] if ca is not None else None
iso = brute_iso(n, a, b, ca, cb)
dist = kwl_distinguishes((n, a), (n, b), 2, ca, cb)
assert not (iso and dist)
if n <= 5:
k = max(2, n - 1)
assert kwl_distinguishes((n, a), (n, b), k, ca, cb) == (not iso)
checked += 1
assert checked == 800
def test_regular_pairs() -> None:
c8 = (8, [(i, (i + 1) % 8) for i in range(8)])
two_c4 = (8, [(i, (i + 1) % 4) for i in range(4)] + [(4 + i, 4 + (i + 1) % 4) for i in range(4)])
cube = (8, [(a, a ^ b) for a in range(8) for b in (1, 2, 4) if a < a ^ b])
two_k4 = (8, [(a, b) for a in range(8) for b in range(a + 1, 8) if a // 4 == b // 4])
assert kwl_distinguishes(c8, two_c4, 2)
assert kwl_distinguishes(cube, two_k4, 2)
assert not kwl_distinguishes(cube, cube, 2)
What the next level needs
- the tests did not run in this build
Read the full page about k-dimensional Weisfeiler-Leman algorithm