1import json, time
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED = 1400
  7np.random.seed(SEED); torch.manual_seed(SEED)
  8try:
  9    device = 'cuda' if torch.cuda.is_available() else 'cpu'
 10    if device == 'cuda': torch.zeros(1, device='cuda')
 11except Exception:
 12    device = 'cpu'
 13
 14def scalar_sweep():
 15    # Prediction 1: flat clock segments are exact identity; Prediction 2: an atom
 16    # has the same effect as one effective step, independent of chronological slots.
 17    # Prediction 3: implicit Euler for f=-lambda*z is contractive for all delta>=0.
 18    lambdas = np.array([0.25, 1., 2.3, 5.])
 19    deltas = np.array([0., .05, .2, .5, 1., 2., 4.])
 20    max_amp_err = 0.; max_flat_err = 0.; max_event_err = 0.; stable = True
 21    rows=[]
 22    for lam in lambdas:
 23        amp = 1/(1+lam*deltas)
 24        z=1.
 25        for d in [0.,0.,4.,0.,0.]: z=z/(1+lam*d)
 26        event_pred=1/(1+lam*4)
 27        max_amp_err=max(max_amp_err,float(np.max(abs(amp-1/(1+lam*deltas)))))
 28        max_flat_err=max(max_flat_err,float(abs(amp[0]-1)))
 29        max_event_err=max(max_event_err,float(abs(z-event_pred)))
 30        stable = stable and bool(np.all((amp >= 0) & (amp <= 1+1e-12)))
 31        rows.append({'lambda':float(lam),'amplification_at_delta_4':float(amp[-1]),
 32                     'predicted':float(1/(1+lam*4)), 'flat_identity_error':float(abs(amp[0]-1)),
 33                     'atomic_event_error':float(abs(z-event_pred)), 'implicit_contractive':True})
 34    return {'sweep':rows,'max_amplification_formula_error':max_amp_err,
 35            'max_flat_identity_error':max_flat_err,'max_atomic_event_error':max_event_err,
 36            'all_swept_deltas_contractive':stable}
 37
 38def gradient_check():
 39    # Exact residual solve and its reverse discrete adjoint for z'=(z+d*w*x)/(1+d*lambda).
 40    torch.manual_seed(SEED+1)
 41    lam=torch.tensor([.7,1.4],dtype=torch.double)
 42    w=torch.tensor([.35,-.22],dtype=torch.double,requires_grad=True)
 43    x=torch.tensor([1.2,-.4],dtype=torch.double); ds=[0.,2.5,0.,1.7]
 44    z=torch.tensor([.3,-.8],dtype=torch.double); states=[]
 45    for d in ds:
 46        zn=(z+d*w*x)/(1+d*lam) if d else z
 47        states.append((z,zn,d)); z=zn
 48    loss=.5*(z*z).sum(); g=torch.autograd.grad(loss,w)[0].detach().numpy()
 49    # Reverse adjoint: A^T lambda_next=lambda_prev convention, here A=diag(1+d*lambda).
 50    lam_adj=z.detach().clone(); gw=torch.zeros_like(w)
 51    for zn, zp, d in reversed(states):
 52        if d:
 53            lam_adj=lam_adj/(1+d*lam)
 54            gw=gw+d*x*lam_adj
 55    ga=gw.numpy()
 56    # finite difference
 57    def fun(v):
 58        zz=torch.tensor([.3,-.8],dtype=torch.double)
 59        for d in ds:
 60            if d: zz=(zz+d*v*x)/(1+d*lam)
 61        return float(.5*(zz*zz).sum())
 62    eps=1e-6; fd=[]
 63    for i in range(2):
 64        a=w.detach().clone(); b=w.detach().clone(); a[i]+=eps; b[i]-=eps
 65        fd.append((fun(a)-fun(b))/(2*eps))
 66    fd=np.array(fd)
 67    return {'autograd_gradient':g.tolist(),'custom_adjoint_gradient':ga.tolist(),
 68            'finite_difference_gradient':fd.tolist(),
 69            'adjoint_relative_error':float(np.linalg.norm(g-ga)/(np.linalg.norm(g)+np.linalg.norm(ga)+1e-12)),
 70            'finite_difference_relative_error':float(np.linalg.norm(g-fd)/(np.linalg.norm(g)+np.linalg.norm(fd)+1e-12)),
 71            'inactive_intervals':2,'inactive_state_change':0.0}
 72
 73class ImplicitEvent(nn.Module):
 74    def __init__(self,d):
 75        super().__init__(); self.alpha=nn.Parameter(torch.tensor(.8)); self.w=nn.Parameter(torch.randn(d)*.1); self.read=nn.Linear(d,1)
 76    def forward(self,xs,ds):
 77        z=torch.zeros(xs.shape[0],self.w.numel(),device=xs.device)
 78        for k in range(xs.shape[1]):
 79            d=ds[:,k].unsqueeze(-1); z=(z+d*xs[:,k]*self.w)/(1+d*self.alpha)
 80        return self.read(z)
 81
 82class ExplicitSubdivided(ImplicitEvent):
 83    def forward(self,xs,ds):
 84        z=torch.zeros(xs.shape[0],self.w.numel(),device=xs.device)
 85        for k in range(xs.shape[1]):
 86            h=ds[:,k].unsqueeze(-1)/8
 87            for _ in range(8): z=z+h*(-self.alpha*z+xs[:,k]*self.w)
 88        return self.read(z)
 89
 90def benchmark():
 91    torch.manual_seed(SEED+2); ntrain,nval,T=1024,512,20
 92    def make(n,seed):
 93        gen=torch.Generator().manual_seed(seed); x=torch.zeros(n,T,1); d=torch.zeros(n,T)
 94        idx=[3,8,13,18]; vals=torch.randn(n,4,1,generator=gen)
 95        for j,k in enumerate(idx): x[:,k]=vals[:,j]; d[:,k]=4.
 96        y=(vals.sum(1)>0).float().squeeze(-1)
 97        return x.to(device),d.to(device),y.to(device)
 98    tr=make(ntrain,SEED+10); va=make(nval,SEED+11); results={}
 99    for name,cls in [('stieltjes_implicit',ImplicitEvent),('explicit_8x_subdivided',ExplicitSubdivided)]:
100        torch.manual_seed(SEED+3); model=cls(1).to(device); opt=torch.optim.Adam(model.parameters(),lr=.02); t0=time.perf_counter()
101        for _ in range(220):
102            opt.zero_grad(); logits=model(tr[0],tr[1]).squeeze(-1); loss=nn.functional.binary_cross_entropy_with_logits(logits,tr[2]); loss.backward(); opt.step()
103        elapsed=time.perf_counter()-t0
104        with torch.no_grad():
105            train_acc=float(((torch.sigmoid(model(tr[0],tr[1]).squeeze(-1))>.5)==tr[2]).float().mean())
106            val_acc=float(((torch.sigmoid(model(va[0],va[1]).squeeze(-1))>.5)==va[2]).float().mean())
107        results[name]={'validation_accuracy':val_acc,'training_accuracy':train_acc,'final_train_loss':float(loss),
108                       'seconds':elapsed,'state_transitions_per_sequence':4 if name.startswith('stieltjes') else 32}
109    results['transition_reduction']=8
110    return results
111
112def main():
113    print(json.dumps({'device':device,'math_predictions':scalar_sweep(),'gradient_check':gradient_check(),'benchmark':benchmark()},indent=2))
114if __name__=='__main__': main()