1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6import torch.nn.functional as F
  7
  8SEED = 3038
  9random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
 10torch.set_num_threads(min(12, torch.get_num_threads()))
 11DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 12try:
 13    if DEVICE == "cuda":
 14        torch.cuda.set_device(0)
 15        torch.zeros(1, device="cuda")
 16except Exception:
 17    DEVICE = "cpu"
 18
 19# Discrete Duffing oscillator. The target is position, while the learned state
 20# contains position and velocity, so long-horizon state propagation is nontrivial.
 21def duffing(x, u=0.0, dt=0.08):
 22    q, v = x[..., 0], x[..., 1]
 23    a = -0.8*q - 0.25*v - 0.15*q**3 + u
 24    return torch.stack((q + dt*v, v + dt*a), -1)
 25
 26def make_data(n, T=140, noise=0.015):
 27    xs = torch.zeros(n, T+1, 2)
 28    xs[:, 0] = torch.randn(n, 2) * torch.tensor([0.45, 0.20])
 29    for k in range(T):
 30        # Mild known excitation makes the identification problem better conditioned.
 31        u = 0.12 * torch.sin(torch.tensor(0.11*k))
 32        xs[:, k+1] = duffing(xs[:, k], u) + noise*torch.randn(n, 2)
 33    return xs
 34
 35class Dynamics(nn.Module):
 36    def __init__(self, hidden=32):
 37        super().__init__()
 38        self.net = nn.Sequential(nn.Linear(2, hidden), nn.Tanh(), nn.Linear(hidden, hidden), nn.Tanh(), nn.Linear(hidden, 2))
 39    def forward(self, x):
 40        return x + 0.08*torch.tanh(self.net(x))
 41
 42class Storage(nn.Module):
 43    def __init__(self, hidden=24, delta=0.01):
 44        super().__init__(); self.delta = delta
 45        self.g = nn.Sequential(nn.Linear(2, hidden), nn.Tanh(), nn.Linear(hidden, hidden), nn.Tanh(), nn.Linear(hidden, 1))
 46    def forward(self, x):
 47        return self.g(x).square().squeeze(-1) + self.delta*x.square().sum(-1)
 48
 49def rollout(model, x0, H):
 50    out=[x0]; x=x0
 51    for _ in range(H):
 52        x=model(x); out.append(x)
 53    return torch.stack(out, 1)
 54
 55def train(train_x, regularized, steps=260, horizon=15):
 56    model=Dynamics().to(DEVICE)
 57    storage=Storage().to(DEVICE) if regularized else None
 58    params=list(model.parameters()) + ([] if storage is None else list(storage.parameters()))
 59    opt=torch.optim.Adam(params, lr=2e-3)
 60    ema=None; rho=0.08; lam=0.35
 61    n=train_x.shape[0]; rng=np.random.default_rng(SEED + (1 if regularized else 0))
 62    for step in range(steps):
 63        idx=torch.tensor(rng.integers(0,n,32), device=DEVICE)
 64        k=torch.tensor(rng.integers(0, train_x.shape[1]-horizon, 32), device=DEVICE)
 65        # Gather independent starting points and targets for a short rollout.
 66        x0=train_x[idx,k].to(DEVICE)
 67        target=torch.stack([train_x[idx,k+j+1,0].to(DEVICE) for j in range(horizon)], 1)
 68        pred=rollout(model,x0,horizon)[:,1:]
 69        losses=(pred[:,:,0]-target).square()
 70        ordinary=losses.mean()
 71        if ema is None: ema=ordinary.detach()
 72        else: ema=0.995*ema + 0.005*ordinary.detach()
 73        loss=ordinary
 74        if regularized:
 75            V=torch.stack([storage(pred[:,j]) for j in range(horizon+0)],1)
 76            V0=storage(x0)
 77            Vprev=torch.cat([V0[:,None], V[:,:-1]],1)
 78            residual=V-Vprev + rho*losses - ordinary.detach() + ema.detach()
 79            diss=F.relu(residual).square().mean()
 80            loss=ordinary+lam*diss
 81        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(params, 5.0); opt.step()
 82    return model, storage
 83
 84@torch.no_grad()
 85def evaluate(model, storage, test_x, H=100):
 86    x0=test_x[:,0].to(DEVICE)
 87    pred=rollout(model,x0,H)[:,1:]
 88    target=test_x[:,1:H+1].to(DEVICE)
 89    err=(pred[:,:,0]-target[:,:,0]).square()
 90    # Evaluate residual against the measured per-step loss and a fixed reference
 91    # l*=0. This is the direct strict-dissipativity residual for this experiment.
 92    if storage is not None:
 93        vv=torch.cat([storage(x0)[:,None], storage(pred)],1)
 94        residual=vv[:,1:]-vv[:,:-1]+0.08*err
 95        pos=F.relu(residual)
 96        frac=(residual>0).float().mean().item()
 97        mean_res=residual.mean().item(); cum_res=residual.sum(1).mean().item()
 98    else: frac=mean_res=cum_res=float('nan')
 99    return dict(mse10=err[:,:10].mean().item(), mse100=err.mean().item(), cum100=err.sum(1).mean().item(), positive_fraction=frac, mean_residual=mean_res, cumulative_residual=cum_res)
100
101def math_check():
102    # Direct numerical verification: sum of V_{k+1}-V_k telescopes exactly;
103    # when residual <= 0, the proposed cumulative bound follows.
104    rng=np.random.default_rng(SEED)
105    B,N=16,40
106    V=np.cumsum(np.abs(rng.normal(size=(B,N+1))),axis=1)
107    ell=np.abs(rng.normal(size=(B,N)))+0.2
108    lstar=np.zeros_like(ell); d2=np.abs(rng.normal(size=(B,N)))
109    rho=0.1
110    r=V[:,1:]-V[:,:-1]+rho*d2-ell+lstar
111    identity=np.max(np.abs((V[:,1:]-V[:,:-1]).sum(1)-(V[:,-1]-V[:,0])))
112    # Construct certified samples by setting ell equal to the RHS plus margin.
113    ell_cert=V[:,1:]-V[:,:-1]+rho*d2+0.5
114    lhs=(rho*d2).sum(1)
115    rhs=(ell_cert).sum(1)+V[:,0]-V[:,-1]
116    bound_violation=np.max(lhs-rhs)
117    return dict(max_telescoping_error=float(identity), max_certified_bound_violation=float(bound_violation), random_mean_positive_residual=float(np.maximum(r,0).mean()))
118
119def main():
120    train_x=make_data(96); test_x=make_data(24)
121    check=math_check()
122    train_x=train_x.to(DEVICE); test_x=test_x.to(DEVICE)
123    base,_=train(train_x,False); diss,store=train(train_x,True)
124    result={'device':DEVICE,'seed':SEED,'math_check':check,
125            'baseline':evaluate(base, None, test_x), 'dissipative':evaluate(diss,store,test_x)}
126    print(json.dumps(result, indent=2))
127    Path('results.json').write_text(json.dumps(result, indent=2))
128
129if __name__=='__main__': main()