1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5
  6SEED = 2843
  7np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  8torch.set_num_threads(4)
  9device = 'cuda' if torch.cuda.is_available() else 'cpu'
 10
 11def jacobian(W, U, us, b, x0, C):
 12    n = W.shape[0]
 13    P = torch.eye(n, dtype=W.dtype, device=W.device)
 14    x = x0
 15    blocks = [C @ P]
 16    masks = []
 17    for u in us:
 18        z = W @ x + U @ u + b
 19        d = (z > 0).to(W.dtype)
 20        masks.append(d)
 21        x = torch.relu(z)
 22        P = torch.diag(d) @ W @ P
 23        blocks.append(C @ P)
 24    return torch.cat(blocks, dim=0), masks
 25
 26def metrics(J, eps=1e-8):
 27    s = torch.linalg.svdvals(J)
 28    gram = J.T @ J + eps * torch.eye(J.shape[1], device=J.device)
 29    return {'rank': int((s > 1e-6).sum()), 'smin': float(s[-1]),
 30            'smax': float(s[0]), 'logdet': float(torch.linalg.slogdet(gram)[1]),
 31            'singular_values': s.detach().cpu().numpy().tolist()}
 32
 33def autograd_check():
 34    n, m, T = 6, 3, 3
 35    g = torch.Generator(device=device).manual_seed(SEED + 99)
 36    W = torch.randn(n,n,generator=g,device=device)/math.sqrt(n)
 37    U = torch.randn(n,2,generator=g,device=device)
 38    b = torch.randn(n,generator=g,device=device)
 39    x0 = torch.randn(n,generator=g,device=device,requires_grad=True)
 40    us = [torch.randn(2,generator=g,device=device) for _ in range(T)]
 41    C = torch.randn(m,n,generator=g,device=device)
 42    J,_ = jacobian(W,U,us,b,x0,C)
 43    def obs(q):
 44        x=q; ys=[C@x]
 45        for u in us:
 46            x=torch.relu(W@x+U@u+b); ys.append(C@x)
 47        return torch.cat(ys)
 48    A=torch.autograd.functional.jacobian(obs,x0)
 49    return float((J-A).abs().max())
 50
 51def setup(seed, n=16):
 52    g=torch.Generator(device=device).manual_seed(seed)
 53    W=torch.randn(n,n,generator=g,device=device)/math.sqrt(n)
 54    U=torch.zeros(n,1,device=device)
 55    b=torch.ones(n,device=device)*0.7
 56    x0=torch.randn(n,generator=g,device=device)
 57    return W,U,b,x0
 58
 59def Cselect(m,n):
 60    C=torch.zeros(m,n,device=device)
 61    C[torch.arange(m),torch.arange(m)]=1
 62    return C
 63
 64def boundary_sweep():
 65    n=16; W,U,b,x0=setup(SEED); u=[torch.zeros(1,device=device)]
 66    out=[]
 67    for m in [2,4,6,7,8,9,12,16]:
 68        q=metrics(jacobian(W,U,u,b,x0,Cselect(m,n))[0])
 69        out.append({'m':m,'counting_max_rank':min(n,2*m),'rank':q['rank'],'smin':q['smin'],'logdet':q['logdet']})
 70    return out
 71
 72def horizon_sweep():
 73    n,m=16,4; W,U,b,x0=setup(SEED+1); us=[torch.zeros(1,device=device) for _ in range(8)]
 74    out=[]
 75    for T in range(9):
 76        q=metrics(jacobian(W,U,us[:T],b,x0,Cselect(m,n))[0])
 77        out.append({'T':T,'counting_max_rank':min(n,m*(T+1)),'rank':q['rank'],'smin':q['smin']})
 78    return out
 79
 80def regularizer_trial(m, regularized, steps=300):
 81    n=16; W,U,b,x0=setup(SEED+10); W.requires_grad_(); b.requires_grad_()
 82    us=[torch.zeros(1,device=device)]
 83    C=Cselect(m,n)
 84    initial=metrics(jacobian(W,U,us,b,x0,C)[0])
 85    if regularized:
 86        opt=torch.optim.Adam([W,b],lr=0.01)
 87        for _ in range(steps):
 88            opt.zero_grad()
 89            J,_=jacobian(W,U,us,b,x0,C)
 90            G=J.T@J+1e-3*torch.eye(n,device=device)
 91            loss=-torch.linalg.slogdet(G)[1]+0.01*(W.square().mean()+b.square().mean())
 92            loss.backward(); opt.step()
 93    final=metrics(jacobian(W,U,us,b,x0,C)[0])
 94    return {'initial':initial,'final':final}
 95
 96def main():
 97    result={'seed':SEED,'device':device,
 98      'autograd_max_abs_error':autograd_check(),
 99      'predictions':{
100       'P1':'At T=1, rank(J)<=2m; m<8 cannot have rank 16, while m=8 is only a necessary counting threshold.',
101       'P2':'At m=4, rank 16 is first possible at T=3 because 4(T+1)>=16.',
102       'P3':'A log-det regularizer should increase finite-horizon Gramian logdet, but cannot overcome structural rank loss.'},
103      'boundary_T1':boundary_sweep(),'horizon_m4':horizon_sweep(),
104      'regularizer_m9':regularizer_trial(9,True),
105      'fixed_control_m9':regularizer_trial(9,False)}
106    Path('results.json').write_text(json.dumps(result,indent=2))
107    print(json.dumps(result,indent=2))
108
109if __name__=='__main__': main()