1import json
  2from pathlib import Path
  3import numpy as np
  4from scipy.fft import dctn, idctn
  5
  6RNG = np.random.default_rng(7)
  7
  8def steering(n, u):
  9    x = np.arange(n) - (n - 1) / 2
 10    return np.exp(1j * np.pi * x * u) / np.sqrt(n)
 11
 12def steering_deriv(n, u):
 13    x = np.arange(n) - (n - 1) / 2
 14    return 1j * np.pi * x * steering(n, u)
 15
 16def atom(nr, nt, ur, ut):
 17    return np.outer(steering(nr, ur), steering(nt, ut).conj())
 18
 19def relerr(a, b):
 20    return float(np.linalg.norm(a - b) / max(np.linalg.norm(a), 1e-15))
 21
 22def channel(nr, nt, k, rng):
 23    H = np.zeros((nr, nt), complex)
 24    paths = []
 25    for _ in range(k):
 26        ur, ut = rng.uniform(-.85, .85, 2)
 27        g = (rng.normal() + 1j * rng.normal()) / np.sqrt(2 * k)
 28        H += g * atom(nr, nt, ur, ut)
 29        paths.append((ur, ut, g))
 30    return H, paths
 31
 32def ls_reconstruct(H, atoms, lam=1e-8):
 33    A = np.stack([x.reshape(-1) for x in atoms], axis=1)
 34    y = H.reshape(-1)
 35    g = np.linalg.solve(A.conj().T @ A + lam * np.eye(len(atoms)), A.conj().T @ y)
 36    return (A @ g).reshape(H.shape), g
 37
 38def analytic_greedy(H, k, grid):
 39    residual = H.copy(); atoms = []; chosen = []
 40    nr, nt = H.shape
 41    for _ in range(k):
 42        best, score = None, -1.0
 43        for ur in grid:
 44            ar = steering(nr, ur)
 45            for ut in grid:
 46                a = np.outer(ar, steering(nt, ut).conj())
 47                s = abs(np.vdot(a, residual))
 48                if s > score:
 49                    score, best = s, (ur, ut, a)
 50        chosen.append(best[:2]); atoms.append(best[2])
 51        rec, _ = ls_reconstruct(H, atoms)
 52        residual = H - rec
 53    rec, g = ls_reconstruct(H, atoms)
 54    return rec, chosen, g
 55
 56def dct_codec(H, k):
 57    # Separate real/imag orthonormal 2-D DCT, retaining largest magnitudes.
 58    coeff = dctn(H.real, norm='ortho') + 1j * dctn(H.imag, norm='ortho')
 59    flat = np.abs(coeff).ravel()
 60    ind = np.argpartition(flat, -k)[-k:]
 61    keep = np.zeros(flat.size, bool); keep[ind] = True
 62    sparse = np.where(keep, coeff.ravel(), 0).reshape(coeff.shape)
 63    return idctn(sparse.real, norm='ortho') + 1j * idctn(sparse.imag, norm='ortho')
 64
 65def taylor_sweep():
 66    # Prediction: first-order error is O(delta^2), with relative error/delta^2
 67    # approaching ||a''||/(2||a||), and is only stable for |delta|*pi*N/2 << 1.
 68    n = 32; u = .17
 69    ds = np.array([1e-4, 3e-4, 1e-3, 3e-3, 1e-2, 3e-2])
 70    errs = []
 71    for d in ds:
 72        exact = steering(n, u+d)
 73        approx = steering(n,u) + d*steering_deriv(n,u)
 74        errs.append(np.linalg.norm(exact-approx)/np.linalg.norm(exact))
 75    slope = np.polyfit(np.log(ds), np.log(errs), 1)[0]
 76    # The quadratic coefficient from the exact second derivative.
 77    x=np.arange(n)-(n-1)/2
 78    second = -(np.pi*x)**2*steering(n,u)
 79    pred_coeff=np.linalg.norm(second)/(2*np.linalg.norm(steering(n,u)))
 80    return {'deltas':ds.tolist(),'errors':np.array(errs).tolist(),'loglog_slope':float(slope),
 81            'predicted_slope':2.0,'quadratic_coeff_predicted':float(pred_coeff),
 82            'quadratic_coeff_observed':float(np.mean(np.array(errs[:3])/ds[:3]**2))}
 83
 84def rank_sweep():
 85    # Prediction: exact K-path channels have zero noiseless residual at K atoms;
 86    # with fewer atoms, residual decreases monotonically as K increases.
 87    rng=np.random.default_rng(11); H, paths=channel(32,32,4,rng)
 88    exact=[]
 89    for kk in range(1,5):
 90        rec,_=ls_reconstruct(H,[atom(32,32,p[0],p[1]) for p in paths[:kk]])
 91        exact.append(relerr(H,rec))
 92    return {'K':list(range(1,5)),'relative_errors':exact,'predicted_at_Kstar':0.0,
 93            'observed_at_Kstar':exact[-1]}
 94
 95def ridge_sweep():
 96    # Prediction: nearly duplicate atoms make unregularized LS ill-conditioned;
 97    # ridge reduces coefficient norm and remains numerically stable.
 98    nr=nt=32; u=.2; eps=np.array([1e-1,1e-2,1e-3,1e-4]); rows=[]
 99    rng=np.random.default_rng(12); H=atom(nr,nt,u,u)+(rng.normal(size=(nr,nt))+1j*rng.normal(size=(nr,nt)))*1e-3
100    for e in eps:
101        A=np.stack([atom(nr,nt,u,u).ravel(),atom(nr,nt,u+e,u+e).ravel()],1)
102        cond=np.linalg.cond(A.conj().T@A)
103        vals=[]
104        for lam in [0,1e-6,1e-3]:
105            try:
106                g=np.linalg.solve(A.conj().T@A+lam*np.eye(2),A.conj().T@H.ravel())
107                vals.append(float(np.linalg.norm(g)))
108            except np.linalg.LinAlgError: vals.append(float('inf'))
109        rows.append({'separation':float(e),'condition':float(cond),'coef_norm_lambda_0_1e-6_1e-3':vals})
110    slope = float(np.polyfit(np.log(eps), np.log([r['condition'] for r in rows]), 1)[0])
111    return {'rows': rows, 'condition_loglog_slope': slope, 'predicted_slope': -2.0}
112
113def codec_compare():
114    rng=np.random.default_rng(21); grid=np.linspace(-.9,.9,17)
115    rows=[]
116    for k in [1,2,3,4]:
117        ae=[]; de=[]
118        for _ in range(10):
119            H,_=channel(32,32,3,rng)
120            ae.append(relerr(H,analytic_greedy(H,k,grid)[0]))
121            de.append(relerr(H,dct_codec(H, k*4)))
122        rows.append({'atoms':k,'analytic_nmse':float(np.mean(np.array(ae)**2)),
123                     'dct_coefficients':4*k,'dct_nmse':float(np.mean(np.array(de)**2))})
124    # Transfer: same continuous physical paths, reconstruct at larger dimensions without decoder retraining.
125    H32,p=channel(32,32,3,np.random.default_rng(22)); H48=np.zeros((48,48),complex)
126    for ur,ut,g in p: H48 += g*atom(48,48,ur,ut)
127    transfer=relerr(H48,sum(g*atom(48,48,ur,ut) for ur,ut,g in p))
128    return rows, transfer
129
130def main():
131    out={'taylor':taylor_sweep(),'rank':rank_sweep(),'ridge':ridge_sweep()}
132    out['codec_rows'],out['transfer_relative_error']=codec_compare()
133    Path('results.json').write_text(json.dumps(out,indent=2))
134    print(json.dumps(out,indent=2))
135
136if __name__=='__main__': main()