1import json
  2import math
  3import numpy as np
  4
  5
  6def weighted_box_projection(z, lo, hi, m):
  7    """Exact argmin over a box of 0.5*(y-z)^T M^{-1}(y-z)."""
  8    z = np.asarray(z, dtype=float)
  9    return np.minimum(hi, np.maximum(lo, z))
 10
 11
 12class StateMetricProjected:
 13    def __init__(self, dim, alpha=0.1, eta=1.0, beta2=0.9,
 14                 metric_beta=0.5, m_min=1e-3, m_max=10.0,
 15                 lo=-2.0, hi=2.0, eps=1e-12):
 16        self.alpha, self.eta, self.beta2 = alpha, eta, beta2
 17        self.metric_beta = metric_beta
 18        self.m_min, self.m_max = m_min, m_max
 19        self.v = np.zeros(dim)
 20        self.lo, self.hi, self.eps = lo, hi, eps
 21
 22    def step(self, theta, grad):
 23        self.v = self.beta2 * self.v + (1.0 - self.beta2) * grad * grad
 24        t = self.v / max(1.0 - self.beta2, self.eps)  # bias correction
 25        logm = self.metric_beta * np.log(t + self.eps)
 26        m = np.exp(np.clip(logm, np.log(self.m_min), np.log(self.m_max)))
 27        z = theta - self.alpha * m * grad
 28        p = weighted_box_projection(z, self.lo, self.hi, m)
 29        new_theta = theta + self.eta * (p - theta)
 30        return new_theta, m
 31
 32
 33def projected_sgd_step(theta, grad, lr, lo=-2.0, hi=2.0):
 34    return np.clip(theta - lr * grad, lo, hi)
 35
 36
 37def projected_adam_step(theta, grad, state, lr=0.1, beta1=.9, beta2=.999,
 38                        lo=-2.0, hi=2.0, eps=1e-8):
 39    state[0] = beta1 * state[0] + (1-beta1) * grad
 40    state[1] = beta2 * state[1] + (1-beta2) * grad * grad
 41    state[2] += 1
 42    mh = state[0] / (1-beta1**state[2])
 43    vh = state[1] / (1-beta2**state[2])
 44    return np.clip(theta - lr * mh / (np.sqrt(vh)+eps), lo, hi)
 45
 46
 47def verify_projection():
 48    rng = np.random.default_rng(7)
 49    z = rng.normal(size=8)
 50    lo, hi = -0.4, 0.6
 51    m = np.exp(rng.normal(size=8))
 52    p = weighted_box_projection(z, lo, hi, m)
 53    # KKT sign check for the separable weighted projection.
 54    stationarity = (p-z)/m
 55    ok = True
 56    for i in range(len(p)):
 57        if lo < p[i] < hi:
 58            ok &= abs(stationarity[i]) < 1e-12
 59        elif p[i] <= lo:
 60            ok &= stationarity[i] >= -1e-12
 61        elif p[i] >= hi:
 62            ok &= stationarity[i] <= 1e-12
 63    return {"max_box_violation": float(max(np.max(lo-p), np.max(p-hi), 0)),
 64            "kkt_sign_check": bool(ok)}
 65
 66
 67def contraction_sweep():
 68    # For x' = x + eta(P(x-alpha*m*h*x)-x), interior factor is 1-eta*alpha*m*h.
 69    h, eta = 37.0, 1.0
 70    alphas = np.array([.001, .01, .02, .03, .04, .05, .06, .07, .08])
 71    observed, predicted = [], []
 72    for alpha in alphas:
 73        x = 0.73
 74        vals = []
 75        for _ in range(30):
 76            vals.append(abs(x))
 77            x = x + eta * ((x-alpha*h*x)-x)
 78        # geometric ratio over early nonzero iterates
 79        ratios = np.abs(np.asarray(vals[1:]) / np.asarray(vals[:-1]))
 80        observed.append(float(np.median(ratios)))
 81        predicted.append(float(abs(1-eta*alpha*h)))
 82    # Stability transition predicted at alpha*h=2; classify by a 100-step bounded run.
 83    boundary_pred = 2.0 / (eta*h)
 84    boundary_obs = None
 85    for alpha in np.linspace(.045, .06, 301):
 86        x = .73
 87        for _ in range(100):
 88            x = x + eta*((x-alpha*h*x)-x)
 89        if abs(x) > 1e3:
 90            boundary_obs = float(alpha)
 91            break
 92    return {"h": h, "eta": eta, "alphas": alphas.tolist(),
 93            "predicted_contraction": predicted, "observed_contraction": observed,
 94            "stability_boundary_predicted": boundary_pred,
 95            "stability_boundary_observed_first_divergence": boundary_obs,
 96            "contraction_max_abs_error": float(np.max(np.abs(np.array(observed)-np.array(predicted))))}
 97
 98
 99def metric_scaling_sweep():
100    # Prediction: for an interior coordinate, the per-step fractional decrease is
101    # alpha*m*h (until the sign-flip boundary), hence gain is linear in m.
102    h, alpha, eta = 23.0, .01, 1.0
103    ms = np.array([.1, .2, .4, .8, 1.2])
104    observed, predicted = [], []
105    for m in ms:
106        x = .91
107        x1 = x + eta*((x-alpha*m*h*x)-x)
108        observed.append(float((x-x1)/x))
109        predicted.append(float(eta*alpha*m*h))
110    return {"h": h, "alpha": alpha, "metrics": ms.tolist(),
111            "predicted_fractional_decrease": predicted,
112            "observed_fractional_decrease": observed,
113            "linearity_max_abs_error": float(np.max(np.abs(np.array(observed)-np.array(predicted))))}
114
115
116def quadratic_comparison():
117    rng = np.random.default_rng(123)
118    H = np.array([1.0, 10000.0])
119    theta0 = np.array([1.0, 1.0])
120    lo, hi = -2.0, 2.0
121    steps = 180
122    def grad(x): return H*x
123    def loss(x): return .5*float(np.sum(H*x*x))
124    records = {}
125    # Stable but deliberately conservative SGD for the ill-conditioned direction.
126    x = theta0.copy(); ls=[]; viol=0
127    for _ in range(steps):
128        g=grad(x); x=projected_sgd_step(x,g,1.5e-4,lo,hi); ls.append(loss(x)); viol=max(viol,float(max(np.max(lo-x),np.max(x-hi),0)))
129    records['projected_sgd']={'loss_curve':ls,'final_loss':ls[-1],'violations':viol}
130    x=theta0.copy(); ls=[]; viol=0; st=[np.zeros(2),np.zeros(2),0]
131    for _ in range(steps):
132        x=projected_adam_step(x,grad(x),st,lr=.1,lo=lo,hi=hi); ls.append(loss(x)); viol=max(viol,float(max(np.max(lo-x),np.max(x-hi),0)))
133    records['projected_adam']={'loss_curve':ls,'final_loss':ls[-1],'violations':viol}
134    x=theta0.copy(); ls=[]; viol=0; opt=StateMetricProjected(2,alpha=.03,eta=1,beta2=.9,metric_beta=.5,m_min=.01,m_max=10,lo=lo,hi=hi)
135    metric_trace=[]
136    for _ in range(steps):
137        g=grad(x); x,m=opt.step(x,g); metric_trace.append(m.copy()); ls.append(loss(x)); viol=max(viol,float(max(np.max(lo-x),np.max(x-hi),0)))
138    records['state_metric_projected']={'loss_curve':ls,'final_loss':ls[-1],'violations':viol,'final_metric':metric_trace[-1].tolist()}
139    return {k:{kk:vv for kk,vv in v.items() if kk!='loss_curve'} for k,v in records.items()}, records
140
141
142def main():
143    proj=verify_projection()
144    contraction=contraction_sweep()
145    scaling=metric_scaling_sweep()
146    summary, curves=quadratic_comparison()
147    out={'projection_check':proj,'prediction_sweeps':{'contraction':contraction,'metric_scaling':scaling},'quadratic_summary':summary,
148         'loss_curves':{k:v['loss_curve'] for k,v in curves.items()}}
149    with open('results.json','w') as f: json.dump(out,f,indent=2)
150    print(json.dumps({'projection_check':proj,'prediction_sweeps':out['prediction_sweeps'],'quadratic_summary':summary},indent=2))
151
152if __name__=='__main__': main()