1import json, math, time, random
  2from pathlib import Path
  3import numpy as np
  4
  5try:
  6    import torch
  7    import torch.nn.functional as F
  8except Exception:
  9    torch = None
 10
 11
 12def cyclic_design(v, base):
 13    blocks = [sorted({(x + t) % v for x in base}) for t in range(v)]
 14    return np.asarray(blocks, dtype=np.int64)
 15
 16
 17def make_designs():
 18    # Cyclic (7,3,1) and (13,4,1) projective-plane designs.
 19    return {
 20        "Fano_7": cyclic_design(7, [0, 1, 3]),
 21        "projective_13": cyclic_design(13, [0, 1, 3, 9]),
 22    }
 23
 24
 25def incidence(blocks, v):
 26    M = np.zeros((v, len(blocks)), dtype=np.int64)
 27    for j, b in enumerate(blocks):
 28        M[b, j] = 1
 29    return M
 30
 31
 32def math_check(name, blocks):
 33    v = int(blocks.max()) + 1
 34    M = incidence(blocks, v)
 35    G = M @ M.T
 36    r = int(M.sum(1)[0])
 37    off = G[~np.eye(v, dtype=bool)]
 38    lam = int(off[0])
 39    pair_vals = off.tolist()
 40    # Z(x,y) is the column-vector membership difference.
 41    norm_err = []
 42    for x in range(v):
 43        for y in range(x + 1, v):
 44            z = M[x] - M[y]
 45            norm_err.append(int(z @ z))
 46    return {
 47        "name": name, "v": v, "b": len(blocks), "k": len(blocks[0]),
 48        "r": r, "lambda": lam, "row_min": int(M.sum(1).min()),
 49        "row_max": int(M.sum(1).max()), "pair_min": int(min(pair_vals)),
 50        "pair_max": int(max(pair_vals)), "pair_variance": float(np.var(pair_vals)),
 51        "mmT_identity_maxerr": int(np.max(np.abs(G - ((r-lam)*np.eye(v, dtype=int)+lam*np.ones((v,v), dtype=int))))),
 52        "predicted_z_norm_sq": 2*(r-lam), "observed_z_norm_sq_min": min(norm_err),
 53        "observed_z_norm_sq_max": max(norm_err), "z_norm_sq_mean": float(np.mean(norm_err)),
 54    }
 55
 56
 57def random_blocks(v, k, b, seed=0):
 58    rng = np.random.default_rng(seed)
 59    return np.asarray([rng.choice(v, k, replace=False) for _ in range(b)], dtype=np.int64)
 60
 61
 62def contiguous_blocks(v, k, b):
 63    return np.asarray([[(i + t) % v for t in range(k)] for i in range(b)], dtype=np.int64)
 64
 65
 66def torch_attention(x, blocks, heads=2):
 67    # One-head equivalent batched block attention, with scatter-and-average.
 68    n, d = x.shape
 69    q, k, val = x, x, x
 70    y = torch.zeros_like(x)
 71    counts = torch.zeros(n, device=x.device)
 72    scale = math.sqrt(d)
 73    for idx_np in blocks:
 74        idx = torch.as_tensor(idx_np, device=x.device, dtype=torch.long)
 75        s = (q[idx] @ k[idx].T) / scale
 76        a = F.softmax(s, dim=-1)
 77        y.index_add_(0, idx, a @ val[idx])
 78        counts.index_add_(0, idx, torch.ones(len(idx), device=x.device))
 79    return y / counts[:, None]
 80
 81
 82def dense_attention(x):
 83    return F.softmax(x @ x.T / math.sqrt(x.shape[1]), dim=-1) @ x
 84
 85
 86def benchmark(seed=123):
 87    random.seed(seed); np.random.seed(seed)
 88    if torch is None:
 89        return {"error": "torch unavailable"}
 90    try:
 91        device = "cuda" if torch.cuda.is_available() else "cpu"
 92        torch.manual_seed(seed)
 93        if device == "cuda": torch.cuda.manual_seed_all(seed)
 94        # Genuine balanced designs are benchmarked on their complete point sets.
 95        # This avoids incorrectly claiming that disjoint copies preserve global coverage.
 96        cases_by_n = {
 97            7: {"dense": None, "contiguous": contiguous_blocks(7, 3, 7),
 98                "random": random_blocks(7, 3, 7, seed), "design": make_designs()["Fano_7"]},
 99            13: {"dense": None, "contiguous": contiguous_blocks(13, 4, 13),
100                 "random": random_blocks(13, 4, 13, seed), "design": make_designs()["projective_13"]},
101        }
102        out = {"device": device, "results": {}}
103        for n, cases in cases_by_n.items():
104            d = 32; k = len(cases["design"][0]); x = torch.randn(n, d, device=device)
105            out["results"][str(n)] = {}
106            for name, blocks in cases.items():
107                if name == "dense":
108                    fn = lambda: dense_attention(x)
109                    flops = n*n*d*2; pair_var = 0.0; row_var = 0.0
110                else:
111                    fn = lambda blocks=blocks: torch_attention(x, blocks)
112                    flops = len(blocks)*k*k*d*2
113                    M = incidence(blocks, n); co = M @ M.T
114                    pair_var = float(np.var(co[~np.eye(n, dtype=bool)]))
115                    row_var = float(np.var(M.sum(1)))
116                fn()
117                if device == "cuda": torch.cuda.synchronize()
118                ts=[]
119                for _ in range(5):
120                    t=time.perf_counter(); fn()
121                    if device == "cuda": torch.cuda.synchronize()
122                    ts.append(time.perf_counter()-t)
123                out["results"][str(n)][name] = {
124                    "median_ms": 1000*float(np.median(ts)),
125                    "attention_score_flops": int(flops),
126                    "pair_coverage_variance": pair_var, "row_degree_variance": row_var}
127        return out
128    except Exception as e:
129        return {"error": repr(e)}
130
131
132def main():
133    designs = make_designs()
134    checks = [math_check(name, blocks) for name, blocks in designs.items()]
135    # Prediction sweep: for every exact design, all pair norms must equal 2(r-lambda).
136    # A second prediction is invariance under scaling the incidence matrix by s copies:
137    # r, lambda, and ||Z||^2 scale by s, while pair variance remains zero.
138    base = designs["Fano_7"]
139    sweep=[]
140    for copies in [1,2,4,8]:
141        b=np.tile(base, (copies,1))
142        c=math_check(f"Fano copies={copies}", b)
143        c["predicted_r"] = 3*copies
144        c["predicted_lambda"] = 1*copies
145        c["predicted_z_norm_sq"] = 2*(3-1)*copies
146        sweep.append(c)
147    report={"math_checks": checks, "scaling_sweep": sweep, "benchmark": benchmark()}
148    Path("results.json").write_text(json.dumps(report, indent=2))
149    print(json.dumps(report, indent=2))
150
151if __name__ == "__main__": main()