scripts/common.py

"""Shared utilities for the reproduction of arXiv:2601.20180.

Everything is CPU-only, deterministic given the seed, and pure numpy/sympy.
"""

import json
import os
import numpy as np

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
OUT = os.path.join(ROOT, "outputs")
FIGS = os.path.join(ROOT, "figs")
os.makedirs(OUT, exist_ok=True)
os.makedirs(FIGS, exist_ok=True)


def dump(name, obj):
    """Write a JSON artifact (json.dump only -- never the Write tool)."""
    path = os.path.join(OUT, name)
    with open(path, "w") as f:
        json.dump(obj, f, indent=2, sort_keys=True, default=_default)
    print("wrote", path)
    return path


def _default(o):
    if isinstance(o, (np.floating,)):
        return float(o)
    if isinstance(o, (np.integer,)):
        return int(o)
    if isinstance(o, np.ndarray):
        return o.tolist()
    raise TypeError(type(o))


# ---------------------------------------------------------------- box helpers
def box_vi_gap(xstar, v, lo=0.0, hi=1.0):
    """max_{x in [lo,hi]^d} <xstar - x, v>  (i.e. the VI / stability gap).

    x* is an eps-approximate VI solution  <=>  box_vi_gap(x*, F(x*)) <= eps.
    Closed form on a box, exact up to floating point.
    """
    xstar = np.asarray(xstar, dtype=float)
    v = np.asarray(v, dtype=float)
    return float(np.sum(np.maximum((xstar - lo) * v, (xstar - hi) * v)))


def box_proj(x, lo=0.0, hi=1.0):
    return np.clip(x, lo, hi)


def random_affine_vi(d, rng):
    """A in R^{d x d} with ||A||_1 <= 1 and ||A||_inf <= 1, plus b."""
    A = rng.normal(size=(d, d))
    s = max(np.abs(A).sum(axis=0).max(), np.abs(A).sum(axis=1).max())
    A = A / s
    b = rng.uniform(-1.0, 1.0, size=d)
    return A, b


def norm1(A):
    return float(np.abs(A).sum(axis=0).max())


def norminf(A):
    return float(np.abs(A).sum(axis=1).max())