code/wavecurriculum_run/interpret_kl_spread.py

"""Turn the printed val_acc / kl* / spread* into a statement about uniformity.

kl* and spread* both measure non-uniformity inside the candidate set but they
weight it very differently, so the pair of them pins down more than either
alone. spread* is bounded in [0,1] and barely moves when one candidate is
crushed; kl* = KL(Uniform(S) || q) diverges as any candidate goes to zero.
A single skew cannot produce a high spread* and a high kl* at once, so the
observed pair implies a minority of cells with a near-zeroed candidate.
"""
import numpy as np
from scipy.optimize import brentq

OUT = 0.0039      # val_acc == out* == 1 - mass*
KL = 0.5138       # kl*
SPREAD = 0.9156   # spread*


def skew_two(x, k):
    """q with one candidate at x, the rest sharing (1-x) uniformly."""
    q = np.full(k, (1.0 - x) / (k - 1))
    q[0] = x
    return q


def kl_u(q):
    k = len(q)
    return float(np.sum((1.0 / k) * np.log((1.0 / k) / np.maximum(q, 1e-300))))


def spread_of(q):
    k = len(q)
    h = float(-np.sum(q * np.log(np.maximum(q, 1e-300))))
    return h / np.log(k)


print(f"current eval: val_acc(out*) {OUT}   kl* {KL}   spread* {SPREAD}")
print(f"\nidentity check: excess* = log(1/mass*) + kl*")
leak = np.log(1.0 / (1.0 - OUT))
print(f"  leak = log(1/{1 - OUT:.4f}) = {leak:.4f}")
print(f"  leak + kl* = {leak + KL:.4f}   (log prints excess* ~ 0.518)")
print(f"  -> {leak / (leak + KL) * 100:.1f}% of excess is leakage, "
      f"{KL / (leak + KL) * 100:.1f}% is non-uniformity inside S")

print(f"\nreading 1: if EVERY cell had the same skew, kl*={KL} means")
for k in (2, 3, 4):
    lo, hi = 1e-12, 1.0 / k - 1e-12
    x = brentq(lambda z: kl_u(skew_two(z, k)) - KL, lo, hi)
    q = skew_two(x, k)
    print(f"  |S|={k}: one candidate at {x:.4f}, others {q[1]:.3f} each "
          f"-> spread would be {spread_of(q):.3f}")
print(f"  but measured spread* is {SPREAD}, far higher. So the cells are NOT")
print("  all uniformly skewed: most are near-flat and a few are crushed.")

print(f"\nreading 2: two populations. fraction f of cells have one candidate")
print(f"  pushed to eps, the rest ({1:.0%}-f) are exactly uniform.")
print(f"  {'|S|':>4}  {'f (crushed cells)':>18}  {'eps':>10}  "
      f"{'that candidate gets':>20}")
for k in (2, 3, 4):
    def resid(le):
        eps = np.exp(le)
        q = skew_two(eps, k)
        sB, kB = spread_of(q), kl_u(q)
        f = (1.0 - SPREAD) / max(1.0 - sB, 1e-12)
        return f * kB - KL
    le = brentq(resid, np.log(1e-14), np.log(1.0 / k - 1e-9))
    eps = np.exp(le)
    q = skew_two(eps, k)
    f = (1.0 - SPREAD) / (1.0 - spread_of(q))
    print(f"  {k:>4}  {f * 100:>17.1f}%  {eps:>10.2e}  "
          f"{1 / eps:>17.0f}x less than uniform")

print(f"\nreading 3: effective number of candidates actually used")
for k, share in ((2, 0.42), (3, 0.29), (4, 0.16), (5, 0.13)):
    print(f"  |S|={k} (~{share:.0%} of multi cells): model spreads over "
          f"{k ** SPREAD:.2f} of {k} candidates")