code/geobasin_bed.py

"""geobasin_bed.py — the CE-replacement geometric arm, reconstructed and
properly tested. #TAG:geobasin #TAG:ce_replacement #TAG:loss_campaign

RECONSTRUCTION PROVENANCE (2026-07-25): the Oct '25 arm survives VERBATIM in
the installed package — geofractal.losses.pure_geometric (PureGeometricLoss:
attraction (1-s_y)^2 + repulsion sum(s_c^2) + margin hinge + range clamp;
GeometricPrototypeLoss; HierarchicalGeometricLoss) and the GBC 4-factor
compatibility head in geofractal/model/experiment_geometric_basin.py.
L-031's "term forms NOT recorded" and L-032's "unbuilt" are both CLOSED.
Losses are IMPORTED from the package, never rewritten (reuse law).

THE TEST (the ce_fixedcode lesson: isolate the LOSS by holding the head
identical): all arms share one cosine-anchor score head on the certified
addr_msl64 read — s = normalize(feats); scores_c = (cos(s, A_c)+1)/2,
A: (256 classes x 256), ~param-matched to the ce head (65,536 vs 65,792).
CE arms consume cos*10 as logits (fixed scale, disclosed).

ARMS: geo_ce_scores (CE on the SAME head - the decisive control) | geo_pure
(verbatim, learned anchors) | geo_pure_frozen (L-108 cell) | geo_pure_norep
(attraction+range ONLY - the absolute-only, doctrine-clean variant; repulsion
+margin are roster-comparative terms) | geo_proto (verbatim, own projector -
extra params disclosed) | geo_hier (nibble hierarchy 16x16 - bytes' natural
coarse structure) | geo_hybrid (0.5 CE + 0.5 PureGeometric).

PREREG (3 seeds; baselines ce 2.4769/acc .505): P1 the Oct'25 trade was -12%
relative accuracy - geo_pure within -12% of geo_ce_scores matches history,
parity overturns it, acc<0.30 refutes viability. P2 (L-108): learned anchors
show a collapse signature vs frozen; falsifier: learned > frozen by >2 pts.
P3 (absolute-beats-relative): norep >= pure - noise; falsifier: pure beats
norep by >2 pts (the comparative terms would be load-bearing - a scope
amendment to the law). P4: hybrid bpb within 0.15 of geo_ce_scores.

Run:  python tools/geobasin_bed.py --arm <name> --seed N  |  --list
"""
import json
import math
import os
import sys
import time
import zlib

import torch
import torch.nn as nn
import torch.nn.functional as F


def _root():
    d = os.path.abspath(os.getcwd())
    while True:
        if os.path.exists(os.path.join(d, "MANIFEST.md")):
            return d
        p = os.path.dirname(d)
        if p == d:
            return os.getcwd()
        d = p

ROOT = _root()
if os.path.join(ROOT, "tools") not in sys.path:
    sys.path.insert(0, os.path.join(ROOT, "tools"))

import ar_differentiation_bed as bed                             # noqa: E402
from geofractal.losses.pure_geometric import (                   # noqa: E402
    GeometricPrototypeLoss, HierarchicalGeometricLoss, PureGeometricLoss)
from loss_forms import prim_ce                                   # noqa: E402

torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
DEV = "cuda" if torch.cuda.is_available() else "cpu"
if DEV == "cuda":
    torch.cuda.set_per_process_memory_fraction(0.73)

DATA_ROOT = os.environ.get("GEOLIP_DATA", "./data")
RUNS = os.path.join(ROOT, "tools", "geobasin_runs")
CE_SCALE = 10.0

def seed_for(name):
    return zlib.crc32(name.encode()) & 0x7FFFFFFF

ARMS = ("geo_ce_scores", "geo_pure", "geo_pure_frozen", "geo_pure_norep",
        "geo_proto", "geo_hier", "geo_hybrid")


class FastHierarchical(HierarchicalGeometricLoss):
    """Vectorizes ONLY the per-sample fine_mask loop (a 60x Python-loop
    slowdown at B*T=8192); every formula identical — parity-asserted against
    the verbatim original at build time."""
    def forward(self, compatibility_scores, labels, mixed_labels=None,
                lam=None):
        if mixed_labels is not None:
            return super().forward(compatibility_scores, labels,
                                   mixed_labels, lam)
        B = compatibility_scores.shape[0]
        dev = compatibility_scores.device
        sl = labels // self.subclasses_per_super
        scores_r = compatibility_scores.view(B, self.num_superclasses,
                                             self.subclasses_per_super)
        ss = scores_r.sum(dim=2)
        coarse_correct = ss[torch.arange(B, device=dev), sl]
        ct = float(self.subclasses_per_super)
        coarse_loss = ((coarse_correct - ct) / ct).pow(2).mean()
        cmask = torch.ones_like(ss)
        cmask[torch.arange(B, device=dev), sl] = 0
        coarse_rep = (ss * cmask).pow(2).sum(dim=1).mean()
        fine_correct = compatibility_scores[torch.arange(B, device=dev), labels]
        fine_loss = (1.0 - fine_correct).pow(2).mean()
        # vectorized fine_mask: ones on the label's superclass block, zero at
        # the label column (identical to the original's per-sample loop)
        cols = (sl * self.subclasses_per_super).unsqueeze(1) + \
            torch.arange(self.subclasses_per_super, device=dev).unsqueeze(0)
        fmask = torch.zeros_like(compatibility_scores)
        fmask.scatter_(1, cols, 1.0)
        fmask[torch.arange(B, device=dev), labels] = 0
        fine_rep = (compatibility_scores * fmask).pow(2).sum(dim=1).mean()
        consistency = F.relu(fine_correct * 2 - coarse_correct).mean()
        w_c = torch.sigmoid(self.coarse_weight)
        w_f = torch.sigmoid(self.fine_weight)
        return (w_c * (coarse_loss + 0.3 * coarse_rep)
                + w_f * (fine_loss + 0.3 * fine_rep) + 0.2 * consistency)


class ScoreHead(nn.Module):
    """The shared cosine-anchor basin head. scores in [0,1] per class."""
    def __init__(self, dim=256, classes=256, frozen=False, gen=None):
        super().__init__()
        A = torch.randn(classes, dim, generator=gen)
        if frozen:
            self.register_buffer("A", F.normalize(A, dim=-1))
        else:
            self.A = nn.Parameter(A)

    def cos(self, feats):
        return F.normalize(feats, dim=-1) @ F.normalize(self.A, dim=-1).t()

    def forward(self, feats):
        return (self.cos(feats) + 1) / 2


def build(arm, seed):
    torch.manual_seed(seed_for(f"geobasin:{arm}:{seed}"))
    lm = bed.ByteLM("addr_msl64").to(DEV)
    lm.head = nn.Identity()                       # forward -> 256-d msl feats
    gh = torch.Generator().manual_seed(seed_for(f"geobasin-head:{seed}"))
    head = ScoreHead(frozen=(arm == "geo_pure_frozen"), gen=gh).to(DEV)
    aux = None
    if arm == "geo_proto":
        aux = GeometricPrototypeLoss(num_classes=256, prototype_dim=64).to(DEV)
    elif arm == "geo_hier":
        aux = FastHierarchical(num_classes=256, num_superclasses=16).to(DEV)
        # parity vs the VERBATIM original on a random batch (reference-check
        # pattern): the vectorization must be mathematics-identical
        ref = HierarchicalGeometricLoss(num_classes=256,
                                        num_superclasses=16).to(DEV)
        ref.load_state_dict(aux.state_dict())
        gpar = torch.Generator().manual_seed(seed_for("hier-parity"))
        sc = torch.rand(16, 256, generator=gpar).to(DEV)
        yy = torch.randint(0, 256, (16,), generator=gpar).to(DEV)
        assert torch.allclose(aux(sc, yy), ref(sc, yy), atol=1e-6), \
            "FastHierarchical diverged from the verbatim original"
    elif arm != "geo_ce_scores":
        aux = PureGeometricLoss()                 # stateless
    params = list(lm.parameters()) + list(head.parameters())
    if aux is not None:
        params += list(aux.parameters())
    return lm, head, aux, [p for p in params if p.requires_grad]


def loss_of(arm, head, aux, feats, y):
    scores = head(feats).reshape(-1, 256)
    yy = y.reshape(-1)
    if arm == "geo_ce_scores":
        return prim_ce((scores * 2 - 1).reshape(*y.shape, 256) * CE_SCALE, y).mean()
    if arm == "geo_hybrid":
        ce = prim_ce((scores * 2 - 1).reshape(*y.shape, 256) * CE_SCALE, y).mean()
        return 0.5 * ce + 0.5 * PureGeometricLoss()(scores, yy)
    if arm == "geo_pure_norep":
        B = scores.shape[0]
        correct = scores[torch.arange(B, device=scores.device), yy]
        attraction = (1.0 - correct).pow(2).mean()          # verbatim term 1
        rng = F.relu(scores - 1.0).pow(2).mean() \
            + F.relu(-scores).pow(2).mean()                 # verbatim term 4
        return attraction + 0.1 * rng                       # comparative terms DROPPED
    return aux(scores, yy)                                  # verbatim package losses


@torch.no_grad()
def evaluate(head, lm, va, g):
    lm.eval()
    tot_ce, tot_ok, n = 0.0, 0, 0
    for _ in range(8):
        x, y = bed._batch(va, 32, 256, DEV, g)
        cos = head.cos(lm(x))
        lp = F.log_softmax(cos * CE_SCALE, -1)
        tot_ce += float(-lp.gather(-1, y.unsqueeze(-1)).sum())
        tot_ok += int((cos.argmax(-1) == y).sum())
        n += y.numel()
    x, _ = bed._batch(va, 8, 256, DEV, g)
    _ = lm(x)
    v = lm.head_addr.vitals(lm.head_proj(lm._last_h).view(
        *lm._last_h.shape[:-1], lm.n_slots, 4))
    lm.train()
    return (tot_ce / n) / math.log(2), tot_ok / n, v


def train_arm(arm, seed, steps=2000):
    assert arm in ARMS, f"unknown arm {arm}"
    os.makedirs(RUNS, exist_ok=True)
    tr, va = bed._wikitext_bytes(DATA_ROOT)
    lm, head, aux, params = build(arm, seed)
    g = torch.Generator().manual_seed(seed_for(f"geobasin-data:{arm}:{seed}"))
    ge = torch.Generator().manual_seed(seed_for("geobasin-eval"))
    opt = torch.optim.Adam(params, lr=3e-4, weight_decay=0.0)
    t0 = time.time()
    # anchor-collapse gauge: pairwise |cos| spread of A at start vs end
    def anchor_spread():
        An = F.normalize(head.A.detach(), dim=-1)
        pc = (An @ An.t()).abs()
        off = pc[~torch.eye(256, dtype=torch.bool, device=pc.device)]
        return float(off.mean())
    sp0 = anchor_spread()
    for step in range(steps):
        x, y = bed._batch(tr, 32, 256, DEV, g)
        opt.zero_grad(set_to_none=True)
        L = loss_of(arm, head, aux, lm(x), y)
        L.backward()
        opt.step()
        if step == 10 and DEV == "cuda":
            print(f"[{arm} s{seed}] step10 loss {float(L.detach()):.4f} "
                  f"peak {torch.cuda.max_memory_allocated()/2**30:.2f}GB",
                  flush=True)
    bpb, acc, vit = evaluate(head, lm, va, ge)
    rec = {"arm": arm, "seed": seed, "steps": steps, "bpb": round(bpb, 4),
           "decoded_acc": round(acc, 4),
           "anchor_abs_cos_mean_init": round(sp0, 4),
           "anchor_abs_cos_mean_final": round(anchor_spread(), 4),
           "vitals": vit, "wall_s": round(time.time() - t0, 1),
           "n_params_extra": sum(p.numel() for p in aux.parameters())
           if aux is not None and any(True for _ in aux.parameters()) else 0}
    out = os.path.join(RUNS, f"{arm}_s{seed}_t{steps}.jsonl")
    with open(out, "a", encoding="utf-8") as f:
        f.write(json.dumps(rec) + "\n")
    print(f"[DONE {arm} s{seed}] bpb {bpb:.4f} acc {acc:.4f} "
          f"anchors |cos| {sp0:.3f}->{rec['anchor_abs_cos_mean_final']:.3f} "
          f"({rec['wall_s']}s)", flush=True)
    return rec


if __name__ == "__main__":
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("--arm", default=None)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--steps", type=int, default=2000)
    ap.add_argument("--list", action="store_true")
    a, _ = ap.parse_known_args()
    if a.list or not a.arm:
        print("arms:", " ".join(ARMS))
        sys.exit(0)
    train_arm(a.arm, a.seed, a.steps)