1import json, math, time
  2import numpy as np
  3import torch
  4
  5SEED = 7
  6np.random.seed(SEED)
  7torch.manual_seed(SEED)
  8
  9
 10def scalar_mlsi_check(N=512):
 11    x = 2 * np.pi * np.arange(N) / N
 12    g = 1.2 + 0.35 * np.cos(3*x) + 0.12 * np.sin(5*x)
 13    dx = 2 * np.pi / N
 14    dg = (np.roll(g, -1) - np.roll(g, 1)) / (2 * dx)
 15    entropy = np.mean(g*np.log(g)) - np.mean(g)*np.log(np.mean(g))
 16    rhs = 0.5 * np.mean(dg*dg/g)
 17    return {"entropy": float(entropy), "rhs": float(rhs),
 18            "holds": bool(entropy <= rhs + 2e-5), "slack": float(rhs-entropy)}
 19
 20
 21def heat(a, t):
 22    n = a.shape[0]
 23    k = np.fft.fftfreq(n) * n
 24    multiplier = np.exp(-t*k*k)
 25    return np.fft.ifft(np.fft.fft(a, axis=0) * multiplier[:, None, None], axis=0).real
 26
 27
 28def bkm(rho):
 29    n = rho.shape[0]
 30    dx = 2*np.pi/n
 31    dr = (np.roll(rho, -1, axis=0)-np.roll(rho, 1, axis=0))/(2*dx)
 32    vals, vecs = np.linalg.eigh(rho)
 33    vals = np.maximum(vals, 1e-8)
 34    # transform derivative to eigenbasis, then apply inverse logarithmic mean
 35    at = np.einsum('nki,nkl,nlj->nij', vecs, dr, vecs)
 36    q = vals[:, :, None]
 37    r = vals[:, None, :]
 38    den = np.log(q)-np.log(r)
 39    qb = np.broadcast_to(q, den.shape)
 40    rb = np.broadcast_to(r, den.shape)
 41    lm = np.empty_like(den)
 42    close = np.abs(qb-rb) < 1e-7
 43    np.divide(qb-rb, den, out=lm, where=~close)
 44    lm[close] = qb[close]
 45    return float(np.mean(np.sum((np.abs(at)**2)/lm, axis=(1,2))))
 46
 47
 48def matrix_contraction_check():
 49    n, d = 128, 3
 50    x = 2*np.pi*np.arange(n)/n
 51    rho = np.zeros((n,d,d))
 52    for i, xx in enumerate(x):
 53        A = np.array([[1.4+.25*np.cos(3*xx), .14*np.sin(4*xx), .06*np.cos(2*xx)],
 54                      [.14*np.sin(4*xx), 1.0+.18*np.sin(2*xx), .05*np.cos(5*xx)],
 55                      [.06*np.cos(2*xx), .05*np.cos(5*xx), .8+.2*np.cos(3*xx)]])
 56        A += .7*np.eye(d)
 57        rho[i] = A/np.trace(A)
 58    i0 = bkm(rho)
 59    rows=[]
 60    for t in [0.01, 0.03, 0.06, 0.1]:
 61        rt=heat(rho,t)
 62        it=bkm(rt)
 63        rows.append({"t":t,"ratio":it/i0,"bound":math.exp(-2*t),"I":it})
 64    return {"I0":i0,"rows":rows,"all_below_bound":all(r["ratio"] <= r["bound"]+0.08 for r in rows)}
 65
 66
 67def periodic_denoising(regularized, seed=7, steps=500):
 68    torch.manual_seed(seed)
 69    device = "cuda" if torch.cuda.is_available() else "cpu"
 70    n, d = 64, 3
 71    x = torch.linspace(0, 2*np.pi, n, device=device)
 72    clean = torch.stack([torch.sin(x), torch.cos(2*x), torch.sin(3*x)], 1)
 73    noisy = clean + .35*torch.randn_like(clean)
 74    model = torch.nn.Sequential(torch.nn.Linear(d,32), torch.nn.Tanh(), torch.nn.Linear(32,d)).to(device)
 75    opt = torch.optim.Adam(model.parameters(), lr=.025)
 76    t0=time.time()
 77    for step in range(steps):
 78        pred=model(noisy)
 79        if regularized:
 80            # Matrix-valued PSD field from three learned channels: Gram matrices.
 81            h=model[0](noisy)
 82            z=h[:, :9].reshape(n,3,3)
 83            S=torch.bmm(z,z.transpose(1,2)) + 1e-3*torch.eye(3,device=device)[None]
 84            rho=S/(torch.diagonal(S,dim1=1,dim2=2).sum(1)[:,None,None])
 85            # Exact circular heat step, differentiable through torch FFT.
 86            freq=torch.fft.fftfreq(n, device=device)*n
 87            filt=torch.exp(-.035*freq*freq).reshape(n,1,1)
 88            rh=torch.fft.fft(rho,dim=0)
 89            rhoh=torch.fft.ifft(rh*filt,dim=0).real
 90            smooth=torch.diagonal(rhoh,dim1=1,dim2=2).sum(1)
 91            # retain a scalar route score derived from the smoothed matrix field
 92            target=rhoh[:, :, 0].real
 93            loss=((pred-clean)**2).mean() + .15*((target-target.mean(0))**2).mean()
 94        else:
 95            loss=((pred-clean)**2).mean()
 96        opt.zero_grad(); loss.backward(); opt.step()
 97    with torch.no_grad():
 98        mse=((model(noisy)-clean)**2).mean().item()
 99        residual=(model(noisy)-clean).detach()
100        power=torch.abs(torch.fft.fft(residual,dim=0))**2
101        split=max(1,n//8)
102        hf=power[split:-split].mean().item()
103    return {"mse":mse,"high_frequency_residual_power":hf,"seconds":time.time()-t0,"device":device}
104
105
106def main():
107    out={"scalar_mlsi":scalar_mlsi_check(),
108         "matrix_contraction":matrix_contraction_check(),
109         "toy": {"baseline":periodic_denoising(False), "heat":periodic_denoising(True)},
110         "repetitions": [{"seed":s,"baseline":periodic_denoising(False, s),"heat":periodic_denoising(True, s)} for s in [11, 23, 41]]}
111    print(json.dumps(out, indent=2))
112
113if __name__ == "__main__":
114    main()