1import json, math, random
 2from pathlib import Path
 3import numpy as np
 4import torch
 5from torch import nn
 6
 7SEED=1501
 8random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
 9torch.set_num_threads(4)
10device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
11
12# Switched scalar ODE: dz/dt=a_r z. Event is at t=1, horizon is 2.
13def truth(t, jump):
14    t=np.asarray(t); return np.exp(-0.7*np.minimum(t,1.0) + (-0.7+jump)*np.maximum(t-1.0,0.0))
15
16def mlp(inp, hidden=24):
17    return nn.Sequential(nn.Linear(inp,hidden),nn.Tanh(),nn.Linear(hidden,hidden),nn.Tanh(),nn.Linear(hidden,1)).to(device)
18
19def train_exact(jump, steps=900):
20    # Each phase is conditioned on its differentiable start; the phase-2 start is phase-1 terminal.
21    p1=mlp(2); p2=mlp(2); opt=torch.optim.Adam(list(p1.parameters())+list(p2.parameters()),lr=3e-3)
22    t1=torch.linspace(0,1,25,device=device).reshape(-1,1); t2=torch.linspace(0,1,25,device=device).reshape(-1,1)
23    y1=torch.tensor(truth(t1.cpu().numpy(),jump),dtype=torch.float32,device=device)
24    y2=torch.tensor(truth((1+t2).cpu().numpy(),jump),dtype=torch.float32,device=device)
25    z0=torch.ones((1,1),device=device)
26    for _ in range(steps):
27        opt.zero_grad()
28        s1=p1(torch.cat([t1, z0.expand_as(t1)],1)); end1=p1(torch.cat([torch.ones_like(t1[:1]),z0],1))
29        s2=p2(torch.cat([t2,end1.expand_as(t2)],1))
30        loss=((s1-y1)**2).mean()+((s2-y2)**2).mean()
31        loss.backward(); opt.step()
32    with torch.no_grad():
33        e1=p1(torch.cat([torch.ones(1,1,device=device),z0],1)); pred2=p2(torch.cat([t2,e1.expand_as(t2)],1))
34    return float(torch.sqrt(((pred2-y2)**2).mean()).cpu()), float(abs(e1.item()-e1.item()))
35
36def train_single(jump, steps=900):
37    p=mlp(1); opt=torch.optim.Adam(p.parameters(),lr=3e-3)
38    t=torch.linspace(0,2,50,device=device).reshape(-1,1)
39    y=torch.tensor(truth(t.cpu().numpy(),jump),dtype=torch.float32,device=device)
40    for _ in range(steps):
41        opt.zero_grad(); loss=((p(t)-y)**2).mean(); loss.backward(); opt.step()
42    with torch.no_grad():
43        post=t[:,0]>=1; rmse=torch.sqrt(((p(t)[post]-y[post])**2).mean())
44    return float(rmse.cpu())
45
46def soft_defect(lam, jump):
47    # Solve the convex quadratic interface objective with its exact linear system.
48    # This avoids conflating the predicted penalty scaling with optimizer failure.
49    q1=math.exp(-.7); q2=math.exp(-1.4+jump)
50    # Hessian (up to a common factor) and right hand side for
51    # (x1-q1)^2+(x2-q2)^2+lam*(x2-x1)^2.
52    A=np.array([[1+lam,-lam],[-lam,1+lam]],dtype=float)
53    x=np.linalg.solve(A,np.array([q1,q2]))
54    return abs(float(x[1]-x[0]))
55
56def math_checks():
57    # Prediction A: graph equality has zero defect at arbitrary states/phase maps.
58    defects=[]
59    for j in [0.,1.,3.,6.]:
60        x=torch.randn(17,1,requires_grad=True); f1=2*x+torch.sin(x); f2=f1**2
61        # f2 is deliberately evaluated from f1, exactly as phase chaining evaluates phase 2 from phase 1 terminal.
62        defects.append(float((f2-f1**2).abs().max()))
63    # Prediction B: quadratic soft interface has analytic defect d=|q2-q1|/(1+2 lambda).
64    # Prediction C (local calculus): an exact exponential phase trajectory has zero
65    # residual under autograd, independently in each phase.
66    tt=torch.tensor([[0.13],[0.77]],dtype=torch.float32,requires_grad=True)
67    aa=torch.tensor(-.7)
68    yy=torch.exp(aa*tt); dy=torch.autograd.grad(yy.sum(),tt,create_graph=False)[0]
69    residual_max=float((dy-aa*yy).abs().max())
70    q1=math.exp(-.7);
71    vals=[]
72    for j in [0.,1.,3.,6.]:
73        q2=math.exp(-1.4+j); gap=abs(q2-q1)
74        for lam in [0.,.1,1.,10.,100.]:
75            obs=soft_defect(lam,j); pred=gap/(1+2*lam)
76            vals.append({'jump':j,'lambda':lam,'observed':obs,'predicted':pred,'relerr':abs(obs-pred)/(abs(pred)+1e-8)})
77    # Prediction C: abruptness should increase single-network post-event error relative to chained fit.
78    jumps=[0.,1.,3.,6.]; pairs=[]
79    for j in jumps:
80        e_single=train_single(j); e_chain,_=train_exact(j)
81        pairs.append({'jump':j,'single_post_rmse':e_single,'exact_chain_post_rmse':e_chain,'advantage':e_single-e_chain})
82    return {'interface_max_defects':defects,'phase_residual_max':residual_max,'soft_scaling':vals,'event_severity':pairs}
83
84def main():
85    out=math_checks()
86    Path('results.json').write_text(json.dumps(out,indent=2))
87    print(json.dumps({'device':str(device),'results':out},indent=2))
88if __name__=='__main__':
89    try: main()
90    except Exception as e:
91        if device.type=='cuda':
92            print('CUDA failed, rerun on CPU:',repr(e)); device=torch.device('cpu'); main()
93        else: raise