bench.py

#!/usr/bin/env python3
"""
Long-range recall bench for the bit-native machine — a controlled capability probe.

Each example is a bit sequence:

    [TYPE bit] [shared BODY of length L] [BOUNDARY 111] [ANSWER] [STOP 000]

The two classes (type 0 / type 1) for a given body use the SAME body bits and differ ONLY
in the leading TYPE bit and the ANSWER it dictates (type 0 -> "11", type 1 -> "00").

SCOPE (verified): for L >= R-3 the answer-position R-bit register is byte-identical across the
two classes, so a memoryless predictor cannot beat chance -- the clean recall regime. For
L <= R-4 the TYPE bit still sits inside the answer-position register; that small-L regime is
NOT clean recall (limited by global label conflict, not memory) and is shown only for context.
The memory horizon is exactly L = R + h - 4 (the -4 = 1 TYPE bit + 3 boundary bits).

Sweeping L produces a MEMORY CURVE that separates:
  * register : horizon R-4     (chance for L >= R-3)
  * shift+h  : horizon R+h-4   (a wider window -- collapses once the gap exceeds it)
  * fold+h   : fixed-width xor-compression (no clean horizon; hurts Hamming smoothness)

RULE-SCRAMBLE CONTROL (--scramble): the type->answer map is randomised per body, so no
transferable rule exists. A held-out cell above chance in the intact arm but at chance in the
scrambled arm proves genuine body-invariant rule transfer, not address coincidence.

Designed capability test (like the copy/parity tasks used to validate LSTMs), generated by a
documented, seeded rule -- not arbitrary mock data. BODY avoids three consecutive 1s AND a
leading [1,1], so TYPE=1 ++ body never makes a spurious 111.
"""

import argparse
import random
import statistics

import blm

BOUNDARY = [1, 1, 1]
STOP = [0, 0, 0]
ANSWER = {0: [1, 1], 1: [0, 0]}


def gen_body(L, rng):
    """L random bits, never three 1s in a row AND never a leading [1,1] (so TYPE=1 ++ body
    cannot create a spurious 111 boundary)."""
    body, run = [], 0
    for i in range(L):
        if run >= 2 or (i == 1 and body[0] == 1):
            b = 0
        else:
            b = rng.randint(0, 1)
        run = run + 1 if b == 1 else 0
        body.append(b)
    return body


def sequence(type_bit, body, answer):
    return [type_bit] + body + BOUNDARY + answer + STOP


def dataset(L, K, seed, scramble=False):
    """K shared bodies; each emitted with type 0 and type 1 -> 2K examples.
    scramble: randomise the type->answer map per body (control with no transferable rule)."""
    rng = random.Random(seed)
    items = []
    for k in range(K):
        body = gen_body(L, rng)
        flip = rng.randint(0, 1) if scramble else 0
        for t in (0, 1):
            answer = ANSWER[t ^ flip]                  # scramble breaks answer=f(type) globally
            items.append({"seq": sequence(t, body, answer), "type": t, "answer": answer,
                          "body_id": k, "ans_start": 1 + L + len(BOUNDARY)})
    return items


def params(R, addr, h, seed, epochs, weights="uniform"):
    class A:
        pass
    a = A()
    a.mode = "frozen"; a.addr = addr; a.R = R; a.hist = h; a.epochs = epochs; a.seed = seed
    a.beta = 2.0; a.move_prob = 0.6; a.anneal = 0.999; a.max_move = 1
    a.alloc_radius = 1; a.push_radius = 2; a.weights = weights
    return blm.default_params(a)


def run_once(R, addr, h, L, K, seed, epochs, holdout=False, scramble=False, weights="uniform"):
    items = dataset(L, K, seed, scramble)
    train_items = test_items = items
    if holdout:                                        # generalise to UNSEEN bodies
        n_test = max(1, K // 4)
        test_bodies = set(range(K - n_test, K))
        train_items = [it for it in items if it["body_id"] not in test_bodies]
        test_items = [it for it in items if it["body_id"] in test_bodies]

    p = params(R, addr, h, seed, epochs, weights)
    pooled = []
    for it in train_items:
        pooled += blm.make_pairs("".join(map(str, it["seq"])), R, addr, h)
    m = blm.Machine(p)
    m.train(pooled)

    correct = 0
    for it in test_items:
        gen = m.generate_primed(it["seq"][:it["ans_start"]], len(it["answer"]))
        correct += (gen == it["answer"])
    return correct / len(test_items)


def memory_curve(R, K, seed, epochs, Ls, holdout):
    print(f"=== Long-range recall: memory curve  (R={R}, K={K} bodies, "
          f"{'HELD-OUT bodies' if holdout else 'in-sample'}, seed={seed}) ===")
    print("answer-accuracy vs gap L  (chance = 0.50; horizon: register=R-4, shift/fold=R+h-4)\n")
    modes = [("register", 0), ("shift", 4), ("fold", 4)]
    head = "  L  | " + " | ".join(f"{a + ('+'+str(h) if a!='register' else ''):>10}" for a, h in modes)
    print(head); print("-" * len(head))
    for L in Ls:
        cells = []
        for addr, h in modes:
            accs = [run_once(R, addr, h, L, K, seed + s, epochs, holdout) for s in range(3)]
            cells.append(f"{statistics.mean(accs):>10.2f}")
        flag = "  <- TYPE inside register" if (1 + L + 3) <= R else ""
        print(f"  {L:>2} | " + " | ".join(cells) + flag)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--R", type=int, default=6)
    ap.add_argument("--K", type=int, default=8)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--epochs", type=int, default=200)
    ap.add_argument("--holdout", action="store_true", help="test on unseen bodies")
    ap.add_argument("--Ls", type=int, nargs="+", default=[1, 2, 4, 6, 8, 10, 12])
    args = ap.parse_args()
    memory_curve(args.R, args.K, args.seed, args.epochs, args.Ls, args.holdout)


if __name__ == "__main__":
    main()