Resolution-aware operator data budget / resolution_budget_experiment.py

Running benchmark…

Raw ⬇ ZIP
  1import json, math, time
  2from pathlib import Path
  3import numpy as np
  4from scipy.interpolate import CubicSpline
  5from sklearn.metrics import mean_squared_error
  6
  7SEED = 469
  8rng = np.random.default_rng(SEED)
  9
 10
 11def fill_distance_1d(points, grid=None):
 12    """Numerical fill distance on [0,1] for a finite point set."""
 13    points = np.sort(np.asarray(points))
 14    if grid is None:
 15        grid = np.linspace(0.0, 1.0, 20001)
 16    return np.max(np.min(np.abs(grid[:, None] - points[None, :]), axis=1))
 17
 18
 19def verify_math():
 20    # Uniform discretizations have h approximately 1/(2m), while random
 21    # points have noticeably poorer coverage at the same cardinality.
 22    ms = np.array([4, 8, 16, 32, 64])
 23    uniform_h = np.array([fill_distance_1d(np.linspace(0, 1, int(m))) for m in ms])
 24    random_h = np.array([
 25        fill_distance_1d(np.sort(rng.uniform(0, 1, int(m)))) for m in ms
 26    ])
 27    slope = np.polyfit(np.log(ms), np.log(uniform_h), 1)[0]
 28    expected = 1.0 / (2.0 * (ms - 1))
 29    return {
 30        "m": ms.tolist(),
 31        "uniform_fill_distance": uniform_h.tolist(),
 32        "random_fill_distance": random_h.tolist(),
 33        "uniform_theory": expected.tolist(),
 34        "loglog_slope": float(slope),
 35        "max_uniform_theory_error": float(np.max(np.abs(uniform_h - expected))),
 36    }
 37
 38
 39def features(z):
 40    # Compact input class: z in [-1,1]^d, with bounded Fourier features.
 41    z = np.asarray(z)
 42    return np.concatenate([z, np.sin(np.pi*z), np.cos(np.pi*z)], axis=1)
 43
 44
 45def operator(z, x):
 46    """Smooth nonlinear coefficient-to-function operator on [0,1]."""
 47    z = np.asarray(z)
 48    x = np.asarray(x)
 49    # Coefficients decay slowly enough that refinement reveals new content.
 50    out = np.zeros((len(z), len(x)))
 51    for k in range(1, 13):
 52        phase = (k * x)[None, :]
 53        coeff = (np.sin((k + 0.35) * z[:, 0])
 54                 + 0.65 * np.cos((k + 1.2) * z[:, 1])
 55                 + 0.35 * z[:, 2] * z[:, 3]) / (k ** 1.15)
 56        out += coeff[:, None] * np.sin(2 * np.pi * phase)
 57    # nonlinear low-frequency component
 58    out += 0.30 * (z[:, 0] * z[:, 1])[:, None] * np.cos(2 * np.pi * x)[None, :]
 59    return out
 60
 61
 62def rbf_predict(xtr, ytr, xte, length=1.35, ridge=2e-3):
 63    # Small exact kernel regression is the offline operator A_off.
 64    def kernel(a, b):
 65        d2 = np.sum((a[:, None, :] - b[None, :, :]) ** 2, axis=2)
 66        return np.exp(-d2 / (2 * length * length))
 67    K = kernel(xtr, xtr)
 68    alpha = np.linalg.solve(K + ridge * np.eye(len(K)), ytr)
 69    return kernel(xte, xtr) @ alpha
 70
 71
 72def relative_l2(pred, truth):
 73    return float(np.linalg.norm(pred - truth) / np.linalg.norm(truth))
 74
 75
 76def run_experiment():
 77    # Verify the claimed fill-distance scaling before any learning experiment.
 78    math_check = verify_math()
 79    local = np.random.default_rng(SEED + 1)
 80    d = 4
 81    ntrain_max = 128
 82    ntest = 256
 83    ztrain = local.uniform(-1, 1, size=(ntrain_max, d))
 84    ztest = local.uniform(-1, 1, size=(ntest, d))
 85    ftrain = features(ztrain)
 86    ftest = features(ztest)
 87    xfine = np.linspace(0, 1, 513)
 88    truth_test = operator(ztest, xfine)
 89    # A fixed fine grid is the reference function for all reconstruction errors.
 90    resolutions = [16, 32, 64]
 91    kappas = [0.5, 1.0, 1.5]
 92    C = 1.0
 93    records = []
 94    for m in resolutions:
 95        x = np.linspace(0, 1, m)
 96        # Exact output observations and online cubic reconstruction form oracle.
 97        exact = operator(ztest, x)
 98        oracle = np.empty_like(truth_test)
 99        for j in range(ntest):
100            oracle[j] = CubicSpline(x, exact[j], bc_type='natural')(xfine)
101        oracle_err = relative_l2(oracle, truth_test)
102        # Same computational protocol for every kappa.
103        for kappa in kappas:
104            N = min(ntrain_max, int(math.ceil(C * m ** kappa)))
105            t0 = time.perf_counter()
106            pred_obs = rbf_predict(ftrain[:N], operator(ztrain[:N], x), ftest)
107            pred = np.empty_like(truth_test)
108            for j in range(ntest):
109                pred[j] = CubicSpline(x, pred_obs[j], bc_type='natural')(xfine)
110            wall = time.perf_counter() - t0
111            learned_err = relative_l2(pred, truth_test)
112            feature_h = np.max(np.min(
113                np.sum((ftest[:, None, :] - ftrain[None, :N, :]) ** 2, axis=2), axis=1
114            ) ** 0.5)
115            records.append({
116                "m": m, "kappa": kappa, "N": N,
117                "oracle_error": oracle_err,
118                "learned_error": learned_err,
119                "learned_oracle_gap": learned_err - oracle_err,
120                "feature_fill_proxy": float(feature_h),
121                "kernel_solve_flops_proxy": int(N ** 3),
122                "wall_seconds": wall,
123            })
124    # Fixed-data baseline uses N equal to the largest budget at each m only for
125    # the comparison requested by the idea: N=16 does not grow with resolution.
126    fixed = []
127    for m in resolutions:
128        N = 16
129        x = np.linspace(0, 1, m)
130        t0 = time.perf_counter()
131        obs = rbf_predict(ftrain[:N], operator(ztrain[:N], x), ftest)
132        pred = np.empty_like(truth_test)
133        for j in range(ntest):
134            pred[j] = CubicSpline(x, obs[j], bc_type='natural')(xfine)
135        fixed.append({"m": m, "N": N, "error": relative_l2(pred, truth_test),
136                      "wall_seconds": time.perf_counter() - t0})
137    return {"math_check": math_check, "records": records, "fixed_N_baseline": fixed}
138
139
140if __name__ == "__main__":
141    result = run_experiment()
142    Path("results.json").write_text(json.dumps(result, indent=2))
143    print(json.dumps(result, indent=2))