1import json, random, time
  2from itertools import combinations
  3import numpy as np
  4
  5# Explicit orthogonality graph: vertices are representatives of nonzero {-1,0,1}^3 vectors.
  6V = np.array([(0,0,1),(0,1,-1),(0,1,0),(0,1,1),(1,-1,-1),(1,-1,0),
  7              (1,-1,1),(1,0,-1),(1,0,0),(1,0,1),(1,1,-1),(1,1,0),(1,1,1)], dtype=np.float32)
  8N = len(V)
  9E = [(i,j) for i in range(N) for j in range(i) if np.dot(V[i], V[j]) == 0]
 10
 11# Small exact backtracking coloring check.
 12def colorable(k):
 13    adj = [set() for _ in range(N)]
 14    for i,j in E: adj[i].add(j); adj[j].add(i)
 15    c = [-1]*N
 16    def go(done):
 17        if done == N: return True
 18        u = max((i for i in range(N) if c[i] < 0), key=lambda i: len(adj[i]))
 19        for col in range(k):
 20            if all(c[w] != col for w in adj[u]):
 21                c[u] = col
 22                if go(done+1): return True
 23                c[u] = -1
 24        return False
 25    return go(0)
 26
 27def graph_check():
 28    gram = V @ V.T
 29    edge_vals = np.array([gram[i,j] for i,j in E])
 30    # Unit-normalized vectors are the orthogonal-rank witness.
 31    U = V / np.linalg.norm(V, axis=1, keepdims=True)
 32    ug = U @ U.T
 33    violation = max(abs(ug[i,j]) for i,j in E)
 34    clique = max((len(C) for r in range(1,N+1) for C in combinations(range(N),r)
 35                  if all(((a,b) in E or (b,a) in E) for a,b in combinations(C,2))), default=0)
 36    return {'vertices':N, 'edges':len(E), 'explicit_dimension':3,
 37            'edge_dot_max_raw':float(max(abs(edge_vals))),
 38            'edge_violation_normalized':float(violation),
 39            'clique_lower_bound_xi':clique,
 40            'colorable_3':colorable(3), 'colorable_4':colorable(4)}
 41
 42# Dataset: retain the first event through distractors and decide whether the final
 43# event is exclusive with it. This is a graph-defined relational language.
 44def make_data(num, T, seed):
 45    rng = np.random.default_rng(seed)
 46    # A learnable graph-walk-style promise language: the first event is the
 47    # context, fixed neutral distractors follow, and the final event is queried.
 48    # The target remains exactly edge compatibility with the initial event.
 49    x = np.zeros((num,T), dtype=np.int64)
 50    x[:,0] = rng.integers(0,N,size=num)
 51    x[:,-1] = rng.integers(0,N,size=num)
 52    edge = np.zeros((N,N), dtype=np.float32)
 53    for i,j in E: edge[i,j]=edge[j,i]=1
 54    y = 1.0 - edge[x[:,0], x[:,-1]]
 55    return x, y.astype(np.float32)
 56
 57
 58import torch
 59import torch.nn as nn
 60
 61def train_one(name, width, orth_lambda, seed, device):
 62    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
 63    class Recognizer(nn.Module):
 64        def __init__(self):
 65            super().__init__()
 66            self.emb = nn.Embedding(N,width)
 67            self.gru = nn.GRU(width,width,batch_first=True)
 68            self.out = nn.Linear(width,1)
 69        def forward(self,x):
 70            h,_ = self.gru(self.emb(x))
 71            return self.out(h[:,-1]).squeeze(-1)
 72        def orth_loss(self):
 73            u = self.emb.weight/(self.emb.weight.norm(dim=1,keepdim=True)+1e-8)
 74            dots = (u @ u.T)
 75            return sum(dots[i,j]**2 for i,j in E) / len(E)
 76    try:
 77        model=Recognizer().to(device)
 78        opt=torch.optim.Adam(model.parameters(),lr=3e-3)
 79        xb,yb=make_data(4000,16,seed+100)
 80        xt,yt=make_data(2000,64,seed+200)
 81        xb=torch.tensor(xb,device=device); yb=torch.tensor(yb,device=device)
 82        xt=torch.tensor(xt,device=device); yt=torch.tensor(yt,device=device)
 83        loss_fn=nn.BCEWithLogitsLoss()
 84        t0=time.time()
 85        for epoch in range(90):
 86            perm=torch.randperm(len(xb),device=device)
 87            for ix in perm.split(128):
 88                loss=loss_fn(model(xb[ix]),yb[ix])
 89                if orth_lambda: loss=loss+orth_lambda*model.orth_loss()
 90                opt.zero_grad(); loss.backward(); opt.step()
 91        with torch.no_grad():
 92            model.eval()
 93            def score(x,y):
 94                p=(torch.sigmoid(model(x))>.5).float()
 95                return float((p==y).float().mean()), float(loss_fn(model(x),y))
 96            a,l=score(xt,yt); tr,_=score(xb,yb)
 97            u=model.emb.weight/(model.emb.weight.norm(dim=1,keepdim=True)+1e-8)
 98            gv=max(abs(float((u[i]@u[j]).cpu())) for i,j in E)
 99        return {'name':name,'width':width,'train_acc':tr,'long_acc':a,'long_loss':l,
100                'edge_violation':gv,'seconds':time.time()-t0,'params':sum(p.numel() for p in model.parameters())}
101    except Exception as e:
102        raise RuntimeError(f'{name}: {e}') from e
103
104def main():
105    check=graph_check()
106    # CUDA is attempted but any runtime/device problem falls back to CPU.
107    device='cuda' if torch.cuda.is_available() else 'cpu'
108    try:
109        if device=='cuda': torch.zeros(1,device='cuda')
110    except Exception: device='cpu'
111    results=[]
112    # Standard width-4 GRU (classical color lower bound), width-3 controls,
113    # and width-3 with the proposed edge orthogonality penalty. If CUDA has
114    # any allocation/cuDNN failure, rerun the complete matched comparison CPU.
115    configs=[('gru_width_chi',4,0.0),('gru_width_xi',3,0.0),('orthogonal_memory',3,3.0)]
116    try:
117        for args in configs: results.append(train_one(*args,seed=7,device=device))
118    except Exception as exc:
119        if device != 'cpu':
120            device='cpu'; results=[]
121            for args in configs: results.append(train_one(*args,seed=7,device=device))
122        else: raise
123    out={'device':device,'graph_check':check,'results':results,
124         'note':'xi is only upper-bounded by the explicit 3D representation; clique size supplies a matching lower bound here.'}
125    with open('results.json','w') as f: json.dump(out,f,indent=2)
126    print(json.dumps(out,indent=2))
127if __name__=='__main__': main()