scripts/exp_nonquadratic.py

"""Non-quadratic targets: needed for claim 6 (and reinforcing claim 4).

Why: on a quadratic V every scheme here is a linear-Gaussian recursion and the
first-order weak error vanishes identically, so LMC's KL bias is anomalously
O(h^2) instead of its generic O(h).  Any ULMC-vs-LMC comparison run only on
Gaussians is therefore biased *against* the paper.  We redo it on a genuinely
non-quadratic ridge-separable target.

Part A  bias-floor law of ULMC / LMC / composite-LMC on
        f(x) = (alpha/2) x^2 + log cosh(x)      (beta = alpha+1, H = alpha+1)
Part B  the same on the pure quadratic (the degenerate control)
Part C  head-to-head iteration complexity on the d-dimensional ridge-separable
        target: m stiff coordinates with the f above, d-m flat coordinates with
        (alpha/2)x^2 (exact Gaussian).  tr(H) = m(alpha+1) + (d-m) alpha << beta d.
Part D  claim 4: alpha = 0 EXACTLY, V(x) = sum_i log cosh(c_i x_i).
        inf grad^2 V = 0, H = diag(c_i^2), tr(H) = sum c_i^2.
        Dimension-freeness at alpha=0 holds iff the per-coordinate bias is
        proportional to c^2, so that the total depends on tr(H) and not on d.
"""

import numpy as np

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

NX, NP, NSTEP = 512, 192, 420
res = {}


def make_pot(alpha):
    return (
        lambda x: alpha * x**2 / 2 + np.log(np.cosh(x)),
        lambda x: alpha * x + np.tanh(x),
        lambda x: alpha + 1.0 / np.cosh(x) ** 2,
    )


def ulmc_floor_grid(f, df, ddf, beta, h, Lx=9.0, nstep=NSTEP):
    g = C.gamma_of(beta)
    gr = G.Grid2DPullback(Lx, NX, 7.0, NP, 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 // 6)
    return float(kl[-1]), [float(v) for v in kl]


def od_floor_grid(scheme, f, df, beta, alpha, h, Lx=9.0, nx=4096, nstep=4000):
    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 // 6)
    return float(kl[-1]), [float(v) for v in kl]


# ---------------------------------------------------------------- Part A / B
ALPHA = 0.05
BETA = 1.0 + ALPHA
f, df, ddf = make_pot(ALPHA)
fq, dfq, ddfq = (
    lambda x: ALPHA * x**2 / 2,
    lambda x: ALPHA * x,
    lambda x: ALPHA * np.ones_like(x),
)

for tag, (ff, dff, ddff), bb in (
    ("nonquadratic", (f, df, ddf), BETA),
    ("quadratic_control", (fq, dfq, ddfq), ALPHA),
):
    blk = {}
    hs_u = np.array([0.17, 0.14, 0.115, 0.095, 0.078, 0.064])
    hs_u = hs_u[hs_u <= 1.0 / C.gamma_of(bb)]
    fl = [ulmc_floor_grid(ff, dff, ddff, bb, float(h))[0] for h in hs_u]
    s, r2 = C.fit_exponent(hs_u, fl)
    blk["ulmc"] = {"h": hs_u.tolist(), "floor": fl, "h_exponent": s, "r2": r2}
    hs_o = np.array([0.4, 0.28, 0.2, 0.14, 0.1, 0.07])
    for sch in ("lmc", "composite"):
        fl = [od_floor_grid(sch, ff, dff, bb, ALPHA, float(h))[0] for h in hs_o]
        s, r2 = C.fit_exponent(hs_o, fl)
        blk[sch] = {"h": hs_o.tolist(), "floor": fl, "h_exponent": s, "r2": r2}
    res[f"bias_law_{tag}"] = blk
    print(tag, {k: round(v["h_exponent"], 3) for k, v in blk.items()}, flush=True)

# ---------------------------------------------------------------- Part C
# ridge-separable, m stiff non-quadratic coordinates + (d-m) flat quadratic
M_RIDGE, D = 4, 400
NA = 3000


def stiff_kl_curve(scheme, h, nsteps=NA):
    if scheme == "ulmc":
        kl, curve = ulmc_floor_grid(f, df, ddf, BETA, h, nstep=NSTEP)
        return np.array(curve), kl
    _, curve = od_floor_grid(scheme, f, df, BETA, ALPHA, h, nstep=nsteps)
    return np.array(curve), curve[-1]


NGRID = U.n_grid(int(4e9), 1.04)
rows = []
for eps in (0.10, 0.07, 0.05, 0.035, 0.025, 0.018, 0.013):
    row = {"eps": float(eps)}
    for scheme, hs in (
        ("ulmc", np.geomspace(0.064, 1 / C.gamma_of(BETA), 8)),
        ("lmc", np.geomspace(0.004, 0.5, 14)),
        ("composite", np.geomspace(0.004, 0.5, 14)),
    ):
        best = None
        for h in hs:
            _, plateau = stiff_kl_curve(scheme, float(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, float(h), C.gamma_of(BETA))
            elif scheme == "lmc":
                mp = U.lmc_maps(a, float(h))
            else:
                mp = U.composite_lmc_maps(a, float(h), ALPHA)
            flat = U.kl_curve(mp, S0, m0, a, mult, NGRID)
            tot = flat + M_RIDGE * plateau  # stiff coords are at their floor
            n = U.first_below(NGRID, tot, eps * eps)
            if n and (best is None or n < best[0]):
                best = (n, float(h), float(plateau))
        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,
        }
    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")},
        "speedup",
        round(row.get("speedup_ulmc_over_composite", 0), 2),
        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(M_RIDGE * BETA + (D - M_RIDGE) * ALPHA),
    "beta_times_d": float(BETA * D),
    "ulmc_beats_composite_at_small_eps": bool(
        ok and ok[-1]["speedup_ulmc_over_composite"] > 1
    ),
    "speedups": [(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}_predicted_eps_exponent"] = (
        -1.0 if sch == "ulmc" else -2.0
    )


# ---------------------------------------------------------------- Part D
# alpha = 0 EXACTLY:  V(x) = sum_i log cosh(c_i x_i)
def logcosh(c):
    return (
        lambda x: np.log(np.cosh(c * x)),
        lambda x: c * np.tanh(c * x),
        lambda x: c * c / np.cosh(c * x) ** 2,
    )


rowsD = []
for c in (0.7, 0.85, 1.0, 1.2, 1.45, 1.7):
    fc, dfc, ddfc = logcosh(c)
    beta_c = c * c  # sup grad^2 V ; inf = 0  => alpha = 0
    hmax = 1.0 / C.gamma_of(beta_c)
    hh = min(0.12, 0.95 * hmax)
    fl, _ = ulmc_floor_grid(fc, dfc, ddfc, beta_c, hh, Lx=max(9.0, 9.0 / c))
    rowsD.append(
        {
            "c": c,
            "trH_per_coord_c2": c * c,
            "beta": beta_c,
            "h": hh,
            "gamma_h": float(C.gamma_of(beta_c) * hh),
            "floor": fl,
            "floor_over_c2": fl / c**2,
        }
    )
# to isolate the c-dependence, hold the dimensionless gamma*h fixed
sc, r2c = C.fit_exponent([r["c"] for r in rowsD], [r["floor"] for r in rowsD])
res["alpha0_logcosh"] = {
    "rows": rowsD,
    "c_exponent_of_bias": sc,
    "r2": r2c,
    "predicted_c_exponent_for_dimension_freeness": 2.0,
    "note": "bias ~ c^2 = per-coordinate tr(H) contribution => total bias is a "
    "function of tr(H)=sum c_i^2 with no explicit d, at alpha = 0 exactly",
    "gamma_h_held_fixed": [r["gamma_h"] for r in rowsD],
}
print(
    "alpha=0 log-cosh: c-exponent of ULMC bias =",
    round(sc, 3),
    "(2.0 => dimension-free)",
    flush=True,
)

# mixed spectrum: does the total bias track tr(H) rather than d?
mix = []
for d, cs in (
    (2, [1.0, 1.0]),
    (5, [1.0, 0.5, 0.5, 0.5, 0.5]),
    (10, [1.0] + [np.sqrt(1.0 / 9)] * 9),
    (26, [1.0] + [np.sqrt(1.0 / 25)] * 25),
):
    beta_c = 1.0
    hh = 0.12
    tot, trH = 0.0, 0.0
    for c in cs:
        fc, dfc, ddfc = logcosh(c)
        fl, _ = ulmc_floor_grid(fc, dfc, ddfc, beta_c, hh, Lx=max(9.0, 9.0 / c))
        tot += fl
        trH += c * c
    mix.append(
        {"d": d, "trH": trH, "total_bias": tot, "h": hh, "bias_over_trH": tot / trH}
    )
res["alpha0_dimension_free_mixed_spectra"] = {
    "rows": mix,
    "bias_over_trH_spread": float(
        max(r["bias_over_trH"] for r in mix) / min(r["bias_over_trH"] for r in mix)
    ),
    "d_range": [mix[0]["d"], mix[-1]["d"]],
}
print(
    "alpha=0 mixed spectra bias/trH:",
    [round(r["bias_over_trH"], 8) for r in mix],
    "over d =",
    [r["d"] for r in mix],
    flush=True,
)

C.dump("nonquadratic", res)