1import json, math, random, time
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6
  7SEED=3015
  8random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  9torch.set_num_threads(4)
 10device = "cuda" if torch.cuda.is_available() else "cpu"
 11try:
 12    if device == "cuda": torch.cuda.set_device(0)
 13except Exception:
 14    device = "cpu"
 15DT=torch.float32
 16NPER=12
 17ANGLE=2*math.pi/NPER
 18ROT=torch.tensor([[math.cos(ANGLE),-math.sin(ANGLE)],[math.sin(ANGLE),math.cos(ANGLE)]],dtype=DT)
 19TRUE_A=.86
 20
 21def true_map(z):
 22    # radial contraction toward the unit circle, followed by a fixed rotation
 23    r=torch.sqrt((z*z).sum(-1,keepdim=True)+1e-9)
 24    rn=1+TRUE_A*(r-1)
 25    return (rn/r)*z @ ROT.to(z.device).T
 26
 27def make_data(ntraj=18, length=90, noise=.035):
 28    xs=[]; ys=[]
 29    for _ in range(ntraj):
 30        z=torch.tensor([1+np.random.uniform(-.18,.18),0.],dtype=DT)
 31        phase=np.random.uniform(0,2*math.pi)
 32        z=torch.tensor([math.cos(phase),math.sin(phase)],dtype=DT)*float(np.random.uniform(.82,1.18))
 33        for k in range(length):
 34            xn=z+noise*torch.randn(2)
 35            zn=true_map(z)
 36            yn=zn+noise*torch.randn(2)
 37            xs.append(xn); ys.append(yn); z=zn
 38    return torch.stack(xs).to(device),torch.stack(ys).to(device)
 39
 40class Dynamics(nn.Module):
 41    def __init__(self):
 42        super().__init__()
 43        self.net=nn.Sequential(nn.Linear(2,32),nn.Tanh(),nn.Linear(32,32),nn.Tanh(),nn.Linear(32,2))
 44    def forward(self,z): return self.net(z)
 45
 46def monodromy(model,z):
 47    # Exact autodiff Jacobian of the N-step return map at the section point.
 48    zz=z
 49    for _ in range(NPER): zz=model(zz)
 50    rows=[]
 51    for j in range(2): rows.append(torch.autograd.grad(zz[0,j],z,create_graph=True,retain_graph=True)[0][0])
 52    return torch.stack(rows),zz
 53
 54def train(kind, x, y, steps=1800):
 55    model=Dynamics().to(device)
 56    opt=torch.optim.Adam(model.parameters(),lr=2e-3)
 57    rng=np.random.default_rng(SEED+ (0 if kind=='baseline' else 1))
 58    trace=[]
 59    for step in range(steps):
 60        ix=torch.tensor(rng.integers(0,len(x),128),device=device)
 61        pred=model(x[ix]); loss=((pred-y[ix])**2).mean()
 62        if kind=='idea':
 63            # Section h(z)=y=0, with a known 2pi symmetry return (alpha=2pi).
 64            # Start exactly on the section and require P(z)=z.
 65            z0=torch.tensor([[1.,0.]],device=device,requires_grad=True)
 66            J,ret=monodromy(model,z0)
 67            retloss=((ret-z0)**2).mean()
 68            eig=torch.linalg.eigvals(J)
 69            rho=torch.abs(eig).max()
 70            floq=torch.relu(rho-(1-.03))**2
 71            loss=loss+1.2*retloss+0.25*floq
 72        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5.0); opt.step()
 73        if step in (0,steps//2,steps-1): trace.append(float(loss.detach().cpu()))
 74    with torch.no_grad():
 75        z0=torch.tensor([[1.,0.]],device=device)
 76    z0.requires_grad_(True)
 77    J,ret=monodromy(model,z0)
 78    rho=float(torch.abs(torch.linalg.eigvals(J)).max().detach().cpu())
 79    return model, rho, float(((ret-z0)**2).mean().detach().cpu()), trace
 80
 81def rollout_error(model, nperiods=15):
 82    # Compare a perturbed radius trajectory with the true dynamics; report error at event returns.
 83    z=torch.tensor([[.90,0.]],device=device)
 84    truth=z.clone(); errs=[]
 85    with torch.no_grad():
 86        for _ in range(nperiods):
 87            for _ in range(NPER): z=model(z); truth=true_map(truth)
 88            errs.append(float(torch.linalg.norm(z-truth).cpu()))
 89    return errs
 90
 91def main():
 92    t=time.time(); x,y=make_data()
 93    results={'device':device,'seed':SEED,'period_steps':NPER,'true_radial_floquet':TRUE_A}
 94    alltr={}
 95    for kind in ('baseline','idea'):
 96        model,rho,ret,trace=train(kind,x,y)
 97        errs=rollout_error(model)
 98        results[kind]={'rho':rho,'return_mse':ret,'event_errors':errs,'loss_trace':trace}
 99        alltr[kind]=model.state_dict()
100    # Independent core math sanity check: finite perturbation ratios for the analytic map.
101    z=torch.tensor([[1.001,0.]],dtype=DT); q=torch.tensor([[1.,0.]],dtype=DT)
102    ratios=[]
103    for n in range(1,6):
104        for _ in range(NPER): z=true_map(z); q=true_map(q)
105        ratios.append(float((torch.linalg.norm(z-q)/.001).item()))
106    results['analytic_perturbation_ratios']=ratios
107    results['runtime_sec']=time.time()-t
108    Path('results.json').write_text(json.dumps(results,indent=2))
109    print(json.dumps(results,indent=2))
110
111if __name__=='__main__': main()