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))
 10SWEEP_SEEDS = tuple(range(4))
 11EPOCHS = 18
 12BATCH = 128
 13# Union is shared: baseline is evaluated at every lr considered by the idea.
 14LRS = [1e-3, 2e-3, 3e-3]
 15
 16
 17def seed_all(seed):
 18    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 19    if torch.cuda.is_available():
 20        try: torch.cuda.manual_seed_all(seed)
 21        except Exception: pass
 22
 23
 24class VectorCell(nn.Module):
 25    """Unconstrained recurrent vector field, used only as a matched reference."""
 26    def __init__(self, hidden=64):
 27        super().__init__()
 28        self.inp = nn.Linear(3, hidden)
 29        self.field = nn.Sequential(nn.Linear(hidden, hidden), nn.Tanh(),
 30                                   nn.Linear(hidden, hidden))
 31        self.head = nn.Linear(hidden, 1)
 32
 33    def forward(self, x):
 34        seq = x.view(x.shape[0], -1, 3)
 35        z = torch.tanh(self.inp(seq[:, 0]))
 36        for k in range(seq.shape[1]):
 37            u = torch.tanh(self.inp(seq[:, k]))
 38            z = z + 0.12 * self.field(z + u)
 39        return self.head(z)
 40
 41
 42class EnergyCell(nn.Module):
 43    """Same recurrent scaffold, but the state update is -grad_z E(z,input)."""
 44    def __init__(self, hidden=64, mu=0.03):
 45        super().__init__()
 46        self.inp = nn.Linear(3, hidden)
 47        self.energy_net = nn.Sequential(nn.Linear(hidden, hidden), nn.Tanh(),
 48                                        nn.Linear(hidden, 1))
 49        self.head = nn.Linear(hidden, 1)
 50        self.mu = mu
 51
 52    def energy(self, z, u):
 53        # Conditioning is fixed during each state update; mu gives coercivity.
 54        return self.energy_net(torch.tanh(z + u)).squeeze(-1) + self.mu * (z*z).sum(-1) / 2
 55
 56    def forward(self, x, collect=False):
 57        seq = x.view(x.shape[0], -1, 3)
 58        z = torch.tanh(self.inp(seq[:, 0]))
 59        energies, grads, states = [], [], []
 60        for k in range(seq.shape[1]):
 61            u = torch.tanh(self.inp(seq[:, k]))
 62            z = z.requires_grad_(True)
 63            e = self.energy(z, u)
 64            g = torch.autograd.grad(e.sum(), z, create_graph=True)[0]
 65            z = z - 0.12 * g
 66            if collect:
 67                energies.append(e.detach()); grads.append(g.detach()); states.append(z.detach())
 68        out = self.head(z)
 69        if collect:
 70            return out, energies, grads, states
 71        return out
 72
 73
 74def train_energy(ds, epochs, lr, seed, collect=False):
 75    seed_all(seed)
 76    model = EnergyCell().to('cuda' if torch.cuda.is_available() else 'cpu')
 77    # Explicit fallback mirrors train_model's robust CUDA->CPU behavior.
 78    devices = ['cuda', 'cpu'] if torch.cuda.is_available() else ['cpu']
 79    last = None
 80    for device in devices:
 81        try:
 82            model = model.to(device)
 83            x, y = ds['xtr'].to(device), ds['ytr'].to(device)
 84            opt = torch.optim.Adam(model.parameters(), lr=lr)
 85            for _ in range(epochs):
 86                model.train(); perm = torch.randperm(len(x), device=device)
 87                for i in range(0, len(x), BATCH):
 88                    idx = perm[i:i+BATCH]
 89                    loss = ((model(x[idx]) - y[idx]) ** 2).mean()
 90                    opt.zero_grad(); loss.backward(); opt.step()
 91            model.eval()
 92            with torch.enable_grad():
 93                pred = model(ds['xte'].to(device))
 94                metric = float(((pred - ds['yte'].to(device)) ** 2).mean().detach().cpu())
 95            return metric, model, device
 96        except RuntimeError as exc:
 97            last = exc
 98            if device == 'cuda':
 99                torch.cuda.empty_cache()
100                continue
101            raise
102    raise last
103
104
105def idea_metric(cfg, seed, keep=False):
106    ds = get_dataset('dynamics', seed)
107    metric, model, device = train_energy(ds, EPOCHS, cfg['lr'], seed, keep)
108    if keep: return metric, model, device, ds
109    return metric
110
111
112def signature(model, ds, device):
113    model.eval(); x = ds['xte'][:64].to(device)
114    with torch.enable_grad():
115        _, es, gs, zs = model(x, collect=True)
116    # Re-test the learned model's actual discrete dissipation, not an analytic toy.
117    e = torch.stack(es, 1).mean(0).detach().cpu().numpy()
118    gn = torch.stack([g.norm(dim=1) for g in gs], 1).mean(0).detach().cpu().numpy()
119    diffs = np.diff(e)
120    return {'predicted': 'energy should not increase under sufficiently small Euler steps',
121            'observed_energy_nonincreasing_fraction': float(np.mean(diffs <= 1e-7)),
122            'observed_energy_first_last': [float(e[0]), float(e[-1])],
123            'observed_grad_norm_first_last': [float(gn[0]), float(gn[-1])],
124            'predicted_boundary_eta_L': 2.0,
125            'observed_local_energy_boundary': 'not estimated (learned Hessian unavailable in budget)',
126            'confirmed': bool(np.all(diffs <= 1e-7) and gn[-1] < gn[0])}
127
128
129def main():
130    # Baseline sweep uses the standard bench train_model and the full lr union.
131    def base_fn(cfg):
132        def run(seed):
133            seed_all(seed); ds = get_dataset('dynamics', seed)
134            net = VectorCell(hidden=64).to('cpu')
135            _, metric, _ = train_model(net, ds, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH, log=lambda *_: None)
136            return metric
137        return run
138    base = sweep_baseline(base_fn, [{'lr': v} for v in LRS], seeds=SWEEP_SEEDS)
139    # Evaluate each idea setting on the same tuning seeds, then full 8 for best.
140    idea_sweep = []
141    for lr in LRS:
142        r = evaluate(lambda s, lr=lr: idea_metric({'lr': lr}, s), seeds=SWEEP_SEEDS)
143        idea_sweep.append({'cfg': {'lr': lr}, 'mean': r['mean']})
144    best_lr = min(idea_sweep, key=lambda q: q['mean'])['cfg']['lr']
145    idea = evaluate(lambda s: idea_metric({'lr': best_lr}, s), seeds=SEEDS)
146    # Signature from one of the trained benchmark models, seed 0.
147    _, trained, dev, ds0 = idea_metric({'lr': best_lr}, 0, keep=True)
148    sig = signature(trained, ds0, dev)
149    report = make_report('dynamics', 'rnn_small',
150        {'best_cfg': base['best_cfg'], 'sweep': base['sweep'], 'full': base['full']},
151        idea, {'idea_sweep': idea_sweep, **sig})
152    report['protocol_notes'] = {'track_choice': 'dynamics matches stability/control/Lyapunov structure',
153        'shared_lr_union': LRS, 'epochs': EPOCHS, 'batch': BATCH,
154        'baseline_architecture': 'matched VectorCell recurrent scaffold (vector field)',
155        'idea_architecture': 'matched recurrent hidden width with scalar energy gradient'}
156    with open('bench_report.json', 'w') as f: json.dump(report, f, indent=2)
157    print(json.dumps(report, indent=2))
158
159if __name__ == '__main__': main()