exp014_gd/read_codebook.py

"""read_codebook.py — projective reading of cultivated aleph codebooks.
Antipodal-collapse extraction on trained codebooks + projective statistics on
RP^(D-1) — applied to the exp012 AR-bed specimens.

Recipe per the Polygonal Omega article (geometric-tri-band-ft2): collapse = (row_i - row_j)/2 normalized for each MUTUAL-STRONGEST
pair with cos < -0.9 — "a deterministic tensor operation," not clustering.
Projective metric ALWAYS arccos|<a,b>| (metric-alignment rule, reading-voids-ft1).
D=4 scope is the validated regime (D=5 walked back; axis count grows with D).

Readouts per specimen:
  pairs / n_axes / unpaired          — antipodal structure
  proj_angle mean vs uniform baseline, deviation  — near-uniform RP^(D-1)?
  drift from home + binding fraction @0.29154     — cultivation record
  erank of the axis set                            — spectral occupancy
  verdict: PROJECTIVE-CLEAN (|dev|<0.05, util>0.95, secondary pairs<=3) /
           -MOSTLY / STRUCTURED / DEGENERATE (per Polygonal Omega thresholds)

Usage (terminal): python read_codebook.py <ckpt_or_dir> [more paths...]
Colab: paste geolip_vitals.py cell first (optional), then this file, then
       read_all(r"/content/data/ar_ckpts").
"""
from __future__ import annotations
import math
import sys
import torch
import torch.nn.functional as F

BINDING = 0.29154


@torch.no_grad()
def antipodal_collapse(codebook: torch.Tensor, thresh: float = -0.9) -> dict:
    """Mutual-strongest antipodal pairing + collapse to axes on RP^(D-1)."""
    A = F.normalize(codebook.float(), dim=-1)
    K = A.shape[0]
    cos = A @ A.T
    cos.fill_diagonal_(2.0)                      # exclude self from minima
    nearest_neg = cos.argmin(dim=-1)             # most-antipodal partner
    pairs = []
    used = set()
    for i in range(K):
        j = int(nearest_neg[i])
        if i < j and int(nearest_neg[j]) == i and cos[i, j] < thresh:
            pairs.append((i, j))
            used.update((i, j))
    axes = [F.normalize((A[i] - A[j]) / 2.0, dim=-1) for i, j in pairs]
    axes += [A[i] for i in range(K) if i not in used]   # unpaired rows as axes
    axes = torch.stack(axes) if axes else A[:0]
    # sign-canon onto RP: first nonzero coordinate positive
    for r in range(axes.shape[0]):
        nz = torch.nonzero(axes[r].abs() > 1e-8)
        if nz.numel() and axes[r, nz[0, 0]] < 0:
            axes[r] = -axes[r]
    return {"pairs": len(pairs), "n_axes": axes.shape[0],
            "unpaired": K - 2 * len(pairs), "axes": axes}


@torch.no_grad()
def projective_stats(axes: torch.Tensor, n_baseline: int = 20000,
                     seed: int = 0) -> dict:
    """Mean projective angle arccos|<a,b>| vs a uniform-RP baseline at same (n, D)."""
    n, D = axes.shape
    if n < 2:
        return {"proj_angle_mean": None, "uniform_baseline": None,
                "deviation": None, "erank": None}
    def mean_angle(rows):
        c = (rows @ rows.T).abs().clamp(max=1.0)
        iu = torch.triu_indices(rows.shape[0], rows.shape[0], offset=1)
        return torch.arccos(c[iu[0], iu[1]]).mean().item()
    obs = mean_angle(axes)
    g = torch.Generator().manual_seed(seed)
    base_angles = []
    m = max(2, n)
    for _ in range(max(1, n_baseline // max(1, m * (m - 1) // 2))):
        r = F.normalize(torch.randn(m, D, generator=g), dim=-1)
        base_angles.append(mean_angle(r))
    base = sum(base_angles) / len(base_angles)
    s = torch.linalg.svdvals(axes)
    p = (s / s.sum().clamp_min(1e-12))
    erank = float(torch.exp(-(p.clamp_min(1e-12) * p.clamp_min(1e-12).log()).sum()))
    return {"proj_angle_mean": round(obs, 4), "uniform_baseline": round(base, 4),
            "deviation": round(obs - base, 4), "erank": round(erank, 3)}


@torch.no_grad()
def read_specimen(path: str) -> dict:
    ck = torch.load(path, map_location="cpu", weights_only=True)
    out = {"file": path.split("\\")[-1].split("/")[-1],
           "arm": ck.get("arm"), "seed": ck.get("seed"),
           "steps": ck.get("steps"), "val_bpb": round(ck.get("val_bpb", -1), 4)}
    if "state_dict" in ck:            # full specimen checkpoint
        sd = ck["state_dict"]
        books = {k[:-len(".codebook")]: sd[k] for k in sd
                 if k.endswith("addr.codebook") or k.endswith("head_addr.codebook")}
        homes = {k[:-len(".home")]: sd[k] for k in sd if k.endswith(".home")}
    else:                             # bare genome dict (exp014+ champion files):
        # books under flat/root/branch* keys; *_proj entries are projections
        books = {k: v for k, v in ck.items()
                 if torch.is_tensor(v) and v.ndim == 2
                 and (k in ("flat", "root") or k.startswith("branch"))}
        homes = {}
    reads = {}
    for name, cb in books.items():
        col = antipodal_collapse(cb)
        stats = projective_stats(col["axes"])
        home = homes.get(name)
        drift = None
        binding = None
        if home is not None and home.shape == cb.shape:
            a = F.normalize(cb.float(), dim=-1)
            b = F.normalize(home.float(), dim=-1)
            dr = torch.arccos((a * b).sum(-1).clamp(-1, 1))
            drift = round(dr.mean().item(), 4)
            binding = round(((dr - BINDING).abs() <= 0.05).float().mean().item(), 4)
        util = col["n_axes"] / cb.shape[0]
        dev = stats["deviation"]
        if dev is not None and abs(dev) < 0.05 and util > 0.95 and col["pairs"] <= 3:
            verdict = "PROJECTIVE-CLEAN"
        elif dev is not None and abs(dev) < 0.05:
            verdict = "PROJECTIVE-MOSTLY"
        elif dev is not None and dev > 0.05:
            verdict = "STRUCTURED(repulsive)"
        else:
            verdict = "DEGENERATE/CLUMPED" if dev is not None else "TOO-FEW-AXES"
        reads[name] = {
            "pairs": col["pairs"], "n_axes": col["n_axes"], **stats,
            "drift": drift, "binding_frac": binding, "verdict": verdict}
    out["codebooks"] = reads
    return out


def read_all(root: str) -> list:
    import glob, os
    results = []
    for p in sorted(glob.glob(os.path.join(root, "*.pt"))):
        r = read_specimen(p)
        print(r, flush=True)
        results.append(r)
    return results


def _in_notebook() -> bool:
    try:
        get_ipython()  # type: ignore[name-defined]  # noqa: F821
        return True
    except NameError:
        return False


if __name__ == "__main__":
    if _in_notebook():
        print("Notebook mode: call read_all(r'<data_root>/ar_ckpts') in the next cell.")
    else:
        args = [a for a in sys.argv[1:] if not a.startswith("-")]
        if not args:
            print("usage: python read_codebook.py <ckpt_or_dir> [...]")
        for a in args:
            import os
            read_all(a) if os.path.isdir(a) else print(read_specimen(a))