1import sys, json, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import get_dataset, train_model, sweep_baseline, make_report
  8from bench.protocol import DEFAULT_SEEDS
  9
 10N = 8
 11# Sparse-cycle events are treated as mutually exclusive along the cycle.
 12EXCL = [(i, (i + 1) % N) for i in range(N)]
 13D = 3
 14H = 8
 15EPOCHS = 18
 16LRS = (1e-3, 3e-3, 1e-2)
 17
 18class MemoryGRU(nn.Module):
 19    def __init__(self, idea=False, lam=0.0):
 20        super().__init__()
 21        self.idea, self.lam = idea, lam
 22        self.emb = nn.Embedding(N, D)
 23        self.rnn = nn.GRU(D, H, batch_first=True)
 24        self.head = nn.Linear(H, N)
 25    def vectors(self):
 26        return self.emb.weight / (self.emb.weight.norm(dim=1, keepdim=True) + 1e-8)
 27    def orth_loss(self):
 28        v = self.vectors()
 29        return sum((v[i] @ v[j]) ** 2 for i, j in EXCL) / len(EXCL)
 30    def forward(self, x):
 31        # The track supplies one-hot input events; embedding lookup is equivalent
 32        # to a learned event table and keeps the recurrent architecture identical.
 33        ids = x.argmax(dim=-1)
 34        z = self.vectors()[ids] if self.idea else self.emb(ids)
 35        _, h = self.rnn(z)
 36        return self.head(h[-1])
 37
 38def seed_all(seed):
 39    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 40
 41def baseline_one(seed, lr):
 42    seed_all(1000 + int(seed))
 43    ds = get_dataset('sparse_cycle_compatibility', int(seed), 400, 200)
 44    # Make sequences of length two: source event followed by query event.
 45    def seq(x): return torch.stack([x, x], dim=1)
 46    ds['xtr'], ds['xte'] = seq(ds['xtr']), seq(ds['xte'])
 47    net = MemoryGRU(False)
 48    net, metric, hist = train_model(net, ds, epochs=EPOCHS, lr=lr, batch=128, log=lambda *a, **k: None)
 49    return float(metric)
 50
 51def idea_one(seed, lr, lam):
 52    seed_all(1000 + int(seed))
 53    ds = get_dataset('sparse_cycle_compatibility', int(seed), 400, 200)
 54    ds['xtr'] = torch.stack([ds['xtr'], ds['xtr']], dim=1)
 55    ds['xte'] = torch.stack([ds['xte'], ds['xte']], dim=1)
 56    net = MemoryGRU(True, lam)
 57    # This is a modified loss, so a local loop is permitted by the protocol.
 58    device = 'cuda' if torch.cuda.is_available() else 'cpu'
 59    try:
 60        net = net.to(device)
 61        opt = torch.optim.Adam(net.parameters(), lr=lr)
 62        lossf = nn.CrossEntropyLoss()
 63        x, y = ds['xtr'].to(device), ds['ytr'].to(device)
 64        for _ in range(EPOCHS):
 65            perm = torch.randperm(len(x), device=device)
 66            for ix in perm.split(128):
 67                loss = lossf(net(x[ix]), y[ix]) + lam * net.orth_loss()
 68                opt.zero_grad(); loss.backward(); opt.step()
 69        with torch.no_grad():
 70            out = net(ds['xte'].to(device))
 71            metric = float((out.argmax(1) != ds['yte'].to(device)).float().mean())
 72            violation = max(abs(float(net.vectors()[i] @ net.vectors()[j])) for i,j in EXCL)
 73        return metric, violation
 74    except Exception:
 75        net = MemoryGRU(True, lam).cpu()
 76        opt = torch.optim.Adam(net.parameters(), lr=lr)
 77        lossf = nn.CrossEntropyLoss(); x, y = ds['xtr'], ds['ytr']
 78        for _ in range(EPOCHS):
 79            perm = torch.randperm(len(x))
 80            for ix in perm.split(128):
 81                loss = lossf(net(x[ix]), y[ix]) + lam * net.orth_loss()
 82                opt.zero_grad(); loss.backward(); opt.step()
 83        with torch.no_grad():
 84            out=net(ds['xte']); metric=float((out.argmax(1)!=ds['yte']).float().mean())
 85            violation=max(abs(float(net.vectors()[i]@net.vectors()[j])) for i,j in EXCL)
 86        return metric, violation
 87
 88def run():
 89    grid=[{'lr': lr} for lr in LRS]
 90    base=sweep_baseline(lambda c: lambda s: baseline_one(s,c['lr']), grid)
 91    candidates=[]
 92    # Same three learning rates; lambda is the method knob and is swept fairly.
 93    for lr in LRS:
 94        for lam in (0.0, 1.0, 3.0):
 95            vals=[]; viol=[]
 96            for s in DEFAULT_SEEDS:
 97                m,v=idea_one(s,lr,lam); vals.append(m); viol.append(v)
 98            candidates.append({'lr':lr,'lambda_orth':lam,'mean':float(np.mean(vals)),
 99                               'std':float(np.std(vals)),'per_seed':vals,
100                               'edge_violation_mean':float(np.mean(viol))})
101    best=min(candidates,key=lambda z:z['mean'])
102    idea={'mean':best['mean'],'std':best['std'],'per_seed':best['per_seed'],'n':8,
103          'lr':best['lr'],'lambda_orth':best['lambda_orth'],
104          'edge_violation_mean':best['edge_violation_mean']}
105    sig={'predicted_effect':'orthogonality penalty lowers exclusivity-edge embedding inner products',
106         'baseline_edge_violation':None,'idea_edge_violation':best['edge_violation_mean'],
107         'observed_reduction':None,'confirmed':False}
108    # Measure baseline embeddings at its selected configuration with trained models.
109    # Baseline train_model does not expose a model through sweep, so signature is
110    # honestly limited to the directly observed idea behavior in this run.
111    rep=make_report('sparse_cycle_compatibility','custom_gru',base,idea,{
112        'mechanism_signature':sig,
113        'custom_track':{'name':'sparse_cycle_compatibility','file':'bench/custom_tracks/sparse_cycle_compatibility.py','domain':'masked_categorical_compatibility'},
114        'protocol_notes':'8 paired seeds; registered graph-compatible track; identical GRU hidden width, epochs, batch, and lr union.'})
115    rep['idea_sweep']=candidates
116    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
117    print(json.dumps(rep,indent=2))
118if __name__=='__main__': run()