Merge pull request #51 from robin-janssen/iterative-eval

Add iterative comparative evaluation and corresponding plots

Author: robin-janssen

codes/benchmark/bench_fcts.py

@@ -483,6 +483,8 @@ def evaluate_iterative_predictions(
         "mean_squared_error": mse,
         "mean_absolute_error": mae,
         "absolute_errors": abs_errors,
+        "absolute_errors_log": abs_errors_log,
+        "iteration_interval": iter_interval,
     }
 
 
@@ -1145,6 +1147,9 @@ def compare_models(metrics: dict, config: dict):
     if config["losses"]:
         compare_main_losses(metrics, config)
 
+    if config["iterative"]:
+        compare_iterative(metrics, config)
+
     if config["gradients"]:
         compare_gradients(metrics, config)
 
@@ -1274,6 +1279,44 @@ def compare_errors(metrics: dict[str, dict], config: dict) -> None:
         plot_error_distribution_comparative(log_errors, config, mode="deltadex")
 
 
+def compare_iterative(metrics: dict[str, dict], config: dict) -> None:
+    """
+    Compare the iterative prediction errors of different surrogate models.
+
+    Args:
+        metrics (dict[str, dict]): dictionary containing the benchmark metrics for each surrogate model.
+        config (dict): Configuration dictionary.
+
+    Returns:
+        None
+    """
+    iterative_errors = {}
+    mean_iterative_errors = {}
+    median_iterative_errors = {}
+
+    for surrogate, surrogate_metrics in metrics.items():
+        if "iterative" in surrogate_metrics:
+            iterative_errors[surrogate] = surrogate_metrics["iterative"][
+                "absolute_errors_log"
+            ]
+            mean_iterative_errors[surrogate] = np.mean(
+                iterative_errors[surrogate], axis=(0, 2)
+            )
+            median_iterative_errors[surrogate] = np.median(
+                iterative_errors[surrogate], axis=(0, 2)
+            )
+
+    plot_errors_over_time(
+        mean_iterative_errors,
+        median_iterative_errors,
+        surrogate_metrics["timesteps"],
+        config,
+        mode="iterative",
+        iter_interval=surrogate_metrics["iterative"]["iteration_interval"],
+    )
+    plot_error_distribution_comparative(iterative_errors, config, mode="iterative")
+
+
 def compare_inference_time(
     metrics: dict[str, dict], config: dict, save: bool = True
 ) -> None:

codes/benchmark/bench_plots.py

@@ -1593,45 +1593,55 @@ def plot_errors_over_time(
     config: dict,
     save: bool = True,
     show_title: bool = True,
-    mode: str = "relative",  # "relative" or "deltadex"
+    mode: str = "relative",  # "relative", "deltadex", or "iterative"
+    iter_interval: int | None = None,
 ) -> None:
     """
-    Plot errors over time for different surrogate models (relative or Δdex).
+    Plot errors over time for different surrogate models (relative, Δdex, or iterative Δdex).
 
     Args:
         mean_errors (dict): Mean errors for each surrogate.
         median_errors (dict): Median errors for each surrogate.
         timesteps (np.ndarray): Array of timesteps.
         config (dict): Configuration dictionary.
-        save (bool): Whether to save the figure.
-        show_title (bool): Whether to add a title.
-        mode (str): "relative" (percentage errors) or "deltadex" (log-space abs. errors).
+                save (bool): Whether to save the figure.
+                show_title (bool): Whether to add a title.
+                mode (str):
+                        - "relative": percentage errors (y-axis log scale)
+                        - "deltadex": log-space absolute errors (Δdex)
+                        - "iterative": like "deltadex" but also draws dashed vertical lines at every
+                            n-th timestep to indicate the iterative retrigger interval.
+                iter_interval (int | None): Interval for vertical guide lines when mode == "iterative".
     """
     plt.figure(figsize=(6, 4))
     colors = plt.cm.viridis(np.linspace(0, 0.95, len(mean_errors)))
     linestyles = ["-", "--"]
 
+    # Support both dict inputs and array-like median_errors
     for i, surrogate in enumerate(mean_errors.keys()):
-        mean_val = np.mean(mean_errors[surrogate])
-        median_val = np.mean(median_errors[surrogate])
+        mean_series = mean_errors[surrogate]
+        median_series = median_errors[surrogate]
+
+        mean_val = float(np.mean(mean_series))
+        median_val = float(np.mean(median_series))
 
         if mode == "relative":
             mean_label = f"{surrogate}\nMean = {mean_val * 100:.2f}%"
             median_label = f"{surrogate}\nMedian = {median_val * 100:.2f}%"
-        else:  # deltadex
+        else:  # deltadex or iterative
             mean_label = f"{surrogate}\nMean = {mean_val:.3f} dex"
             median_label = f"{surrogate}\nMedian = {median_val:.3f} dex"
 
         plt.plot(
             timesteps,
-            mean_errors[surrogate],
+            mean_series,
             label=mean_label,
             color=colors[i],
             linestyle=linestyles[0],
         )
         plt.plot(
             timesteps,
-            median_errors[surrogate],
+            median_series,
             label=median_label,
             color=colors[i],
             linestyle=linestyles[1],
@@ -1644,10 +1654,23 @@ def plot_errors_over_time(
         plt.yscale("log")
         fname = "accuracy_rel_errors_time_models.png"
         title = "Comparison of Relative Errors Over Time"
-    else:
+    elif mode == "deltadex":
         plt.ylabel(r"Log-MAE ($\Delta dex$)")
         fname = "accuracy_delta_dex_time.png"
         title = "Comparison of Δdex Errors Over Time"
+    elif mode == "iterative":
+        # Single backslash inside raw string to render the LaTeX Delta properly
+        plt.ylabel(r"Log-MAE ($\Delta dex$)")
+        fname = "iterative_delta_dex_time.png"
+        title = "Comparison of Δdex Errors Over Time for Iterative Predictions"
+        # Add subtle dashed vertical lines at every n-th timestep if provided and valid
+        if isinstance(iter_interval, int) and iter_interval > 0:
+            # start at iter_interval to avoid drawing a line at the very first x-limit
+            for idx in range(iter_interval, len(timesteps), iter_interval):
+                x = timesteps[idx]
+                plt.axvline(x=x, linestyle="--", color="gray", alpha=0.3, linewidth=0.8)
+    else:
+        raise ValueError(f"Unknown mode: {mode}")
 
     if config["dataset"]["log_timesteps"]:
         plt.xscale("log")
@@ -2252,7 +2275,7 @@ def plot_error_distribution_comparative(
     conf: dict,
     save: bool = True,
     show_title: bool = True,
-    mode: str = "relative",  # "relative" or "deltadex"
+    mode: str = "relative",  # "relative", "deltadex", or "iterative"
 ) -> None:
     """
     Plot comparative error distributions for each surrogate model.
@@ -2262,7 +2285,8 @@ def plot_error_distribution_comparative(
         conf (dict): Configuration dictionary.
         save (bool): Whether to save the figure.
         show_title (bool): Whether to add a title.
-        mode (str): "relative" (unitless %) or "deltadex" (log-space abs. errors).
+        mode (str): "relative" (unitless %), "deltadex" (log-space abs. errors), or
+            "iterative" (same as deltadex plotting, different title/filename for iterative context).
     """
     model_names = list(errors.keys())
     num_models = len(model_names)
@@ -2321,10 +2345,17 @@ def plot_error_distribution_comparative(
         xlabel = "Relative Error Magnitude"
         title = "Distribution of Surrogate Relative Errors"
         fname = "accuracy_error_dist_relative.png"
-    else:
+    elif mode == "deltadex":
         xlabel = r"Log-MAE ($\Delta dex$)"
         title = "Distribution of Surrogate Δdex Errors"
         fname = "accuracy_error_dist_deltadex.png"
+    elif mode == "iterative":
+        # Plot identical to deltadex but labeled for iterative evaluation context
+        xlabel = r"Log-MAE ($\Delta dex$)"
+        title = "Distribution of Surrogate Relative Errors for Iterative Prediction"
+        fname = "iterative_error_dist_deltadex.png"
+    else:
+        raise ValueError(f"Unknown mode: {mode}")
 
     plt.xlabel(xlabel)
     plt.ylabel("Smoothed Histogram Count")

codes/benchmark/bench_utils.py

@@ -415,6 +415,7 @@ def clean_metrics(metrics: dict, conf: dict) -> dict:
 
     if conf["iterative"]:
         write_metrics["iterative"].pop("absolute_errors", None)
+        write_metrics["iterative"].pop("absolute_errors_log", None)
     if conf["gradients"]:
         write_metrics["gradients"].pop("gradients", None)
         write_metrics["gradients"].pop("max_counts", None)

test/test_model_comparison.py

@@ -12,6 +12,7 @@ def record_calls(monkeypatch):
     calls = []
     names = [
         "compare_errors",
+        "compare_iterative",
         "compare_main_losses",
         "compare_gradients",
         "compare_inference_time",
@@ -68,6 +69,10 @@ def make_dummy_metrics():
                 "correlation_metrics": None,
                 "weighted_diff": None,
             },
+            "iterative": {
+                "iteration_interval": 10,
+                "absolute_errors_log": None,
+            },
         }
     }
 
@@ -79,6 +84,7 @@ def make_dummy_metrics():
         (
             {
                 "losses": True,
+                "iterative": True,
                 "gradients": True,
                 "timing": True,
                 "interpolation": {"enabled": True},
@@ -90,6 +96,7 @@ def make_dummy_metrics():
             [
                 "compare_errors",
                 "compare_main_losses",
+                "compare_iterative",
                 "compare_gradients",
                 "compare_inference_time",
                 "compare_interpolation",
@@ -105,6 +112,7 @@ def make_dummy_metrics():
         (
             {
                 "losses": False,
+                "iterative": False,
                 "gradients": False,
                 "timing": False,
                 "interpolation": {"enabled": False},
@@ -122,6 +130,7 @@ def make_dummy_metrics():
         (
             {
                 "losses": True,
+                "iterative": False,
                 "gradients": False,
                 "timing": False,
                 "interpolation": {"enabled": False},
@@ -143,6 +152,7 @@ def test_compare_models_branching(record_calls, flags, expected_sequence):
         "training_id": "test",
         "devices": ["cpu"],  # for compare_main_losses
         "losses": flags["losses"],
+        "iterative": flags["iterative"],
         "gradients": flags["gradients"],
         "timing": flags["timing"],
         "interpolation": flags["interpolation"],