1import json, math, random
 2from pathlib import Path
 3import numpy as np
 4from scipy.optimize import minimize
 5
 6# J_{a,b}(u)=(u^2-1)^2+a*u^2+b*u, u in [-2,2].
 7# At b=0: N=2 for a<2, N=1 for a>=2. Around a<2, r=sqrt(1-a/2),
 8# H(r)=4(2-a), so the positive-well displacement is |b|/[4(2-a)] + O(b^2).
 9def objective(u, a, b):
10    x = float(np.asarray(u).reshape(-1)[0])
11    return (x*x-1.0)**2 + a*x*x + b*x
12
13def grad(x, a, b):
14    x = float(x)
15    return 4*x**3 + 2*(a-2)*x + b
16
17def hess(x, a, b):
18    x = float(x)
19    return 12*x*x + 2*(a-2)
20
21def audit(a, b, starts, eps=1e-3, grad_tol=2e-6, hess_tol=2e-6):
22    ends=[]
23    vals=[]
24    for x0 in starts:
25        r=minimize(lambda z: objective(z,a,b), [float(x0)], jac=lambda z: np.array([grad(z[0],a,b)]),
26                   bounds=[(-2,2)], method='L-BFGS-B', options={'ftol':1e-14,'gtol':1e-10,'maxiter':300})
27        x=float(r.x[0])
28        if abs(grad(x,a,b)) <= grad_tol and hess(x,a,b) >= -hess_tol:
29            ends.append(x); vals.append(objective(x,a,b))
30    clusters=[]
31    for x,v in sorted(zip(ends,vals)):
32        if not clusters or abs(x-clusters[-1]['u']) > eps:
33            clusters.append({'u':x,'J':v,'members':1})
34        else:
35            c=clusters[-1]; c['u']=(c['u']*c['members']+x)/(c['members']+1); c['J']=min(c['J'],v); c['members']+=1
36    best=min(clusters,key=lambda c:c['J']) if clusters else {'u':float('nan'),'J':float('nan')}
37    return {'N':len(clusters),'best_u':best['u'],'best_J':best['J'],'clusters':clusters}
38
39def main():
40    np.random.seed(7); random.seed(7)
41    starts=np.linspace(-1.95,1.95,81)
42    # Prediction 1: pitchfork count transition at a=2.
43    count_rows=[]
44    for a in np.linspace(0,3,13):
45        out=audit(float(a),0.0,starts)
46        predicted=2 if a < 2 else 1
47        count_rows.append({'a':float(a),'N_observed':out['N'],'N_predicted':predicted})
48    # Prediction 2: displacement is linear in |b| with slope 1/[4(2-a)].
49    a=1.0; ref=audit(a,0.0,starts)['best_u']
50    drift_rows=[]
51    for b in np.linspace(0,0.16,9):
52        out=audit(a,float(b),starts)
53        # Track positive well, which is the best well for b>=0; reference is +1.
54        observed=abs(out['best_u']-ref)
55        predicted=abs(b)/(4*(2-a))
56        drift_rows.append({'b':float(b),'e_observed':observed,'e_predicted_linear':predicted})
57    # Prediction 3: larger curvature gap (2-a) suppresses drift; slope sweep.
58    slope_rows=[]
59    for aa in [0.5,1.0,1.5]:
60        rr=audit(aa,0.0,starts); r=rr['best_u']; bs=np.array([0.02,0.04,0.06])
61        es=[]
62        for bb in bs: es.append(abs(audit(aa,float(bb),starts)['best_u']-r))
63        fit=float(np.dot(bs,es)/np.dot(bs,bs))
64        slope_rows.append({'a':aa,'slope_observed':fit,'slope_predicted':1/(4*(2-aa))})
65    # Drift trajectory: nearly identical prediction loss, but surrogate decision drifts.
66    traj=[]
67    for t,b in enumerate(np.linspace(0,0.18,19)):
68        out=audit(1.0,float(b),starts)
69        e=abs(out['best_u']-ref)
70        val_loss=1e-4*(1.0-b/0.18) # mildly improving but flat observational metric
71        traj.append({'t':t,'b':float(b),'validation_loss':float(val_loss),'e':float(e),'N':out['N'],'J':float(out['best_J'])})
72    tau_u=0.02; tau_J=0.03; nmax=2
73    accepted=[r for r in traj if r['e']<=tau_u and abs(r['J']-traj[0]['J'])<=tau_J and r['N']<=nmax]
74    # validation-only chooses last (slightly improving) checkpoint; audit chooses latest accepted.
75    baseline=traj[-1]; audited=accepted[-1]
76    result={'seed':7,'count_sweep':count_rows,'drift_sweep':drift_rows,'curvature_sweep':slope_rows,
77            'trajectory':traj,'selection':{'baseline_validation_only':baseline,'audit_selected':audited,
78             'tau_u':tau_u,'tau_J':tau_J,'N_max':nmax},
79            'summary':{'count_all_correct':all(x['N_observed']==x['N_predicted'] for x in count_rows),
80             'max_drift_linear_abs_error':max(abs(x['e_observed']-x['e_predicted_linear']) for x in drift_rows),
81             'slope_relative_errors':[abs(x['slope_observed']-x['slope_predicted'])/x['slope_predicted'] for x in slope_rows]}}
82    Path('results.json').write_text(json.dumps(result,indent=2))
83    print(json.dumps(result['summary'],indent=2))
84    print('selection',json.dumps(result['selection'],indent=2))
85
86if __name__=='__main__': main()