scripts/setup_problem.py

"""Shared logistic-regression testbed used by claims 1, 2, 3, 6.

Logistic regression satisfies Assumptions (A)-(C) of the paper with *closed-form*
constants: convex (A); L_n-smooth with L_n=||x_n||^2/4 and third-derivative bound
M_n = ||x_n||^3 / (6 sqrt 3) since sup_z |sigma''(z)| = 1/(6 sqrt 3) (B);
mu-strongly convex with mu = lambda_min(Gamma)/N thanks to the Gaussian prior (C).
It is therefore the cleanest model in which every constant of Theorem 4.1 /
Corollary 4.6 can actually be evaluated (for Poisson regression M_n = infinity
unless the parameter space is bounded, as the paper notes).
"""

import numpy as np
import sys

sys.path.insert(0, "/home/ubuntu/samuel/sgmcmc-uq-repro/scripts")
from common import Logistic, NoiseModel


def make(
    N=500,
    D=3,
    seed=20260725,
    gamma=2.0,
    design="gaussian",
    df=3.0,
    rank_deficient=False,
):
    rng = np.random.default_rng(seed)
    if design == "gaussian":
        X = rng.standard_normal((N, D))
    elif design == "heavy":  # Student-t covariates: inflates M, tau4
        X = rng.standard_t(df, size=(N, D)) / np.sqrt(df / (df - 2))
    elif design == "binary":
        X = rng.integers(0, 2, size=(N, D)).astype(float) * 2 - 1
    else:
        raise ValueError(design)
    if rank_deficient:
        X[:, -1] = X[:, 0]  # exactly collinear -> mu_hat = 0 if gamma = 0
    th_star = rng.standard_normal(D)
    y = (rng.random(N) < 1.0 / (1.0 + np.exp(-X @ th_star))).astype(float)
    Gam = gamma * np.eye(D)
    m = Logistic(X, y, Gam)
    that = m.map_estimate()
    nm = NoiseModel(m.grad_n(that), m.hess_n(that), that, Gam, N)
    return m, that, nm, m.constants(that)