Include train/val/test predictions in scenario charts and standardize y_true/y_pred colors

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

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:

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]

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]