generate_star.py

#!/usr/bin/env python3
"""Generate a B-branch star reachability dataset in the rbs format.

Generalises generate_2arm_star.py from a hardcoded 2 arms to any branch count B.

Two disconnected, structurally identical B-branch stars:
  - reachable component   (root R):  branches 0..B-1; one leaf is `target`.
  - unreachable component (root R2): branches 0..B-1; one leaf is `neg_target`.

The reachable frontier at hop k is therefore B nodes (one per branch), so the
balance score in graph_metrics -- ce_score = exp(log|F| - CE) = |F| * GM(p) --
is 1.0 only when all B branches hold probability 1/B. Nothing in the metric or
the curriculum assumes |F| == 2, so B > 2 works with no training-code change.

Node budget: 2 components * B branches * L nodes, plus the two roots, so
    n_nodes = 2 + 2*B*L
must fit the tokenizer (node id == token id, NUM_NODES ids available). That is
the binding constraint: B*L <= (NUM_NODES - 2) / 2, i.e. 49 at NUM_NODES=100.

Emits both flavors per split (identical graphs, differing only in frontier fields):
  *_fo_coconut.json : neighbor_k[k] = [target-branch node]        (unique path)
  *_fo_bfs.json     : neighbor_k[k] = [node from every branch]    (full frontier)

Example:
  python generate_star.py --B 4 --L 10 --n_train 100000 \
      --out_prefix data/star_4arm_L
"""
import json, random, argparse

try:
    from stokenizer import NUM_NODES as VOCAB_NODES  # single source of truth
except Exception:
    VOCAB_NODES = 100


def _branch(node_iter, root, L):
    """One chain of L nodes hanging off `root`."""
    nodes, edges, prev = [], [], root
    for _ in range(L):
        x = next(node_iter)
        edges.append([prev, x])
        prev = x
        nodes.append(x)
    return nodes, edges


def make_star(B, L, rng, node_pool=None):
    """Raw B-branch double-star structure (flavor-independent).

    node_pool: size of the id pool to sample labels from (ids 0..node_pool-1).
    A dense pool matters: sampling from the full 100-id vocab when a graph only
    uses ~30 ids makes individual ids rare and empirically breaks stage-0
    learning. Default keeps the original dense ratio (nodes needed + 5).
    """
    n_nodes = 2 + 2 * B * L
    if node_pool is None:
        node_pool = min(n_nodes + 5, VOCAB_NODES)
    assert n_nodes <= node_pool <= VOCAB_NODES, (
        f"B={B} L={L} needs {n_nodes} nodes; pool={node_pool} vocab={VOCAB_NODES}"
    )
    ids = rng.sample(range(node_pool), n_nodes)
    it = iter(ids)

    R = next(it)
    pos_branches, edges = [], []
    for _ in range(B):
        nodes, e = _branch(it, R, L)
        pos_branches.append(nodes)
        edges += e

    R2 = next(it)
    neg_branches = []
    for _ in range(B):
        nodes, e = _branch(it, R2, L)
        neg_branches.append(nodes)
        edges += e

    tgt_branch = rng.choice(pos_branches)
    neg_branch = rng.choice(neg_branches)
    rng.shuffle(edges)
    return {
        "R": R, "R2": R2,
        "pos_branches": pos_branches, "neg_branches": neg_branches,
        "tgt_branch": tgt_branch, "neg_branch": neg_branch,
        "edges": edges, "L": L, "B": B,
    }


def to_sample(star, flavor):
    L = star["L"]
    R, R2 = star["R"], star["R2"]
    tgt_branch, neg_branch = star["tgt_branch"], star["neg_branch"]
    target, neg_target = tgt_branch[-1], neg_branch[-1]

    neighbor_k = {"0": [R]}
    neg_neighbor_k = {"0": [R2]}
    for k in range(1, L + 1):
        if flavor == "coconut":                    # unique shortest-path frontier
            neighbor_k[str(k)] = [tgt_branch[k - 1]]
            neg_neighbor_k[str(k)] = [neg_branch[k - 1]]
        else:                                      # bfs: every branch at depth k
            neighbor_k[str(k)] = [b[k - 1] for b in star["pos_branches"]]
            neg_neighbor_k[str(k)] = [b[k - 1] for b in star["neg_branches"]]

    all_nodes = sorted(
        [R, R2]
        + [n for b in star["pos_branches"] for n in b]
        + [n for b in star["neg_branches"] for n in b]
    )
    return {
        "question": "",
        "answer": str(target),
        "steps": [str(x) for x in tgt_branch],     # length L (rbs uses only len())
        "edges": star["edges"],
        "root": R,
        "target": target,
        "neg_target": neg_target,
        "neighbor_k": neighbor_k,
        "neg_root": R2,
        "neg_neighbor_k": neg_neighbor_k,
        "idx_to_symbol": [str(x) for x in all_nodes],
        "difficulty": L,
    }


def gen(B, L, n, rng, seen, node_pool=None):
    out = []
    while len(out) < n:
        star = make_star(B, L, rng, node_pool=node_pool)
        key = (star["R"], star["tgt_branch"][-1], star["neg_branch"][-1],
               frozenset(map(tuple, star["edges"])))
        if key in seen:
            continue
        seen.add(key)
        out.append(star)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--B", type=int, default=2, help="branches per component")
    ap.add_argument("--L", type=int, default=10, help="branch length = reasoning depth")
    ap.add_argument("--n_train", type=int, default=14000)
    ap.add_argument("--n_valid", type=int, default=256)
    ap.add_argument("--n_test", type=int, default=256)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--out_prefix", default=None,
                    help="default: data/star_{B}arm_L")
    ap.add_argument("--node_pool", type=int, default=None,
                    help="sample ids from 0..node_pool-1 (default: nodes+5, dense)")
    a = ap.parse_args()

    n_nodes = 2 + 2 * a.B * a.L
    assert n_nodes <= VOCAB_NODES, (
        f"B={a.B} L={a.L} needs {n_nodes} nodes > vocab {VOCAB_NODES}; "
        f"max B*L is {(VOCAB_NODES - 2) // 2}"
    )
    prefix = a.out_prefix or f"data/star_{a.B}arm_L"
    rng = random.Random(a.seed)
    pool = a.node_pool if a.node_pool is not None else min(n_nodes + 5, VOCAB_NODES)
    seen = set()  # dedup across splits (no leakage)
    for split, n in [("train", a.n_train), ("valid", a.n_valid), ("test", a.n_test)]:
        stars = gen(a.B, a.L, n, rng, seen, node_pool=pool)
        for flavor in ("coconut", "bfs"):
            data = [to_sample(s, flavor) for s in stars]
            path = f"{prefix}{a.L}_{split}_fo_{flavor}.json"
            json.dump(data, open(path, "w"))
            print(f"{path}: {len(data)} samples, B={a.B}, {n_nodes} nodes/sample, "
                  f"{2 * a.B * a.L} edges/sample, frontier={a.B}/hop, node_pool={pool}")


if __name__ == "__main__":
    main()