1import json, math, random
 2from pathlib import Path
 3import numpy as np
 4import torch
 5
 6SEED = 2270
 7random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
 8torch.set_default_dtype(torch.float64)
 9
10# A two-dimensional neural-ODE-like polynomial vector field.
11# At p=(0,0), J=[[0,1],[0,0]], hence J is rank-one nilpotent and ker(J)=span((1,0)).
12# F2=c*x^2 + k*x^3 + d*x*y + e*x^2*y + r*y^2.
13# Along q0=(1,0): a=c and b=d exactly.
14def field(z, c, d, k=0.37, e=-0.21, r=0.13, m=0.8):
15    x, y = z[0], z[1]
16    return torch.stack((y, c*x*x + k*x**3 + m*x**4 + d*x*y + e*x*x*y + r*y*y))
17
18def jacobian_at(z, c, d, k=0.37, e=-0.21, r=0.13, m=0.8):
19    z = z.detach().clone().requires_grad_(True)
20    return torch.autograd.functional.jacobian(lambda zz: field(zz, c, d, k, e, r, m), z)
21
22def monitor(c, d, h):
23    p = torch.zeros(2)
24    J = jacobian_at(p, c, d)
25    # right singular vector for the smallest singular value
26    _, _, vh = torch.linalg.svd(J)
27    q = vh[-1]
28    q = q / torch.linalg.vector_norm(q)
29    # orient deterministically to make reported values comparable
30    if q[0] < 0: q = -q
31    jp = jacobian_at(p + h*q, c, d)
32    jm = jacobian_at(p - h*q, c, d)
33    a = -0.5 * (torch.linalg.det(jp)-torch.linalg.det(jm))/(2*h)
34    b = (torch.trace(jp)-torch.trace(jm))/(2*h)
35    return float(a), float(b), float(torch.linalg.det(J)), float(torch.trace(J)), q.numpy().tolist()
36
37def math_checks():
38    # Prediction 1: a scales linearly with c; prediction 2: b scales linearly with d.
39    cs = np.linspace(-1.2, 1.2, 9); ds = np.linspace(-1.2, 1.2, 9)
40    av = np.array([monitor(float(c), .41, 1e-3)[0] for c in cs])
41    bv = np.array([monitor(.41, float(d), 1e-3)[1] for d in ds])
42    slope_a, intercept_a = np.polyfit(cs, av, 1)
43    slope_b, intercept_b = np.polyfit(ds, bv, 1)
44    # Prediction 3: symmetric finite differences are stable as h is reduced, including at zero crossings.
45    hs = [1e-1, 5e-2, 2.5e-2, 1.25e-2, 6.25e-3]
46    conv = [monitor(.73, -.62, h)[:2] for h in hs]
47    zero_a = monitor(0., .4, 1e-3)[0]; zero_b = monitor(.4, 0., 1e-3)[1]
48    return {
49        'predictions': {
50            'a(c) predicted slope/intercept': [1.0, 0.0],
51            'a(c) observed slope/intercept': [float(slope_a), float(intercept_a)],
52            'b(d) predicted slope/intercept': [1.0, 0.0],
53            'b(d) observed slope/intercept': [float(slope_b), float(intercept_b)],
54            'zero crossing a at c=0 observed': zero_a,
55            'zero crossing b at d=0 observed': zero_b,
56            'finite_difference predicted a(h),b(h)': [[0.73 + 1.6*h*h, -0.62] for h in hs],
57            'finite_difference observed h=[0.1,...]': conv,
58        },
59        'nilpotent_check': monitor(.73, -.62, 1e-3)[2:]
60    }
61
62def train(use_bt, steps=180, lr=.035):
63    # Same tiny task: recover coefficients c=d=1 from a weak signal. BT term avoids
64    # the degenerate region but does not replace the task objective.
65    c = torch.tensor(0.015, requires_grad=True)
66    d = torch.tensor(-0.02, requires_grad=True)
67    opt = torch.optim.Adam([c, d], lr=lr)
68    eps = .35; lam = 0.8
69    history=[]
70    for step in range(steps):
71        task = (c-1.)**2 + (d-1.)**2
72        a,b,_,_,_ = monitor(float(c.detach()), float(d.detach()), 1e-3)
73        # Rebuild differentiable exact monitor for this controlled family: a=c,b=d.
74        bt = lam*(torch.relu(torch.tensor(eps)-torch.abs(c))**2 + torch.relu(torch.tensor(eps)-torch.abs(d))**2)
75        loss = task + (bt if use_bt else 0.)
76        opt.zero_grad(); loss.backward(); opt.step()
77        if step in (0, 9, 29, 89, 179): history.append((step, float(task), float(a), float(b)))
78    return {'c':float(c), 'd':float(d), 'task_loss':float((c-1.)**2+(d-1.)**2), 'history':history}
79
80def main():
81    out = {'math': math_checks(), 'training': {'baseline':train(False), 'bt_penalty':train(True)}}
82    Path('results.json').write_text(json.dumps(out, indent=2))
83    print(json.dumps(out, indent=2))
84
85if __name__ == '__main__': main()