repro/src/verify_caffnet.py

#!/usr/bin/env python3
"""CAffNet: Hard Constraint-Affine Neural Networks (arXiv:2605.24437). Reproduces Theorem 3.5 / 3.4:
  [0a] Hard constraint satisfaction: the CAffNet parameterization P(x) = f(x) - A^+(A f(x) - b(x)) +
       (I - A^+ A) w(x) makes the output ALWAYS satisfy the affine constraint A(x) y = b(x) EXACTLY
       (0 violation), for ANY base network f and null-space network w.
  [0b] Universal approximation (Thm 3.5): if the unconstrained base class approximates a target, then
       CAffNet approximates the constraint-satisfying target -- approximation error -> 0 as base width
       grows, while the constraint stays satisfied exactly.
Pure numpy MLP; equality + inequality affine constraints. Deterministic seeds.
"""
import numpy as np, json, hashlib

def mlp_init(sizes, rng):
    return [(rng.standard_normal((a, b)) * np.sqrt(2.0 / a), np.zeros(b)) for a, b in zip(sizes[:-1], sizes[1:])]
def mlp_fwd(W, X):
    a = X
    for i, (w, b) in enumerate(W):
        a = a @ w + b
        if i < len(W) - 1: a = np.maximum(a, 0.0)
    return a

def main():
    R = {"claim": "CAffNet_hard_constraint_and_universal_approx", "paper": "arXiv:2605.24437"}
    rng = np.random.default_rng(0)
    N = 400; X = rng.uniform(0, 1, (N, 1))
    m = 3                                                 # 3-dim output

    # affine EQUALITY constraint A y = b(x): sum of outputs must equal g(x) = 1 + x
    A = np.ones((1, m)); Apinv = A.T @ np.linalg.inv(A @ A.T)      # Moore-Penrose pseudoinverse
    def bvec(x): return (1.0 + x)                                  # b(x) in R^1 (per sample)

    def caffnet(fbase, wnull, Xb):
        b = bvec(Xb)                                               # (N,1)
        # P = f - A^+(A f - b) + (I - A^+ A) w
        Af = fbase @ A.T                                           # (N,1)
        proj = fbase - (Af - b) @ Apinv.T                         # enforce equality
        nullw = wnull - (wnull @ A.T) @ Apinv.T                   # (I-A^+A) w
        return proj + nullw

    # ---------- [0a] hard constraint satisfaction for ARBITRARY (random, untrained) base+null nets ----------
    max_viol = 0.0
    for s in range(50):
        r = np.random.default_rng(100 + s)
        fb = mlp_fwd(mlp_init([1, 16, m], r), X); wn = mlp_fwd(mlp_init([1, 16, m], r), X)
        P = caffnet(fb, wn, X)
        viol = np.max(np.abs(P @ A.T - bvec(X)))                  # |A P - b|
        max_viol = max(max_viol, float(viol))
    R["max_constraint_violation_equality"] = float(f"{max_viol:.2e}")
    R["thm3.4_hard_equality_satisfied"] = max_viol < 1e-9

    # inequality constraint y_1 <= 3x (project only when violated) -> 0 violation
    def caffnet_ineq(fbase, Xb):
        cap = 3 * Xb[:, 0]                                        # y_1 <= 3x
        out = fbase.copy(); mask = out[:, 0] > cap
        out[mask, 0] = cap[mask]                                  # project onto feasible (closest)
        return out
    max_viol_ineq = 0.0
    for s in range(50):
        r = np.random.default_rng(200 + s)
        fb = mlp_fwd(mlp_init([1, 16, m], r), X)
        out = caffnet_ineq(fb, X)
        max_viol_ineq = max(max_viol_ineq, float(np.max(np.maximum(out[:, 0] - 3 * X[:, 0], 0))))
    R["max_constraint_violation_inequality"] = float(f"{max_viol_ineq:.2e}")
    R["thm3.4_hard_inequality_satisfied"] = max_viol_ineq < 1e-12

    # ---------- [0b] universal approximation: error -> 0 as base width grows (constraint held exactly) ----------
    # constrained target: y*(x) = softmax-like split of g(x)=1+x among 3 coords with a nonlinear pattern
    def target(x):
        raw = np.stack([np.sin(3 * x[:, 0]), np.cos(2 * x[:, 0]), 0.5 * x[:, 0]], 1)
        raw = np.exp(raw); raw = raw / raw.sum(1, keepdims=True); return raw * (1.0 + x)   # sums to g(x)=1+x
    Yt = target(X)
    # P = (I - A^+A) f + A^+ b   (single base net f; the projection is affine in f). Fit f so P ~ Yt.
    Proj = np.eye(m) - Apinv @ A                          # symmetric idempotent projector onto null(A)
    def fit_and_error(width, epochs=3000):
        r = np.random.default_rng(7)
        W = mlp_init([1, width, width, m], r); lr = 0.05
        Ab = bvec(X) @ Apinv.T                            # A^+ b term (fixed)
        for _ in range(epochs):
            a = X; acts = [X]
            for i, (w, b) in enumerate(W):
                a = a @ w + b
                if i < len(W) - 1: a = np.maximum(a, 0.0)
                acts.append(a)
            P = acts[-1] @ Proj + Ab                       # CAffNet output
            g = 2 * (P - Yt) / N
            delta = g @ Proj.T                             # grad through the projector
            for i in range(len(W) - 1, -1, -1):
                gw = acts[i].T @ delta + 1e-5 * W[i][0]; gb = delta.sum(0)
                if i > 0: delta = (delta @ W[i][0].T) * (acts[i] > 0)
                W[i] = (W[i][0] - lr * gw, W[i][1] - lr * gb)
        a = X
        for i, (w, b) in enumerate(W):
            a = a @ w + b
            if i < len(W) - 1: a = np.maximum(a, 0.0)
        P = a @ Proj + Ab
        err = float(np.sqrt(np.mean((P - Yt) ** 2)))
        viol = float(np.max(np.abs(P @ A.T - bvec(X))))
        return err, viol
    rows = []
    for width in (4, 8, 16, 32, 64):
        err, viol = fit_and_error(width)
        rows.append({"width": width, "approx_rmse": round(err, 4), "constraint_violation": float(f"{viol:.1e}")})
    R["approx_vs_width"] = rows
    R["thm3.5_error_decreases_with_width"] = rows[-1]["approx_rmse"] < rows[0]["approx_rmse"] - 1e-4
    R["constraint_held_during_training"] = all(r["constraint_violation"] < 1e-9 for r in rows)

    R["verdict"] = "supports" if (R["thm3.4_hard_equality_satisfied"] and R["thm3.4_hard_inequality_satisfied"]
                                  and R["thm3.5_error_decreases_with_width"] and R["constraint_held_during_training"]) else "inconclusive"

    print("claim: " + R["claim"])
    print(f"[0a] hard constraint satisfaction (arbitrary nets): equality |Ay-b|={R['max_constraint_violation_equality']} -> {R['thm3.4_hard_equality_satisfied']}; "
          f"inequality max viol={R['max_constraint_violation_inequality']} -> {R['thm3.4_hard_inequality_satisfied']}")
    print(f"[0b] universal approximation (Thm 3.5): RMSE to constrained target vs base width:")
    for r in rows: print(f"      width={r['width']:>2}: approx RMSE={r['approx_rmse']}, constraint violation={r['constraint_violation']}")
    print(f"    error decreases with width: {R['thm3.5_error_decreases_with_width']}; constraint held exactly throughout: {R['constraint_held_during_training']}")
    print(f"verdict: {R['verdict']}")

    def _np(o):
        if isinstance(o, np.bool_): return bool(o)
        if isinstance(o, np.integer): return int(o)
        if isinstance(o, np.floating): return float(o)
        raise TypeError
    import os; os.makedirs("outputs", exist_ok=True)
    open("outputs/caffnet_results.json", "w").write(json.dumps(R, indent=2, default=_np))
    print("RESULTS_SHA256=" + hashlib.sha256(json.dumps(R, sort_keys=True, default=_np).encode()).hexdigest())
    return 0 if R["verdict"] == "supports" else 1

if __name__ == "__main__":
    raise SystemExit(main())