1import json, math
  2from pathlib import Path
  3import numpy as np
  4
  5SEED = 1452
  6
  7def sphere(rng, n, k):
  8    x = rng.normal(size=(k, n))
  9    return x / np.linalg.norm(x, axis=1, keepdims=True)
 10
 11
 12def estimate_jacobian(theta, beta_fn, n, m, delta, bsp, bmin, sigma, rng):
 13    u = sphere(rng, n, bsp)
 14    bp, bm = [], []
 15    for ui in u:
 16        ep = rng.normal(0.0, sigma / math.sqrt(bmin), size=m)
 17        em = rng.normal(0.0, sigma / math.sqrt(bmin), size=m)
 18        bp.append(beta_fn(theta + delta * ui) + ep)
 19        bm.append(beta_fn(theta - delta * ui) + em)
 20    bp, bm = np.asarray(bp), np.asarray(bm)
 21    return (n / bsp) * ((bp - bm) / (2 * delta)).T @ u
 22
 23
 24def mean_mse(theta, beta, J, n, m, delta, bsp, bmin, sigma, reps, seed):
 25    rng = np.random.default_rng(seed)
 26    vals = []
 27    for _ in range(reps):
 28        h = estimate_jacobian(theta, beta, n, m, delta, bsp, bmin, sigma, rng)
 29        vals.append(float(np.sum((h - J) ** 2)))
 30    return float(np.mean(vals)), float(np.std(vals) / math.sqrt(reps))
 31
 32
 33def mechanism_checks():
 34    rng = np.random.default_rng(SEED)
 35    n, m = 8, 3
 36    A = rng.normal(size=(m, n)) / math.sqrt(n)
 37    theta = rng.normal(size=n)
 38    q = rng.normal(size=n)
 39    c = 0.7
 40    beta_lin = lambda x: A @ x
 41    beta_non = lambda x: A @ x + c * (q @ x) ** 3 * np.ones(m)
 42    Jlin = A
 43    Jnon = A + 3 * c * (q @ theta) ** 2 * np.ones((m, 1)) @ q[None, :]
 44    out = {}
 45
 46    # Prediction 1: with a linear response and negligible rollout noise, error falls ~1/b_sp.
 47    vals = []
 48    for k in [2, 4, 8, 16, 32, 64]:
 49        e, se = mean_mse(theta, beta_lin, Jlin, n, m, 0.08, k, 100000, 0.0, 250, SEED + k)
 50        vals.append((k, e, se))
 51    slope = np.polyfit(np.log([x[0] for x in vals]), np.log([x[1] for x in vals]), 1)[0]
 52    out['sphere_scaling'] = {'prediction': 'MSE proportional to 1/b_sp (slope -1)', 'observed': vals, 'loglog_slope': float(slope)}
 53
 54    # Prediction 2: rollout noise contribution falls as 1/(delta^2*b_sp*b_min).
 55    vals = []
 56    for bmin in [4, 16, 64, 256]:
 57        e, se = mean_mse(theta, beta_lin, Jlin, n, m, 0.10, 16, bmin, 0.8, 300, SEED + bmin)
 58        vals.append((bmin, e, se))
 59    slope = np.polyfit(np.log([x[0] for x in vals]), np.log([x[1] for x in vals]), 1)[0]
 60    out['noise_scaling'] = {'prediction': 'noise MSE proportional to 1/b_min', 'observed': vals, 'loglog_slope': float(slope)}
 61
 62    # Prediction 3: central-difference nonlinear bias has MSE proportional to delta^4.
 63    vals = []
 64    for d in [0.01, 0.02, 0.04, 0.08, 0.16]:
 65        e, se = mean_mse(theta, beta_non, Jnon, n, m, d, 2000, 100000, 0.0, 30, SEED + int(d * 10000))
 66        vals.append((d, e, se))
 67    slope = np.polyfit(np.log([x[0] for x in vals]), np.log([x[1] for x in vals]), 1)[0]
 68    out['delta_bias'] = {'prediction': 'finite-difference MSE proportional to delta^4', 'observed': vals, 'loglog_slope': float(slope)}
 69    return out
 70
 71
 72def optimizer_comparison():
 73    # Performative quadratic: ordinary loss has gradient theta-y; deployed mean y=beta(theta)=a theta.
 74    # True objective is 0.5*(theta-a theta)^2 + 0.5*sigma_y^2, whose gradient is (1-a)^2 theta.
 75    rng = np.random.default_rng(SEED + 99)
 76    n, m, a = 6, 2, 0.70
 77    B = np.zeros((m, n)); B[:, :m] = a * np.eye(m)
 78    target = np.ones(m)
 79    theta0 = np.full(n, 1.5)
 80    alpha = 0.18
 81    steps = 80
 82    def beta(x): return B @ x
 83    def run(method):
 84        theta = theta0.copy(); losses=[]
 85        for t in range(steps):
 86            # Current-environment sample mean; loss = .5 ||theta[:m]-y||^2.
 87            ymean = beta(theta)
 88            gtheta = np.zeros(n); gtheta[:m] = theta[:m] - ymean - target
 89            # g_beta for .5||theta-y-beta target||^2 with y=beta+target convention.
 90            # Use a fixed differentiable surrogate: loss(beta,theta)=.5||theta[:m]-beta-target||^2.
 91            gbeta = -(theta[:m] - ymean - target)
 92            if method == 'sgd':
 93                g = gtheta
 94            elif method == 'oracle':
 95                g = gtheta + B.T @ gbeta
 96            else:
 97                h = estimate_jacobian(theta, beta, n, m, 0.04, 16, 128, 0.20, rng)
 98                g = gtheta + h.T @ gbeta
 99            theta -= alpha * g
100            residual = theta[:m] - beta(theta) - target
101            losses.append(0.5 * float(residual @ residual))
102        return losses
103    results = {k: run(k) for k in ['sgd', 'sphere', 'oracle']}
104    return {'final_loss': {k: float(v[-1]) for k,v in results.items()},
105            'loss_at_20': {k: float(v[19]) for k,v in results.items()},
106            'a': a, 'steps': steps, 'alpha': alpha}
107
108
109def main():
110    checks = mechanism_checks()
111    comparison = optimizer_comparison()
112    report = {'seed': SEED, 'checks': checks, 'optimizer_comparison': comparison}
113    Path('results.json').write_text(json.dumps(report, indent=2))
114    print(json.dumps(report, indent=2))
115
116if __name__ == '__main__':
117    main()