1import json, time, random
 2import numpy as np
 3import torch
 4import torch.nn as nn
 5
 6SEED = 388
 7random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
 8torch.set_num_threads(4)
 9device = 'cuda' if torch.cuda.is_available() else 'cpu'
10try:
11    if device == 'cuda': torch.cuda.empty_cache()
12except Exception:
13    device = 'cpu'
14
15def antihermitian(W):
16    # i times a real symmetric matrix is anti-Hermitian
17    S = (W + W.transpose(-1, -2)) / 2
18    return 1j * S.to(torch.complex64)
19
20def target_map(x, h0):
21    n = h0.numel()
22    H1 = antihermitian(torch.arange(n*n, device=x.device, dtype=torch.float32).reshape(n,n)/17)
23    H2 = antihermitian(torch.flip(torch.arange(n*n, device=x.device, dtype=torch.float32), (0,)).reshape(n,n)/19)
24    out = h0[None].expand(x.shape[0], -1).to(torch.complex64)
25    for c, H in [(0.55*torch.sin(x[:,0]), H1), (0.4*torch.cos(x[:,1]), H2)]:
26        E = torch.matrix_exp(c[:,None,None] * H[None])
27        out = torch.bmm(E, out[...,None])[...,0]
28    return out
29
30class StructuredKLU(nn.Module):
31    def __init__(self, d=3, n=4, K=4):
32        super().__init__(); self.d=d; self.n=n; self.K=K
33        self.H = nn.Parameter(torch.randn(K,n,n)*0.12)
34        self.J = nn.Parameter(torch.randn(K,n,n)*0.12)
35        self.psi = nn.ModuleList([nn.ModuleList([nn.Sequential(nn.Linear(1,8),nn.Tanh(),nn.Linear(8,1)) for _ in range(d)]) for _ in range(K)])
36        self.coeff = nn.ModuleList([nn.Sequential(nn.Linear(1,8),nn.Tanh(),nn.Linear(8,2),nn.Tanh()) for _ in range(K)])
37    def forward(self, x, h, return_u=False):
38        B=x.shape[0]; out=h.to(torch.complex64)
39        U=torch.eye(self.n, device=x.device, dtype=torch.complex64)[None].expand(B,-1,-1).clone()
40        for k in range(self.K):
41            s=sum(self.psi[k][j](x[:,j:j+1]) for j in range(self.d))
42            ab=self.coeff[k](s)
43            G=ab[:,0,None,None]*antihermitian(self.H[k])[None] + ab[:,1,None,None]*antihermitian(self.J[k])[None]
44            E=torch.matrix_exp(G)
45            out=torch.bmm(E,out[...,None])[...,0]
46            U=torch.bmm(E,U)
47        return (out,U) if return_u else out
48
49class DenseController(nn.Module):
50    def __init__(self, d=3, n=4, K=4):
51        super().__init__(); self.n=n; self.K=K
52        self.net=nn.Sequential(nn.Linear(d,32),nn.Tanh(),nn.Linear(32,2*K*n*n))
53    def forward(self,x,h,return_u=False):
54        B=x.shape[0]; q=self.net(x).reshape(B,self.K,2,self.n,self.n)
55        out=h.to(torch.complex64); U=torch.eye(self.n,device=x.device,dtype=torch.complex64)[None].expand(B,-1,-1).clone()
56        for k in range(self.K):
57            G=antihermitian(q[:,k,0])+1j*(q[:,k,1]-q[:,k,1].transpose(-1,-2))/2
58            E=torch.matrix_exp(G*0.06)
59            out=torch.bmm(E,out[...,None])[...,0]; U=torch.bmm(E,U)
60        return (out,U) if return_u else out
61
62def params(m): return sum(p.numel() for p in m.parameters() if p.requires_grad)
63
64def main():
65    n,d,K=4,3,4; Ntr,Nte=768,256
66    g=torch.Generator(device='cpu').manual_seed(SEED)
67    xtr=(torch.rand(Ntr,d,generator=g)*2-1).to(device); xte=(torch.rand(Nte,d,generator=g)*2-1).to(device)
68    h0=torch.tensor([1+0j,.2+.3j,-.4+.1j,.5-.2j],device=device)
69    ytr=target_map(xtr,h0).detach(); yte=target_map(xte,h0).detach()
70    # independent mathematical sanity check, including a longer product
71    W=torch.randn(32,4,4,device=device); G=antihermitian(W); E=torch.matrix_exp(G)
72    one=torch.eye(4,device=device,dtype=torch.complex64)
73    single_res=(E.conj().transpose(-1,-2)@E-one).norm().item()
74    P=one
75    for i in range(32): P=torch.matrix_exp(G[i])@P
76    product_res=(P.conj().transpose(-1,-2)@P-one).norm().item()
77    results={'device':device,'math_single_residual':single_res,'math_32_product_residual':product_res}
78    for name, cls in [('KLU',StructuredKLU),('dense',DenseController)]:
79        torch.manual_seed(SEED+ (0 if name=='KLU' else 1))
80        model=cls(d,n,K).to(device); opt=torch.optim.Adam(model.parameters(),lr=3e-3)
81        t0=time.perf_counter(); losses=[]; gradvals=[]
82        for step in range(301):
83            ix=torch.randint(0,Ntr,(64,),device=device)
84            pred=model(xtr[ix],h0.expand(64,-1))
85            loss=((pred-ytr[ix]).abs()**2).mean()
86            opt.zero_grad(); loss.backward(); gradvals.append(float(torch.nn.utils.clip_grad_norm_(model.parameters(),100))); opt.step()
87            if step in (0,100,300): losses.append(float(loss.detach().cpu()))
88        with torch.no_grad():
89            pred,U=model(xte,h0.expand(Nte,-1),return_u=True)
90            test=float(((pred-yte).abs()**2).mean().cpu())
91            norm_drift=float((pred.abs().norm(dim=1)-h0.abs().norm()).abs().mean().cpu())
92            unit=float(((U.conj().transpose(-1,-2)@U-torch.eye(n,device=device,dtype=torch.complex64)).norm(dim=(1,2))).mean().cpu())
93        results[name]={'parameters':params(model),'train_loss_0_100_300':losses,'test_mse':test,'norm_drift':norm_drift,'unitarity_residual':unit,'seconds':time.perf_counter()-t0,'mean_grad_norm':float(np.mean(gradvals))}
94    with open('results.json','w') as f: json.dump(results,f,indent=2)
95    print(json.dumps(results,indent=2))
96if __name__=='__main__': main()