1import json, time
 2import numpy as np
 3import torch
 4from torch import nn
 5from toposcan import grid_tokens, make_transformer
 6
 7SEED = 19
 8np.random.seed(SEED); torch.manual_seed(SEED)
 9DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
10try:
11    if DEVICE == "cuda": torch.cuda.get_device_name(0)
12except Exception:
13    DEVICE = "cpu"
14
15
16def cycle_graph(lengths):
17    edges=[]; off=0
18    for L in lengths:
19        for i in range(L): edges.append((off+i, off+(i+1)%L))
20        off += L
21    return off, edges
22
23
24def make_data(n_each=160):
25    # Every graph has 12 vertices and every vertex degree 2.
26    # Label 0: one 12-cycle (beta1=1); label 1: two 6-cycles (beta1=2).
27    graphs=[]
28    for y, lengths in [(0,[12]), (1,[6,6])]:
29        for _ in range(n_each):
30            n,e=cycle_graph(lengths)
31            # identical filtration keeps topology active in one shared window
32            h=np.full(n, .5, dtype=np.float32)
33            graphs.append((n,e,h,y))
34    rng=np.random.default_rng(SEED); rng.shuffle(graphs)
35    return graphs
36
37
38def features(graphs):
39    degree=[]; topo=[]; labels=[]; prep=0.
40    grid=np.linspace(0.,1.,9)
41    for n,e,h,y in graphs:
42        d=np.zeros(n, dtype=np.float32)
43        for u,v in e: d[u]+=1; d[v]+=1
44        degree.append(np.bincount(d.astype(int), minlength=4).astype(np.float32)/n)
45        t0=time.perf_counter()
46        tok=grid_tokens(n,e,h,grid,m=4,stride=2)
47        prep += time.perf_counter()-t0
48        topo.append(tok); labels.append(y)
49    return np.asarray(degree), topo, np.asarray(labels), prep
50
51
52def train_mlp(xtr,ytr,xte,yte, epochs=100):
53    model=nn.Sequential(nn.Linear(xtr.shape[1],16),nn.ReLU(),nn.Linear(16,2)).to(DEVICE)
54    opt=torch.optim.Adam(model.parameters(),lr=.03)
55    X=torch.tensor(xtr,dtype=torch.float32,device=DEVICE); Y=torch.tensor(ytr,device=DEVICE)
56    for _ in range(epochs):
57        opt.zero_grad(); loss=nn.functional.cross_entropy(model(X),Y); loss.backward(); opt.step()
58    with torch.no_grad():
59        pred=model(torch.tensor(xte,dtype=torch.float32,device=DEVICE)).argmax(1).cpu().numpy()
60    return float(np.mean(pred==yte)), sum(p.numel() for p in model.parameters())
61
62
63def train_topo(ttr,ytr,tte,yte,epochs=100):
64    # Batch padding is unnecessary here: all token sequences have the same length.
65    model=nn.Sequential(make_transformer(width=16,heads=4,layers=1),nn.Linear(16,2)).to(DEVICE)
66    opt=torch.optim.Adam(model.parameters(),lr=.02)
67    X=torch.tensor(np.asarray(ttr),dtype=torch.float32,device=DEVICE); Y=torch.tensor(ytr,device=DEVICE)
68    for _ in range(epochs):
69        opt.zero_grad(); loss=nn.functional.cross_entropy(model(X),Y); loss.backward(); opt.step()
70    with torch.no_grad():
71        pred=model(torch.tensor(np.asarray(tte),dtype=torch.float32,device=DEVICE)).argmax(1).cpu().numpy()
72    return float(np.mean(pred==yte)), sum(p.numel() for p in model.parameters())
73
74
75def main():
76    data=make_data(); cut=round(.7*len(data))
77    tr,te=data[:cut],data[cut:]
78    dtr,ttr,ytr,p1=features(tr); dte,tte,yte,p2=features(te)
79    # perturbation test: filtration shifts do not change the relative topology here;
80    # use a separate score perturbation and recompute tokens.
81    acc_d,par_d=train_mlp(dtr,ytr,dte,yte)
82    acc_t,par_t=train_topo(ttr,ytr,tte,yte)
83    # A cheap direct topological classifier isolates representation power.
84    ztr=np.asarray([a[:,1].max() for a in ttr])[:,None]
85    zte=np.asarray([a[:,1].max() for a in tte])[:,None]
86    acc_direct,par_direct=train_mlp(ztr,ytr,zte,yte)
87    results={"device":DEVICE,"graphs":len(data),"degree_hist_unique":int(np.unique(np.concatenate([dtr,dte]),axis=0).shape[0]),
88             "baseline_degree_accuracy":acc_d,"toposcan_transformer_accuracy":acc_t,
89             "direct_beta1_accuracy":acc_direct,"baseline_params":par_d,"toposcan_params":par_t,
90             "preprocess_seconds_total":p1+p2,"preprocess_seconds_per_graph":(p1+p2)/len(data),
91             "note":"All graphs have identical degree histogram; topology labels are beta1=1 vs 2."}
92    print(json.dumps(results,indent=2))
93
94if __name__=='__main__': main()