Interlevel Betti Token Transformer / toposcan.py

✓✓ Beats tuned baseline

Raw ⬇ ZIP
 1"""Interlevel Betti tokens for graph 1-complexes, with GF(2) homology."""
 2import math
 3import numpy as np
 4
 5
 6def gf2_rank(a):
 7    """Rank over GF(2), using compact XOR elimination."""
 8    a = (np.asarray(a, dtype=np.uint8) & 1).copy()
 9    if a.ndim != 2:
10        raise ValueError("matrix must be 2-D")
11    rows, cols = a.shape
12    rank = 0
13    for col in range(cols):
14        piv = np.flatnonzero(a[rank:, col])
15        if len(piv) == 0:
16            continue
17        p = rank + int(piv[0])
18        if p != rank:
19            a[[rank, p]] = a[[p, rank]]
20        for r in range(rows):
21            if r != rank and a[r, col]:
22                a[r] ^= a[rank]
23        rank += 1
24        if rank == rows:
25            break
26    return rank
27
28
29def boundary_matrices(n, edges):
30    """Return B1 (vertices x edges) and B0 (0 x vertices)."""
31    b1 = np.zeros((n, len(edges)), dtype=np.uint8)
32    for j, (u, v) in enumerate(edges):
33        b1[u, j] = 1
34        b1[v, j] ^= 1
35    return np.zeros((0, n), dtype=np.uint8), b1
36
37
38def betti_graph(n, edges):
39    """Betti numbers of a graph after boundary closure, over GF(2)."""
40    b0, b1 = boundary_matrices(n, edges)
41    r0, r1 = gf2_rank(b0), gf2_rank(b1)
42    beta0 = n - r0 - r1
43    beta1 = len(edges) - r1
44    return int(beta0), int(beta1)
45
46
47def interlevel_betti(n, edges, h, lo, hi):
48    """Upper-star interval complex with boundary closure.
49
50    Edges whose max endpoint filtration value is in [lo, hi] are retained;
51    all their endpoints are then retained to enforce simplicial boundary closure.
52    Isolated vertices whose own value is in the interval are also retained.
53    """
54    h = np.asarray(h, dtype=float)
55    keep_edges = [(u, v) for (u, v) in edges
56                  if lo <= max(h[u], h[v]) <= hi]
57    keep_vertices = {i for i, x in enumerate(h) if lo <= x <= hi}
58    for u, v in keep_edges:
59        keep_vertices.update((u, v))
60    verts = sorted(keep_vertices)
61    remap = {v: i for i, v in enumerate(verts)}
62    redges = [(remap[u], remap[v]) for u, v in keep_edges]
63    return betti_graph(len(verts), redges), (len(verts), len(redges))
64
65
66def grid_tokens(n, edges, h, grid, m, stride):
67    grid = np.asarray(grid, dtype=float)
68    out = []
69    for start in range(0, len(grid) - m, stride):
70        lo, hi = grid[start], grid[start + m]
71        (b0, b1), (nv, ne) = interlevel_betti(n, edges, h, lo, hi)
72        out.append([b0, b1, math.log1p(nv), math.log1p(ne)])
73    return np.asarray(out, dtype=np.float32)
74
75
76# Optional neural integration point: tokens -> positional Transformer -> graph vector.
77def make_transformer(width=32, heads=4, layers=2):
78    import torch
79    import torch.nn as nn
80    class TopoScanTransformer(nn.Module):
81        def __init__(self):
82            super().__init__()
83            self.proj = nn.Linear(4, width)
84            enc = nn.TransformerEncoderLayer(width, heads, 4 * width,
85                                              batch_first=True, dropout=0.0)
86            self.encoder = nn.TransformerEncoder(enc, layers)
87            self.norm = nn.LayerNorm(width)
88        def forward(self, tokens):
89            z = self.proj(tokens)
90            z = self.encoder(z)
91            return self.norm(z.mean(dim=1))
92    return TopoScanTransformer()