1import json, math
  2from pathlib import Path
  3import numpy as np
  4
  5SEED = 938
  6rng = np.random.default_rng(SEED)
  7
  8# Quadratic target regularizer R(x)=lambda*x^2/2. Its proximal map is a*v,
  9# where a=1/(1+gamma*lambda). PnP-PGD uses A=I and f=.5||x-y||^2.
 10def prox_coeff(gamma, lam):
 11    return 1.0 / (1.0 + gamma * lam)
 12
 13def pnp_run(a, gamma, lam, y, steps=60):
 14    x = 0.0
 15    hist = []
 16    for _ in range(steps):
 17        v = x - gamma * (x - y)
 18        x = a * v
 19        F = .5 * (x-y)**2 + .5 * lam*x*x
 20        station = abs((x-y) + lam*x)
 21        hist.append((x, F, station))
 22    return np.asarray(hist)
 23
 24def main():
 25    out = {"seed": SEED}
 26
 27    # ---------------- Stage 1: analytic mechanism verification ----------------
 28    # Prediction 1: prox coefficient is exactly 1/(1+gamma*lambda).
 29    coeff_rows=[]
 30    for gamma in [0.1, 0.3, 0.7, 1.0]:
 31        for lam in [0.0, 0.5, 2.0, 8.0]:
 32            v = rng.normal(size=20000)
 33            # empirical least-squares fit to teacher action
 34            t = prox_coeff(gamma, lam) * v
 35            fitted = float(np.dot(v,t)/np.dot(v,v))
 36            pred = prox_coeff(gamma, lam)
 37            coeff_rows.append({"gamma":gamma,"lambda":lam,"pred":pred,"fit":fitted,"abs_err":abs(fitted-pred)})
 38    max_coeff_err=max(r["abs_err"] for r in coeff_rows)
 39
 40    # Prediction 2: squared mismatch is quadratic in coefficient error, slope 2 on log-log.
 41    v = rng.normal(size=100000)
 42    target_a = prox_coeff(.6, 3.0)
 43    errors=np.logspace(-4,-1,10)
 44    mismatch=[]
 45    for e in errors:
 46        mismatch.append(float(np.mean(((target_a+e)*v-target_a*v)**2)))
 47    slope=float(np.polyfit(np.log(errors),np.log(mismatch),1)[0])
 48    # Also verify lambda=0 gives identity teacher and identity coefficient has zero mismatch.
 49    zero_mismatch=float(np.mean((1.0*v-1.0*v)**2))
 50
 51    # Prediction 3: PnP linear contraction factor q=a(1-gamma), boundary |q|=1.
 52    # With gamma=.5 and positive a, predicted boundary a=2.0. Measure whether
 53    # trajectory grows (unstable) or contracts from x0=1,y=0.
 54    gamma=.5; y=0.0
 55    boundary=1.0/(1.0-gamma)
 56    stability=[]
 57    for a in [0.5, 1.0, 1.8, 1.99, 2.01, 2.2, 3.0]:
 58        h=pnp_run(a,gamma,0.0,y,steps=25)
 59        ratio=abs(h[-1,0])/max(abs(h[0,0]),1e-12)
 60        q=a*(1-gamma)
 61        stability.append({"a":a,"predicted_q":q,"observed_abs_final_over_first":float(ratio),"stable_observed":bool(ratio<1)})
 62    # pnp_run starts x=0, so explicitly use recurrence for stability test from x0=1
 63    stability=[]
 64    for a in [0.5,1.0,1.8,1.99,2.01,2.2,3.0]:
 65        q=a*(1-gamma); vals=np.array([q**k for k in range(25)])
 66        stability.append({"a":a,"predicted_q":q,"observed_final_abs":float(abs(vals[-1])),"stable_observed":bool(abs(vals[-1])<1)})
 67    stable_boundary=max(r["a"] for r in stability if r["stable_observed"])
 68    unstable_boundary=min(r["a"] for r in stability if not r["stable_observed"])
 69    out["math_check"]={"coefficient_max_abs_error":max_coeff_err,
 70        "coefficient_rows":coeff_rows,"mismatch_loglog_slope":slope,
 71        "mismatch_quadratic_prediction":2.0,"lambda_zero_mismatch":zero_mismatch,
 72        "stability_predicted_boundary":boundary,"largest_stable_grid_a":stable_boundary,
 73        "smallest_unstable_grid_a":unstable_boundary,"stability_rows":stability}
 74
 75    # ---------------- Stage 2: tiny practical adaptation experiment ----------------
 76    # Equal-capacity denoiser D_a(v)=a*v. Target regularizer is known to teacher,
 77    # while ordinary clean-MSE sees noisy clean-image pairs at another noise level.
 78    rng2=np.random.default_rng(SEED+1)
 79    n=4000; lam=2.0; gamma=.6; a_star=prox_coeff(gamma,lam)
 80    target_var=.25; sigma_mse=.35
 81    clean=rng2.normal(0,np.sqrt(target_var),n)
 82    mse_inputs=clean+rng2.normal(0,sigma_mse,n)
 83    mse_a=float(np.dot(mse_inputs,clean)/np.dot(mse_inputs,mse_inputs))
 84    # Generate actual PnP intermediate states from random target observations.
 85    ys=clean+rng2.normal(0,.45,n)
 86    states=[]; teachers=[]
 87    for y0 in ys:
 88        x0=0.0
 89        for k in range(12):
 90            v0=x0-gamma*(x0-y0)
 91            states.append(v0); teachers.append(a_star*v0)
 92            x0=a_star*v0
 93    states=np.asarray(states); teachers=np.asarray(teachers)
 94    prox_a=float(np.dot(states,teachers)/np.dot(states,states))
 95    # A mixed objective with equal weighting in normalized scalar regression.
 96    mixed_a=float(np.dot(states,teachers)+np.dot(mse_inputs,clean))/(
 97        np.dot(states,states)+np.dot(mse_inputs,mse_inputs))
 98    # Evaluate reconstruction of held-out y. Exact target optimum is y/(1+lambda).
 99    test_y=rng2.normal(0,1.0,1000)
100    rows=[]
101    for name,a in [("clean_mse",mse_a),("proximal_match",prox_a),("mixed",mixed_a),("ideal",a_star)]:
102        losses=[]; errs=[]; stations=[]
103        for y0 in test_y:
104            h=pnp_run(a,gamma,lam,float(y0),steps=30)
105            losses.append(h[-1,1]); stations.append(h[-1,2]); errs.append((h[-1,0]-y0/(1+lam))**2)
106        rows.append({"method":name,"a":a,"final_F":float(np.mean(losses)),
107                     "final_stationarity":float(np.mean(stations)),
108                     "reconstruction_MSE_to_optimum":float(np.mean(errs)),
109                     "prox_mismatch_on_states":float(np.mean((a*states-teachers)**2))})
110    out["adaptation"]={"a_star":a_star,"clean_mse_coefficient":mse_a,
111        "proximal_coefficient":prox_a,"mixed_coefficient":mixed_a,"results":rows}
112    out["interpretation"]={
113      "mechanism_confirmed": bool(max_coeff_err<2e-3 and abs(slope-2)<0.03 and zero_mismatch<1e-20 and stable_boundary < boundary < unstable_boundary),
114      "practical_prox_better_than_clean": bool(rows[1]["final_F"] < rows[0]["final_F"] and rows[1]["final_stationarity"] < rows[0]["final_stationarity"]),
115      "note":"The toy teacher is exact and quadratic; this isolates the proposed proximal-mismatch mechanism, not image-denoising generalization."}
116    Path("results.json").write_text(json.dumps(out,indent=2))
117    print(json.dumps({"math_check":out["math_check"],"adaptation":out["adaptation"],"interpretation":out["interpretation"]},indent=2))
118
119if __name__ == '__main__': main()