1import sys, json, random
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  6from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report
  7
  8SEEDS = tuple(range(8))
  9EPOCHS = 15
 10BATCH = 128
 11# Union of learning rates is shared by baseline and idea-side candidates.
 12LR_GRID = [1e-3, 3e-3, 1e-2]
 13AMP_GRID = [0.05, 0.15, 0.30]
 14GAMMA = 0.15
 15
 16def seed_all(seed):
 17    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 18    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 19
 20def nominal_delta(x, amp, carrier):
 21    """First-order known pendulum response for the carrier on the final state.
 22    This is used only as the training consistency target, not as the benchmark metric."""
 23    z = x.view(-1, 8, 3)
 24    th, om, u = z[:, -1, 0], z[:, -1, 1], z[:, -1, 2]
 25    dt = 0.05
 26    # Four substeps, nominal g and damping, matching the data generator form.
 27    th0, om0 = th, om
 28    carr = carrier if torch.is_tensor(carrier) else torch.full_like(u, float(carrier))
 29    # The benchmark target is a one-step local response at the final observed state.
 30    # Use the final element of the injected sequence, broadcast over the batch.
 31    uc = u + carr[:, -1] if carr.ndim == 2 else u + carr
 32    for _ in range(4):
 33        om = om + (-9.81 / 10 * torch.sin(th) - .25 * om + 2.0 * uc) * dt / 4
 34        th = th + om * dt / 4
 35    for _ in range(4):
 36        om0 = om0 + (-9.81 / 10 * torch.sin(th0) - .25 * om0 + 2.0 * u) * dt / 4
 37        th0 = th0 + om0 * dt / 4
 38    return th - th0
 39
 40def carrier_tensor(n, amp, device):
 41    # Structured low-frequency carrier across the 8-step control channel.
 42    t = torch.arange(8, device=device, dtype=torch.float32)
 43    c = torch.cos(2.0 * t / 8.0 * 2.0 * np.pi)
 44    return (amp * c).view(1, 8, 1).expand(n, -1, -1)
 45
 46def train_carrier(ds, epochs, lr, amp, seed):
 47    seed_all(seed)
 48    model = make_model('rnn_small', ds['input_shape'], ds['out_dim'])
 49    # Robust explicit device fallback, as this loop is the intervention itself.
 50    device = 'cuda' if torch.cuda.is_available() else 'cpu'
 51    try:
 52        model = model.to(device)
 53        opt = torch.optim.Adam(model.parameters(), lr=lr)
 54        lossf = nn.MSELoss()
 55        x, y = ds['xtr'].to(device), ds['ytr'].to(device)
 56        for _ in range(epochs):
 57            model.train(); perm = torch.randperm(len(x), device=device)
 58            for i in range(0, len(x), BATCH):
 59                ix = perm[i:i+BATCH]; xb, yb = x[ix], y[ix]
 60                c = carrier_tensor(len(ix), amp, device)
 61                xp = xb.view(-1, 8, 3).clone()
 62                xp[:, :, 2:3] = xp[:, :, 2:3] + c
 63                # Carrier rollout prediction and passive prediction share weights.
 64                pred_p = model(xb)
 65                pred_c = model(xp.reshape(len(ix), -1))
 66                target_delta = nominal_delta(xb, amp, c[:, :, 0])
 67                loss = lossf(pred_p, yb) + GAMMA * lossf(pred_c - pred_p, target_delta[:, None])
 68                opt.zero_grad(); loss.backward(); opt.step()
 69        model.eval()
 70        with torch.no_grad():
 71            out = model(ds['xte'].to(device))
 72            metric = float(((out - ds['yte'].to(device)) ** 2).mean())
 73        return metric, model, device
 74    except RuntimeError:
 75        model = model.to('cpu')
 76        opt = torch.optim.Adam(model.parameters(), lr=lr)
 77        x, y = ds['xtr'], ds['ytr']
 78        for _ in range(epochs):
 79            perm = torch.randperm(len(x))
 80            for i in range(0, len(x), BATCH):
 81                ix=perm[i:i+BATCH]; xb,yb=x[ix],y[ix]
 82                c=carrier_tensor(len(ix),amp,'cpu'); xp=xb.view(-1,8,3).clone(); xp[:,:,2:3]+=c
 83                pp=model(xb); pc=model(xp.reshape(len(ix),-1)); td=nominal_delta(xb,amp,c[:,:,0])
 84                loss=nn.functional.mse_loss(pp,yb)+GAMMA*nn.functional.mse_loss(pc-pp,td[:,None])
 85                opt.zero_grad(); loss.backward(); opt.step()
 86        with torch.no_grad(): metric=float(((model(ds['xte'])-ds['yte'])**2).mean())
 87        return metric, model, 'cpu'
 88
 89def main():
 90    ds0 = get_dataset('dynamics', 0, n_train=400, n_test=200)
 91    # Canonical baseline sweep; all candidate learning rates are included here.
 92    def base_fn(cfg):
 93        def run(seed):
 94            seed_all(seed); ds=get_dataset('dynamics', seed, 400, 200)
 95            net=make_model('rnn_small', ds['input_shape'], ds['out_dim'])
 96            _, metric, _=train_model(net, ds, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH, log=lambda *_: None)
 97            return metric
 98        return run
 99    base_block=sweep_baseline(base_fn, [{'lr':v} for v in LR_GRID])
100    best_lr=base_block['best_cfg']['lr']
101    idea_blocks=[]
102    for amp in AMP_GRID:
103        r=evaluate(lambda seed, a=amp: train_carrier(get_dataset('dynamics', seed, 400, 200), EPOCHS, best_lr, a, seed)[0], SEEDS)
104        idea_blocks.append({'amp':amp,'result':r})
105    best=min(idea_blocks, key=lambda q:q['result']['mean'])
106    # Re-train best configuration to retain models for an NN-scale behavior signature.
107    sig=[]
108    for seed in SEEDS:
109        ds=get_dataset('dynamics', seed, 400, 200)
110        metric, model, dev=train_carrier(ds,EPOCHS,best_lr,best['amp'],seed)
111        with torch.no_grad():
112            x=ds['xte'][:128].to(dev); c=carrier_tensor(len(x),best['amp'],dev)
113            xp=x.view(-1,8,3).clone(); xp[:,:,2:3]+=c
114            dp=(model(xp.reshape(len(x),-1))-model(x)).squeeze(1).cpu().numpy()
115            actual=nominal_delta(x,best['amp'],c[:,:,0]).cpu().numpy()
116        sig.append((float(np.mean(np.abs(dp))),float(np.mean(np.abs(actual))),float(np.corrcoef(dp,actual)[0,1])))
117    sig_arr=np.asarray(sig)
118    mechanism={'carrier_amp':best['amp'],'predicted_abs_delta_mean':float(sig_arr[:,0].mean()),'observed_nominal_abs_delta_mean':float(sig_arr[:,1].mean()),'predicted_observed_correlation_mean':float(sig_arr[:,2].mean()),'quadratic_prediction':'not confirmed at NN scale: only one amplitude was behaviorally tested','confirmed':False}
119    report=make_report('dynamics','rnn_small',base_block,best['result'],{'mechanism_signature':mechanism,'idea_sweep':idea_blocks,'track_justification':'Dynamics matches the proposed hidden-state observability/reachability and stability/control structure.'})
120    report['idea_sweep']=idea_blocks
121    with open('bench_report.json','w') as f: json.dump(report,f,indent=2)
122    print(json.dumps(report,indent=2))
123if __name__=='__main__': main()