1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4from scipy.stats import norm
  5from sklearn.linear_model import Lasso
  6import torch
  7import torch.nn as nn
  8
  9SEED=1639
 10np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
 11try:
 12    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 13    # small tensors only; fall back if CUDA initialization fails
 14    torch.zeros(1, device=device)
 15except Exception:
 16    device=torch.device('cpu')
 17
 18# True oscillator has a known linear base and an unobserved, slowly varying forcing.
 19# The observed acceleration residual is predominantly -0.8*x^3 + 0.2*u.
 20def true_acc(x,v,u,h):
 21    return -x - .8*x**3 + .2*u + .18*h
 22
 23def make_data(ntraj=80, steps=40, dt=.04, noise=.015, seed=SEED):
 24    rng=np.random.default_rng(seed); rows=[]; test=[]
 25    for q in range(ntraj):
 26        x,v,h=rng.uniform(-1.0,1.0),rng.uniform(-.7,.7),rng.uniform(-.4,.4)
 27        for k in range(steps):
 28            u=np.sin(.07*k+q*.31)
 29            a=true_acc(x,v,u,h)
 30            # observed derivative target, with measurement/derivative noise
 31            rows.append([x,v,u,a + rng.normal(0,noise)])
 32            x=x+dt*v; v=v+dt*a; h=.995*h + .025*np.sin(.11*k+q)
 33    # independent trajectories for rollout evaluation
 34    for q in range(40):
 35        x,v,h=rng.uniform(-1.0,1.0),rng.uniform(-.7,.7),rng.uniform(-.4,.4)
 36        traj=[]
 37        for k in range(100):
 38            u=np.sin(.07*k+(q+100)*.31); traj.append((x,v,u))
 39            a=true_acc(x,v,u,h); x=x+dt*v; v=v+dt*a; h=.995*h+.025*np.sin(.11*k+q+100)
 40        test.append(np.asarray(traj))
 41    return np.asarray(rows), test
 42
 43def features(z):
 44    x,v,u=z[:,0],z[:,1],z[:,2]
 45    return np.column_stack([x,x**3,u,x*u,v,v**3])
 46
 47class Residual(nn.Module):
 48    def __init__(self):
 49        super().__init__(); self.net=nn.Sequential(nn.Linear(3,32),nn.Tanh(),nn.Linear(32,32),nn.Tanh(),nn.Linear(32,1))
 50    def forward(self,x): return self.net(x).squeeze(-1)
 51
 52def fit_neural(X, y, epochs=350):
 53    m=Residual().to(device); opt=torch.optim.Adam(m.parameters(),lr=.008)
 54    tx=torch.tensor(X,dtype=torch.float32,device=device); ty=torch.tensor(y,dtype=torch.float32,device=device)
 55    for _ in range(epochs):
 56        opt.zero_grad(); loss=((m(tx)-ty)**2).mean(); loss.backward(); opt.step()
 57    return m
 58
 59def sparse_project(P, target, alpha=.002, tau=2.5):
 60    # elastic-net-like L1 projection; standardize only for optimization, then unscale.
 61    mu=P.mean(0); sd=P.std(0)+1e-8; Q=(P-mu)/sd
 62    reg=Lasso(alpha=alpha,fit_intercept=True,max_iter=10000).fit(Q,target)
 63    b=reg.coef_/sd
 64    # OLS refit on active terms, with empirical covariance and t-statistics.
 65    active=np.flatnonzero(np.abs(b)>1e-5)
 66    selected=[]; bref=np.zeros(P.shape[1]); t=np.zeros(P.shape[1])
 67    if len(active):
 68        A=np.column_stack([np.ones(len(P)),P[:,active]])
 69        coef=np.linalg.lstsq(A,target,rcond=None)[0]
 70        resid=target-A@coef; dof=max(1,len(P)-A.shape[1]); s2=(resid@resid)/dof
 71        cov=s2*np.linalg.pinv(A.T@A); se=np.sqrt(np.maximum(np.diag(cov),1e-14))
 72        for i,j in enumerate(active):
 73            bref[j]=coef[i+1]; t[j]=coef[i+1]/se[i+1]
 74            if abs(t[j])>tau: selected.append(j)
 75    return bref, selected, t
 76
 77def acceleration_symbolic(z,b): return features(np.asarray(z).reshape(1,3))[0]@b
 78
 79def rollout(model_kind,b=None,m=None,dt=.04,seed=999):
 80    rng=np.random.default_rng(seed); errs=[]; one=[]
 81    for q in range(40):
 82        x,v,h=rng.uniform(-1,1),rng.uniform(-.7,.7),rng.uniform(-.4,.4); xp,vp=x,v
 83        for k in range(100):
 84            u=np.sin(.07*k+(q+100)*.31); a=true_acc(x,v,u,h)
 85            if model_kind=='neural':
 86                with torch.no_grad(): ar=float(m(torch.tensor([[x,v,u]],dtype=torch.float32,device=device)).cpu())
 87            else: ar=float(features(np.array([[x,v,u]]))[0]@b)
 88            # both models share known base; idea replaces neural correction by sparse field
 89            ap=-x+ar; at=a
 90            x=x+dt*v; v=v+dt*ap; h=.995*h+.025*np.sin(.11*k+q+100)
 91            xp=xp+dt*v # only used as a benign numerical reference
 92            one.append((ap-at)**2); errs.append((ap-at)**2)
 93    return float(np.sqrt(np.mean(one))), float(np.sqrt(np.mean(errs)))
 94
 95def mechanism_sweep():
 96    # OLS theory predicts t(signal) proportional to sqrt(N); null t is N(0,1),
 97    # so a two-sided threshold tau has acceptance 2*Phi(-tau).
 98    rng=np.random.default_rng(44); out=[]
 99    beta=.8; sigma=.25; tau=2.5
100    for n in [40,80,160,320,640]:
101        ts=[]; null_accept=[]
102        for rep in range(180):
103            x=rng.normal(size=n); X=np.column_stack([x,x**3,x**2])
104            y=beta*x**3+rng.normal(0,sigma,n)
105            c=np.linalg.lstsq(np.column_stack([np.ones(n),X]),y,rcond=None)[0]
106            e=y-np.column_stack([np.ones(n),X])@c; s2=e@e/(n-4); cov=s2*np.linalg.pinv(np.column_stack([np.ones(n),X]).T@np.column_stack([np.ones(n),X]))
107            ts.append(c[2]/math.sqrt(cov[2,2])); null_accept.append(abs(c[3]/math.sqrt(cov[3,3]))>tau)
108        out.append({'n':n,'signal_t_mean':float(np.mean(ts)),'sqrt_n_scaled':float(np.mean(ts)/math.sqrt(n)),'null_accept_rate':float(np.mean(null_accept))})
109    ratio=out[-1]['signal_t_mean']/out[0]['signal_t_mean']
110    slope=np.polyfit(np.log([x['n'] for x in out]),np.log([x['signal_t_mean'] for x in out]),1)[0]
111    return {'predicted_signal_t_ratio_640_over_40':math.sqrt(640/40),'observed_signal_t_ratio_640_over_40':float(ratio),'predicted_t_scaling_exponent':0.5,'observed_t_scaling_exponent':float(slope),'predicted_null_accept_rate_tau_2.5':2*norm.sf(2.5),'observed_null_accept_rate_mean':float(np.mean([x['null_accept_rate'] for x in out])),'observed':out}
112
113def main():
114    data,test=make_data(); X=data[:,:3]; y=data[:,3]
115    # Train on derivative residual relative to known base F_obs acceleration=-x.
116    residual_target=y + X[:,0]
117    m=fit_neural(X,residual_target)
118    with torch.no_grad(): pred=m(torch.tensor(X,dtype=torch.float32,device=device)).cpu().numpy()
119    P=features(X); b,sel,t=sparse_project(P,pred)
120    # Refit accepted coefficients directly to noisy physical residual for fair final field.
121    if sel:
122        A=np.column_stack([np.ones(len(P)),P[:,sel]]); co=np.linalg.lstsq(A,residual_target,rcond=None)[0]; bf=np.zeros(6); bf[sel]=co[1:]; b=bf
123    # Sparse regression directly on the measured residual is the non-neural baseline.
124    b0,sel0,t0=sparse_project(P,residual_target)
125    bs=np.zeros(6)
126    if sel0:
127        A0=np.column_stack([np.ones(len(P)),P[:,sel0]])
128        co0=np.linalg.lstsq(A0,residual_target,rcond=None)[0]
129        bs[sel0]=co0[1:]
130    nr=rollout('neural',m=m); sr=rollout('symbolic',b=b); sparse_r=rollout('symbolic',b=bs)
131    result={'device':str(device),'n_train':len(X),'selected_terms':[ ['x','x3','u','xu','v','v3'][j] for j in sel], 'coefficients':b.tolist(),'t_statistics':t.tolist(), 'neural_train_rmse':float(np.sqrt(np.mean((pred-residual_target)**2))), 'rollout_rmse_acceleration':{'neural':nr[1],'sparse_only':sparse_r[1],'residual_to_symbolic':sr[1]}, 'sparse_only_terms':[['x','x3','u','xu','v','v3'][j] for j in sel0], 'mechanism_sweep':mechanism_sweep()}
132    Path('results.json').write_text(json.dumps(result,indent=2))
133    print(json.dumps(result,indent=2))
134if __name__=='__main__': main()