import os, sys, json, math import numpy as np import torch import torch.nn as nn import torch.nn.functional as F sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import get_dataset, train_model, make_report, sweep_baseline DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' EPOCHS = 20 NTR, NTE = 400, 200 class BaseDynamics(nn.Module): def __init__(self, idea=False, hidden=64, dt=0.05, channels=2): super().__init__() self.idea = idea; self.dt = dt self.rnn = nn.GRU(3, hidden, batch_first=True) self.head = nn.Linear(hidden, 1) # Equal-sized projection used by both systems; transport has its # mechanism heads, while the baseline uses an unconstrained residual. self.mix = nn.Linear(hidden, hidden) if idea: self.rho_head = nn.Linear(hidden, channels) self.pi_head = nn.Linear(hidden, channels) self.edge_head = nn.Linear(3 * hidden, channels) self.transport_proj = nn.Linear(channels, hidden) def latent(self, x): seq = x.view(x.shape[0], -1, 3) try: hs, _ = self.rnn(seq) except RuntimeError: old = torch.backends.cudnn.enabled; torch.backends.cudnn.enabled = False try: hs, _ = self.rnn(seq) finally: torch.backends.cudnn.enabled = old return hs def forward(self, x): h = self.latent(x) if not self.idea: # Ordinary unconstrained residual graph mixing on the 8 time nodes. z = self.mix(h) h = h + self.dt * (z.mean(1, keepdim=True) - z) else: rho = F.softplus(self.rho_head(h)) + 1e-5 pi = F.softplus(self.pi_head(h)) + 1e-5 pi = pi / pi.sum(1, keepdim=True) * rho.sum(1, keepdim=True).detach() # Complete graph, symmetric feature-dependent conductance. a = h.unsqueeze(2).expand(-1,-1,h.size(1),-1) b = h.unsqueeze(1).expand(-1,h.size(1),-1,-1) e = torch.cat([a,b,(a-b).abs()], -1) cij = F.softplus(self.edge_head(e)) + 1e-5 cij = (cij + cij.transpose(1,2)) / 2 q = rho / pi dr = (cij * (q.unsqueeze(1) - q.unsqueeze(2))).sum(2) rho_new = rho + self.dt * dr # dt is deliberately small relative to the learned positive rates. h = h + self.transport_proj(torch.log(rho_new + 1e-6)) return self.head(h[:, -1]) @torch.no_grad() def signature(self, x): if not self.idea: return None h = self.latent(x) rho = F.softplus(self.rho_head(h)) + 1e-5 pi = F.softplus(self.pi_head(h)) + 1e-5 pi = pi / pi.sum(1, keepdim=True) * rho.sum(1, keepdim=True) a = h.unsqueeze(2).expand(-1,-1,h.size(1),-1) b = h.unsqueeze(1).expand(-1,h.size(1),-1,-1) cij = F.softplus(self.edge_head(torch.cat([a,b,(a-b).abs()],-1)))+1e-5 cij=(cij+cij.transpose(1,2))/2 q=rho/pi; dr=(cij*(q.unsqueeze(1)-q.unsqueeze(2))).sum(2) rn=rho+self.dt*dr mass_err=((rn.sum(1)-rho.sum(1)).abs()/(rho.sum(1)+1e-8)).mean().item() def energy(r): return (r*torch.log(r/pi)).sum((1,2)) dF=(energy(rn)-energy(rho)).mean().item() mu=torch.log(q) qa,qb=q.unsqueeze(1),q.unsqueeze(2) mua,mub=mu.unsqueeze(1),mu.unsqueeze(2) lm=torch.where((mua-mub).abs()<1e-7, (qa+qb)/2, (qa-qb)/(mua-mub)) pred=-(cij*lm*(mua-mub)**2).sum((1,2)).mean().item()/2 # Euler finite-step energy change should agree with dt*derivative up to a # conservative 25% discretization tolerance for this trained NN state. rel=abs(dF-self.dt*pred)/(abs(self.dt*pred)+1e-12) return {'mass_conservation_rel_error':mass_err,'observed_energy_change':dF, 'predicted_energy_derivative':pred,'derivative_relative_error':rel, 'energy_dissipates':bool(dF<=1e-7), 'confirmed': bool(mass_err < 1e-5 and dF <= 0 and rel < 0.25)} def run_one(seed, idea, lr, dt=0.05, return_net=False): torch.manual_seed(seed); np.random.seed(seed) ds=get_dataset('dynamics', seed, n_train=NTR, n_test=NTE) net=BaseDynamics(idea=idea, dt=dt) net, metric, hist=train_model(net, ds, epochs=EPOCHS, lr=lr, batch=128) if return_net: x=ds['xte'][:64].to(next(net.parameters()).device) return metric, net.signature(x) if net is not None else None return metric def main(): # Shared union of tested learning rates. Baseline sweep has same candidate lr set; # dt is the transport step and is fixed a priori for the final matched comparison. grid=[{'lr':1e-3},{'lr':3e-3},{'lr':1e-2}] base=sweep_baseline(lambda c: lambda s: run_one(s,False,c['lr']), grid) best_lr=base['best_cfg']['lr'] idea_grid=[best_lr/3, best_lr, best_lr*3] # All idea lrs are also explicitly evaluated on baseline side (parity). union=sorted(set([1e-3,3e-3,1e-2]+idea_grid)) base_union=[] for lr in union: if lr not in [g['lr'] for g in grid]: r=[run_one(s,False,lr) for s in range(8)] base_union.append({'cfg':{'lr':lr},'full':{'mean':float(np.mean(r)),'std':float(np.std(r)),'per_seed':r,'n':8}}) # Choose best transport setting on the same four-seed sweep budget. itried=[] for lr in idea_grid: vals=[run_one(s,True,lr) for s in range(4)] itried.append({'cfg':{'lr':lr,'dt':0.05},'mean':float(np.mean(vals))}) best_i=min(itried,key=lambda z:z['mean'])['cfg'] idea_vals=[run_one(s,True,best_i['lr'],best_i['dt']) for s in range(8)] idea={'mean':float(np.mean(idea_vals)),'std':float(np.std(idea_vals)),'per_seed':idea_vals,'n':8, 'sweep':itried,'best_cfg':best_i} # Signature measured on trained benchmark models, one per seed. sigs=[] for s in range(8): sigs.append(run_one(s,True,best_i['lr'],best_i['dt'],True)[1]) sig={k:float(np.mean([x[k] for x in sigs])) for k in sigs[0] if isinstance(sigs[0][k],(int,float))} sig['confirmed']=all(x['confirmed'] for x in sigs) base['union_extra']=base_union report=make_report('dynamics','rnn_small',base,idea,extra=sig) 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} with open('bench_report.json','w') as f: json.dump(report,f,indent=2) print(json.dumps(report,indent=2)) if __name__=='__main__': main()