1import json
 2from pathlib import Path
 3import numpy as np
 4import math
 5import torch
 6from torch import nn
 7
 8SEED = 17
 9np.random.seed(SEED)
10torch.manual_seed(SEED)
11device = "cuda" if torch.cuda.is_available() else "cpu"
12try:
13    if device == "cuda": torch.cuda.set_device(0)
14except Exception:
15    device = "cpu"
16
17# Exact underdamped oscillator labels, vectorized.
18# q'' + .7 q' + 1.5 q = 0, q(0)=p, q'(0)=v.
19def true_flow(x, t=1.0):
20    p, v = x[:, 0], x[:, 1]
21    alpha = 0.35
22    omega = (1.5 - alpha * alpha) ** 0.5
23    e = math.exp(-alpha * t)
24    c, s = math.cos(omega * t), math.sin(omega * t)
25    q = e * (p * c + (v + alpha * p) * s / omega)
26    qv = e * (v * c - (alpha * v + 1.5 * p) * s / omega)
27    return torch.stack((q, qv), dim=1)
28
29class Field(nn.Module):
30    def __init__(self):
31        super().__init__()
32        self.net = nn.Sequential(nn.Linear(2, 24), nn.Tanh(), nn.Linear(24, 24), nn.Tanh(), nn.Linear(24, 2))
33    def forward(self, z): return self.net(z)
34
35def rollout(model, z, factors=None, steps=16, T=1.0):
36    dt = T / steps
37    for i in range(steps):
38        fac = 1.0 if factors is None else factors[:, i:i+1]
39        z = z + dt * fac * model(z)
40    return z
41
42def train(use_ss):
43    torch.manual_seed(SEED)
44    model = Field().to(device)
45    opt = torch.optim.Adam(model.parameters(), lr=4e-3)
46    for _ in range(180):
47        x = torch.rand(64, 2, device=device) * 2 - 1
48        y = true_flow(x)
49        nominal = rollout(model, x)
50        loss = ((nominal - y) ** 2).mean()
51        if use_ss:
52            eps = 0.08
53            factors = 1 + (2 * torch.rand(64, 16, device=device) - 1) * eps
54            pseudo = rollout(model, x, factors)
55            # rho is deliberately nonzero as in the proposed hinge loss.
56            ss = torch.relu(torch.linalg.vector_norm(pseudo - nominal, dim=1) - .08).pow(2).mean()
57            loss = loss + 2.0 * ss
58        opt.zero_grad(); loss.backward(); opt.step()
59    with torch.no_grad():
60        x = torch.rand(256, 2, device=device) * 2 - 1
61        y = true_flow(x)
62        nominal = rollout(model, x)
63        factors = 1 + (2 * torch.rand(256, 16, device=device) - 1) * .08
64        pseudo = rollout(model, x, factors)
65        return {"clean_rmse": float(((nominal-y)**2).mean().sqrt()),
66                "timing_perturbed_rmse": float(((pseudo-y)**2).mean().sqrt()),
67                "mean_standard_tracking_error": float(torch.linalg.vector_norm(pseudo-nominal,dim=1).mean())}
68
69if __name__ == "__main__":
70    try:
71        result = {"device": device, "baseline": train(False), "standard_shadowing": train(True)}
72    except Exception as e:
73        torch.manual_seed(SEED); device = "cpu"
74        result = {"device": "cpu", "error_fallback": repr(e), "baseline": train(False), "standard_shadowing": train(True)}
75    Path("neural_ode_results.json").write_text(json.dumps(result, indent=2))
76    print(json.dumps(result, indent=2))