1import json, math, time
  2import numpy as np
  3
  4
  5def make_features(z, omega, kappa):
  6    z = np.asarray(z)
  7    return np.exp(1j * kappa * (z @ omega.T)) / np.sqrt(len(omega))
  8
  9
 10def projector_from_sources(sources, omega, kappa, noise=0.0, rng=None):
 11    A = make_features(sources, omega, kappa).T
 12    if noise:
 13        if rng is None: rng = np.random.default_rng(0)
 14        A = A + noise * (rng.normal(size=A.shape) + 1j*rng.normal(size=A.shape)) / np.sqrt(2*len(omega))
 15    U, _, _ = np.linalg.svd(A, full_matrices=False)
 16    return U[:, :len(sources)]
 17
 18
 19def score_grad(z, U, omega, kappa):
 20    """q=1-||U^*phi||^2 and its exact real spatial gradient."""
 21    phi = make_features(np.asarray(z)[None, :], omega, kappa)[0]
 22    Pphi = U @ (U.conj().T @ phi)
 23    q = float(1.0 - np.vdot(phi, Pphi).real)
 24    dphi = (1j * kappa * omega) * phi[:, None]
 25    grad = -2.0 * np.real(np.conj(dphi).T @ Pphi)
 26    return q, grad
 27
 28
 29def project(z, lo, hi):
 30    return np.minimum(np.maximum(z, lo), hi)
 31
 32
 33def refine(z0, U, omega, kappa, lo, hi, steps=40, eta=0.1, tol=1e-7):
 34    z = np.array(z0, dtype=float)
 35    used = 0
 36    for t in range(steps):
 37        q, g = score_grad(z, U, omega, kappa)
 38        h = eta / (kappa*kappa)
 39        # Backtracking makes the stated fixed scaling safe under finite samples.
 40        while True:
 41            zn = project(z - h*g, lo, hi)
 42            qn, _ = score_grad(zn, U, omega, kappa)
 43            if qn <= q + 1e-12 or h < 1e-8/(kappa*kappa): break
 44            h *= 0.5
 45        z = zn; used += 1
 46        if np.linalg.norm(g) < tol: break
 47    return z, score_grad(z, U, omega, kappa)[0], used
 48
 49
 50def nms(points, values, radius):
 51    order = np.argsort(values)
 52    kept = []
 53    for i in order:
 54        if all(np.linalg.norm(points[i]-points[j]) > radius for j in kept): kept.append(i)
 55    return kept
 56
 57
 58def nearest_errors(found, sources):
 59    return [float(np.min(np.linalg.norm(np.asarray(found)-s, axis=1))) if len(found) else float('inf') for s in sources]
 60
 61
 62def main():
 63    rng = np.random.default_rng(3138)
 64    d, M, kappa = 2, 96, 18.0
 65    lo, hi = np.zeros(d), np.ones(d)
 66    sources = np.array([[.23,.31],[.68,.72],[.79,.25]])
 67    omega = rng.normal(size=(M,d)); omega /= np.linalg.norm(omega, axis=1, keepdims=True)
 68    U = projector_from_sources(sources, omega, kappa, noise=0.015, rng=rng)
 69
 70    # Stage 1: exact gradient versus central differences, and h~k^-2 stability.
 71    ztest = np.array([.41,.57]); q,g = score_grad(ztest,U,omega,kappa)
 72    eps=1e-5; fd=[]
 73    for j in range(d):
 74        a=ztest.copy(); b=ztest.copy(); a[j]+=eps; b[j]-=eps
 75        fd.append((score_grad(a,U,omega,kappa)[0]-score_grad(b,U,omega,kappa)[0])/(2*eps))
 76    grad_rel_err=float(np.linalg.norm(g-np.array(fd))/(np.linalg.norm(g)+1e-12))
 77    stability=[]
 78    for kk in [kappa/2,kappa,2*kappa]:
 79        # same dimensionless eta, measure average one-step decrease from random points
 80        drops=[]
 81        for z in rng.uniform(.05,.95,size=(100,d)):
 82            qq,gg=score_grad(z,U,omega,kk) # intentionally U fixed: tests step scaling only
 83            zn=project(z-.1*gg/(kk*kk),lo,hi)
 84            drops.append(score_grad(zn,U,omega,kk)[0]-qq)
 85        stability.append(float(np.mean(np.array(drops)<=1e-10)))
 86
 87    # A grid at c_g/kappa; coarse stage keeps only low-score cells and NMSes them.
 88    spacing=0.8/kappa
 89    axes=[np.arange(0,1+1e-9,spacing) for _ in range(d)]
 90    grid=np.array(np.meshgrid(*axes,indexing='ij')).reshape(d,-1).T
 91    vals=np.array([score_grad(z,U,omega,kappa)[0] for z in grid])
 92    labels=np.min(np.linalg.norm(grid[:,None,:]-sources[None,:,:],axis=2),axis=1)<0.65/kappa
 93    pos, neg = vals[labels], vals[~labels]
 94    tau=float(.5*(np.max(pos)+np.min(neg))) if np.max(pos)<np.min(neg) else float(np.quantile(vals,.12))
 95    accepted=np.flatnonzero(vals<=tau)
 96    survivors=nms(grid[accepted], vals[accepted], spacing/2)
 97    starts_coarse=grid[accepted][survivors]
 98
 99    def run(starts):
100        out=[]; steps=0
101        t=time.perf_counter()
102        for z in starts:
103            x,v,n=refine(z,U,omega,kappa,lo,hi)
104            out.append(x); steps+=n
105        elapsed=time.perf_counter()-t
106        keep=nms(np.array(out),np.array([score_grad(x,U,omega,kappa)[0] for x in out]),.55/kappa) if out else []
107        return np.array(out), len(keep), steps, elapsed
108    all_out, all_unique, all_steps, all_time=run(grid)
109    coarse_out, coarse_unique, coarse_steps, coarse_time=run(starts_coarse)
110    result={
111      'config':{'kappa':kappa,'M':M,'grid_points':len(grid),'sources':len(sources),'grid_spacing':spacing,'threshold':tau},
112      'math_check':{'gradient_relative_error':grad_rel_err,'decrease_fraction_eta_over_k2':[{'kappa':x,'fraction':y} for x,y in zip([kappa/2,kappa,2*kappa],stability)]},
113      'baseline_exhaustive':{'starts':len(grid),'refinement_steps':all_steps,'unique_outputs':all_unique,'source_errors':nearest_errors(all_out, sources),'seconds':all_time},
114      'idea_coarse_to_fine':{'accepted':len(accepted),'survivor_starts':len(starts_coarse),'refinement_steps':coarse_steps,'unique_outputs':coarse_unique,'source_errors':nearest_errors(coarse_out,sources),'seconds':coarse_time},
115    }
116    print(json.dumps(result, indent=2))
117
118if __name__=='__main__': main()