scripts/core.py

"""
Core independent re-implementation of the definitions and constructions of

  "The Interplay Between Interpolation and Aggregation in Regression:
   Optimal Sample Complexity"  (arXiv 2605.29819, OpenReview qXlovWytwg)

Everything here is written from the *statements* of Definitions 3.1, 3.2, 3.3,
3.4, 3.6, 3.7, 3.9 and the *constructions* sketched in the proof overviews of
Theorems 3.5, 3.8, 3.10, 3.11, 3.12.  No code from the authors was used
(none is released).

CPU only.  numpy + itertools.
"""

from __future__ import annotations

import itertools
import math
from dataclasses import dataclass
from typing import Callable, Iterable, Sequence

import numpy as np

# ----------------------------------------------------------------------------
# Finite hypothesis classes: a class is an (|H| x |X|) matrix of values in [0,1]
# ----------------------------------------------------------------------------


@dataclass
class FiniteClass:
    """H subset of [0,1]^X for a finite domain X = {0,...,m-1}."""

    values: np.ndarray  # shape (|H|, |X|), entries in [0,1]
    name: str = ""

    @property
    def n_hyp(self) -> int:
        return self.values.shape[0]

    @property
    def n_pts(self) -> int:
        return self.values.shape[1]


# ----------------------------------------------------------------------------
# Definition 3.1 (gamma-graph shattering) and Definition 3.2 (gamma-graph dim)
# ----------------------------------------------------------------------------


def _relation_matrix(values: np.ndarray, witness: int, gamma: float) -> np.ndarray:
    """R[h, j] = 0 if h(x_j) == witness(x_j)
                 1 if |h(x_j) - witness(x_j)| > gamma
                 2 otherwise (h unusable at x_j for shattering with this witness)

    This is a *verbatim* transcription of Definition 3.1's two bullet points.
    """
    w = values[witness]
    diff = np.abs(values - w[None, :])
    R = np.full(values.shape, 2, dtype=np.int8)
    R[diff == 0.0] = 0
    R[diff > gamma] = 1
    return R


def is_gamma_graph_shattered(
    cls: FiniteClass, cols: Sequence[int], gamma: float, witness: int
) -> bool:
    """Definition 3.1: is the sequence (x_j)_{j in cols} gamma-graph shattered
    by H with witness h_{witness}?"""
    R = _relation_matrix(cls.values, witness, gamma)[:, list(cols)]
    d = len(cols)
    usable = ~(R == 2).any(axis=1)
    if not usable.any():
        return False
    pats = R[usable].astype(np.int64)
    weights = (1 << np.arange(d)).astype(np.int64)
    codes = pats @ weights
    return len(np.unique(codes)) == (1 << d)


def gamma_graph_dim(
    cls: FiniteClass,
    gamma: float,
    max_d: int | None = None,
    return_witness: bool = False,
    witnesses=None,
):
    """Definition 3.2, computed EXHAUSTIVELY.

    gamma-graph shattering is downward closed in the point set (drop a
    coordinate from every b), so we enumerate level-wise: only supersets of an
    already-shattered set can be shattered.  For each witness the shattered
    families are represented with bitmasks, so a single shattering test is one
    vectorised numpy pass over H.
    """
    m = cls.n_pts
    cap = m if max_d is None else min(m, max_d)
    wit_iter = range(cls.n_hyp) if witnesses is None else list(witnesses)
    best, best_cert = 0, None

    for w in wit_iter:
        R = _relation_matrix(cls.values, w, gamma)
        bits = (1 << np.arange(m)).astype(np.int64)
        vmask = ((R != 2).astype(np.int64) * bits).sum(axis=1)   # usable columns
        fmask = ((R == 1).astype(np.int64) * bits).sum(axis=1)   # 'gamma-far' columns

        def shattered(cmask: int, d: int) -> bool:
            sel = (vmask & cmask) == cmask
            if not sel.any():
                return False
            codes = np.unique(fmask[sel] & cmask)
            return len(codes) == (1 << d)

        level = [(int(bits[j]), (j,)) for j in range(m) if shattered(int(bits[j]), 1)]
        if level and best < 1:
            best, best_cert = 1, (w, list(level[0][1]))
        d = 1
        while level and d < cap:
            nxt = []
            for cmask, cols in level:
                for j in range(cols[-1] + 1, m):
                    nm = cmask | int(bits[j])
                    if shattered(nm, d + 1):
                        nxt.append((nm, cols + (j,)))
            if not nxt:
                break
            d += 1
            level = nxt
            if d > best:
                best, best_cert = d, (w, list(level[0][1]))
        if best >= cap:
            break

    if return_witness:
        return best, best_cert
    return best


# ----------------------------------------------------------------------------
# Definition 3.9 (gamma-OIG dimension) -- one-inclusion (hyper)graph machinery
# ----------------------------------------------------------------------------


def one_inclusion_graph(cls: FiniteClass, cols: Sequence[int]):
    """Vertices V = H|_S (distinct restrictions).  Hyperedge (f, i) collects all
    v in V agreeing with f on every coordinate except i.  Returns
    (vertices ndarray, list of (direction i, tuple of vertex indices))."""
    sub = cls.values[:, list(cols)]
    verts = np.unique(sub, axis=0)
    n = len(cols)
    edges = []
    for i in range(n):
        keys = {}
        for vi, v in enumerate(verts):
            key = tuple(np.delete(v, i))
            keys.setdefault(key, []).append(vi)
        for _, group in keys.items():
            edges.append((i, tuple(group)))
    return verts, edges


def min_max_outdegree(
    verts: np.ndarray, edges, gamma: float, brute_cap: int = 2_000_000
):
    """min over orientations sigma of  max_v |{i : |sigma(e_{v,i})_i - v_i| > gamma}|.

    Singleton edges are self-oriented and contribute 0.  For the remaining
    (multi-vertex) edges we brute force all orientations when the search space
    is small, otherwise we return the value achieved by the paper's explicit
    'orient towards the smallest value in the direction of the edge' rule and
    flag it as an upper bound."""
    multi = [(i, g) for (i, g) in edges if len(g) > 1]
    nv = len(verts)
    if not multi:
        return 0, True

    space = 1
    for _, g in multi:
        space *= len(g)
        if space > brute_cap:
            break

    def outdeg_for(choice):
        deg = np.zeros(nv, dtype=int)
        for (i, g), tgt in zip(multi, choice):
            tv = verts[tgt][i]
            for v in g:
                if abs(tv - verts[v][i]) > gamma:
                    deg[v] += 1
        return int(deg.max()) if nv else 0

    if space <= brute_cap:
        best = None
        for choice in itertools.product(*[g for _, g in multi]):
            val = outdeg_for(choice)
            if best is None or val < best:
                best = val
                if best == 0:
                    break
        return best, True

    # paper's rule: orient each edge to the vertex with the smallest value in
    # the direction of the edge
    choice = [min(g, key=lambda v: verts[v][i]) for (i, g) in multi]
    return outdeg_for(choice), False


def gamma_oig_dim(cls: FiniteClass, gamma: float, max_k: int = 5, exact: bool = True):
    """Definition 3.9.  Returns (dimension estimate, per-k detail).

    k counts iff there EXISTS a point set S of size k and a finite subgraph for
    which EVERY orientation has max out-degree strictly greater than k/3.
    We evaluate the full induced one-inclusion graph and (for small instances)
    every vertex-induced subgraph, which is the family of subgraphs the paper's
    argument ranges over."""
    detail = {}
    dim = 0
    m = cls.n_pts
    for k in range(1, min(max_k, m) + 1):
        worst = 0
        exact_all = True
        for cols in itertools.combinations(range(m), k):
            verts, edges = one_inclusion_graph(cls, cols)
            val, is_exact = min_max_outdegree(verts, edges, gamma)
            exact_all = exact_all and is_exact
            worst = max(worst, val)
            if exact and len(verts) <= 9:
                # also scan vertex-induced subgraphs
                idx = list(range(len(verts)))
                for r in range(2, len(verts) + 1):
                    for keep in itertools.combinations(idx, r):
                        sub = FiniteClass(verts[list(keep)][:, :], "sub")
                        v2, e2 = one_inclusion_graph(sub, range(k))
                        val2, ex2 = min_max_outdegree(v2, e2, gamma)
                        exact_all = exact_all and ex2
                        worst = max(worst, val2)
        counts = worst > k / 3.0
        detail[k] = {
            "min_max_outdegree": worst,
            "threshold_k_over_3": k / 3.0,
            "counts": bool(counts),
            "exact": bool(exact_all),
        }
        if counts:
            dim = k
    return dim, detail


# ----------------------------------------------------------------------------
# Aggregation rules  (Definitions 3.4 and 3.7)
# ----------------------------------------------------------------------------


def rule_median(z: np.ndarray) -> float:
    return float(np.median(z))


def rule_min(z: np.ndarray) -> float:
    return float(np.min(z))


def rule_max(z: np.ndarray) -> float:
    return float(np.max(z))


def rule_mean(z: np.ndarray) -> float:
    return float(np.mean(z))


def rule_midrange(z: np.ndarray) -> float:
    return 0.5 * (float(np.min(z)) + float(np.max(z)))


def make_order_statistic(q: float) -> Callable[[np.ndarray], float]:
    def r(z: np.ndarray) -> float:
        s = np.sort(z)
        return float(s[min(len(s) - 1, int(q * len(s)))])

    return r


def is_proper_rule(rule, rng, trials=2000, m_max=7, odd_only=False) -> bool:
    """Definition 3.4 checked empirically.  `odd_only` restricts the arity to odd
    m, which is the regime in which the paper calls the median proper."""
    for _ in range(trials):
        m = int(rng.integers(1, m_max + 1))
        if odd_only and m % 2 == 0:
            m += 1
        z = rng.random(m)
        out = rule(z)
        if not np.any(np.isclose(out, z)):
            return False
    return True


def is_interpolating_rule(rule, rng, trials=2000, m_max=7) -> bool:
    """Definition 3.7 checked empirically."""
    for _ in range(trials):
        m = int(rng.integers(1, m_max + 1))
        z = rng.random(m)
        out = rule(z)
        if out < z.min() - 1e-12 or out > z.max() + 1e-12:
            return False
    return True


# ----------------------------------------------------------------------------
# Cutoff loss (Eq. 1)
# ----------------------------------------------------------------------------


def cutoff_loss(
    pred: np.ndarray, labels: np.ndarray, masses: np.ndarray, gamma: float
) -> float:
    """L^gamma_D(h) = P_{(x,y)~D}[|h(x)-y| > gamma] for a finitely supported D."""
    return float(masses[np.abs(pred - labels) > gamma].sum())



def _unique_values(n: int, gamma: float) -> np.ndarray:
    """n distinct values gamma_A in (gamma, 1], all within a window of width
    < gamma of each other.

    The paper only requires gamma_{A} in (gamma, 1] and *unique*.  Keeping the
    window narrower than gamma additionally guarantees that two distinct unique
    values are never gamma-far from each other, which is the regime the proofs
    of Theorems 3.8 / 3.10 / 3.12 implicitly work in (there, agreement at a
    unique value is the ONLY way another hypothesis can match the witness off
    its zero set).  Widening the window is exercised as a sensitivity probe.
    """
    lo = gamma + 0.05 * gamma
    hi = min(1.0, lo + 0.5 * gamma)
    if n == 1:
        return np.array([lo])
    return np.linspace(lo, hi, n)


# ----------------------------------------------------------------------------
# Construction A -- Theorem 3.5 hard instance
#   shattered set x_1..x_d with witness h; mass 1 - a*eps on x_1,
#   a*eps/(d-1) on each of x_2..x_d.  Worst-case interpolator returns the
#   hypothesis that is gamma-far from the witness on every unobserved point.
# ----------------------------------------------------------------------------


def thm35_class(d: int, gamma: float) -> FiniteClass:
    """Smallest class realising gamma-graph dimension exactly d: the witness is
    all-zero on d points and for each b in {0,1}^d there is h_b that is 0 where
    b_i = 0 and (gamma+delta) where b_i = 1."""
    hi = min(1.0, gamma + (1.0 - gamma) / 2.0)
    rows = []
    for b in itertools.product([0, 1], repeat=d):
        rows.append([hi if bi else 0.0 for bi in b])
    return FiniteClass(np.array(rows, dtype=float), f"thm35_d{d}")


def thm35_masses(d: int, eps: float, a: float = 2.0) -> np.ndarray:
    """1 - a*eps on the heavy point, a*eps/(d-1) on each of the other d-1."""
    p = np.empty(d)
    p[0] = 1.0 - a * eps
    p[1:] = a * eps / (d - 1)
    return p


def thm35_expected_loss_exact(d: int, eps: float, n: int, a: float = 2.0) -> float:
    """Exact E_S[L^gamma_D(A'(S))] lower bound for ANY interpolator-based
    aggregation with a proper rule: the predictor is gamma-far from the witness
    on every point missing from S, so the loss is at least the mass of the
    missing light points.  (The heavy point is essentially always observed; we
    include it exactly anyway.)"""
    p = thm35_masses(d, eps, a)
    return float(np.sum(p * (1.0 - p) ** n))


# ----------------------------------------------------------------------------
# Construction B -- Theorem 3.8 "Cantor-like" class
#   H = {h_A : A subset of [ku], |A| = d},  h_A = 0 on A, unique gamma_A > gamma
#   elsewhere.
# ----------------------------------------------------------------------------


def thm38_class(
    ku: int, d: int, gamma: float
) -> tuple[FiniteClass, list[tuple[int, ...]]]:
    subsets = list(itertools.combinations(range(ku), d))
    n_sub = len(subsets)
    rows = np.zeros((n_sub, ku))
    # unique value gamma_A in (gamma, 1] for each A
    uniq = _unique_values(n_sub, gamma)
    for i, A in enumerate(subsets):
        rows[i, :] = uniq[i]
        for x in A:
            rows[i, x] = 0.0
    return FiniteClass(rows, f"thm38_ku{ku}_d{d}"), subsets


# ----------------------------------------------------------------------------
# Construction C -- Theorem 3.10 split-space class (one block k = i^2)
#   H_k = {h_{k,A} : A subset [k], |A| = sqrt(k)}, 0 on A, unique value elsewhere
# ----------------------------------------------------------------------------


def thm310_block_class(
    k: int, gamma: float
) -> tuple[FiniteClass, list[tuple[int, ...]]]:
    s = int(round(math.isqrt(k)))
    assert s * s == k, "k must be a perfect square"
    subsets = list(itertools.combinations(range(k), s))
    rows = np.zeros((len(subsets), k))
    uniq = _unique_values(len(subsets), gamma)
    for i, A in enumerate(subsets):
        rows[i, :] = uniq[i]
        for x in A:
            rows[i, x] = 0.0
    return FiniteClass(rows, f"thm310_k{k}"), subsets


# ----------------------------------------------------------------------------
# Construction D -- Theorem 3.12 class
#   block k: H_k = {h_{k,A}: A subset [k], |A| = d-1}, h_{k,A} = 0 on [k]\A,
#   unique value gamma_{k,A} > gamma on A (and on every other block).
# ----------------------------------------------------------------------------


def thm312_block_class(
    k: int, d: int, gamma: float
) -> tuple[FiniteClass, list[tuple[int, ...]]]:
    subsets = list(itertools.combinations(range(k), d - 1))
    rows = np.zeros((len(subsets), k))
    uniq = _unique_values(len(subsets), gamma)
    for i, A in enumerate(subsets):
        rows[i, :] = 0.0
        for x in A:
            rows[i, x] = uniq[i]
    return FiniteClass(rows, f"thm312_k{k}_d{d}"), subsets


def thm312_class(ks, d: int, gamma: float):
    """The FULL Theorem 3.12 class over a split input space X = union_k X_k.

    h_{k,A} (A subset of [k], |A| = d-1) is 0 on {(k,x) : x in [k]\\A}, equal to
    the unique value gamma_{k,A} in (gamma,1] on {(k,x) : x in A}, and equal to
    that same unique value on EVERY point of every other block.  The extra
    blocks push the gamma-graph dimension from d-1 up to exactly d."""
    ks = list(ks)
    offsets, tot = {}, 0
    for k in ks:
        offsets[k] = tot
        tot += k
    specs = []
    for k in ks:
        for A in itertools.combinations(range(k), d - 1):
            specs.append((k, A))
    rows = np.zeros((len(specs), tot))
    uniq = _unique_values(len(specs), gamma)
    for i, (k, A) in enumerate(specs):
        rows[i, :] = uniq[i]  # unique value everywhere ...
        off = offsets[k]
        for x in range(k):  # ... except zeros on [k]\A inside block k
            if x not in A:
                rows[i, off + x] = 0.0
    return FiniteClass(rows, f"thm312_full_d{d}"), specs, offsets


def thm310_class(ks, gamma: float):
    """The FULL Theorem 3.10 class over the split space X = union_{k=i^2} X_k:
    h_{k,A} (A subset of [k], |A| = sqrt(k)) is 0 on {(k,x): x in A} and equal
    to a unique value gamma_{k,A} in (gamma,1] on every other point of X."""
    ks = list(ks)
    offsets, tot = {}, 0
    for k in ks:
        offsets[k] = tot
        tot += k
    specs = []
    for k in ks:
        s = math.isqrt(k)
        assert s * s == k
        for A in itertools.combinations(range(k), s):
            specs.append((k, A))
    rows = np.zeros((len(specs), tot))
    uniq = _unique_values(len(specs), gamma)
    for i, (k, A) in enumerate(specs):
        rows[i, :] = uniq[i]
        off = offsets[k]
        for x in A:
            rows[i, off + x] = 0.0
    return FiniteClass(rows, "thm310_full"), specs, offsets


# ----------------------------------------------------------------------------
# helpers
# ----------------------------------------------------------------------------


def fit_loglog_slope(xs: Iterable[float], ys: Iterable[float]) -> tuple[float, float]:
    x = np.log(np.asarray(list(xs), dtype=float))
    y = np.log(np.asarray(list(ys), dtype=float))
    A = np.vstack([x, np.ones_like(x)]).T
    coef, *_ = np.linalg.lstsq(A, y, rcond=None)
    resid = y - A @ coef
    ss_res = float((resid**2).sum())
    ss_tot = float(((y - y.mean()) ** 2).sum())
    r2 = 1.0 - ss_res / ss_tot if ss_tot > 0 else float("nan")
    return float(coef[0]), r2


def dump_json(path, obj):
    import json

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

    with open(path, "w") as f:
        json.dump(obj, f, indent=2, default=default)
    return path