import json, math, time import numpy as np import torch import torch.nn as nn SEED = 2904 torch.manual_seed(SEED); np.random.seed(SEED) torch.set_default_dtype(torch.float64) DEVICE = "cuda" if torch.cuda.is_available() else "cpu" try: if DEVICE == "cuda": torch.cuda.empty_cache() except Exception: DEVICE = "cpu" class MLP(nn.Module): def __init__(self, inp, out, width=32): super().__init__() self.net = nn.Sequential(nn.Linear(inp,width), nn.Tanh(), nn.Linear(width,width), nn.Tanh(), nn.Linear(width,out)) def forward(self,x): return self.net(x) class PHField(nn.Module): def __init__(self, d=2, eps=0.03): super().__init__(); self.d=d; self.eps=eps self.h=MLP(d,1); self.a=MLP(d,d*d); self.l=MLP(d,d*d) self.qraw=nn.Parameter(torch.eye(d)*0.7) def energy(self,z): q=self.qraw @ self.qraw.T + 0.05*torch.eye(self.d, device=z.device) return torch.nn.functional.softplus(self.h(z).squeeze(-1)) + .5*torch.sum((z@q)*z,dim=-1) def matrices(self,z): n=z.shape[0]; A=self.a(z).reshape(n,self.d,self.d); L=self.l(z).reshape(n,self.d,self.d) J=A-A.transpose(1,2); R=L@L.transpose(1,2)+self.eps*torch.eye(self.d,device=z.device) return J,R def forward(self,z): zz=z.detach().requires_grad_(True) H=self.energy(zz); g=torch.autograd.grad(H.sum(),zz,create_graph=self.training)[0] J,R=self.matrices(zz) return torch.bmm((J-R),g.unsqueeze(-1)).squeeze(-1) def derivative_terms(self,z): zz=z.detach().requires_grad_(True); H=self.energy(zz); g=torch.autograd.grad(H.sum(),zz)[0] J,R=self.matrices(zz); dz=torch.bmm((J-R),g.unsqueeze(-1)).squeeze(-1) return H.detach(), g.detach(), J.detach(), R.detach(), (g*dz).sum(1).detach(), -(torch.bmm(R,g.unsqueeze(-1)).squeeze(-1)*g).sum(1).detach() class Unconstrained(nn.Module): def __init__(self,d=2): super().__init__(); self.net=MLP(d,d) def forward(self,z): return self.net(z) def rk4(fun,x,dt): k1=fun(x); k2=fun(x+.5*dt*k1); k3=fun(x+.5*dt*k2); k4=fun(x+dt*k3) return x+dt*(k1+2*k2+2*k3+k4)/6 def oscillator(x): # nonlinear Hamiltonian oscillator with linear damping q,p=x[...,0],x[...,1] return torch.stack((p, -q-0.15*q**3-0.12*p),-1) def make_data(n=32, steps=16, dt=.08): x=torch.randn(n,2)*1.1; xs=[] with torch.no_grad(): for _ in range(steps+1): xs.append(x.clone()); x=rk4(oscillator,x,dt) return torch.stack(xs,1) def train_model(model, data, epochs=100, dt=.08): opt=torch.optim.Adam(model.parameters(),lr=3e-3); t0=time.perf_counter() for _ in range(epochs): # random one-step minibatch, same data and objective b=torch.randint(data.shape[0],(min(64,data.shape[0]),)); k=torch.randint(data.shape[1]-1,(len(b),)) x=data[b,k]; target=data[b,k+1] pred=rk4(model,x,dt) loss=((pred-target)**2).mean(); opt.zero_grad(); loss.backward(); opt.step() return time.perf_counter()-t0 def rollout(model,x,steps,dt): out=[x.clone()] # PH evaluation still needs input autograd for grad_z H; do not wrap in no_grad. was_training = model.training model.eval() with torch.enable_grad(): for _ in range(steps): out.append(rk4(model,out[-1],dt)) if was_training: model.train() return torch.stack(out,1) def main(): d=4; z=torch.randn(200,d) # Structural prediction 1: J antisymmetry is exactly zero for every scale. scales=[0.0,.25,.5,1.,2.,4.] skew=[]; psd=[]; drift=[]; predicted=[] ph=PHField(d=d).eval() with torch.no_grad(): for s in scales: A=torch.randn(200,d,d)*s; L=torch.randn(200,d,d)*s J=A-A.transpose(1,2); R=L@L.transpose(1,2)+ph.eps*torch.eye(d) skew.append(float((J+J.transpose(1,2)).abs().max())) psd.append(float(torch.linalg.eigvalsh(R).amin())) # use fixed random g; predicted dH/dt has affine lambda dependence g=torch.randn(200,d); val=-(torch.bmm(R,g.unsqueeze(-1)).squeeze(-1)*g).sum(1) drift.append(float(val.mean())) predicted.append(float(-ph.eps*(g*g).sum(1).mean()-s*s*((torch.bmm((torch.randn(200,d,d)*0+L),g.unsqueeze(-1)).squeeze(-1))**2).sum(1).mean())) # Cleaner dissipation sweep: fixed L0,g gives exact prediction slope in lambda^2. L0=torch.randn(200,d,d); g0=torch.randn(200,d); eps=.03; lam=torch.tensor(scales) observed=[]; theory=[] base=eps*(g0*g0).sum(1).mean(); coeff=(torch.bmm(L0,g0.unsqueeze(-1)).squeeze(-1)**2).sum(1).mean() for s in scales: R=s*s*(L0@L0.transpose(1,2))+eps*torch.eye(d) observed.append(float(-(torch.bmm(R,g0.unsqueeze(-1)).squeeze(-1)*g0).sum(1).mean())) theory.append(float(-base-s*s*coeff)) # Mini supervised trajectory fit. data=make_data(); ph2=PHField(2,eps=.03); uc=Unconstrained(2) ph_time=train_model(ph2,data); uc_time=train_model(uc,data) test=make_data(n=32,steps=100)[0:] ph_pred=rollout(ph2,test[:,0],100,.08); uc_pred=rollout(uc,test[:,0],100,.08) target=test ph_err=float(torch.sqrt(((ph_pred-target)**2).mean())); uc_err=float(torch.sqrt(((uc_pred-target)**2).mean())) # Unforced energy check over random states and finite RK4 trajectory. H,g,J,R,dh,formula=ph.derivative_terms(z) x=z[:1]; energies=[] with torch.enable_grad(): for _ in range(100): energies.append(float(ph.energy(x)[0].detach())); x=rk4(ph,x,.02) energy_increases=sum(np.diff(energies)>1e-10) result={"device":DEVICE,"structural":{"max_skew_residual_by_scale":dict(zip(scales,skew)),"min_R_eigenvalue_by_scale":dict(zip(scales,psd)),"max_energy_formula_residual":float((dh-formula).abs().max()),"energy_derivative_max":float(dh.max()),"continuous_energy_nonpositive":bool(dh.max()<=1e-10),"rk4_energy_increases":int(energy_increases)},"dissipation_sweep":{"lambda":scales,"observed_dHdt":observed,"predicted_dHdt":theory,"max_abs_prediction_error":float(max(abs(a-b) for a,b in zip(observed,theory))),"slope_observed_vs_lambda2":float(np.polyfit(np.array(scales)**2,observed,1)[0]),"slope_predicted":float(-coeff)},"mini_experiment":{"ph_rmse":ph_err,"unconstrained_rmse":uc_err,"ph_train_seconds":ph_time,"unconstrained_train_seconds":uc_time,"steps":100,"ph_params":sum(p.numel() for p in ph2.parameters()),"unconstrained_params":sum(p.numel() for p in uc.parameters())}} print(json.dumps(result,indent=2)) if __name__=='__main__': main()