1import json, math, random, time
  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)
  9device = "cuda" if torch.cuda.is_available() else "cpu"
 10# A conservative fallback: all computations fit easily on CPU or a small GPU.
 11try:
 12    if device == "cuda":
 13        torch.cuda.empty_cache()
 14except Exception:
 15    device = "cpu"
 16
 17dtype = torch.float32
 18omega = torch.tensor([1.0, math.sqrt(2.0)], device=device, dtype=dtype)
 19
 20# K(theta) = [cos(theta1), sin(theta1), cos(theta2), sin(theta2)].
 21# This is a finite Fourier embedding; coefficients are trainable in the experiment.
 22class FourierTorus(nn.Module):
 23    def __init__(self):
 24        super().__init__()
 25        # columns: cosine/sine coefficients for each phase and state coordinate
 26        self.a = nn.Parameter(torch.tensor([[1.,0.],[0.,1.],[0.,0.],[0.,0.]], device=device))
 27        self.b = nn.Parameter(torch.tensor([[0.,0.],[1.,0.],[0.,0.],[0.,1.]], device=device))
 28    def forward(self, theta):
 29        # theta [N,2], output [N,4]
 30        return torch.cos(theta) @ self.a.T + torch.sin(theta) @ self.b.T
 31    def tangent(self, theta):
 32        # dK/dtheta_j * omega_j, exact Fourier differentiation
 33        return (-torch.sin(theta) * omega) @ self.a.T + (torch.cos(theta) * omega) @ self.b.T
 34
 35class Field(nn.Module):
 36    def __init__(self):
 37        super().__init__()
 38        self.net = nn.Sequential(nn.Linear(4, 48), nn.Tanh(), nn.Linear(48, 48), nn.Tanh(), nn.Linear(48,4))
 39    def forward(self, x): return self.net(x)
 40
 41def true_field(x):
 42    y = torch.zeros_like(x)
 43    y[:,0] = -omega[0]*x[:,1]; y[:,1] = omega[0]*x[:,0]
 44    y[:,2] = -omega[1]*x[:,3]; y[:,3] = omega[1]*x[:,2]
 45    return y
 46
 47def rk4(model, x, dt):
 48    k1=model(x); k2=model(x+dt*k1/2); k3=model(x+dt*k2/2); k4=model(x+dt*k3)
 49    return x + dt*(k1+2*k2+2*k3+k4)/6
 50
 51def make_data(n=192, noise=0.025):
 52    # Short observations: only a fraction of the long quasiperiodic beat.
 53    t = torch.rand(n, device=device)*7.0
 54    phases = torch.stack([0.4 + t, 1.1 + math.sqrt(2.0)*t], dim=1)
 55    x = torch.stack([torch.cos(phases[:,0]), torch.sin(phases[:,0]), torch.cos(phases[:,1]), torch.sin(phases[:,1])], dim=1)
 56    y = true_field(x) + noise*torch.randn_like(x)
 57    return x, y
 58
 59def train(use_torus, seed=SEED, steps=1400):
 60    torch.manual_seed(seed)
 61    x, y = make_data()
 62    field, torus = Field().to(device), FourierTorus().to(device)
 63    opt = torch.optim.Adam(list(field.parameters()) + (list(torus.parameters()) if use_torus else []), lr=3e-3)
 64    L=12
 65    grid = torch.stack(torch.meshgrid(torch.arange(L, device=device)*2*math.pi/L,
 66                                      torch.arange(L, device=device)*2*math.pi/L, indexing='ij'), dim=-1).reshape(-1,2)
 67    for step in range(steps):
 68        opt.zero_grad()
 69        pred = field(x)
 70        loss = ((pred-y)**2).mean()
 71        if use_torus:
 72            K = torus(grid); T = torus.tangent(grid)
 73            colloc = ((field(K)-T)**2).mean()
 74            # Gauge: preserve the nonconstant Fourier embedding amplitudes and avoid collapse.
 75            gauge = (torus.a.pow(2).sum(dim=0)-1).pow(2).mean() + (torus.b.pow(2).sum(dim=0)-1).pow(2).mean()
 76            loss = loss + 0.8*colloc + 0.03*gauge
 77        loss.backward(); opt.step()
 78    # Evaluate on a long trajectory with exact starting point.
 79    t = torch.linspace(0, 80, 1601, device=device)
 80    ph = torch.stack([0.4+t, 1.1+math.sqrt(2.0)*t], dim=1)
 81    truth = torch.stack([torch.cos(ph[:,0]),torch.sin(ph[:,0]),torch.cos(ph[:,1]),torch.sin(ph[:,1])], dim=1)
 82    z=truth[0:1].clone(); traj=[z]
 83    with torch.no_grad():
 84        for _ in range(1600):
 85            z=rk4(field,z,0.05); traj.append(z)
 86        predtraj=torch.cat(traj)
 87        err=torch.sqrt(((predtraj-truth)**2).sum(dim=1))
 88        radii=torch.stack([torch.sqrt(predtraj[:,0]**2+predtraj[:,1]**2), torch.sqrt(predtraj[:,2]**2+predtraj[:,3]**2)],dim=1)
 89        one_step=((field(x)-true_field(x))**2).mean().sqrt().item()
 90        # Collocation residual on the same fixed grid.
 91        cres=((field(torus(grid))-torus.tangent(grid))**2).mean().sqrt().item() if use_torus else float('nan')
 92    return {'long_rmse': float(torch.sqrt((predtraj-truth).pow(2).mean()).item()),
 93            'final_error': float(err[-1].item()), 'max_radius_dev': float((radii-1).abs().max().item()),
 94            'one_step_field_rmse':one_step, 'collocation_rmse':cres}
 95
 96def derivative_check():
 97    torch.manual_seed(3); tor=FourierTorus().to(device)
 98    th=torch.rand(11,2,device=device)*2*math.pi
 99    eps=1e-3
100    direction=torch.tensor([0.37,-0.52],device=device)
101    fd=(tor(th+eps*direction)-tor(th-eps*direction))/(2*eps)
102    # Check the actual Fourier derivative independently, using arbitrary direction.
103    actual=(-torch.sin(th)*direction)@tor.a.T + (torch.cos(th)*direction)@tor.b.T
104    return float((fd-actual).abs().max().item()), float((tor.tangent(th)-((-torch.sin(th)*omega)@tor.a.T+(torch.cos(th)*omega)@tor.b.T)).abs().max().item())
105
106if __name__ == '__main__':
107    t0=time.time(); fd_err, tangent_err=derivative_check()
108    base=[train(False, SEED+i) for i in range(3)]
109    idea=[train(True, SEED+i) for i in range(3)]
110    def avg(rows):
111        return {k: float(np.mean([r[k] for r in rows])) for k in rows[0]}
112    out={'device':device,'derivative_fd_max_abs':fd_err,'analytic_tangent_self_check':tangent_err,
113         'baseline':avg(base),'idea':avg(idea),'baseline_runs':base,'idea_runs':idea,
114         'seconds':time.time()-t0}
115    print(json.dumps(out, indent=2)); open('results.json','w').write(json.dumps(out, indent=2, allow_nan=True))