1import os, json, math, random
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5
  6SEED=17
  7random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  8torch.set_num_threads(4)
  9try:
 10    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 11    if device.type=='cuda': torch.zeros(1,device=device)
 12except Exception:
 13    device=torch.device('cpu')
 14
 15def sym_eigs(M): return torch.linalg.eigvalsh((M+M.transpose(-1,-2))/2)
 16
 17# Core numerical verification: direct finite-step LMI and energy ratio.
 18def exact_check():
 19    h=.08; alpha=.35
 20    # A is stable in a non-Euclidean metric; P changes with theta.
 21    P0=torch.tensor([[2.0,.35],[.35,1.0]])
 22    P1=P0.clone()
 23    # Construct A from a contracting similarity transform, then test the actual discrete condition.
 24    Q=torch.tensor([[-1.4, .25],[-.10,-.9]])
 25    L=torch.linalg.cholesky(P0)
 26    A=torch.linalg.solve(L.T, Q @ L.T) # gives A^T P + P A = L(Q^T+Q)L^T
 27    J=torch.eye(2)+h*A
 28    S=J.T@P1@J-math.exp(-2*alpha*h)*P0
 29    # Generalized normalized matrix via symmetric inverse square root
 30    ew,U=torch.linalg.eigh(P0); Pinv=(U*ew.rsqrt())@U.T
 31    lam=torch.linalg.eigvalsh(Pinv@S@Pinv).max().item()
 32    # simulate perturbations and compare exact metric energies
 33    d=torch.tensor([1.0,-.7]); v0=(d@P0@d).item(); d1=J@d; v1=(d1@P1@d1).item()
 34    ratio=v1/v0; bound=math.exp(-2*alpha*h)
 35    # A deliberately noncontracting control in same P.
 36    Jbad=torch.eye(2)+h*(A+1.8*torch.eye(2)); Sbad=Jbad.T@P1@Jbad-bound*P0
 37    lambad=torch.linalg.eigvalsh(Pinv@Sbad@Pinv).max().item()
 38    return {'device':str(device),'h':h,'alpha':alpha,'finite_lmi_lambda':lam,
 39            'energy_ratio':ratio,'bound':bound,'bad_lambda':lambad,
 40            'lmi_satisfied':bool(lam<=1e-7),'energy_satisfied':bool(ratio<=bound)}
 41
 42class Dyn(nn.Module):
 43    def __init__(self, hidden=32):
 44        super().__init__()
 45        self.net=nn.Sequential(nn.Linear(3,hidden),nn.Tanh(),nn.Linear(hidden,hidden),nn.Tanh(),nn.Linear(hidden,2))
 46        self.pnet=nn.Sequential(nn.Linear(1,16),nn.Tanh(),nn.Linear(16,3))
 47    def forward(self,z,theta): return z + .12*self.net(torch.cat([z,theta],-1))
 48    def metric(self,theta):
 49        x=self.pnet(theta); L=torch.zeros(x.shape[0],2,2,device=x.device)
 50        L[:,0,0]=torch.nn.functional.softplus(x[:,0])+0.15
 51        L[:,1,0]=x[:,1]; L[:,1,1]=torch.nn.functional.softplus(x[:,2])+0.15
 52        return L@L.transpose(-1,-2)+.05*torch.eye(2,device=x.device)
 53
 54def target(z,theta):
 55    # Two changing operating regimes, both stable but with different shear/rotation.
 56    th=theta[:,0]
 57    a=0.73+0.10*torch.sin(th); b=.18*torch.cos(th)
 58    M=torch.stack([torch.stack([a,b],-1),torch.stack([-b,a-.08*torch.sin(th)],-1)],-2)
 59    return torch.bmm(M,z.unsqueeze(-1)).squeeze(-1)
 60
 61def jacobian_batch(model,z,theta):
 62    # Explicit small-state Jacobians, retaining autograd for the penalty.
 63    outs=[]
 64    for i in range(z.shape[0]):
 65        zi=z[i:i+1].detach().requires_grad_(True); ti=theta[i:i+1]
 66        yi=model(zi,ti).squeeze(0)
 67        rows=[]
 68        for q in range(2): rows.append(torch.autograd.grad(yi[q],zi,create_graph=True,retain_graph=True)[0].squeeze(0))
 69        outs.append(torch.stack(rows))
 70    return torch.stack(outs)
 71
 72def lyap_penalty(model,z,theta,alpha=.22,h=.12):
 73    J=jacobian_batch(model,z,theta); P=model.metric(theta)
 74    # P_{k+1}: parameter changes between adjacent operating conditions (same batch shuffled partner).
 75    theta1=torch.roll(theta,1,0); P1=model.metric(theta1)
 76    E=torch.eye(2,device=z.device).expand(z.shape[0],-1,-1)
 77    R=J.transpose(1,2)@P1@J-math.exp(-2*alpha*h)*P
 78    # generalized max eigenvalue, normalized by P^{-1/2}; squared positive part.
 79    ev,U=torch.linalg.eigh(P); Pinv=(U*ev.rsqrt().unsqueeze(-2))@U.transpose(1,2)
 80    lam=torch.linalg.eigvalsh(Pinv@R@Pinv)[:,-1]
 81    return torch.relu(lam).square().mean(), lam.detach()
 82
 83def run(penalized, steps=120):
 84    torch.manual_seed(SEED+int(penalized)); model=Dyn().to(device)
 85    opt=torch.optim.Adam(model.parameters(),lr=2e-3)
 86    for step in range(steps):
 87        z=torch.randn(48,2,device=device); th=(torch.rand(48,1,device=device)*2-1)*math.pi
 88        y=target(z,th)
 89        pred=model(z,th); loss=((pred-y)**2).mean()
 90        if penalized:
 91            lp,_=lyap_penalty(model,z,th); loss=loss+0.8*lp
 92        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5); opt.step()
 93    with torch.no_grad():
 94        z=torch.randn(256,2,device=device); th=(torch.rand(256,1,device=device)*2-1)*math.pi
 95        mse=((model(z,th)-target(z,th))**2).mean().item()
 96    # Need jacobians with gradients enabled, no_grad cannot be used.
 97    z=torch.randn(128,2,device=device); th=(torch.rand(128,1,device=device)*2-1)*math.pi
 98    lp,lam=lyap_penalty(model,z,th)
 99    # independent same-regime P_{k+1}=P(theta), to isolate local contraction.
100    J=jacobian_batch(model,z,th); P=model.metric(th); ev,U=torch.linalg.eigh(P); Pinv=(U*ev.rsqrt().unsqueeze(-2))@U.transpose(1,2)
101    local=torch.linalg.eigvalsh(Pinv@(J.transpose(1,2)@P@J)@Pinv)[:,-1]
102    return {'mse':mse,'penalty_batch':lp.item(),'fraction_violating':(lam>0).float().mean().item(),
103            'mean_generalized_lambda':lam.mean().item(),'max_generalized_lambda':lam.max().item(),
104            'local_contraction_ratio_mean':local.mean().item(),'local_ratio_max':local.max().item()}
105
106if __name__=='__main__':
107    out={'exact_check':exact_check(),'baseline':run(False),'lyapunov':run(True)}
108    print(json.dumps(out,indent=2))
109    with open('results.json','w') as f: json.dump(out,f,indent=2)