1import json, math, random
 2import numpy as np
 3import torch
 4from torch import nn
 5
 6SEED = 3081
 7random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
 8torch.set_num_threads(4)
 9device = 'cuda' if torch.cuda.is_available() else 'cpu'
10try:
11    if device == 'cuda': torch.cuda.get_device_name(0)
12except Exception:
13    device = 'cpu'
14
15# Periodic advection-diffusion: x'(t)=A x(t).
16d = 8
17L = np.zeros((d,d), np.float32); D = np.zeros((d,d), np.float32)
18for i in range(d):
19    L[i,i] = -2; L[i,(i-1)%d] = 1; L[i,(i+1)%d] = 1
20    D[i,(i+1)%d] = .5; D[i,(i-1)%d] = -.5
21A = .7*L + .35*D
22At = torch.tensor(A, device=device)
23def exact(x,t): return x @ torch.matrix_exp(At.T*float(t))
24def rmse(x): return float(torch.sqrt(torch.mean(x*x)).detach().cpu())
25
26def math_check():
27    torch.manual_seed(SEED+1); x=torch.randn(64,d,device=device); gen=x@At.T
28    gen_errors=[]; growth=[]
29    for h in [.2,.1,.05,.025]:
30        gen_errors.append(rmse((x+h*gen-x)/h-gen))
31        B=torch.randn(d,d,device=device); B=B/torch.linalg.matrix_norm(B)*.002
32        ref=x.clone(); bad=x.clone(); errs=[]
33        for _ in range(round(1/h)):
34            ref=ref+h*(ref@At.T); bad=bad+h*(bad@At.T)+bad@B.T; errs.append(rmse(bad-ref))
35        # Measured induced norm gives the discrete Gronwall bound.
36        Q=torch.eye(d,device=device)+h*At+B
37        Lh=float(torch.linalg.matrix_norm(Q,ord=2).detach().cpu())
38        eps=.002*float(torch.linalg.matrix_norm(x,ord=2).detach().cpu())
39        bound=eps*sum(Lh**j for j in range(round(1/h)))
40        growth.append({'h':h,'final_error':errs[-1],'bound':bound,'Lh':Lh})
41    return {'generator_errors':gen_errors,'growth':growth}
42
43class Derivative(nn.Module):
44    def __init__(self):
45        super().__init__(); self.net=nn.Sequential(nn.Linear(d,32),nn.Tanh(),nn.Linear(32,d))
46    def forward(self,x): return self.net(x)
47
48def rollout(model,x,h,K):
49    y=x
50    for _ in range(K): y=y+h*model(y)
51    return y
52
53def train_model(model, tied=True, steps=500):
54    opt=torch.optim.Adam(model.parameters(),lr=3e-3)
55    torch.manual_seed(SEED+4)
56    x=torch.randn(512,d,device=device)
57    for it in range(steps):
58        ix=torch.randint(0,x.shape[0],(64,),device=device); z=x[ix]
59        # Train all models on the same physical horizon t=1.
60        h=.05; K=20; target=exact(z,1.0)
61        if tied:
62            pred=rollout(model,z,h,K)
63            one=(model(z)-((exact(z,.05)-z)/.05)).pow(2).mean()
64        else:
65            y=z
66            for layer in model: y=y+h*layer(y)
67            pred=y; one=0.0
68        loss=((pred-target)**2).mean()+(.1*one if tied else 0)
69        opt.zero_grad(); loss.backward(); opt.step()
70    return model
71
72def main():
73    check=math_check(); torch.manual_seed(SEED+2)
74    test=torch.randn(256,d,device=device)
75    tied=train_model(Derivative().to(device),True)
76    torch.manual_seed(SEED+3)
77    untied=train_model(nn.ModuleList([Derivative() for _ in range(20)]).to(device),False)
78    results={}
79    # Untied baseline is a conventional depth-20 network trained for t=1.
80    with torch.no_grad():
81        yt=rollout(tied,test,.05,20)
82        yu=test.clone()
83        for layer in untied: yu=yu+.05*layer(yu)
84        target=exact(test,1)
85        results['1']={'tied':rmse(yt-target),'untied':rmse(yu-target)}
86        for T in [2,5,10]:
87            results[str(T)]={'tied':rmse(rollout(tied,test,.05,round(T/.05))-exact(test,T))}
88    out={'device':device,'math':check,'params_tied':sum(p.numel() for p in tied.parameters()),'params_untied':sum(p.numel() for p in untied.parameters()),'rollout':results}
89    print(json.dumps(out,indent=2))
90    with open('results.json','w') as f: json.dump(out,f,indent=2)
91if __name__=='__main__': main()
92if __name__=='__main__': main()