arborist/scripts/generate_t3_bound_kat.py

#!/usr/bin/env python3
"""Regenerate the T3-bound calculator known-answer-test fixture.

Writes ``bench/fixtures/t3-bound/known-answer-tests.jsonl`` from the
fixed config list below by running ``t3_bound_bits`` on each and
recording the full output. Run this whenever the calculator's
algorithm changes — and then bump ``CALCULATOR_VERSION`` in
``bench/scripts/t3_bound_calculator.py`` so the fixture's
``calculator_version`` field changes too. The KAT regression test
(``tests/test_t3_bound_calculator.py::test_t3_bound_known_answer_tests``)
pins these values.

Usage::

    python -m scripts.generate_t3_bound_kat        # from repo root
    # or:  python scripts/generate_t3_bound_kat.py

Config list rationale (12 entries):

- ``small-deployment-§7.1`` / ``medium-deployment-§7.2`` /
  ``hardened-deployment-§7.3`` — the worked examples in
  ``docs/soft-hash-channel-t3-bound.md`` §7, under the default
  ``b1_model=max_envelope``.
- ``extreme-low-g`` / ``tight-window-W=100`` — corner regimes.
- ``tightened-c-b1`` / ``all-constants-tight`` — non-default
  C_B* overrides (replay-completeness for the constants path).
- ``b3-floor-regime`` — exercises the B3 floor-at-zero branch.
- ``gradient-fraction-zero`` — the g=0 component-isolation case
  (B1=0; B2+B3 still fire).
- ``b1-effective-control-v1-mode`` / ``b1-fraction-channels-mode``
  / ``b1-aggregate-bias-mode`` — explicit ``b1_model`` pins so the
  non-default models are regression-guarded.
"""

from __future__ import annotations

import json
import pathlib
import sys

# Allow ``python scripts/generate_t3_bound_kat.py`` from repo root.
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1]))

from bench.scripts.t3_bound_calculator import (  # noqa: E402
    CALCULATOR_VERSION,
    t3_bound_bits,
)

_FIXTURE = (
    pathlib.Path(__file__).resolve().parents[1]
    / "bench" / "fixtures" / "t3-bound" / "known-answer-tests.jsonl"
)

# (label, inputs-dict). inputs may include c_b1/c_b2/c_b3/b1_model.
_CONFIGS: list[tuple[str, dict]] = [
    ("small-deployment-§7.1", dict(
        gradient_fraction=0.05, gradient_norm_max=1.0, gradient_noise_stddev=0.1,
        lr_decision_interval=100, lr_grid_size=8, window_length=10000,
        batches_per_epoch=1024, steps_per_epoch=1024)),
    ("medium-deployment-§7.2", dict(
        gradient_fraction=0.05, gradient_norm_max=1.0, gradient_noise_stddev=0.1,
        lr_decision_interval=1000, lr_grid_size=16, window_length=100000,
        batches_per_epoch=8192, steps_per_epoch=8192)),
    ("hardened-deployment-§7.3", dict(
        gradient_fraction=0.01, gradient_norm_max=1.0, gradient_noise_stddev=0.1,
        lr_decision_interval=100, lr_grid_size=4, window_length=10000,
        batches_per_epoch=1024, steps_per_epoch=1024)),
    ("extreme-low-g", dict(
        gradient_fraction=0.001, gradient_norm_max=1.0, gradient_noise_stddev=0.5,
        lr_decision_interval=500, lr_grid_size=2, window_length=5000,
        batches_per_epoch=256, steps_per_epoch=256)),
    ("tight-window-W=100", dict(
        gradient_fraction=0.05, gradient_norm_max=2.0, gradient_noise_stddev=0.1,
        lr_decision_interval=50, lr_grid_size=4, window_length=100,
        batches_per_epoch=64, steps_per_epoch=64)),
    ("tightened-c-b1", dict(
        gradient_fraction=0.05, gradient_norm_max=1.0, gradient_noise_stddev=0.1,
        lr_decision_interval=100, lr_grid_size=8, window_length=10000,
        batches_per_epoch=1024, steps_per_epoch=1024, c_b1=0.5)),
    ("all-constants-tight", dict(
        gradient_fraction=0.05, gradient_norm_max=1.0, gradient_noise_stddev=0.1,
        lr_decision_interval=100, lr_grid_size=8, window_length=10000,
        batches_per_epoch=1024, steps_per_epoch=1024, c_b1=0.1, c_b2=0.1, c_b3=0.1)),
    ("b3-floor-regime", dict(
        gradient_fraction=0.05, gradient_norm_max=100.0, gradient_noise_stddev=0.001,
        lr_decision_interval=100, lr_grid_size=8, window_length=10000,
        batches_per_epoch=4, steps_per_epoch=4)),
    ("gradient-fraction-zero", dict(
        gradient_fraction=0.0, gradient_norm_max=1.0, gradient_noise_stddev=0.1,
        lr_decision_interval=100, lr_grid_size=8, window_length=10000,
        batches_per_epoch=1024, steps_per_epoch=1024)),
    ("b1-effective-control-v1-mode", dict(
        gradient_fraction=0.05, gradient_norm_max=1.0, gradient_noise_stddev=0.1,
        lr_decision_interval=100, lr_grid_size=8, window_length=10000,
        batches_per_epoch=1024, steps_per_epoch=1024, b1_model="effective_control_v1")),
    ("b1-fraction-channels-mode", dict(
        gradient_fraction=0.05, gradient_norm_max=1.0, gradient_noise_stddev=0.1,
        lr_decision_interval=100, lr_grid_size=8, window_length=10000,
        batches_per_epoch=1024, steps_per_epoch=1024, b1_model="fraction_channels")),
    ("b1-aggregate-bias-mode", dict(
        gradient_fraction=0.05, gradient_norm_max=1.0, gradient_noise_stddev=0.1,
        lr_decision_interval=100, lr_grid_size=8, window_length=10000,
        batches_per_epoch=1024, steps_per_epoch=1024, b1_model="aggregate_bias")),
]


def build_lines() -> list[str]:
    lines = [
        f"# T3 bound calculator known-answer tests — version {CALCULATOR_VERSION}",
        "# Regenerated by scripts/generate_t3_bound_kat.py. Algorithm change",
        "# MUST bump CALCULATOR_VERSION and re-run this script; old runs",
        "# replay against the old fixture (do not delete history).",
    ]
    for label, inp in _CONFIGS:
        r = t3_bound_bits(**inp)
        lines.append(json.dumps({
            "label": label,
            "calculator_version": r["calculator_version"],
            "b1_model": r["b1_model"],
            "inputs": r["inputs"],
            "expected_total": r["I_window_bits_upper_bound"],
            "expected_b1": r["B1_contribution"],
            "expected_b1_selected": r["b1_selected"],
            "expected_b1_fraction_channels": r["B1_fraction_channels"],
            "expected_b1_aggregate_bias": r["B1_aggregate_bias"],
            "expected_b1_effective_control_v1": r["B1_effective_control_v1"],
            "expected_b2": r["B2_contribution"],
            "expected_b3": r["B3_contribution"],
            "expected_snr_grad": r["snr_grad"],
            "expected_snr_per_channel": r["snr_per_channel"],
            "expected_certification_status": r["certification_status"],
        }, ensure_ascii=False))
    return lines


def main() -> int:
    lines = build_lines()
    _FIXTURE.parent.mkdir(parents=True, exist_ok=True)
    _FIXTURE.write_text("\n".join(lines) + "\n", encoding="utf-8")
    print(f"wrote {len(_CONFIGS)} KAT entries to {_FIXTURE} (version {CALCULATOR_VERSION})")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())