1import json, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5
  6SEED = 2553
  7torch.set_default_dtype(torch.float64)
  8np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  9
 10# Canonical BT unfolding:
 11# x' = y
 12# y' = mu*x + nu*y + alpha*x^2 + beta*x*y
 13# At mu=nu=0 and z*=0, J=[[0,1],[0,0]], q0=(1,0).
 14def field(z, p):
 15    x, y = z
 16    mu, nu, alpha, beta = p
 17    return torch.stack((y, mu*x + nu*y + alpha*x*x + beta*x*y))
 18
 19def jacobian(z, p):
 20    # Explicit form keeps the z -> J -> invariant graph for higher derivatives.
 21    x, y = z
 22    mu, nu, alpha, beta = p
 23    return torch.stack((torch.stack((0*x, 1+0*x)),
 24                        torch.stack((mu + 2*alpha*x + beta*y, nu + beta*x))))
 25
 26def invariants(J):
 27    e1 = J[0, 0] + J[1, 1]
 28    e2 = J[0, 0]*J[1, 1] - J[0, 1]*J[1, 0]
 29    return e1, e2  # e1, e2; polynomial form supports higher derivatives
 30
 31def monitor(p, q_override=None, z_value=(0., 0.)):
 32    z = torch.tensor(z_value, dtype=p.dtype, requires_grad=True)
 33    J = jacobian(z, p)
 34    e1, e2 = invariants(J)
 35    if q_override is None:
 36        _, _, vh = torch.linalg.svd(J.detach())
 37        q = vh[-1].detach()
 38    else:
 39        q = torch.as_tensor(q_override, dtype=z.dtype)
 40    # For n=2: delta=e2/e0=e2 and tau=e1.
 41    gdelta = torch.autograd.grad(e2, z, retain_graph=True)[0]
 42    gtau = torch.autograd.grad(e1, z)[0]
 43    Ddelta = torch.dot(gdelta, q)
 44    Dtau = torch.dot(gtau, q)
 45    return {
 46        'J': J.detach().numpy().tolist(), 'q': q.numpy().tolist(),
 47        'delta': float(e2.detach()), 'tau': float(e1.detach()),
 48        'a': float((-0.5*Ddelta).detach()), 'b': float(Dtau.detach()),
 49        'e_transverse': 1.0,  # e_(n-2)=e0 for n=2
 50    }
 51
 52def finite_difference(alpha, beta, h=1e-5):
 53    p = torch.tensor([0., 0., alpha, beta])
 54    def vals(x):
 55        r = monitor(p, [1., 0.], (x, 0.))
 56        return r['delta'], r['tau']
 57    dm = (vals(h)[0] - vals(-h)[0])/(2*h)
 58    dt = (vals(h)[1] - vals(-h)[1])/(2*h)
 59    return -0.5*dm, dt
 60
 61def run_target(kind):
 62    torch.manual_seed(SEED)
 63    p = torch.nn.Parameter(torch.tensor([0.65, -0.55, 1.2, -0.8]))
 64    opt = torch.optim.Adam([p], lr=0.05)
 65    for _ in range(250):
 66        mu, nu = p[0], p[1]
 67        if kind == 'monitor':
 68            loss = mu*mu + nu*nu
 69        else:
 70            # Spectral-abscissa stability penalty: stable systems receive zero loss.
 71            disc = nu*nu + 4*mu
 72            rp = torch.where(disc >= 0, (nu + torch.sqrt(torch.clamp(disc, min=0)))/2, nu/2)
 73            loss = torch.relu(rp)**2
 74        opt.zero_grad(); loss.backward(); opt.step()
 75    r = monitor(p.detach(), [1., 0.])
 76    return {'final_p': p.detach().numpy().tolist(), 'delta': r['delta'], 'tau': r['tau'],
 77            'monitor_loss': r['delta']**2 + r['tau']**2}
 78
 79def verify():
 80    # Prediction 1: exact unfolding identities delta=-mu, tau=nu over a grid.
 81    identity = []
 82    for mu in np.linspace(-1, 1,  nine := 9):
 83        for nu in np.linspace(-1, 1, nine):
 84            r = monitor(torch.tensor([mu, nu, 1., 1.]), [1., 0.])
 85            identity.append((mu, nu, r['delta'], r['tau']))
 86    err_delta = max(abs(d + mu) for mu, nu, d, t in identity)
 87    err_tau = max(abs(t - nu) for mu, nu, d, t in identity)
 88
 89    # Prediction 2: a scales as alpha and b scales as beta at the BT point.
 90    coeff = []
 91    for alpha, beta in [(0.3,-0.7), (1.,2.), (-1.5,.25), (2.5,-3.)]:
 92        r = monitor(torch.tensor([0.,0.,alpha,beta]), [1.,0.])
 93        af, bf = finite_difference(alpha, beta)
 94        coeff.append({'alpha':alpha, 'beta':beta, 'a':r['a'], 'b':r['b'],
 95                      'a_fd':af, 'b_fd':bf, 'a_error':abs(r['a']-alpha),
 96                      'b_error':abs(r['b']-beta)})
 97
 98    # Prediction 3: the candidate is simultaneous delta=tau=0 only at (mu,nu)=(0,0).
 99    candidates = []
100    for mu in [-.2, 0., .2]:
101        for nu in [-.2, 0., .2]:
102            r = monitor(torch.tensor([mu,nu,1.,1.]), [1.,0.])
103            candidates.append({'mu':mu, 'nu':nu, 'delta':r['delta'], 'tau':r['tau'],
104                               'distance':abs(r['delta'])+abs(r['tau'])})
105    out = {
106        'prediction_1_unfolding': {'points':len(identity), 'max_abs_delta_minus_neg_mu':err_delta,
107                                   'max_abs_tau_minus_nu':err_tau},
108        'prediction_2_coefficients': coeff,
109        'prediction_3_candidate_grid': candidates,
110        'baseline': run_target('baseline'), 'idea_monitor': run_target('monitor')
111    }
112    return out
113
114def main():
115    result = verify()
116    Path('results.json').write_text(json.dumps(result, indent=2))
117    print(json.dumps(result, indent=2))
118
119if __name__ == '__main__':
120    main()