1import os, sys, json, random
 2import numpy as np
 3import torch
 4import torch.nn as nn
 5
 6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 7from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report
 8
 9SEEDS = tuple(range(8))
10GRID = [
11    {'lr': 0.003, 'epochs': 25, 'steps': 0},
12    {'lr': 0.006, 'epochs': 25, 'steps': 0},
13    {'lr': 0.012, 'epochs': 25, 'steps': 0},
14]
15IDEA_GRID = [
16    {'lr': 0.003, 'epochs': 25, 'steps': 1},
17    {'lr': 0.006, 'epochs': 25, 'steps': 2},
18    {'lr': 0.012, 'epochs': 25, 'steps': 3},
19]
20DT = 0.05
21
22def projection(raw, x, steps):
23    if steps <= 0:
24        return raw
25    q = x.view(x.shape[0], -1, 3)
26    th, om = q[:, -1, 0], q[:, -1, 1]
27    # Kinematic equilibrium F(z;u)=z-(theta+dt*omega), J_F=1.
28    target = (th + DT * om).unsqueeze(-1)
29    z = raw
30    for _ in range(steps):
31        z = z - (z - target)
32    return z
33
34class ProjectedModel(nn.Module):
35    def __init__(self, base, steps):
36        super().__init__(); self.base = base; self.steps = steps
37    def forward(self, x):
38        return projection(self.base(x), x, self.steps)
39
40def seed_all(seed):
41    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
42    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
43
44def run(cfg, idea):
45    def one(seed):
46        seed_all(seed)
47        d = get_dataset('dynamics', seed, n_train=400, n_test=160)
48        base = make_model('rnn_small', d['input_shape'], d['out_dim'])
49        net = ProjectedModel(base, cfg['steps']) if idea else base
50        _, metric, _ = train_model(net, d, epochs=cfg['epochs'], lr=cfg['lr'], batch=128, log=lambda *_: None)
51        return float(metric)
52    return one
53
54def mechanism_signature():
55    rows=[]
56    for seed in SEEDS:
57        seed_all(seed); d=get_dataset('dynamics', seed, n_train=400, n_test=160)
58        net=ProjectedModel(make_model('rnn_small', d['input_shape'], d['out_dim']), 2)
59        net, _, _=train_model(net, d, epochs=25, lr=0.006, batch=128, log=lambda *_: None)
60        try:
61            dev=next(net.parameters()).device
62            x=d['xte'].to(dev)
63            with torch.no_grad():
64                raw=net.base(x); out=net(x)
65                q=x.view(x.shape[0],-1,3)
66                target=(q[:,-1,0]+DT*q[:,-1,1]).unsqueeze(-1)
67                rr=(raw-target).abs().mean().item(); rp=(out-target).abs().mean().item()
68            rows.append((rr,rp))
69        except RuntimeError:
70            rows.append((float('nan'), float('nan')))
71    raw=np.array([r[0] for r in rows]); proj=np.array([r[1] for r in rows])
72    ratio=float(np.nanmean(proj)/(np.nanmean(raw)+1e-12))
73    return {'claim':'implicit equality projection removes the measured equilibrium residual',
74            'trained_model_raw_residual_mean':float(np.nanmean(raw)),
75            'trained_model_projected_residual_mean':float(np.nanmean(proj)),
76            'residual_ratio':ratio, 'confirmed': bool(ratio < 1e-3)}
77
78def main():
79    base=sweep_baseline(lambda c: run(c, False), GRID, seeds=(0,1,2,3))
80    idea_sweep=[]
81    for c in IDEA_GRID:
82        r=evaluate(run(c, True), seeds=(0,1,2,3))
83        idea_sweep.append({'cfg':c,'mean':r['mean']})
84    best=min(idea_sweep,key=lambda x:x['mean'])['cfg']
85    idea=evaluate(run(best, True), seeds=SEEDS)
86    rep=make_report('dynamics','rnn_small',base,idea,{
87        'track_match':'dynamics: controlled pendulum rollout',
88        'idea_config_sweep':idea_sweep, 'best_idea_cfg':best,
89        'mechanism':mechanism_signature()})
90    rep['protocol_notes']='8 paired seeds; baseline and idea share rnn_small, dataset sizes, lr/epoch union, and train_model.'
91    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
92    print(json.dumps(rep,indent=2))
93
94if __name__=='__main__': main()