test: capture mouse generate() golden output (pre-migration)

This commit is contained in:
2026-05-12 00:19:39 +08:00
parent 99f7ae29cb
commit 98daef54ca
2 changed files with 49 additions and 0 deletions

View File

@@ -0,0 +1,49 @@
"""One-shot script to capture golden mouse trajectories from the current torch
implementation. Run BEFORE the migration so we can verify the numpy/ORT rewrite
in Phase 4 produces equivalent output.
Output: tests/unit/data/golden_mouse.npz
"""
from __future__ import annotations
import random
import sys
from pathlib import Path
# Allow running as `uv run python scripts/build_golden_mouse.py` from project root.
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
import numpy as np
import torch
from ai_mouse import generate
CASES: list[tuple[tuple[int, int], tuple[int, int]]] = [
((100, 200), (900, 400)), # horizontal 800px
((500, 500), (500, 100)), # vertical 400px upward
((200, 600), (800, 200)), # 720px diagonal
((100, 100), (130, 110)), # very short 31px
((50, 50), (1500, 900)), # very long 1700px
((400, 300), (500, 300)), # short horizontal 100px
((300, 300), (700, 700)), # 45° diagonal
((600, 400), (200, 100)), # reverse diagonal
]
SEEDS = (0, 1, 2, 3)
def main() -> None:
out: dict[str, np.ndarray] = {}
for case_idx, (start, end) in enumerate(CASES):
for seed in SEEDS:
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
pts = generate(start=start, end=end)
out[f"case{case_idx}_seed{seed}"] = np.array(pts, dtype=np.int64)
out_path = Path("tests/unit/data/golden_mouse.npz")
np.savez_compressed(out_path, **out)
print(f"Wrote {len(out)} golden traces to {out_path}")
if __name__ == "__main__":
main()