1import json, math, time, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5
  6SEED = 1620
  7np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  8if torch.cuda.is_available():
  9    try: torch.cuda.manual_seed_all(SEED)
 10    except Exception: pass
 11try:
 12    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 13    # A tiny allocation makes CUDA failures visible before the experiment.
 14    if device.type == "cuda": torch.zeros(1, device=device)
 15except Exception:
 16    device = torch.device("cpu")
 17
 18
 19def rand_u2(rng):
 20    z = rng.normal(size=(2,2)) + 1j*rng.normal(size=(2,2))
 21    q, r = np.linalg.qr(z)
 22    q = q @ np.diag(np.conj(np.diag(r))/np.abs(np.diag(r)))
 23    return q.astype(np.complex128)
 24
 25
 26def apply_butterfly(x, gates, scale=1.0):
 27    """Apply log2(N) stages of local 2x2 gates to vectors or a matrix."""
 28    n = x.shape[-1]
 29    y = np.array(x, dtype=np.complex128, copy=True)
 30    levels = int(math.log2(n))
 31    for s in range(levels):
 32        half = 2**s
 33        step = 2*half
 34        out = y.copy()
 35        for base in range(0, n, step):
 36            for j in range(half):
 37                inds = [base+j, base+j+half]
 38                out[..., inds] = np.einsum('ab,...b->...a', scale*gates[s][base//step*half+j], y[..., inds])
 39        y = out
 40    return y
 41
 42
 43def full_transform(n, rng, scale=1.0):
 44    gates = [[rand_u2(rng) for _ in range(n//2)] for _ in range(int(math.log2(n)))]
 45    # apply_butterfly expects one gate per pair at every stage
 46    cols = np.eye(n, dtype=np.complex128)
 47    return apply_butterfly(cols, gates, scale=scale), gates
 48
 49
 50def unitary_checks():
 51    rng = np.random.default_rng(SEED)
 52    rows=[]
 53    for n in [4, 8, 16, 32, 64]:
 54        T, gates = full_transform(n, rng)
 55        gram_err = float(np.linalg.norm(T.conj().T @ T - np.eye(n), ord='fro'))
 56        x = rng.normal(size=n)+1j*rng.normal(size=n)
 57        ratio = float(np.linalg.norm(T@x)/np.linalg.norm(x))
 58        rows.append({'N':n, 'levels':int(math.log2(n)), 'unitarity_fro_error':gram_err, 'norm_ratio':ratio,
 59                     'predicted_norm_ratio':1.0})
 60    # Controlled violation: multiplying each local gate by (1+eps) gives exactly
 61    # (1+eps)^levels for every vector, because each path crosses one gate per level.
 62    drift=[]
 63    n=32; x=rng.normal(size=n)+1j*rng.normal(size=n)
 64    # Sweep both perturbation strength and repeated depth. The exact prediction is
 65    # (1+eps)^(log2(N)*repetitions), since each path crosses one gate per level.
 66    for eps in [0.0, 1e-3, 1e-2, 5e-2]:
 67        rng_eps=np.random.default_rng(SEED+1)
 68        Tp, _ = full_transform(n, rng_eps, scale=1+eps)
 69        for depth in [1,4,8]:
 70            observed=float(np.linalg.norm(np.linalg.matrix_power(Tp, depth)@x)/np.linalg.norm(x))
 71            predicted=float((1+eps)**(int(math.log2(n))*depth))
 72            drift.append({'epsilon':eps, 'repetitions':depth, 'observed':observed,
 73                          'predicted':predicted,
 74                          'relative_error':abs(observed-predicted)/(predicted if predicted else 1)})
 75    # Since T is unitary, its conjugate transpose is its exact inverse.
 76    rng_inv=np.random.default_rng(SEED+2); T,_=full_transform(32,rng_inv)
 77    inverse_err=float(np.linalg.norm(T.conj().T@[email protected](32)-np.ones(32)))
 78    return rows, drift, inverse_err
 79
 80
 81class StructuredMixer(torch.nn.Module):
 82    # Four real coordinates parameterize every U(2), using a standard SU(2) chart
 83    # plus a global phase. Gates are shared over samples but not over butterfly pairs.
 84    def __init__(self, n, init=None):
 85        super().__init__(); self.n=n; self.levels=int(math.log2(n))
 86        self.p=torch.nn.Parameter(torch.zeros(self.levels, n//2, 4))
 87        if init is not None:
 88            with torch.no_grad(): self.p.copy_(init)
 89    def gates(self):
 90        a,b,c,d=self.p.unbind(-1)
 91        ca=torch.cos(c); sa=torch.sin(c)
 92        phase=torch.exp(1j*a)
 93        row0=torch.stack([torch.exp(1j*b)*ca, torch.exp(1j*d)*sa], -1)
 94        row1=torch.stack([-torch.exp(-1j*d)*sa, torch.exp(-1j*b)*ca], -1)
 95        return phase[...,None,None]*torch.stack([row0,row1],-2)
 96    def forward(self, x):
 97        y=x
 98        gs=self.gates()
 99        for s in range(self.levels):
100            half=2**s; step=2*half; out=y.clone()
101            for base in range(0,self.n,step):
102                for j in range(half):
103                    ids=[base+j,base+j+half]
104                    out[:,ids]=torch.einsum('ab,nb->na',gs[s,base//step*half+j],y[:,ids])
105            y=out
106        return y
107
108
109def fit_toy():
110    torch.manual_seed(SEED); n=16; samples=256
111    # Generate a target exactly in the proposed topology, making identifiability fair.
112    true=StructuredMixer(n).to(device)
113    with torch.no_grad(): true.p.normal_(0,0.35)
114    x=(torch.randn(samples,n)+1j*torch.randn(samples,n)).to(device)
115    with torch.no_grad(): y=true(x).detach()
116    structured=StructuredMixer(n).to(device)
117    dense=torch.nn.Parameter((0.05*(torch.randn(n,n)+1j*torch.randn(n,n))).to(device))
118    os=torch.optim.Adam(structured.parameters(),lr=0.06)
119    od=torch.optim.Adam([dense],lr=0.06)
120    losses_s=[]; losses_d=[]; t0=time.time()
121    for step in range(301):
122        os.zero_grad(); pred=structured(x); ls=(pred-y).abs().pow(2).mean(); ls.backward(); os.step()
123        od.zero_grad(); pred2=x@dense.T; ld=(pred2-y).abs().pow(2).mean(); ld.backward(); od.step()
124        if step in [0,50,100,200,300]: losses_s.append(float(ls.detach().cpu())); losses_d.append(float(ld.detach().cpu()))
125    elapsed=time.time()-t0
126    return {'N':n,'samples':samples,'steps':300,'structured_params':int(sum(p.numel() for p in structured.parameters())),
127            'dense_complex_params':int(dense.numel()),'structured_final_mse':losses_s[-1],
128            'dense_final_mse':losses_d[-1],'structured_curve':losses_s,'dense_curve':losses_d,
129            'seconds':elapsed,'device':str(device)}
130
131
132def main():
133    checks, drift, inverse_err=unitary_checks()
134    fit=fit_toy()
135    result={'seed':SEED,'predictions':[
136      'For every N=2^k, ||T x||/||x|| = 1 exactly up to floating point.',
137      'T^*T-I is at machine precision and does not grow materially with N.',
138      'If every local gate is scaled by 1+eps, repeated-pass norm drift is (1+eps)^(k*r), k=log2(N).'
139    ],'unitarity_sweep':checks,'controlled_drift':drift,'inverse_reconstruction_error':inverse_err,'toy_fit':fit}
140    Path('results.json').write_text(json.dumps(result,indent=2))
141    print(json.dumps(result,indent=2))
142
143if __name__=='__main__': main()