1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6
  7SEED = 2765
  8np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  9torch.set_num_threads(4)
 10
 11
 12def tube_rollout(Js, r0, ds):
 13    r = np.asarray(r0, dtype=float).copy()
 14    out = [r.copy()]
 15    for J, d in zip(Js, ds):
 16        r = np.abs(np.asarray(J)) @ r + np.asarray(d)
 17        out.append(r.copy())
 18    return np.asarray(out)
 19
 20
 21def toy_verification():
 22    # For J=gamma*lambda, the theory predicts q=1 boundary, slope log(q) above it,
 23    # and stationary radius d/(1-q) below it.
 24    lam = 0.8
 25    gamma_values = np.array([0.50, 0.90, 1.10, 1.30, 1.50])
 26    d = 0.01; H = 400
 27    rows = []
 28    for gamma in gamma_values:
 29        q = gamma * lam
 30        rs = tube_rollout([[[q]]]*H, [0.0], [[d]]*H)[:, 0]
 31        if q < 1:
 32            predicted_plateau = d/(1-q)
 33            observed_plateau = float(np.mean(rs[-10:]))
 34            relerr = abs(observed_plateau-predicted_plateau)/predicted_plateau
 35            slope_pred = 0.0
 36            slope_obs = 0.0
 37        else:
 38            predicted_plateau = float('inf')
 39            # Fit log radius after the transient; asymptotic slope should log(q).
 40            ix = np.arange(150, H+1)
 41            slope_obs = float(np.polyfit(ix, np.log(rs[150:]), 1)[0])
 42            slope_pred = math.log(q)
 43            relerr = abs(slope_obs-slope_pred)/abs(slope_pred)
 44        rows.append(dict(gamma=float(gamma), q=float(q), stable_pred=bool(q<1),
 45                         stable_observed=(bool(abs(slope_obs) <= 1e-10) if q >= 1 else bool(abs(slope_obs) <= 1e-10)), predicted_plateau=predicted_plateau,
 46                         observed_plateau=observed_plateau if q<1 else None,
 47                         predicted_log_slope=slope_pred, observed_log_slope=slope_obs,
 48                         relative_error=float(relerr)))
 49    # Disturbance scaling prediction at q=.72: equilibrium is linear in d.
 50    q=.72; ds=np.array([.002,.005,.01,.02]); final=[]
 51    for di in ds:
 52        rr=tube_rollout([[[q]]]*120,[0],[[di]]*120)[-1,0]
 53        final.append(float(rr))
 54    ratio=np.asarray(final)/ds
 55    scaling_error=float(np.max(np.abs(ratio-1/(1-q))/(1/(1-q))))
 56    # Locate transition empirically among a fine sweep: first q >= 1.
 57    fine=np.linspace(.70,1.50,161); empirical_boundary=float(fine[np.argmax(fine*lam>=1)])
 58    return {'gain_sweep':rows, 'predicted_boundary_gamma':1/lam,
 59            'observed_boundary_gamma':empirical_boundary, 'disturbances':ds.tolist(),
 60            'final_radii':final, 'radius_over_d':ratio.tolist(),
 61            'predicted_radius_over_d':1/(1-q), 'scaling_max_relative_error':scaling_error}
 62
 63
 64class Transition(nn.Module):
 65    def __init__(self, dim=2):
 66        super().__init__()
 67        self.net=nn.Sequential(nn.Linear(dim+1,24),nn.Tanh(),nn.Linear(24,dim))
 68    def forward(self,z,u): return z + 0.18*self.net(torch.cat([z,u],-1))
 69
 70
 71def make_data(n=900, horizon=12):
 72    zs=[]; us=[]; ys=[]
 73    for _ in range(n):
 74        z=np.random.uniform(-1,1,2).astype('float32')
 75        for k in range(horizon):
 76            u=np.random.uniform(-.7,.7,1).astype('float32')
 77            # nonlinear stable plant with transition noise
 78            y=np.array([.88*z[0]+.12*np.tanh(z[1])+ .08*u[0],
 79                        .80*z[1]+.10*np.sin(z[0])- .04*u[0]],dtype='float32')
 80            y += np.random.normal(0,.008,2).astype('float32')
 81            zs.append(z); us.append(u); ys.append(y); z=y
 82    return torch.tensor(np.asarray(zs)),torch.tensor(np.asarray(us)),torch.tensor(np.asarray(ys))
 83
 84
 85def neural_comparison():
 86    z,u,y=make_data(); n=len(z); split=int(.8*n)
 87    # Same initialization and training budget; tube penalty uses exact autograd Jacobian.
 88    def train(tube):
 89        torch.manual_seed(SEED); m=Transition(); opt=torch.optim.Adam(m.parameters(),lr=3e-3)
 90        for epoch in range(28):
 91            ix=torch.randperm(split)[:128]; zb,ub,yb=z[ix],u[ix],y[ix]
 92            pred=m(zb,ub); loss=((pred-yb)**2).mean()
 93            if tube:
 94                # Conservative residual radius: held-out-like fixed upper quantile estimate.
 95                r=torch.full((len(ix),2),.025)
 96                zz=zb.detach().clone().requires_grad_(True)
 97                pp=m(zz,ub)
 98                J=[]
 99                for j in range(2):
100                    grad=torch.autograd.grad(pp[:,j].sum(),zz,create_graph=False,retain_graph=True)[0]
101                    J.append(grad)
102                J=torch.stack(J,1).abs()
103                rnext=torch.bmm(J,r.unsqueeze(-1)).squeeze(-1)+.012
104                # nominal state constraint |z_i| <= 1, robust backoff.
105                robust_violation=torch.relu(torch.abs(pred)+rnext-1.0)
106                loss=loss + .015*rnext.mean() + .20*robust_violation.pow(2).mean()
107            opt.zero_grad(); loss.backward(); opt.step()
108        with torch.no_grad():
109            one=((m(z[split:],u[split:])-y[split:])**2).mean().item()
110        # Free-running 20-step rollout on test starts; compare state constraint violations.
111        viol=[]; mse=[]
112        for start in range(split,min(split+100, n-20),20):
113            zz=z[start:start+1].clone(); truth=z[start:start+20].clone();
114            for k in range(20):
115                uu=u[start+k:start+k+1]
116                zz=m(zz,uu); mse.append(float(((zz-truth[k:k+1])**2).mean()))
117                viol.append(float((torch.abs(zz)>1).any()))
118        return {'one_step_mse':one,'rollout_mse':float(np.mean(mse)),
119                'violation_rate':float(np.mean(viol))}
120    return {'baseline':train(False),'tube_regularized':train(True)}
121
122
123def main():
124    out={'seed':SEED,'toy_verification':toy_verification()}
125    try:
126        out['neural_comparison']=neural_comparison()
127        out['device']='cpu'
128    except Exception as e:
129        out['neural_error']=repr(e)
130    Path('results.json').write_text(json.dumps(out,indent=2))
131    print(json.dumps(out,indent=2))
132
133if __name__=='__main__': main()