diff --git a/scripts/aggregate_report.py b/scripts/aggregate_report.py index 326ec0b..676f30d 100644 --- a/scripts/aggregate_report.py +++ b/scripts/aggregate_report.py @@ -50,6 +50,8 @@ def _is_better(current: float, best: float | None, primary_metric: str) -> bool: def _prepare_preds_for_plot(preds: pd.DataFrame) -> pd.DataFrame: + if "global_index" in preds.columns: + return preds.sort_values("global_index").reset_index(drop=True) if "time_ms" in preds.columns: return preds.sort_values("time_ms").reset_index(drop=True) return preds.reset_index(drop=True) @@ -58,6 +60,16 @@ def _build_timeseries_chart(preds: pd.DataFrame, scenario: str, split_meta: dict sorted_preds = _prepare_preds_for_plot(preds) fig, ax = plt.subplots(figsize=(9, 3)) + if "global_index" in sorted_preds.columns: + x_vals = sorted_preds["global_index"].to_numpy() + x_label = "global row index" + elif "time_ms" in sorted_preds.columns: + x_vals = sorted_preds["time_ms"].to_numpy() + x_label = "timestamp" + else: + x_vals = sorted_preds.index.to_numpy() + x_label = "row index" + if split_meta: tr_s = split_meta.get("train_start_index", 0) tr_e = split_meta.get("train_end_index", -1) @@ -66,31 +78,20 @@ def _build_timeseries_chart(preds: pd.DataFrame, scenario: str, split_meta: dict te_s = split_meta.get("test_start_index", 0) te_e = split_meta.get("test_end_index", len(sorted_preds) - 1) - full_x = np.arange(int(max(te_e + 1, len(sorted_preds)))) - y_true_full = np.full(full_x.shape, np.nan, dtype=float) - y_pred_full = np.full(full_x.shape, np.nan, dtype=float) - test_slice = slice(int(te_s), int(min(te_s + len(sorted_preds), len(full_x)))) - n_fill = test_slice.stop - test_slice.start - y_true_full[test_slice] = sorted_preds["y_true"].values[:n_fill] - y_pred_full[test_slice] = sorted_preds["y_pred"].values[:n_fill] - ax.axvspan(tr_s, tr_e, alpha=0.15, color="#2ca02c", label="Train region") ax.axvspan(va_s, va_e, alpha=0.15, color="#ff7f0e", label="Validation region") ax.axvspan(te_s, te_e, alpha=0.15, color="#1f77b4", label="Test region") ax.axvline(tr_e, color="#2ca02c", linestyle="--", linewidth=1.2) ax.axvline(va_e, color="#ff7f0e", linestyle="--", linewidth=1.2) - ax.plot(full_x, y_true_full, label="y_true (test window)", linewidth=1.4, color="black") - ax.plot(full_x, y_pred_full, label="y_pred (test window)", linewidth=1.1, color="red") - ax.set_xlabel("global row index") - else: - x = sorted_preds["time_ms"] if "time_ms" in sorted_preds.columns else sorted_preds.index - ax.plot(x.values, sorted_preds["y_true"].values, label="y_true", linewidth=1.6) - ax.plot(x.values, sorted_preds["y_pred"].values, label="y_pred", linewidth=1.2) - ax.set_xlabel("timestamp") + ax.plot(x_vals, sorted_preds["y_true"].to_numpy(), label="y_true", linewidth=1.5, color="blue") + ax.plot(x_vals, sorted_preds["y_pred"].to_numpy(), label="y_pred", linewidth=1.2, color="yellow") + ax.set_xlabel(x_label) + ax.set_title(f"{scenario}: y_true vs y_pred across train/validation/test") - ax.set_title(f"{scenario}: predictions with chronological split boundaries") - ax.legend(loc="best", fontsize=8) + handles, labels = ax.get_legend_handles_labels() + dedup = dict(zip(labels, handles)) + ax.legend(dedup.values(), dedup.keys(), loc="best", fontsize=8) img = _fig_to_base64(fig) plt.close(fig) return img diff --git a/scripts/run_scenario.py b/scripts/run_scenario.py index 3b5e954..c28af6c 100644 --- a/scripts/run_scenario.py +++ b/scripts/run_scenario.py @@ -4,11 +4,14 @@ import json import os import subprocess +import tempfile from oran_sim.config import SCENARIOS from datetime import datetime, timezone from pathlib import Path +import pandas as pd + def _iso_now() -> str: return datetime.now(timezone.utc).isoformat() @@ -19,6 +22,22 @@ def run(cmd: list[str]) -> None: subprocess.run(cmd, check=True) +def _combine_split_predictions(split_outputs: list[tuple[str, Path]]) -> pd.DataFrame: + parts: list[pd.DataFrame] = [] + offset = 0 + for split_name, pred_path in split_outputs: + pred_df = pd.read_csv(pred_path).copy() + pred_df["split"] = split_name + pred_df["global_index"] = pred_df["index"].to_numpy() + offset + offset += len(pred_df) + parts.append(pred_df) + + if not parts: + return pd.DataFrame() + + return pd.concat(parts, ignore_index=True) + + def _resolve_dataset_path(scenario: str, dataset: str | None) -> Path: if dataset: return Path(dataset) @@ -114,19 +133,29 @@ def main() -> None: status["seq_len"] = cfg.get("seq_len") print(f"[{scenario}] prediction started", flush=True) - run( - [ - "python", - "-m", - "scripts.predict", - "--model_dir", - str(sdir / "model"), - "--csv", - str(csv.with_name(f"{csv.stem}_test.csv")), - "--output", - str(sdir / "preds.csv"), - ] - ) + split_outputs: list[tuple[str, Path]] = [] + split_names = ["train", "val", "test"] + with tempfile.TemporaryDirectory(prefix=f"{scenario}_preds_") as tmp_dir: + for split_name in split_names: + split_csv = csv.with_name(f"{csv.stem}_{split_name}.csv") + split_pred_path = Path(tmp_dir) / f"preds_{split_name}.csv" + run( + [ + "python", + "-m", + "scripts.predict", + "--model_dir", + str(sdir / "model"), + "--csv", + str(split_csv), + "--output", + str(split_pred_path), + ] + ) + split_outputs.append((split_name, split_pred_path)) + + combined_preds = _combine_split_predictions(split_outputs) + combined_preds.to_csv(sdir / "preds.csv", index=False) print(f"[{scenario}] prediction done", flush=True) status["success"] = True except Exception as exc: diff --git a/tests/test_aggregate_report_sorting.py b/tests/test_aggregate_report_sorting.py index 93007c4..a0d8552 100644 --- a/tests/test_aggregate_report_sorting.py +++ b/tests/test_aggregate_report_sorting.py @@ -16,3 +16,18 @@ def test_prepare_preds_for_plot_sorts_by_time_ms() -> None: assert sorted_preds["time_ms"].tolist() == [10, 20, 30] assert sorted_preds["y_true"].tolist() == [1.0, 2.0, 3.0] + + +def test_prepare_preds_for_plot_prefers_global_index_when_available() -> None: + preds = pd.DataFrame( + { + "global_index": [2, 0, 1], + "time_ms": [100, 300, 200], + "y_true": [3.0, 1.0, 2.0], + "y_pred": [2.8, 1.2, 2.1], + } + ) + + sorted_preds = _prepare_preds_for_plot(preds) + + assert sorted_preds["global_index"].tolist() == [0, 1, 2] diff --git a/tests/test_run_scenario_preds.py b/tests/test_run_scenario_preds.py new file mode 100644 index 0000000..d85d441 --- /dev/null +++ b/tests/test_run_scenario_preds.py @@ -0,0 +1,30 @@ +from pathlib import Path + +import pandas as pd + +from scripts.run_scenario import _combine_split_predictions + + +def test_combine_split_predictions_appends_global_index_and_split(tmp_path: Path) -> None: + train = pd.DataFrame({"index": [0, 1], "y_true": [1.0, 2.0], "y_pred": [1.1, 1.9]}) + val = pd.DataFrame({"index": [0], "y_true": [3.0], "y_pred": [2.9]}) + test = pd.DataFrame({"index": [0, 1], "y_true": [4.0, 5.0], "y_pred": [4.1, 5.2]}) + + train_path = tmp_path / "train.csv" + val_path = tmp_path / "val.csv" + test_path = tmp_path / "test.csv" + + train.to_csv(train_path, index=False) + val.to_csv(val_path, index=False) + test.to_csv(test_path, index=False) + + combined = _combine_split_predictions( + [ + ("train", train_path), + ("val", val_path), + ("test", test_path), + ] + ) + + assert combined["split"].tolist() == ["train", "train", "val", "test", "test"] + assert combined["global_index"].tolist() == [0, 1, 2, 3, 4]