scripts/claim4_float32_probe.py

"""Claim 4 add-on: run the BASELINE dual kernel-SGD in single precision (float32) to
exhibit the actual floating-point overflow of Remark 3.1 (eps = 0.01 -> exp overflows
for z >= 0.887 in float32, z >= 7.098 in float64), and record the first non-finite
iteration.  The proposed semi-dual is run under identical conditions for contrast.
"""

import json
import math
import os
import sys
from multiprocessing import Pool

import numpy as np
from scipy.special import expit

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from claim4_eot import EPS, sample_mu, sample_nu

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
OUT = os.path.join(ROOT, "outputs")


def run(args):
    mode, sigma_sq, C, seed, rho, n_iter, dtype = args
    dt = np.float32 if dtype == "float32" else np.float64
    rng = np.random.default_rng(seed)
    X = sample_mu(rng, n_iter).astype(dt)
    Y = sample_nu(rng, n_iter).astype(dt)
    c = ((X - Y) ** 2).astype(dt)
    coeffs = np.zeros(n_iter, dtype=dt)
    alpha = dt(-c[0] / EPS)
    log_rho = math.log(rho) if rho else None
    first_inf, first_nonfinite = None, None
    max_z = -np.inf
    for k in range(n_iter):
        if k > 0:
            ky = np.exp(-((Y[k] - Y[:k]) ** 2) / dt(sigma_sq))
            v = dt(ky @ coeffs[:k])
        else:
            v = dt(0.0)
        if mode == "dual":
            if k > 0:
                kx = np.exp(-((X[k] - X[:k]) ** 2) / dt(sigma_sq))
                u = dt(kx @ coeffs[:k])
            else:
                u = dt(0.0)
            z = dt(u + v - c[k])  # the quantity of Remark 3.1
            max_z = max(max_z, float(z))
            with np.errstate(over="ignore", invalid="ignore"):
                e = np.exp(dt(z / EPS))
            if not np.isfinite(e) and first_inf is None:
                first_inf = k
            g = dt(1.0) - e
        else:
            z = dt((v - c[k]) / EPS) - alpha
            max_z = max(max_z, float(z * EPS))
            g = dt(1.0 - expit(float(z) + log_rho) / rho)
            alpha = dt(alpha + C * (-EPS * g) / math.sqrt(k + 1))
        coeffs[k] = dt(C * g / math.sqrt(k + 1))
        if not np.isfinite(coeffs[k]) and first_nonfinite is None:
            first_nonfinite = k
            break
    return {
        "mode": mode,
        "dtype": dtype,
        "sigma_sq": sigma_sq,
        "C": C,
        "seed": seed,
        "rho": rho,
        "n_iter": n_iter,
        "first_exp_inf_iter": first_inf,
        "first_nonfinite_coeff_iter": first_nonfinite,
        "max_z": max_z,
        "float32_overflow_z": 0.8872283554077147,
        "float64_overflow_z": 7.09782712893384,
    }


if __name__ == "__main__":
    n = 20000
    tasks = []
    for C in [1e-3, 1e-2, 1e-1, 1.0]:
        for s in range(3):
            tasks.append(("dual", 1.0, C, s, None, n, "float32"))
    for C in [1.0, 10.0]:
        for s in range(3):
            tasks.append(("semidual", 10.0, C, s, 0.1, n, "float32"))
    with Pool(18) as pool:
        recs = pool.map(run, tasks)
    with open(os.path.join(OUT, "claim4_float32_probe.json"), "w") as fh:
        json.dump(recs, fh, indent=2)
    agg = {}
    for r in recs:
        key = f"{r['mode']}_C{r['C']}_{r['dtype']}"
        agg.setdefault(key, {"overflow_runs": 0, "n": 0, "first_inf": [], "max_z": []})
        agg[key]["n"] += 1
        agg[key]["overflow_runs"] += int(r["first_exp_inf_iter"] is not None)
        agg[key]["first_inf"].append(r["first_exp_inf_iter"])
        agg[key]["max_z"].append(round(r["max_z"], 3))
    print(json.dumps(agg, indent=1))