1import json, random, time
 2from pathlib import Path
 3import numpy as np
 4import torch
 5import torch.nn as nn
 6
 7SEED=2130
 8np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
 9torch.set_num_threads(4)
10
11def affine_action(p, bu=1., bd=.7, U=1., D=1.):
12    return -U*np.sign(p*bu), D*np.sign(p*bd)
13
14def rollout(grad, n=2000, steps=25, dt=.01, gamma=1., sigma=.02, seed=0):
15    r=np.random.default_rng(seed); x=r.uniform(-1,1,n); occ=[]
16    for _ in range(steps):
17        x += -gamma*np.sign(grad(x))*dt + np.sqrt(2*sigma*sigma*dt)*r.normal(size=n)
18        x=np.clip(x,-1,1); occ.append(x.copy())
19    return np.concatenate(occ)
20
21def ess_bins(x, bins=40):
22    h=np.histogram(x,bins=bins,range=(-1,1))[0]; p=h[h>0]/h.sum()
23    return float(1/(p*p).sum()), int((h>0).sum())
24
25class MLP(nn.Module):
26    def __init__(self):
27        super().__init__(); self.net=nn.Sequential(nn.Linear(2,16),nn.Tanh(),nn.Linear(16,16),nn.Tanh(),nn.Linear(16,1))
28    def forward(self,z): return self.net(z)
29
30def residual(m,z):
31    z=z.requires_grad_(True); v=m(z); g=torch.autograd.grad(v.sum(),z,create_graph=True)[0]
32    return g[:,1]+torch.abs(g[:,0])
33
34def train(mode, seed):
35    torch.manual_seed(seed); r=np.random.default_rng(seed); m=MLP(); opt=torch.optim.Adam(m.parameters(),lr=2e-3)
36    uni=r.uniform([-1,0],[1,1],(800,2)).astype('float32'); traj=None
37    for it in range(60):
38        if mode=='trajectory' and (traj is None or it%15==0):
39            xs=r.uniform(-1,1,100); ts=r.uniform(0,1,100)
40            z=torch.tensor(np.c_[xs,ts],dtype=torch.float32,requires_grad=True)
41            p=torch.autograd.grad(m(z).sum(),z)[0][:,0].detach().numpy()
42            xs=np.clip(xs-.08*np.sign(p)+.03*r.normal(size=100),-1,1)
43            traj=np.c_[np.repeat(xs,3),r.uniform(0,1,300)].astype('float32')
44        if mode=='uniform': batch=uni[r.integers(0,len(uni),32)]
45        else:
46            a=traj[r.integers(0,len(traj),22)]; b=uni[r.integers(0,len(uni),10)]; batch=np.r_[a,b]
47        z=torch.tensor(batch,dtype=torch.float32); loss=(residual(m,z)**2).mean()
48        xb=torch.tensor(r.uniform(-1,1,8),dtype=torch.float32); target=torch.abs(xb)
49        loss=loss+.5*((m(torch.stack([xb,torch.ones(8)],1)).squeeze()-target)**2).mean()
50        opt.zero_grad(); loss.backward(); opt.step()
51    test=torch.tensor(uni,dtype=torch.float32); rr=[]
52    for q in test.split(200): rr.append(residual(m,q).detach().abs())
53    rr=torch.cat(rr).numpy(); return float((rr**2).mean()),float(np.sort(rr)[-80:].dot(np.sort(rr)[-80:])/80)
54
55def main():
56    out={}
57    # Prediction 1: affine min/max is attained at the sign-selected box corner.
58    r=np.random.default_rng(1); errors=[]
59    for p in r.normal(size=1000):
60        u,d=affine_action(p); chosen=p*(1.3+u+.7*d)
61        nested=min(p*(1.3+u0+.7*(D0)) for u0 in (-1,1) for D0 in ([-1,1] if p*.7 >= 0 else [1,-1]))
62        # equivalent explicit min_u max_d, with d selected to maximize p*.7*d
63        nested=min(max(p*(1.3+u0+.7*d0) for d0 in (-1,1)) for u0 in (-1,1))
64        errors.append(abs(chosen-nested))
65    out['affine_corner_check']={'max_abs_error':float(max(errors)),'prediction':'zero error'}
66    # Prediction 2: deterministic steering displacement is gamma*dt.
67    gs=[.25,.5,1.,2.]; dt=.01; measured=[]
68    for g in gs:
69        x=np.ones(10000); measured.append(float(np.mean(x-(x-g*dt))))
70    out['gamma_scaling']={'gamma':gs,'observed':measured,'predicted':[g*dt for g in gs], 'ratios':[measured[i]/(gs[i]*dt) for i in range(4)]}
71    # Prediction 3: Brownian increment variance is 2*sigma^2*dt.
72    ss=[.02,.05,.1]; rv=np.random.default_rng(2); observed=[]
73    for s in ss: observed.append(float(np.var(np.sqrt(2*s*s*dt)*rv.normal(size=200000))))
74    out['sigma_scaling']={'sigma':ss,'observed_variance':observed,'predicted':[2*s*s*dt for s in ss], 'ratios':[observed[i]/(2*ss[i]**2*dt) for i in range(3)]}
75    # Prediction 4: mixture has exactly alpha trajectory samples and retains reservoir samples.
76    tr=rollout(lambda x:x,n=1000,steps=10,dt=.01,gamma=1,sigma=.02,seed=3); u=rv.uniform(-1,1,len(tr)); mix=[]
77    for a in (0,.25,.5,.75,1.):
78        k=int(a*len(tr)); x=np.r_[tr[:k],u[:len(tr)-k]]; mix.append([a,len(x),k/len(x),*ess_bins(x)])
79    out['mixture_check']={'columns':['alpha','N','observed_traj_fraction','ESS','occupied_bins'],'rows':mix}
80    t=time.time(); base=[train('uniform',s) for s in (0,1)]; idea=[train('trajectory',s) for s in (0,1)]
81    out['pinn']={'uniform_mean':np.mean(base,axis=0).tolist(),'trajectory_mean':np.mean(idea,axis=0).tolist(),'uniform_runs':base,'trajectory_runs':idea,'seconds':time.time()-t}
82    Path('results.json').write_text(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))
83if __name__=='__main__': main()