1import json
  2import math
  3import numpy as np
  4
  5
  6def sqrt_quantile_grid(u, a, n, delta=0.0):
  7    u = np.asarray(u, float); a = np.maximum(np.asarray(a, float), 0) + delta
  8    w = np.sqrt(a)
  9    c = np.zeros_like(u)
 10    c[1:] = np.cumsum((w[1:] + w[:-1]) * np.diff(u) / 2)
 11    targets = np.linspace(0, c[-1], n + 1)
 12    return np.interp(targets, c, u)
 13
 14
 15def piecewise_cost(grid, u, a):
 16    # Dense Gauss-Legendre quadrature per interval avoids coarse-grid bias.
 17    nodes, weights = np.polynomial.legendre.leggauss(12)
 18    lo, hi = grid[:-1], grid[1:]
 19    x = (lo[:, None] + hi[:, None]) / 2 + (hi[:, None] - lo[:, None]) * nodes[None, :] / 2
 20    integ = np.sum(np.interp(x, u, a) * weights[None, :], axis=1) * (hi-lo) / 2
 21    return .5 * np.sum(integ * (hi-lo))
 22
 23
 24def empirical_density(u, temporal_amp=1.0, spatial_amp=0.0, seed=0, batch=512):
 25    """Toy analogue of E[||d_u s||^2 + lambda ||J_x s||_F^2].
 26    The score is s(x,u)=q(u)x + p(u) tanh(x), with x~N(0,1).
 27    Finite differences and exact Jacobian are used as a pilot estimate.
 28    """
 29    rng = np.random.default_rng(seed)
 30    x = rng.normal(size=(batch, 1))
 31    q = 0.35 + 0.7 * u + 1.5 * np.exp(-((u-.70)/.09)**2)
 32    dq = 0.7 + 1.5 * np.exp(-((u-.70)/.09)**2) * (-2*(u-.70)/(.09**2))
 33    p = 0.25 * np.sin(5*u)
 34    dp = 1.25 * np.cos(5*u)
 35    temporal = np.mean((dq*x + dp*np.tanh(x))**2, axis=0)
 36    spatial = np.mean((q + p*(1-np.tanh(x)**2))**2, axis=0)
 37    return temporal_amp * temporal + spatial_amp * spatial
 38
 39
 40def euler_variable_ode(grid, k=18.0):
 41    """y'=-r(u)y, with rapid rate around .7; exact y(1) is known."""
 42    r = lambda z: 0.35 + k*np.exp(-((z-.70)/.075)**2)
 43    y = 1.0
 44    for lo, hi in zip(grid[:-1], grid[1:]):
 45        y *= 1.0 - r((lo+hi)/2)*(hi-lo)
 46    exact = np.exp(-(.35 + k*.075*np.sqrt(np.pi)/2 * (
 47        math.erf((1-.70)/.075) - math.erf((0-.70)/.075))))
 48    return abs(y-exact), y, exact
 49
 50
 51def main():
 52    u = np.linspace(0, 1, 20001)
 53    # Density with one narrow high-error region is the mechanism test.
 54    a = 1.0 + 80.0*np.exp(-((u-.70)/.075)**2)
 55    rows = []
 56    for n in [8, 16, 32, 64, 128]:
 57        uni = np.linspace(0, 1, n+1)
 58        ad = sqrt_quantile_grid(u, a, n, delta=1e-10)
 59        cuni = piecewise_cost(uni, u, a)
 60        cad = piecewise_cost(ad, u, a)
 61        lower = .5*np.trapz(np.sqrt(a), u)**2/n
 62        rows.append({'N': n, 'uniform_cost': cuni, 'adaptive_cost': cad,
 63                     'predicted_lower_bound': lower,
 64                     'adaptive_over_bound': cad/lower,
 65                     'uniform_over_adaptive': cuni/cad,
 66                     'max_equalized_cost_ratio': (np.max(np.interp((ad[:-1]+ad[1:])/2,u,a)*np.diff(ad)**2)/
 67                                                   np.min(np.interp((ad[:-1]+ad[1:])/2,u,a)*np.diff(ad)**2))})
 68
 69    # Prediction 1: adaptive cost tends to lower bound and interval costs equalize.
 70    # Prediction 2: local widths follow inverse square-root density (log correlation).
 71    n = 64; ad = sqrt_quantile_grid(u, a, n); mids=(ad[:-1]+ad[1:])/2
 72    widths=np.diff(ad); amid=np.interp(mids,u,a)
 73    width_pred=1/np.sqrt(amid)
 74    corr=float(np.corrcoef(np.log(widths), np.log(width_pred))[0,1])
 75    # Prediction 3: multiplying density by c changes cost by c but not grid.
 76    scale_rows=[]
 77    ref=sqrt_quantile_grid(u,a,n)
 78    for c in [.25, 1., 4., 16.]:
 79        g=sqrt_quantile_grid(u,c*a,n)
 80        scale_rows.append({'scale':c, 'grid_max_difference':float(np.max(abs(g-ref))),
 81                           'cost_ratio_to_reference':piecewise_cost(g,u,c*a)/piecewise_cost(ref,u,a),
 82                           'predicted_cost_ratio':c})
 83
 84    # Empirical-pilot robustness: noisy density estimates should preserve the allocation.
 85    # Each estimate is positive and smoothed by a small floor, as in the proposal.
 86    robustness = []
 87    true_grid = sqrt_quantile_grid(u, a, 32)
 88    for rel_noise in [0.0, 0.05, 0.10, 0.20, 0.40]:
 89        rng = np.random.default_rng(1000 + int(rel_noise*100))
 90        pilot = np.maximum(a[::200] * (1 + rel_noise*rng.normal(size=len(a[::200]))), 1e-8)
 91        pu = u[::200]
 92        est_grid = sqrt_quantile_grid(pu, pilot, 32, delta=1e-8*np.mean(pilot))
 93        robustness.append({'relative_noise': rel_noise,
 94                           'grid_linf_error': float(np.max(abs(est_grid-true_grid))),
 95                           'cost_ratio_using_true_a': float(piecewise_cost(est_grid,u,a)/piecewise_cost(true_grid,u,a))})
 96
 97    # Mini scheduler experiment: same Euler method, fixed NFE.
 98    ode=[]
 99    for n in [8, 16, 32, 64]:
100        ug=np.linspace(0,1,n+1)
101        # use normalized empirical density; only its shape matters
102        eg=sqrt_quantile_grid(u,a,n)
103        eu=euler_variable_ode(ug); ea=euler_variable_ode(eg)
104        ode.append({'N':n, 'uniform_abs_error':eu[0], 'adaptive_abs_error':ea[0],
105                    'error_ratio_adaptive_over_uniform':ea[0]/eu[0]})
106
107    report={'predictions':[
108        {'claim':'adaptive cost approaches C*=L^2/(2N)', 'observed':rows},
109        {'claim':'width proportional to a^{-1/2}', 'observed_log_correlation':corr,
110         'tolerance':'correlation > 0.98 on this smooth density'},
111        {'claim':'a -> c*a leaves grid invariant and scales cost by c', 'observed':scale_rows}
112    ], 'pilot_robustness': robustness, 'ode_mini_experiment':ode,
113    'density_summary':{'min':float(a.min()),'max':float(a.max()),'peak_u':float(u[a.argmax()])}}
114    print(json.dumps(report, indent=2))
115
116    # Mechanism checks (strict, reproducible).
117    assert rows[-1]['adaptive_over_bound'] < 1.01
118    assert corr > .98
119    assert max(x['grid_max_difference'] for x in scale_rows) < 1e-9
120    assert max(abs(x['cost_ratio_to_reference']-x['predicted_cost_ratio']) for x in scale_rows) < 1e-10
121
122if __name__ == '__main__':
123    main()