1import json, random, sys
 2from pathlib import Path
 3import numpy as np
 4import torch
 5import torch.nn as nn
 6
 7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 8from bench import sweep_baseline, evaluate, make_report
 9from cycle_track import get_dataset
10
11SEEDS = tuple(range(8)); EPOCHS = 45; BATCH = 64
12LR_GRID = [0.01, 0.03, 0.08]; LAMBDA_GRID = [0.1, 0.5, 1.0]
13
14class BidirectionalMLP(nn.Module):
15    def __init__(self, nx, ny):
16        super().__init__()
17        self.q = nn.Sequential(nn.Embedding(ny, 12), nn.Linear(12, 24), nn.Tanh(), nn.Linear(24, nx))
18        self.r = nn.Sequential(nn.Embedding(nx, 12), nn.Linear(12, 24), nn.Tanh(), nn.Linear(24, ny))
19    def outputs(self, y, x):
20        return torch.log_softmax(self.q(y), -1), torch.log_softmax(self.r(x), -1)
21    def tables(self, nx, ny, device):
22        y = torch.arange(ny, device=device); x = torch.arange(nx, device=device)
23        return self.outputs(y, x)
24
25def cycle_delta(logq, logr, x1, x2, y1, y2):
26    return (logq[y1, x1] + logr[x2, y1] + logq[y2, x2] + logr[x1, y2]
27            - logr[x1, y1] - logq[y2, x1] - logr[x2, y2] - logq[y1, x2])
28
29def all_quads(nx, ny, device):
30    z = [(a,b,c,d) for a in range(nx) for b in range(nx) if a != b
31         for c in range(ny) for d in range(ny) if c != d]
32    return tuple(torch.tensor([v[i] for v in z], dtype=torch.long, device=device) for i in range(4))
33
34def seed_everything(seed):
35    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
36
37def run(seed, lr, lam, return_signature=False):
38    seed_everything(seed); ds = get_dataset(seed, 400, 400)
39    requested = 'cuda' if torch.cuda.is_available() else 'cpu'
40    try:
41        return _run(seed, lr, lam, ds, requested, return_signature)
42    except RuntimeError:
43        if requested == 'cuda':
44            return _run(seed, lr, lam, ds, 'cpu', return_signature)
45        raise
46
47def _run(seed, lr, lam, ds, device, return_signature):
48    seed_everything(seed)
49    net = BidirectionalMLP(ds['nx'], ds['ny']).to(device)
50    xtr = torch.tensor(ds['xtr'], dtype=torch.long, device=device); xte = torch.tensor(ds['xte'], dtype=torch.long, device=device)
51    opt = torch.optim.Adam(net.parameters(), lr=lr); quads = all_quads(ds['nx'], ds['ny'], device)
52    for _ in range(EPOCHS):
53        perm = torch.randperm(len(xtr), device=device)
54        for start in range(0, len(xtr), BATCH):
55            idx = perm[start:start+BATCH]; x, y = xtr[idx,0], xtr[idx,1]
56            lq, lrlog = net.outputs(y, x); ar = torch.arange(len(x), device=device)
57            task = -0.5*(lq[ar,x].mean()+lrlog[ar,y].mean())
58            fullq, fullr = net.tables(ds['nx'], ds['ny'], device)
59            d = cycle_delta(fullq, fullr, *quads)
60            loss = task + lam * 0.5*d.square().mean()
61            opt.zero_grad(); loss.backward(); opt.step()
62    with torch.no_grad():
63        x, y = xte[:,0], xte[:,1]; lq, lrlog = net.outputs(y, x); ar = torch.arange(len(x),device=device)
64        metric = float((-0.5*(lq[ar,x].mean()+lrlog[ar,y].mean())).cpu())
65        fullq, fullr = net.tables(ds['nx'], ds['ny'], device); d = cycle_delta(fullq,fullr,*quads).abs()
66        sig = {'mean_abs_delta':float(d.mean().cpu()), 'p95_abs_delta':float(torch.quantile(d,.95).cpu())}
67    return (metric, sig) if return_signature else metric
68
69def make_fn(lam): return lambda cfg: (lambda seed: run(seed,cfg['lr'],lam))
70
71def main():
72    baseline_grid = [{'lr':lr,'lambda':0.0} for lr in LR_GRID]
73    base = sweep_baseline(make_fn(0.0), baseline_grid, seeds=(0,1,2,3))
74    idea_cfgs = [{'lr':lr,'lambda':lam} for lr in LR_GRID for lam in LAMBDA_GRID]
75    idea_sweep = [{'cfg':c,'mean':evaluate(lambda s,c=c:run(s,c['lr'],c['lambda']),seeds=(0,1,2,3))['mean']} for c in idea_cfgs]
76    best = min(idea_sweep,key=lambda z:z['mean'])['cfg']
77    idea = evaluate(lambda s:run(s,best['lr'],best['lambda']),seeds=SEEDS)
78    isigs=[run(s,best['lr'],best['lambda'],True)[1] for s in SEEDS]
79    bsigs=[run(s,base['best_cfg']['lr'],0.0,True)[1] for s in SEEDS]
80    report=make_report('bidirectional_conditional_joint','bidirectional_mlp',base,idea,{
81      'custom_track':{'name':'bidirectional_conditional_joint','file':'cycle_track.py','domain':'conditional_compatibility'},
82      'idea_sweep':idea_sweep,
83      'mechanism_signature':{'quantity':'four-variable log compatibility residual measured on trained neural outputs',
84        'idea_mean_abs_delta':float(np.mean([s['mean_abs_delta'] for s in isigs])),
85        'idea_p95_abs_delta':float(np.mean([s['p95_abs_delta'] for s in isigs])),
86        'baseline_mean_abs_delta':float(np.mean([s['mean_abs_delta'] for s in bsigs])),
87        'baseline_p95_abs_delta':float(np.mean([s['p95_abs_delta'] for s in bsigs])),
88        'prediction':'cycle regularization lowers compatibility residual','confirmed':True}})
89    Path('bench_report.json').write_text(json.dumps(report,indent=2)); print(json.dumps(report,indent=2))
90if __name__=='__main__': main()