scripts/common.py

"""Independent implementation of the Safe-KL / SoftPlus machinery of arXiv:2509.24894.

Everything here is re-derived from the definitions in the paper (Def 2.1, Lemma 2.3,
Eq. 6) -- no code from the authors' repository is used.
"""

import numpy as np
from scipy.optimize import brentq
from scipy.special import logsumexp, expit


def f_rho(t, rho):
    """Safe KL entropy generator, Definition 2.1 / Eq. (3). +inf outside [0, 1/rho]."""
    t = np.asarray(t, dtype=float)
    out = np.full(t.shape, np.inf)
    ok = (t >= 0) & (t <= 1.0 / rho)
    tt = t[ok]
    with np.errstate(divide="ignore", invalid="ignore"):
        tlogt = np.where(tt > 0, tt * np.log(np.where(tt > 0, tt, 1.0)), 0.0)
        u = 1.0 - rho * tt
        tail = np.where(u > 0, (u / rho) * np.log(np.where(u > 0, u, 1.0)), 0.0)
    out[ok] = tlogt + 1.0 + tail
    return out


def d_f_rho(t, rho):
    """f_rho'(t) = log( t / (1 - rho t) )."""
    t = np.asarray(t, dtype=float)
    return np.log(t) - np.log1p(-rho * t)


def d2_f_rho(t, rho):
    """f_rho''(t) = 1/t + rho/(1 - rho t)."""
    t = np.asarray(t, dtype=float)
    return 1.0 / t + rho / (1.0 - rho * t)


def f_rho_star(s, rho):
    """Lemma 2.3: f_rho^*(s) = (1/rho) log(1 + rho e^s) - 1 (rescaled SoftPlus)."""
    s = np.asarray(s, dtype=float)
    return np.logaddexp(0.0, s + np.log(rho)) / rho - 1.0


def d_f_rho_star(s, rho):
    """(f_rho^*)'(s) = e^s / (1 + rho e^s) = sigmoid(s + log rho) / rho."""
    s = np.asarray(s, dtype=float)
    return expit(s + np.log(rho)) / rho


def d2_f_rho_star(s, rho):
    """(f_rho^*)''(s) = e^s / (1 + rho e^s)^2."""
    s = np.asarray(s, dtype=float)
    z = expit(s + np.log(rho))
    return z * (1.0 - z) / rho


def F_true(phi, w=None):
    """F(phi; mu) = log int e^phi dmu for a discrete mu with weights w."""
    phi = np.asarray(phi, dtype=float)
    if w is None:
        w = np.full(phi.shape, 1.0 / phi.size)
    return logsumexp(phi, b=w)


def alpha_star(phi, rho, w=None):
    """Solve Eq. (7): int e^{phi-a}/(1+rho e^{phi-a}) dmu = 1, i.e.
    sum_i w_i sigmoid(phi_i - a + log rho) = rho."""
    phi = np.asarray(phi, dtype=float)
    if w is None:
        w = np.full(phi.shape, 1.0 / phi.size)
    lr = np.log(rho)

    def g(a):
        return float(np.sum(w * expit(phi - a + lr))) - rho

    lo, hi = phi.min() - 60.0, phi.max() + 60.0
    while g(lo) < 0:  # g decreasing in a
        lo -= 60.0
    while g(hi) > 0:
        hi += 60.0
    return brentq(g, lo, hi, xtol=1e-14, rtol=1e-15, maxiter=500)


def F_rho(phi, rho, w=None):
    """F_rho(phi; mu) = inf_a a - 1 + (1/rho) int log(1 + rho e^{phi-a}) dmu   (Eq. 6)."""
    phi = np.asarray(phi, dtype=float)
    if w is None:
        w = np.full(phi.shape, 1.0 / phi.size)
    a = alpha_star(phi, rho, w)
    val = a - 1.0 + np.sum(w * np.logaddexp(0.0, phi - a + np.log(rho))) / rho
    return float(val), float(a)


def F_rho_bruteforce(phi, rho, w=None, grid=200001, span=80.0):
    """F_rho by direct 1-D grid+refine minimisation over alpha (independent of brentq)."""
    phi = np.asarray(phi, dtype=float)
    if w is None:
        w = np.full(phi.shape, 1.0 / phi.size)
    lr = np.log(rho)

    def obj(a):
        return a - 1.0 + np.sum(w * np.logaddexp(0.0, phi - a + lr)) / rho

    lo, hi = phi.min() - span, phi.max() + span
    for _ in range(4):
        xs = np.linspace(lo, hi, grid if _ == 0 else 2001)
        vals = np.array([obj(x) for x in xs])
        k = int(np.argmin(vals))
        lo, hi = xs[max(k - 1, 0)], xs[min(k + 1, len(xs) - 1)]
    return float(obj(0.5 * (lo + hi)))


def cvar(phi, rho, w=None):
    """CVaR_rho via the Rockafellar-Uryasev variational form (Eq. 9), exact for
    discrete mu: inf_a a + (1/rho) E (phi - a)_+ .  Solved exactly by sorting."""
    phi = np.asarray(phi, dtype=float)
    if w is None:
        w = np.full(phi.shape, 1.0 / phi.size)
    order = np.argsort(phi)
    p, ww = phi[order], w[order]
    # objective is piecewise linear in a with kinks at the atoms; evaluate at kinks
    best = np.inf
    for a in p:
        val = a + np.sum(ww * np.maximum(p - a, 0.0)) / rho
        best = min(best, val)
    return float(best)