1import json, math, os
  2import numpy as np
  3
  4# Differentiable equilibrium projection MVP: F(z;u)=z^3+a*z-u.
  5# Its Jacobian is 3z^2+a, so dz*/du=1/(3z*^2+a).
  6
  7def solve_newton(u, a, z0, maxit=30, tol=1e-10):
  8    z=float(z0)
  9    for k in range(maxit):
 10        f=z**3+a*z-u
 11        j=3*z*z+a
 12        if abs(f)<tol: return z,k+1,abs(f)
 13        z -= f/j
 14    return z,maxit,abs(z**3+a*z-u)
 15
 16def solve_gd_residual(u,a,z0,tau,maxit=10000,tol=1e-8):
 17    # z <- z - tau J^T F, as in the proposed least-squares implicit iteration.
 18    z=float(z0)
 19    for k in range(maxit):
 20        f=z**3+a*z-u; j=3*z*z+a
 21        if abs(f)<tol: return z,k+1,abs(f)
 22        z -= tau*j*f
 23    return z,maxit,abs(z**3+a*z-u)
 24
 25def root(u,a): return solve_newton(u,a,0,maxit=100)[0]
 26
 27def main():
 28    rng=np.random.default_rng(7)
 29    rows=[]
 30    # Prediction 1: for linear F=a*z-u, residual contraction is |1-tau*a^2|,
 31    # and the stability boundary is tau*a^2=2.
 32    lin=[]
 33    for a in [0.5,1.,2.]:
 34        for tau in [0.25/a**2, 0.9/a**2, 1.1/a**2, 1.9/a**2, 2.1/a**2]:
 35            u=1.3; z=4.; r0=abs(a*z-u); rs=[]
 36            for k in range(12):
 37                rs.append(abs(a*z-u)); z -= tau*a*(a*z-u)
 38            observed=(rs[-1]/rs[0])**(1/11)
 39            predicted=abs(1-tau*a*a)
 40            lin.append({'a':a,'tau_a2':tau*a*a,'predicted_ratio':predicted,'observed_ratio':observed,'stable_pred':tau*a*a<2,'decreased':rs[-1]<rs[0]})
 41    rows.append(('linear_contraction',lin))
 42
 43    # Prediction 2: implicit sensitivity is inverse Jacobian; finite differences should scale 1/a
 44    sens=[]
 45    for a in [0.1,0.25,0.5,1.,2.,4.]:
 46        u=0.7; eps=1e-5
 47        fd=(root(u+eps,a)-root(u-eps,a))/(2*eps)
 48        pred=1/(3*root(u,a)**2+a)
 49        sens.append({'a':a,'predicted_1_over_J':pred,'finite_difference':fd,'relative_error':abs(fd-pred)/pred})
 50    rows.append(('implicit_sensitivity',sens))
 51
 52    # Controlled conditioning sweep at u=0: z*=0, J=a exactly, so |dz/du|=1/a.
 53    conditioning=[]
 54    for a in [1e-1, 3e-2, 1e-2, 3e-3, 1e-3]:
 55        u=0.0; eps=1e-6
 56        fd=(root(u+eps,a)-root(u-eps,a))/(2*eps)
 57        conditioning.append({'a_sigma_min_J':a,'predicted_gain':1/a,
 58                             'finite_difference_gain':fd,
 59                             'relative_error':abs(fd-1/a)/(1/a)})
 60    rows.append(('conditioning_gain_at_singular_limit',conditioning))
 61
 62    # Verify the custom implicit backward rule against torch autograd finite differences.
 63    try:
 64        import torch
 65        from implicit_layer import implicit_project
 66        u=torch.tensor(0.7,dtype=torch.double,requires_grad=True)
 67        a_t=torch.tensor(0.4,dtype=torch.double,requires_grad=True)
 68        z=implicit_project(u,a_t,steps=30)
 69        (z*z/2).backward()
 70        z0=float(z.detach()); j=3*z0*z0+float(a_t.detach())
 71        expected_u=z0/j; expected_a=-z0*z0/j
 72        implicit_check={'autograd_du':float(u.grad),'expected_du':expected_u,
 73                        'autograd_da':float(a_t.grad),'expected_da':expected_a,
 74                        'max_abs_error':max(abs(float(u.grad)-expected_u),abs(float(a_t.grad)-expected_a))}
 75    except Exception as e:
 76        implicit_check={'error':str(e)}
 77    rows.append(('implicit_backward_check',implicit_check))
 78
 79    # Prediction 3: Newton reaches tolerance from very different predictor errors,
 80    # while penalty-only optimization retains initialization/penalty dependence.
 81    init_sweep=[]
 82    u=1.2; a=.35
 83    for z0 in [-10,-2,0,2,10]:
 84        z,k,r=solve_newton(u,a,z0)
 85        init_sweep.append({'z0':z0,'newton_iters':k,'final_residual':r})
 86    rows.append(('newton_initialization_invariance',init_sweep))
 87
 88    # Small ML comparison: train direct predictor and physics-penalty predictor;
 89    # projected predictor uses same direct network output followed by Newton.
 90    import torch
 91    torch.manual_seed(7)
 92    torch.set_num_threads(4)
 93    n=500; uall=torch.linspace(-2,2,n).unsqueeze(1)
 94    aa=.35
 95    yall=torch.tensor([root(float(u),aa) for u in uall[:,0]]).float().unsqueeze(1)
 96    tr=slice(0,350); te=slice(350,n)
 97    def net(): return torch.nn.Sequential(torch.nn.Linear(1,24),torch.nn.Tanh(),torch.nn.Linear(24,1))
 98    def train(mode, epochs=500):
 99        m=net(); opt=torch.optim.Adam(m.parameters(),lr=.02)
100        for _ in range(epochs):
101            pred=m(uall[tr]); f=pred**3+aa*pred-uall[tr]
102            loss=((pred-yall[tr])**2).mean() if mode=='direct' else (f*f).mean()
103            opt.zero_grad(); loss.backward(); opt.step()
104        with torch.no_grad():
105            p=m(uall[te]); f=p**3+aa*p-uall[te]
106            if mode=='project': pass
107            pred_res=float(f.abs().mean()); mse=float(((p-yall[te])**2).mean())
108        return m,pred_res,mse
109    md,rd,ed=train('direct'); mp,rp,ep=train('penalty')
110    # deterministic projection of penalty network's predictions, fixed Newton solve
111    with torch.no_grad(): raw=mp(uall[te]).numpy().ravel()
112    projected=np.array([solve_newton(float(u),aa,float(z))[0] for u,z in zip(uall[te,0],raw)])
113    up=uall[te,0].numpy(); truth=yall[te,0].numpy()
114    proj_res=float(np.mean(np.abs(projected**3+aa*projected-up)))
115    proj_mse=float(np.mean((projected-truth)**2))
116    rows.append(('mini_experiment',{'direct_test_mean_residual':rd,'penalty_test_mean_residual':rp,'penalty_projected_residual':proj_res,'direct_mse':ed,'penalty_mse':ep,'projected_mse':proj_mse,'train_points':350,'test_points':150}))
117    with open('results.json','w') as f: json.dump(dict(rows),f,indent=2)
118    print(json.dumps(dict(rows),indent=2))
119
120if __name__=='__main__': main()