scripts/exp_nonquadratic2.py

"""Non-quadratic targets, corrected.

Fixes over exp_nonquadratic.py: (i) enough steps to actually reach the stationary
law -- every reported floor carries a convergence check comparing the value at
50 % and 100 % of the horizon; (ii) the stiff ridge coordinate is uniformly
strongly convex so that it relaxes fast, f(y) = (1-a0)/2 y^2 + s*log cosh(y),
grad^2 f in [1, 1+s]; (iii) results cached per (scheme, h) instead of recomputed
inside the eps loop.

Part A  stationary KL bias law of ULMC / LMC / composite on a non-quadratic
        stiff coordinate, and on the quadratic control.
Part C  head-to-head iteration complexity on the d-dimensional ridge-separable
        target with tr(H) << beta d.
"""

import numpy as np

import common as C
import grid1d as G
import ulmc_core as U

res = {}
NGRID = U.n_grid(int(4e9), 1.04)
_cache = {}


def ulmc_floor(f, df, ddf, beta, h, Lx=9.0, nstep=1500, nx=384, npv=160):
    key = ("u", id(f), beta, h, Lx, nstep, nx)
    if key in _cache:
        return _cache[key]
    g = C.gamma_of(beta)
    gr = G.Grid2DPullback(Lx, nx, 7.0, npv, f, df, ddf, g)
    rho0 = np.exp(-(gr.X**2 * beta / 2 + gr.P**2 / 2))
    rho0 /= rho0.sum() * gr.dx * gr.dp
    _, kl = gr.run_ulmc_pb(h, rho0, nstep, record_every=nstep // 2)
    out = (float(kl[-1]), float(abs(kl[-1] - kl[0]) / max(kl[-1], 1e-300)))
    _cache[key] = out
    return out


def od_floor(scheme, f, df, beta, alpha, h, Lx=9.0, nx=4096, nstep=None):
    nstep = nstep or max(3000, int(25.0 / h))
    key = (scheme, id(f), beta, alpha, h, Lx, nx, nstep)
    if key in _cache:
        return _cache[key]
    gr = G.Grid1D(Lx, nx, f, df)
    rho0 = np.exp(-gr.x**2 * beta / 2)
    rho0 /= rho0.sum() * gr.dx
    _, kl = gr.run(scheme, h, rho0, nstep, alpha=alpha, record_every=nstep // 2)
    out = (float(kl[-1]), float(abs(kl[-1] - kl[0]) / max(kl[-1], 1e-300)))
    _cache[key] = out
    return out


# ---------------------------------------------------------------- Part A
ALPHA, SFT = 0.05, 0.5
BETA = 1.0 + SFT  # sup grad^2 f = 1 + s
f = lambda x: (1 - ALPHA) * x**2 / 2 + SFT * np.log(np.cosh(x))
df = lambda x: (1 - ALPHA) * x + SFT * np.tanh(x)
ddf = lambda x: (1 - ALPHA) + SFT / np.cosh(x) ** 2
fq = lambda x: BETA * x**2 / 2
dfq = lambda x: BETA * x
ddfq = lambda x: BETA * np.ones_like(x)

HS_U = np.array([0.13, 0.112, 0.096, 0.082, 0.070])
HS_O = np.array([0.40, 0.28, 0.20, 0.14, 0.10, 0.07])
for tag, (ff, dff, ddff) in (
    ("nonquadratic", (f, df, ddf)),
    ("quadratic_control", (fq, dfq, ddfq)),
):
    blk = {}
    hs = HS_U[HS_U <= 1.0 / C.gamma_of(BETA)]
    vals = [ulmc_floor(ff, dff, ddff, BETA, float(h)) for h in hs]
    s, r2 = C.fit_exponent(hs, [v[0] for v in vals])
    blk["ulmc"] = {
        "h": hs.tolist(),
        "floor": [v[0] for v in vals],
        "converged_rel_change": [v[1] for v in vals],
        "h_exponent": s,
        "r2": r2,
    }
    for sch in ("lmc", "composite"):
        vals = [
            od_floor(sch, ff, dff, BETA, ALPHA, float(h))
            for h in HS_O
        ]
        s, r2 = C.fit_exponent(HS_O, [v[0] for v in vals])
        blk[sch] = {
            "h": HS_O.tolist(),
            "floor": [v[0] for v in vals],
            "converged_rel_change": [v[1] for v in vals],
            "h_exponent": s,
            "r2": r2,
        }
    res[f"bias_law_{tag}"] = blk
    print(
        tag,
        {
            k: (round(v["h_exponent"], 3), round(max(v["converged_rel_change"]), 5))
            for k, v in blk.items()
        },
        flush=True,
    )

# ---------------------------------------------------------------- Part C
M_RIDGE, D = 4, 400
TRH = M_RIDGE * BETA + (D - M_RIDGE) * ALPHA
HSETS = {
    "ulmc": np.geomspace(0.055, 1 / C.gamma_of(BETA), 9),
    "lmc": np.geomspace(0.006, 0.9, 16),
    "composite": np.geomspace(0.006, 0.9, 16),
}
tables = {}
for scheme in ("ulmc", "lmc", "composite"):
    tables[scheme] = []
    for h in HSETS[scheme]:
        h = float(h)
        if scheme == "ulmc":
            plateau, conv = ulmc_floor(f, df, ddf, BETA, h)
        else:
            plateau, conv = od_floor(scheme, f, df, BETA, ALPHA, h)
        a = np.array([ALPHA])
        mult = np.array([float(D - M_RIDGE)])
        S0, m0 = C.init_cold(a, BETA)
        if scheme == "ulmc":
            mp = U.ulmc_maps(a, h, C.gamma_of(BETA))
        elif scheme == "lmc":
            mp = U.lmc_maps(a, h)
        else:
            mp = U.composite_lmc_maps(a, h, ALPHA)
        curve = U.kl_curve(mp, S0, m0, a, mult, NGRID) + M_RIDGE * plateau
        tables[scheme].append((h, curve, plateau, conv))
    print("built", scheme, flush=True)

rows = []
for eps in (0.12, 0.09, 0.065, 0.048, 0.035, 0.026, 0.019, 0.014):
    row = {"eps": float(eps)}
    for scheme in ("ulmc", "lmc", "composite"):
        best = None
        for h, curve, plateau, conv in tables[scheme]:
            n = U.first_below(NGRID, curve, eps * eps)
            if n and (best is None or n < best[0]):
                best = (n, h, plateau, conv)
        row[scheme] = {
            "N_eps": best[0] if best else None,
            "h_star": best[1] if best else None,
            "stiff_floor": best[2] if best else None,
            "floor_rel_change_over_last_half": best[3] if best else None,
        }
    if row["ulmc"]["N_eps"] and row["composite"]["N_eps"]:
        row["speedup_ulmc_over_composite"] = (
            row["composite"]["N_eps"] / row["ulmc"]["N_eps"]
        )
        row["speedup_ulmc_over_lmc"] = row["lmc"]["N_eps"] / row["ulmc"]["N_eps"]
    rows.append(row)
    print(
        "eps",
        eps,
        {k: row[k]["N_eps"] for k in ("ulmc", "lmc", "composite")},
        "ULMC/composite speedup",
        round(row.get("speedup_ulmc_over_composite", 0), 3),
        flush=True,
    )

ok = [r for r in rows if r.get("speedup_ulmc_over_composite")]
res["head_to_head_nonquadratic"] = {
    "rows": rows,
    "m_ridge": M_RIDGE,
    "d": D,
    "alpha": ALPHA,
    "beta": BETA,
    "trH": float(TRH),
    "beta_times_d": float(BETA * D),
    "trH_over_beta_d": float(TRH / (BETA * D)),
    "ulmc_beats_composite_anywhere": bool(
        any(r["speedup_ulmc_over_composite"] > 1 for r in ok)
    ),
    "best_speedup": float(max(r["speedup_ulmc_over_composite"] for r in ok)),
    "best_speedup_eps": float(
        max(ok, key=lambda r: r["speedup_ulmc_over_composite"])["eps"]
    ),
    "speedup_trend": [(r["eps"], r["speedup_ulmc_over_composite"]) for r in ok],
}
for sch in ("ulmc", "lmc", "composite"):
    s, r2 = C.fit_exponent(
        [r["eps"] for r in rows if r[sch]["N_eps"]],
        [r[sch]["N_eps"] for r in rows if r[sch]["N_eps"]],
    )
    res["head_to_head_nonquadratic"][f"{sch}_eps_exponent"] = s
    res["head_to_head_nonquadratic"][f"{sch}_r2"] = r2
res["head_to_head_nonquadratic"]["table1_eps_exponents"] = {
    "ulmc_paper": -1.0,
    "composite_freund": -2.0,
    "lmc": -2.0,
}


C.dump("nonquadratic", res)