1import os, sys, json, math
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5import torch.nn.functional as F
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import get_dataset, train_model, make_report, sweep_baseline
  8
  9DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 10EPOCHS = 20
 11NTR, NTE = 400, 200
 12
 13class BaseDynamics(nn.Module):
 14    def __init__(self, idea=False, hidden=64, dt=0.05, channels=2):
 15        super().__init__()
 16        self.idea = idea; self.dt = dt
 17        self.rnn = nn.GRU(3, hidden, batch_first=True)
 18        self.head = nn.Linear(hidden, 1)
 19        # Equal-sized projection used by both systems; transport has its
 20        # mechanism heads, while the baseline uses an unconstrained residual.
 21        self.mix = nn.Linear(hidden, hidden)
 22        if idea:
 23            self.rho_head = nn.Linear(hidden, channels)
 24            self.pi_head = nn.Linear(hidden, channels)
 25            self.edge_head = nn.Linear(3 * hidden, channels)
 26            self.transport_proj = nn.Linear(channels, hidden)
 27
 28    def latent(self, x):
 29        seq = x.view(x.shape[0], -1, 3)
 30        try:
 31            hs, _ = self.rnn(seq)
 32        except RuntimeError:
 33            old = torch.backends.cudnn.enabled; torch.backends.cudnn.enabled = False
 34            try: hs, _ = self.rnn(seq)
 35            finally: torch.backends.cudnn.enabled = old
 36        return hs
 37
 38    def forward(self, x):
 39        h = self.latent(x)
 40        if not self.idea:
 41            # Ordinary unconstrained residual graph mixing on the 8 time nodes.
 42            z = self.mix(h)
 43            h = h + self.dt * (z.mean(1, keepdim=True) - z)
 44        else:
 45            rho = F.softplus(self.rho_head(h)) + 1e-5
 46            pi = F.softplus(self.pi_head(h)) + 1e-5
 47            pi = pi / pi.sum(1, keepdim=True) * rho.sum(1, keepdim=True).detach()
 48            # Complete graph, symmetric feature-dependent conductance.
 49            a = h.unsqueeze(2).expand(-1,-1,h.size(1),-1)
 50            b = h.unsqueeze(1).expand(-1,h.size(1),-1,-1)
 51            e = torch.cat([a,b,(a-b).abs()], -1)
 52            cij = F.softplus(self.edge_head(e)) + 1e-5
 53            cij = (cij + cij.transpose(1,2)) / 2
 54            q = rho / pi
 55            dr = (cij * (q.unsqueeze(1) - q.unsqueeze(2))).sum(2)
 56            rho_new = rho + self.dt * dr
 57            # dt is deliberately small relative to the learned positive rates.
 58            h = h + self.transport_proj(torch.log(rho_new + 1e-6))
 59        return self.head(h[:, -1])
 60
 61    @torch.no_grad()
 62    def signature(self, x):
 63        if not self.idea: return None
 64        h = self.latent(x)
 65        rho = F.softplus(self.rho_head(h)) + 1e-5
 66        pi = F.softplus(self.pi_head(h)) + 1e-5
 67        pi = pi / pi.sum(1, keepdim=True) * rho.sum(1, keepdim=True)
 68        a = h.unsqueeze(2).expand(-1,-1,h.size(1),-1)
 69        b = h.unsqueeze(1).expand(-1,h.size(1),-1,-1)
 70        cij = F.softplus(self.edge_head(torch.cat([a,b,(a-b).abs()],-1)))+1e-5
 71        cij=(cij+cij.transpose(1,2))/2
 72        q=rho/pi; dr=(cij*(q.unsqueeze(1)-q.unsqueeze(2))).sum(2)
 73        rn=rho+self.dt*dr
 74        mass_err=((rn.sum(1)-rho.sum(1)).abs()/(rho.sum(1)+1e-8)).mean().item()
 75        def energy(r): return (r*torch.log(r/pi)).sum((1,2))
 76        dF=(energy(rn)-energy(rho)).mean().item()
 77        mu=torch.log(q)
 78        qa,qb=q.unsqueeze(1),q.unsqueeze(2)
 79        mua,mub=mu.unsqueeze(1),mu.unsqueeze(2)
 80        lm=torch.where((mua-mub).abs()<1e-7, (qa+qb)/2, (qa-qb)/(mua-mub))
 81        pred=-(cij*lm*(mua-mub)**2).sum((1,2)).mean().item()/2
 82        # Euler finite-step energy change should agree with dt*derivative up to a
 83        # conservative 25% discretization tolerance for this trained NN state.
 84        rel=abs(dF-self.dt*pred)/(abs(self.dt*pred)+1e-12)
 85        return {'mass_conservation_rel_error':mass_err,'observed_energy_change':dF,
 86                'predicted_energy_derivative':pred,'derivative_relative_error':rel,
 87                'energy_dissipates':bool(dF<=1e-7),
 88                'confirmed': bool(mass_err < 1e-5 and dF <= 0 and rel < 0.25)}
 89
 90def run_one(seed, idea, lr, dt=0.05, return_net=False):
 91    torch.manual_seed(seed); np.random.seed(seed)
 92    ds=get_dataset('dynamics', seed, n_train=NTR, n_test=NTE)
 93    net=BaseDynamics(idea=idea, dt=dt)
 94    net, metric, hist=train_model(net, ds, epochs=EPOCHS, lr=lr, batch=128)
 95    if return_net:
 96        x=ds['xte'][:64].to(next(net.parameters()).device)
 97        return metric, net.signature(x) if net is not None else None
 98    return metric
 99
100def main():
101    # Shared union of tested learning rates. Baseline sweep has same candidate lr set;
102    # dt is the transport step and is fixed a priori for the final matched comparison.
103    grid=[{'lr':1e-3},{'lr':3e-3},{'lr':1e-2}]
104    base=sweep_baseline(lambda c: lambda s: run_one(s,False,c['lr']), grid)
105    best_lr=base['best_cfg']['lr']
106    idea_grid=[best_lr/3, best_lr, best_lr*3]
107    # All idea lrs are also explicitly evaluated on baseline side (parity).
108    union=sorted(set([1e-3,3e-3,1e-2]+idea_grid))
109    base_union=[]
110    for lr in union:
111        if lr not in [g['lr'] for g in grid]:
112            r=[run_one(s,False,lr) for s in range(8)]
113            base_union.append({'cfg':{'lr':lr},'full':{'mean':float(np.mean(r)),'std':float(np.std(r)),'per_seed':r,'n':8}})
114    # Choose best transport setting on the same four-seed sweep budget.
115    itried=[]
116    for lr in idea_grid:
117        vals=[run_one(s,True,lr) for s in range(4)]
118        itried.append({'cfg':{'lr':lr,'dt':0.05},'mean':float(np.mean(vals))})
119    best_i=min(itried,key=lambda z:z['mean'])['cfg']
120    idea_vals=[run_one(s,True,best_i['lr'],best_i['dt']) for s in range(8)]
121    idea={'mean':float(np.mean(idea_vals)),'std':float(np.std(idea_vals)),'per_seed':idea_vals,'n':8,
122          'sweep':itried,'best_cfg':best_i}
123    # Signature measured on trained benchmark models, one per seed.
124    sigs=[]
125    for s in range(8): sigs.append(run_one(s,True,best_i['lr'],best_i['dt'],True)[1])
126    sig={k:float(np.mean([x[k] for x in sigs])) for k in sigs[0] if isinstance(sigs[0][k],(int,float))}
127    sig['confirmed']=all(x['confirmed'] for x in sigs)
128    base['union_extra']=base_union
129    report=make_report('dynamics','rnn_small',base,idea,extra=sig)
130    report['protocol_notes']={'structural_match':'controlled pendulum rollout tests stability of recurrent dynamical states','paired_seeds':8,'epochs':EPOCHS,'samples_per_seed':[NTR,NTE],'device':DEVICE}
131    with open('bench_report.json','w') as f: json.dump(report,f,indent=2)
132    print(json.dumps(report,indent=2))
133if __name__=='__main__': main()