1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4
  5SEED = 2218
  6np.random.seed(SEED); random.seed(SEED)
  7try:
  8    import torch
  9    torch.manual_seed(SEED)
 10    if torch.cuda.is_available(): torch.cuda.manual_seed_all(SEED)
 11except Exception:
 12    torch = None
 13
 14
 15def weighted_scores(J, groups, noise=1.0):
 16    """Return tr(J_g'WJ_g) and Frobenius leave-one-group residuals."""
 17    A = np.asarray(J, float) / noise
 18    traces=[]; residuals=[]
 19    blocks=[A[:, g] if isinstance(g, slice) else A[:, list(g)] for g in groups]
 20    for i, block in enumerate(blocks):
 21        traces.append(float(np.sum(block * block)))
 22        others=[blocks[j] for j in range(len(blocks)) if j != i]
 23        if not others:
 24            residuals.append(float(np.sum(block * block))); continue
 25        other=np.concatenate(others, axis=1)
 26        q,_=np.linalg.qr(other, mode='reduced')
 27        proj=q @ (q.T @ block)
 28        residuals.append(float(np.sum((block-proj)**2)))
 29    return np.array(traces), np.array(residuals)
 30
 31
 32def mechanism_check():
 33    # Three quantitative predictions directly implied by I=J'WJ and projection.
 34    rng=np.random.default_rng(SEED)
 35    base=rng.normal(size=(30,3)); base[:,1]=base[:,0] # exact redundant pair
 36    groups=[slice(0,1),slice(1,2),slice(2,3)]
 37    # Prediction 1: scaling sensitivity by alpha scales information and residual by alpha^2.
 38    alphas=np.array([.25,.5,1.,2.,4.])
 39    tr=[]; rr=[]
 40    for a in alphas:
 41        t,r=weighted_scores(a*base, groups)
 42        tr.append(t[2]); rr.append(r[2])
 43    tr=np.array(tr); rr=np.array(rr)
 44    slope_I=np.polyfit(np.log(alphas), np.log(tr), 1)[0]
 45    slope_R=np.polyfit(np.log(alphas), np.log(rr), 1)[0]
 46    # Prediction 2: exact collinearity has zero incremental rank; residual rises as sin^2(angle).
 47    angles=np.array([0., .1, .2, .4, .8, 1.2])
 48    residual_angle=[]
 49    x=base[:,0]
 50    # Use one fixed orthogonal direction so only the prescribed angle changes.
 51    z=rng.normal(size=x.size); z-=z@x/(x@x)*x; z/=np.linalg.norm(z)
 52    for ang in angles:
 53        J=np.column_stack([x, math.cos(ang)*x+math.sin(ang)*z, base[:,2]])
 54        residual_angle.append(weighted_scores(J, groups)[1][1])
 55    residual_angle=np.array(residual_angle)
 56    ratio=np.divide(residual_angle[1:], np.sin(angles[1:])**2, out=np.zeros(5), where=np.sin(angles[1:])!=0)
 57    cv=float(np.std(ratio)/np.mean(ratio))
 58    # Prediction 3: inverse noise std weighting makes information scale as 1/sigma^2.
 59    sigmas=np.array([.5,1.,2.,4.]); infos=[]
 60    for s in sigmas: infos.append(weighted_scores(base, groups, noise=s)[0][2])
 61    noise_slope=np.polyfit(np.log(sigmas), np.log(infos), 1)[0]
 62    return {
 63        'alpha_values':alphas.tolist(), 'information_values':tr.tolist(), 'residual_values':rr.tolist(),
 64        'predicted_alpha_exponent':2.0, 'observed_information_exponent':float(slope_I),
 65        'observed_residual_exponent':float(slope_R),
 66        'angles_rad':angles.tolist(), 'angle_residuals':residual_angle.tolist(),
 67        'predicted_angle_law':'r = C sin(angle)^2', 'angle_law_ratio_cv':cv,
 68        'noise_sigmas':sigmas.tolist(), 'noise_information':infos,
 69        'predicted_noise_exponent':-2.0, 'observed_noise_exponent':float(noise_slope),
 70        'pass': bool(abs(slope_I-2)<.03 and abs(slope_R-2)<.03 and cv<.03 and abs(noise_slope+2)<.03)
 71    }
 72
 73
 74class UnitODE(torch.nn.Module):
 75    def __init__(self, width):
 76        super().__init__(); self.width=width
 77        self.w=torch.nn.Parameter(torch.randn(width)*.5)
 78        self.b=torch.nn.Parameter(torch.randn(width)*.1)
 79        self.a=torch.nn.Parameter(torch.randn(width)*.5/width**.5)
 80    def rhs(self,x): return (self.a*torch.tanh(self.w*x+self.b)).sum(-1,keepdim=True)
 81    def forward(self,x0,dt,steps):
 82        x=x0
 83        out=[x]
 84        for _ in range(steps):
 85            k1=self.rhs(x); k2=self.rhs(x+dt*k1/2); k3=self.rhs(x+dt*k2/2); k4=self.rhs(x+dt*k3)
 86            x=x+dt*(k1+2*k2+2*k3+k4)/6; out.append(x)
 87        return torch.cat(out,dim=1)
 88
 89
 90def train(width, x0, target, epochs=350, keep=None, device='cpu'):
 91    model=UnitODE(width).to(device)
 92    opt=torch.optim.Adam(model.parameters(),lr=.025)
 93    dt=float(2.5/(target.shape[1]-1))
 94    x0=x0.to(device); target=target.to(device)
 95    for _ in range(epochs):
 96        opt.zero_grad(); pred=model(x0,dt,target.shape[1]-1)
 97        loss=((pred-target)**2).mean(); loss.backward(); opt.step()
 98    return model
 99
100
101def unit_sensitivity(model, x0, times, device):
102    # Observed-output sensitivity columns using finite differences of actual trajectories.
103    # This is equivalent to local S propagated through the fixed-step ODE integrator.
104    model.eval(); x0=x0.to(next(model.parameters()).device); dt=float(2.5/(len(times)-1)); base=model(x0,dt,len(times)-1).detach().cpu().numpy().ravel()
105    cols=[]
106    with torch.no_grad():
107        for j in range(model.width):
108            vals=[]
109            for p in [model.w,model.b,model.a]:
110                old=float(p[j]); eps=1e-4*max(1.,abs(old)); p[j]=old+eps
111                vals.append((model(x0,dt,len(times)-1).detach().cpu().numpy().ravel()-base)/eps)
112                p[j]=old
113            cols.append(np.stack(vals,axis=1))
114    return np.stack(cols,axis=1), base
115
116
117def prune_and_score(kind, width=16, retain=8, device='cpu'):
118    t=np.linspace(0,2.5,31); x0=torch.tensor([[.7]],dtype=torch.float32); true=torch.zeros(1,len(t)); true[0,0]=.7
119    # stable nonlinear target generated by known scalar dynamics, integrated with RK4 externally
120    x=.7
121    for k in range(30):
122        f=lambda q: -0.8*q+0.35*q*q
123        k1=f(x); k2=f(x+(t[1]-t[0])*k1/2); k3=f(x+(t[1]-t[0])*k2/2); k4=f(x+(t[1]-t[0])*k3)
124        x += (t[1]-t[0])*(k1+2*k2+2*k3+k4)/6; true[0,k+1]=x
125    model=train(width,x0,true,device=device)
126    J,_=unit_sensitivity(model,x0,t,device)
127    # J shape observations, units, 3 params; group score sums parameter columns
128    groups=[slice(3*j,3*j+3) for j in range(width)]
129    Jflat=J.reshape(J.shape[0],-1)
130    traces,res=weighted_scores(Jflat,groups)
131    if kind=='sensitivity': chosen=np.argsort(traces+res)[:width-retain]
132    elif kind=='random': chosen=np.random.default_rng(SEED).choice(width, width-retain, replace=False)
133    elif kind=='magnitude': chosen=np.argsort(np.array([abs(model.a[j].item())+abs(model.w[j].item())+abs(model.b[j].item()) for j in range(width)]))[:width-retain]
134    else: chosen=np.arange(width-retain)
135    mask=np.ones(width,dtype=bool); mask[chosen]=False
136    # Evaluate the selected compact subnetwork, then retrain it at the same budget.
137    compact=UnitODE(retain).to(device)
138    with torch.no_grad():
139        keep=np.flatnonzero(mask)
140        compact.w.copy_(model.w[keep]); compact.b.copy_(model.b[keep]); compact.a.copy_(model.a[keep])
141    opt=torch.optim.Adam(compact.parameters(),lr=.025)
142    xd=x0.to(device); yd=true.to(device)
143    for _ in range(350):
144        opt.zero_grad(); pred0=compact(xd,2.5/30,30); loss=((pred0-yd)**2).mean(); loss.backward(); opt.step()
145    with torch.no_grad(): pred=compact(xd,2.5/30,30).cpu().numpy()[0]
146    mse=float(np.mean((pred-true.numpy()[0])**2)); longerr=float(abs(pred[-1]-true.numpy()[0,-1]))
147    return {'retained_units':retain,'parameters':3*retain,'mse':mse,'terminal_error':longerr,
148            'removed_mean_information':float(np.mean(traces[chosen])),'removed_mean_residual':float(np.mean(res[chosen]))}
149
150
151def main():
152    device='cuda' if torch is not None and torch.cuda.is_available() else 'cpu'
153    try:
154        mech=mechanism_check(); exp={k:prune_and_score(k,device=device) for k in ['sensitivity','magnitude','random']}
155    except Exception as e:
156        device='cpu'; mech=mechanism_check(); exp={k:prune_and_score(k,device=device) for k in ['sensitivity','magnitude','random']}
157        exp['error_fallback']=repr(e)
158    out={'seed':SEED,'device':device,'mechanism':mech,'pruning_experiment':exp}
159    Path('results.json').write_text(json.dumps(out,indent=2))
160    print(json.dumps(out,indent=2))
161if __name__=='__main__': main()