release/v1/tests/test_make_figures.py

from __future__ import annotations

import importlib.util
import sys
from pathlib import Path
from typing import Any

import pytest


SCRIPT_PATH = Path(__file__).resolve().parents[1] / "scripts" / "make_figures.py"
SPEC = importlib.util.spec_from_file_location("chebyshev_make_figures", SCRIPT_PATH)
assert SPEC is not None and SPEC.loader is not None
FIGURES = importlib.util.module_from_spec(SPEC)
sys.modules[SPEC.name] = FIGURES
SPEC.loader.exec_module(FIGURES)


COORDINATES = (
    (-1.0, -0.5),
    (1.0, -0.5),
    (-1.0, 0.5),
    (1.0, 0.5),
)


def episode_rows(
    method: str,
    seed: int,
    panel_id: str,
    values: tuple[float, float, float, float],
    *,
    protocol: str | None = None,
) -> list[dict[str, str]]:
    rows = []
    for (angle, velocity), value in zip(COORDINATES, values, strict=True):
        row = {
            "status": "success",
            "environment": FIGURES.PENDULUM_ENVIRONMENT,
            "batch_id": "claim5-batch",
            "claim_id": "claim5",
            "panel_id": panel_id,
            "method": method,
            "seed": str(seed),
            "initial_angle": str(angle),
            "initial_angular_velocity": str(velocity),
            "return": str(value),
        }
        if protocol is not None:
            row["evaluation_protocol"] = protocol
        rows.append(row)
    return rows


def pendulum_tables() -> dict[str, list[dict[str, str]]]:
    candidate_seed_11 = (0.0, 0.0, 0.0, 0.0)
    candidate_seed_22 = (1.0, 2.0, 3.0, 4.0)
    episodes = [
        *episode_rows(
            "ch6_ars_pendulum",
            11,
            "pendulum_ch6_ars_seed_grid_v1",
            candidate_seed_11,
        ),
        *episode_rows(
            "ch6_ars_pendulum",
            22,
            "pendulum_ch6_ars_seed_grid_v1",
            candidate_seed_22,
        ),
        *episode_rows(
            "ch6_ars_pendulum_selected",
            22,
            "pendulum_primary_comparison_v1",
            candidate_seed_22,
            protocol="author_exact_polynomial_helper",
        ),
        *episode_rows(
            "ars_baseline_pendulum_released",
            0,
            "pendulum_primary_comparison_v1",
            (-2.0, -2.0, -2.0, -2.0),
            protocol="author_exact_first_observation_bypasses_wrappers",
        ),
        *episode_rows(
            "ars_baseline_pendulum_released_corrected_reset",
            0,
            "pendulum_baseline_reset_sensitivity_v1",
            (-1.0, -1.0, -1.0, -1.0),
            protocol="corrected_first_observation_through_vecnormalize",
        ),
    ]
    return {
        "episodes.csv": episodes,
        "selections.csv": [
            {
                "status": "success",
                "batch_id": "claim5-batch",
                "claim_id": "claim5",
                "method": "ch6_ars_pendulum",
                "candidate_count": "2",
                "selected_seed": "22",
                "selection_and_reporting_grid_reused": "true",
            }
        ],
    }


def test_pendulum_figure_uses_authoritative_selection_and_responsive_2x2(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    captured: dict[str, Any] = {}

    def capture_result(
        staging: Path,
        key: str,
        figure: Any,
        rows: list[dict[str, Any]],
        columns: tuple[str, ...],
        details: dict[str, Any],
    ) -> Any:
        captured.update(
            {
                "staging": staging,
                "key": key,
                "figure": figure,
                "rows": rows,
                "columns": columns,
                "details": details,
            }
        )
        return FIGURES.FigureResult(key=key, status="generated", details=details)

    monkeypatch.setattr(FIGURES, "generated_result", capture_result)
    result = FIGURES.build_pendulum_figure(tmp_path, pendulum_tables())

    assert result.status == "generated"
    assert captured["figure"].layout.width is None
    assert len(captured["figure"].layout.annotations) == 4
    assert len(captured["figure"].data) == 4
    assert captured["details"]["declared_selected_seed"] == 22
    assert captured["details"]["recomputed_selected_seed"] == 22
    assert captured["details"]["selected_copy_matches_candidate_grid"] is True
    assert captured["details"]["corrected_minus_author_exact_mean"] == pytest.approx(
        1.0
    )
    assert len(captured["rows"]) == 16
    assert not any(row["method"] == "ch6_ars_pendulum" for row in captured["rows"])
    assert "evaluation_protocol" in captured["columns"]
    assert "left_protocol" in captured["columns"]
    assert "right_protocol" in captured["columns"]


def test_pendulum_figure_rejects_selection_seed_mismatch(tmp_path: Path) -> None:
    tables = pendulum_tables()
    tables["selections.csv"][0]["selected_seed"] = "11"

    with pytest.raises(
        FIGURES.ValidationError, match="disagrees with complete-grid argmax"
    ):
        FIGURES.build_pendulum_figure(tmp_path, tables)


def test_pendulum_figure_rejects_selected_copy_grid_mismatch(tmp_path: Path) -> None:
    tables = pendulum_tables()
    selected_copy = next(
        row
        for row in tables["episodes.csv"]
        if row["method"] == "ch6_ars_pendulum_selected"
    )
    selected_copy["return"] = "999.0"

    with pytest.raises(
        FIGURES.ValidationError,
        match="selected-policy grid differs from candidate seed",
    ):
        FIGURES.build_pendulum_figure(tmp_path, tables)