diff --git a/README.md b/README.md index e5db941..1aec29e 100644 --- a/README.md +++ b/README.md @@ -132,6 +132,15 @@ python3 scripts/plot_agent_communication.py \ --events outputs/events_20260309_030340.jsonl \ --dialogs outputs/llm_routes_20260309_030340.dialogs.csv +# One-row-per-agent round timeline with departure, arrival, and route-change highlights +python3 scripts/plot_agent_round_timeline.py --run-id 20260309_030340 + +# Or pass explicit files +python3 scripts/plot_agent_round_timeline.py \ + --events outputs/events_20260309_030340.jsonl \ + --replay outputs/llm_routes_20260309_030340.jsonl \ + --metrics outputs/run_metrics_20260309_030340.json + # Compare multiple completed runs or sweep outputs python3 scripts/plot_experiment_comparison.py \ --results-json outputs/experiments/experiment_results.json diff --git a/scripts/__init__.py b/scripts/__init__.py new file mode 100644 index 0000000..7646535 --- /dev/null +++ b/scripts/__init__.py @@ -0,0 +1 @@ +"""Plotting and utility scripts for post-run analysis.""" diff --git a/scripts/plot_agent_round_timeline.py b/scripts/plot_agent_round_timeline.py new file mode 100644 index 0000000..d758c75 --- /dev/null +++ b/scripts/plot_agent_round_timeline.py @@ -0,0 +1,298 @@ +#!/usr/bin/env python3 +"""Plot a round-based agent timeline with departure, arrival, and route-change overlays.""" + +from __future__ import annotations + +import argparse +from pathlib import Path +from typing import Any + +try: + from scripts._plot_common import load_json, load_jsonl, require_matplotlib, resolve_input +except ModuleNotFoundError: + from _plot_common import load_json, load_jsonl, require_matplotlib, resolve_input + + +def _parse_args() -> argparse.Namespace: + """Parse CLI arguments for the round-timeline plot.""" + parser = argparse.ArgumentParser( + description="Plot one row per agent from departure round to arrival round, " + "with route-change rounds highlighted." + ) + parser.add_argument( + "--run-id", + help="Timestamp token such as 20260309_030340. Used to resolve matching outputs files.", + ) + parser.add_argument( + "--events", + help="Path to an events_*.jsonl file. Defaults to the newest outputs/events_*.jsonl.", + ) + parser.add_argument( + "--replay", + help="Path to an llm_routes_*.jsonl file. Defaults to the newest outputs/llm_routes_*.jsonl.", + ) + parser.add_argument( + "--metrics", + help="Path to a run_metrics_*.json file. Defaults to the newest outputs/run_metrics_*.json.", + ) + parser.add_argument( + "--out", + help="Output PNG path. Defaults to .round_timeline.png.", + ) + parser.add_argument( + "--show", + action="store_true", + help="Open the figure window in addition to saving the PNG.", + ) + parser.add_argument( + "--include-no-departure", + action="store_true", + help="Also show agents without a departure_release event, starting from their first route change.", + ) + return parser.parse_args() + + +def _round_table(event_rows: list[dict[str, Any]]) -> list[tuple[int, float]]: + """Extract and sort the `(round, sim_t_s)` table from event rows.""" + rounds = [] + for rec in event_rows: + if rec.get("event") != "decision_round_start": + continue + if rec.get("round") is None or rec.get("sim_t_s") is None: + continue + rounds.append((int(rec["round"]), float(rec["sim_t_s"]))) + rounds = sorted(set(rounds), key=lambda item: item[0]) + if not rounds: + raise SystemExit("No decision_round_start events found; cannot build round timeline.") + return rounds + + +def _round_for_time(t: float, rounds: list[tuple[int, float]]) -> int: + """Return the latest decision round whose time is <= ``t``.""" + selected = rounds[0][0] + for round_idx, round_t in rounds: + if round_t <= float(t) + 1e-9: + selected = round_idx + else: + break + return selected + + +def _departure_times(event_rows: list[dict[str, Any]]) -> dict[str, float]: + """Collect the first recorded departure time for each agent.""" + out: dict[str, float] = {} + for rec in event_rows: + if rec.get("event") != "departure_release": + continue + vid = rec.get("veh_id") + sim_t = rec.get("sim_t_s") + if vid is None or sim_t is None: + continue + out.setdefault(str(vid), float(sim_t)) + return out + + +def _route_change_times(replay_rows: list[dict[str, Any]]) -> dict[str, list[float]]: + """Collect route-change timestamps per agent from the replay log.""" + out: dict[str, list[float]] = {} + for rec in replay_rows: + if rec.get("event") != "route_change": + continue + vid = rec.get("veh_id") + sim_t = rec.get("time_s") + if vid is None or sim_t is None: + continue + out.setdefault(str(vid), []).append(float(sim_t)) + for vid in out: + out[vid] = sorted(set(out[vid])) + return out + + +def _timeline_rows( + event_rows: list[dict[str, Any]], + replay_rows: list[dict[str, Any]], + metrics: dict[str, Any], + *, + include_no_departure: bool, +) -> tuple[list[dict[str, Any]], int]: + """Build per-agent timeline rows from departures, travel times, and route changes. + + Returns: + A tuple `(rows, final_round)` where `rows` contains one dict per agent with + `start_round`, `end_round`, `change_rounds`, and a `status` label. + """ + rounds = _round_table(event_rows) + final_round = rounds[-1][0] + departures = _departure_times(event_rows) + route_changes = _route_change_times(replay_rows) + travel_times = metrics.get("average_travel_time", {}).get("per_agent", {}) or {} + + all_agent_ids = set(departures.keys()) + if include_no_departure: + all_agent_ids.update(route_changes.keys()) + + rows: list[dict[str, Any]] = [] + for vid in sorted(all_agent_ids): + depart_time = departures.get(vid) + change_times = route_changes.get(vid, []) + + if depart_time is None: + if not include_no_departure or not change_times: + continue + start_round = _round_for_time(change_times[0], rounds) + status = "no_departure_event" + else: + start_round = _round_for_time(depart_time, rounds) + status = "completed" if vid in travel_times else "incomplete" + + if vid in travel_times and depart_time is not None: + arrival_time = float(depart_time) + float(travel_times[vid]) + end_round = _round_for_time(arrival_time, rounds) + status = "completed" + else: + end_round = final_round + + end_round = max(end_round, start_round) + change_rounds = sorted({_round_for_time(t, rounds) for t in change_times if _round_for_time(t, rounds) >= start_round}) + + rows.append({ + "veh_id": vid, + "start_round": start_round, + "end_round": end_round, + "change_rounds": change_rounds, + "status": status, + }) + + rows.sort(key=lambda row: (row["start_round"], row["veh_id"])) + return rows, final_round + + +def plot_agent_round_timeline( + *, + events_path: Path, + replay_path: Path, + metrics_path: Path, + out_path: Path, + show: bool, + include_no_departure: bool, +) -> None: + """Render the round-based agent timeline figure and save it to disk.""" + plt = require_matplotlib() + from matplotlib.patches import Patch + + event_rows = load_jsonl(events_path) + replay_rows = load_jsonl(replay_path) + metrics = load_json(metrics_path) + timeline_rows, final_round = _timeline_rows( + event_rows, + replay_rows, + metrics, + include_no_departure=include_no_departure, + ) + if not timeline_rows: + raise SystemExit("No agent timeline rows could be constructed from the provided artifacts.") + + fig_h = max(6.0, 0.32 * len(timeline_rows) + 2.0) + fig, ax = plt.subplots(figsize=(14, fig_h)) + fig.suptitle( + f"Agent Round Timeline\n{events_path.name} | {replay_path.name} | {metrics_path.name}", + fontsize=14, + ) + + yticks = [] + ylabels = [] + base_colors = { + "completed": "#4C78A8", + "incomplete": "#999999", + "no_departure_event": "#BBBBBB", + } + + for idx, row in enumerate(timeline_rows): + y = idx + yticks.append(y) + ylabels.append(row["veh_id"]) + start = float(row["start_round"]) - 0.5 + width = float(row["end_round"] - row["start_round"] + 1) + color = base_colors.get(row["status"], "#4C78A8") + hatch = "//" if row["status"] != "completed" else None + ax.broken_barh( + [(start, width)], + (y - 0.35, 0.7), + facecolors=color, + edgecolors="black", + linewidth=0.4, + hatch=hatch, + alpha=0.9, + ) + change_segments = [(float(round_idx) - 0.5, 1.0) for round_idx in row["change_rounds"]] + if change_segments: + ax.broken_barh( + change_segments, + (y - 0.35, 0.7), + facecolors="#F58518", + edgecolors="#C04B00", + linewidth=0.4, + ) + + ax.set_xlim(0.5, final_round + 0.5) + ax.set_ylim(-1, len(timeline_rows)) + ax.set_xlabel("Decision Round") + ax.set_ylabel("Agent") + ax.set_yticks(yticks) + ax.set_yticklabels(ylabels, fontsize=8) + ax.grid(axis="x", linestyle=":", alpha=0.4) + + ax.legend( + handles=[ + Patch(facecolor="#4C78A8", edgecolor="black", label="Active interval"), + Patch(facecolor="#F58518", edgecolor="#C04B00", label="Route/destination change round"), + Patch(facecolor="#999999", edgecolor="black", hatch="//", label="Still active at run end / inferred"), + ], + loc="upper right", + ) + + fig.tight_layout(rect=(0, 0, 1, 0.97)) + fig.savefig(out_path, dpi=160, bbox_inches="tight") + print(f"[PLOT] events={events_path}") + print(f"[PLOT] replay={replay_path}") + print(f"[PLOT] metrics={metrics_path}") + print(f"[PLOT] output={out_path}") + if show: + plt.show() + plt.close(fig) + + +def main() -> None: + """CLI entry point for generating the round-timeline plot.""" + args = _parse_args() + if args.run_id: + run_id = str(args.run_id) + events_default = f"outputs/events_{run_id}.jsonl" + replay_default = f"outputs/llm_routes_{run_id}.jsonl" + metrics_default = f"outputs/run_metrics_{run_id}.json" + else: + events_default = "outputs/events_*.jsonl" + replay_default = "outputs/llm_routes_*.jsonl" + metrics_default = "outputs/run_metrics_*.json" + + events_path = resolve_input(args.events, events_default) + replay_path = resolve_input(args.replay, replay_default) + metrics_path = resolve_input(args.metrics, metrics_default) + out_path = ( + Path(args.out) + if args.out + else events_path.with_suffix("").with_name(f"{events_path.with_suffix('').name}.round_timeline.png") + ) + out_path.parent.mkdir(parents=True, exist_ok=True) + plot_agent_round_timeline( + events_path=events_path, + replay_path=replay_path, + metrics_path=metrics_path, + out_path=out_path, + show=args.show, + include_no_departure=args.include_no_departure, + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/plot_all_run_artifacts.py b/scripts/plot_all_run_artifacts.py index 68ad3ad..e6003ef 100644 --- a/scripts/plot_all_run_artifacts.py +++ b/scripts/plot_all_run_artifacts.py @@ -10,18 +10,21 @@ try: from scripts._plot_common import newest_file from scripts.plot_agent_communication import plot_agent_communication + from scripts.plot_agent_round_timeline import plot_agent_round_timeline from scripts.plot_departure_timeline import plot_timeline from scripts.plot_experiment_comparison import load_cases, plot_experiment_comparison from scripts.plot_run_metrics import plot_metrics_dashboard except ModuleNotFoundError: from _plot_common import newest_file from plot_agent_communication import plot_agent_communication + from plot_agent_round_timeline import plot_agent_round_timeline from plot_departure_timeline import plot_timeline from plot_experiment_comparison import load_cases, plot_experiment_comparison from plot_run_metrics import plot_metrics_dashboard def _parse_args() -> argparse.Namespace: + """Parse CLI arguments for the aggregate plotting wrapper.""" parser = argparse.ArgumentParser( description="Generate the standard dashboard, timeline, comparison, and communication plots for one run." ) @@ -45,6 +48,7 @@ def _parse_args() -> argparse.Namespace: def _maybe_path(path_arg: str | None) -> Path | None: + """Validate an optional explicit file path and return it as a `Path`.""" if not path_arg: return None path = Path(path_arg) @@ -54,6 +58,7 @@ def _maybe_path(path_arg: str | None) -> Path | None: def _resolve_run_id(args: argparse.Namespace) -> str: + """Resolve the run ID from CLI args or the newest events file.""" if args.run_id: return str(args.run_id) for path_arg in (args.events, args.metrics, args.replay, args.dialogs): @@ -67,6 +72,7 @@ def _resolve_run_id(args: argparse.Namespace) -> str: def _resolve_paths(args: argparse.Namespace, run_id: str) -> dict[str, Path | None]: + """Resolve all input artifact paths for one run.""" metrics = _maybe_path(args.metrics) events = _maybe_path(args.events) replay = _maybe_path(args.replay) @@ -94,6 +100,7 @@ def _resolve_paths(args: argparse.Namespace, run_id: str) -> dict[str, Path | No def main() -> None: + """CLI entry point for generating the standard set of run figures.""" args = _parse_args() run_id = _resolve_run_id(args) paths = _resolve_paths(args, run_id) @@ -129,6 +136,15 @@ def main() -> None: show=args.show, top_n=args.top_n, ) + if replay_path is not None: + plot_agent_round_timeline( + events_path=events_path, + replay_path=replay_path, + metrics_path=metrics_path, + out_path=out_dir / "agent_round_timeline.png", + show=args.show, + include_no_departure=False, + ) comparison_source: Path | None = None if args.results_json: results_path = Path(args.results_json) diff --git a/tests/test_plot_agent_round_timeline.py b/tests/test_plot_agent_round_timeline.py new file mode 100644 index 0000000..ff1b02f --- /dev/null +++ b/tests/test_plot_agent_round_timeline.py @@ -0,0 +1,75 @@ +"""Unit tests for scripts.plot_agent_round_timeline.""" + +from scripts.plot_agent_round_timeline import ( + _round_for_time, + _round_table, + _timeline_rows, +) + + +class TestRoundTable: + def test_extracts_and_sorts_rounds(self): + rows = [ + {"event": "decision_round_start", "round": 2, "sim_t_s": 20.0}, + {"event": "ignored", "round": 99, "sim_t_s": 99.0}, + {"event": "decision_round_start", "round": 1, "sim_t_s": 10.0}, + ] + assert _round_table(rows) == [(1, 10.0), (2, 20.0)] + + +class TestRoundForTime: + def test_maps_to_latest_round_not_exceeding_time(self): + rounds = [(1, 10.0), (2, 20.0), (3, 30.0)] + assert _round_for_time(9.0, rounds) == 1 + assert _round_for_time(20.0, rounds) == 2 + assert _round_for_time(29.9, rounds) == 2 + assert _round_for_time(31.0, rounds) == 3 + + +class TestTimelineRows: + def _event_rows(self): + return [ + {"event": "decision_round_start", "round": 1, "sim_t_s": 10.0}, + {"event": "decision_round_start", "round": 2, "sim_t_s": 20.0}, + {"event": "decision_round_start", "round": 3, "sim_t_s": 30.0}, + {"event": "decision_round_start", "round": 4, "sim_t_s": 40.0}, + {"event": "departure_release", "veh_id": "veh_a", "sim_t_s": 20.0}, + {"event": "departure_release", "veh_id": "veh_b", "sim_t_s": 20.0}, + ] + + def test_completed_agent_uses_departure_plus_travel_time(self): + rows, final_round = _timeline_rows( + self._event_rows(), + [{"event": "route_change", "veh_id": "veh_a", "time_s": 30.0}], + {"average_travel_time": {"per_agent": {"veh_a": 15.0}}}, + include_no_departure=False, + ) + by_id = {row["veh_id"]: row for row in rows} + assert final_round == 4 + assert by_id["veh_a"]["start_round"] == 2 + assert by_id["veh_a"]["end_round"] == 3 + assert by_id["veh_a"]["change_rounds"] == [3] + assert by_id["veh_a"]["status"] == "completed" + + def test_incomplete_agent_extends_to_final_round(self): + rows, _ = _timeline_rows( + self._event_rows(), + [], + {"average_travel_time": {"per_agent": {}}}, + include_no_departure=False, + ) + by_id = {row["veh_id"]: row for row in rows} + assert by_id["veh_a"]["end_round"] == 4 + assert by_id["veh_a"]["status"] == "incomplete" + + def test_include_no_departure_uses_first_route_change_round(self): + rows, _ = _timeline_rows( + self._event_rows(), + [{"event": "route_change", "veh_id": "veh_c", "time_s": 30.0}], + {"average_travel_time": {"per_agent": {}}}, + include_no_departure=True, + ) + by_id = {row["veh_id"]: row for row in rows} + assert by_id["veh_c"]["start_round"] == 3 + assert by_id["veh_c"]["end_round"] == 4 + assert by_id["veh_c"]["status"] == "no_departure_event" diff --git a/tests/test_plot_all_run_artifacts.py b/tests/test_plot_all_run_artifacts.py new file mode 100644 index 0000000..b5c0f0b --- /dev/null +++ b/tests/test_plot_all_run_artifacts.py @@ -0,0 +1,72 @@ +"""Unit tests for scripts.plot_all_run_artifacts.""" + +from argparse import Namespace +from pathlib import Path + +from scripts.plot_all_run_artifacts import _resolve_paths, _resolve_run_id + + +class TestResolveRunId: + def test_prefers_explicit_run_id(self): + args = Namespace( + run_id="20260309_030340", + events=None, + metrics=None, + replay=None, + dialogs=None, + ) + assert _resolve_run_id(args) == "20260309_030340" + + def test_extracts_run_id_from_explicit_path(self): + args = Namespace( + run_id=None, + events="outputs/events_20260309_030340.jsonl", + metrics=None, + replay=None, + dialogs=None, + ) + assert _resolve_run_id(args) == "20260309_030340" + + +class TestResolvePaths: + def test_prefers_matching_run_id_files(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + out = Path("outputs") + out.mkdir() + (out / "run_metrics_20260309_030340.json").write_text("{}", encoding="utf-8") + (out / "events_20260309_030340.jsonl").write_text("", encoding="utf-8") + (out / "llm_routes_20260309_030340.jsonl").write_text("", encoding="utf-8") + (out / "llm_routes_20260309_030340.dialogs.csv").write_text( + "step,time_s,veh_id,control_mode,model,system_prompt,user_prompt,response_text,parsed_json,error\n", + encoding="utf-8", + ) + args = Namespace( + metrics=None, + events=None, + replay=None, + dialogs=None, + ) + paths = _resolve_paths(args, "20260309_030340") + assert paths["metrics"] == out / "run_metrics_20260309_030340.json" + assert paths["events"] == out / "events_20260309_030340.jsonl" + assert paths["replay"] == out / "llm_routes_20260309_030340.jsonl" + assert paths["dialogs"] == out / "llm_routes_20260309_030340.dialogs.csv" + + def test_missing_replay_returns_none(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + out = Path("outputs") + out.mkdir() + (out / "run_metrics_20260309_030340.json").write_text("{}", encoding="utf-8") + (out / "events_20260309_030340.jsonl").write_text("", encoding="utf-8") + (out / "llm_routes_20260309_030340.dialogs.csv").write_text( + "step,time_s,veh_id,control_mode,model,system_prompt,user_prompt,response_text,parsed_json,error\n", + encoding="utf-8", + ) + args = Namespace( + metrics=None, + events=None, + replay=None, + dialogs=None, + ) + paths = _resolve_paths(args, "20260309_030340") + assert paths["replay"] is None