Merge changes from paper plots branch back to main branch while ignoring files which are not required

codes/benchmark/bench_fcts.py

@@ -1,3 +1,4 @@
+import os
 from contextlib import redirect_stdout
 from typing import Any
 
@@ -378,7 +379,6 @@ def evaluate_iterative_predictions(
     # container for the piecewise predictions; seed t=0 with ground truth so errors
     # are computed only on actual predictions for t>=1 while keeping shape intact
     iterative_preds = np.zeros_like(targets)
-    iterative_preds[:, 0, :] = targets[:, 0, :]
 
     # number of chunks
     n_chunks = (n_timesteps + iter_interval - 1) // iter_interval
@@ -433,14 +433,17 @@ def evaluate_iterative_predictions(
         )
         # We predict steps 1..(chunk_len-1) relative to the provided init state (index 0).
         # Map these to global indices [start+1 .. end] inclusively.
+        if i == 0:
+            iterative_preds[:, start : end + 1, :] = preds_chunk[:, : model.n_timesteps, :].detach().cpu().numpy()
         iterative_preds[:, start + 1 : end + 1, :] = (
             preds_chunk[:, 1 : model.n_timesteps, :].detach().cpu().numpy()
         )
 
     iterative_preds_log = model.denormalize(iterative_preds, leave_log=True)
+    full_preds_log = model.denormalize(full_preds, leave_log=True)
     targets_log = model.denormalize(targets, leave_log=True)
     iterative_preds = model.denormalize(iterative_preds)
-    full_preds = model.denormalize(full_preds.detach().cpu().numpy())
+    full_preds_real = model.denormalize(full_preds.detach().cpu().numpy())
     targets = model.denormalize(targets)
 
     # compute error metrics
@@ -466,7 +469,7 @@ def evaluate_iterative_predictions(
         surr_name,
         conf,
         iterative_preds,
-        full_preds,
+        full_preds_real,
         targets,
         timesteps,
         iter_interval=iter_interval,
@@ -1583,7 +1586,7 @@ def compare_UQ(all_metrics: dict, config: dict) -> None:
         ensemble_errors,
         ensemble_std,
         config,
-        flag_fractions=(0.01, 0.05, 0.10, 0.20, 0.30, 0.40, 0.50),
+        flag_fractions=(0, 0.025, 0.05, 0.10, 0.20, 0.30, 0.40, 0.50),
         save=True,
         show_title=True,
     )

codes/benchmark/bench_plots.py

@@ -1661,6 +1661,7 @@ def plot_errors_over_time(
     elif mode == "iterative":
         # Single backslash inside raw string to render the LaTeX Delta properly
         plt.ylabel(r"Log-MAE ($\Delta dex$)")
+        plt.ylim(bottom=0, top=min(np.max(list(mean_errors.values())) * 1.1, 5))
         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
@@ -2035,8 +2036,6 @@ def inference_time_bar_plot(
     # Calculate the upper y-limit to provide space for text
     max_bar = max(means[i] + stds[i] for i in range(len(means)))
     # min_bar = min(means[i] - stds[i] for i in range(len(means)))
-    # Temp!
-    # ax.set_ylim(min_bar * 0.3, max_bar * 2)  # Set limits with some padding
     ax.set_ylim(0, max_bar * 1.2)  # Set limits with some padding
 
     # Add inference time as text to the bars using the format_time function
@@ -2053,8 +2052,6 @@ def inference_time_bar_plot(
 
     ax.set_xlabel("Surrogate Model")
     ax.set_ylabel("Mean Inference Time per Run")
-    # Temp!
-    # ax.set_yscale("log")
     if show_title:
         ax.set_title("Surrogate Mean Inference Time Comparison")
 
@@ -2507,11 +2504,16 @@ def plot_catastrophic_detection_curves(
 
             xs, ys = [], []
             for f in flag_fractions:
-                unc_thr = np.percentile(u, 100.0 * (1.0 - float(f)))
-                flagged = u >= unc_thr
-                recall = (flagged & is_cat).sum() / n_cat if n_cat > 0 else 0.0
-                xs.append(100.0 * flagged.mean())
-                ys.append(100.0 * recall)
+                if f <= 0.0:
+                    xs.append(0.0)
+                    ys.append(0.0)
+                    recall = 0.0
+                else:
+                    unc_thr = np.percentile(u, 100.0 * (1.0 - float(f)))
+                    flagged = u >= unc_thr
+                    recall = (flagged & is_cat).sum() / n_cat if n_cat > 0 else 0.0
+                    xs.append(100.0 * flagged.mean())
+                    ys.append(100.0 * recall)
 
             ax.plot(
                 xs,
@@ -2988,8 +2990,6 @@ def rel_errors_and_uq(
 
     ax1.set_xlabel("Time")
     ax1.set_xlim(timesteps[0], timesteps[-1])
-    # Temp!
-    # ax1.set_ylim(3e-4, 1)
     ax1.set_ylabel("Relative Error")
     ax1.set_yscale("log")
     ax1.set_title("Comparison of Relative Errors Over Time")
@@ -3020,8 +3020,6 @@ def rel_errors_and_uq(
 
     ax2.set_xlabel("Time")
     ax2.set_xlim(timesteps[0], timesteps[-1])
-    # Temp!
-    # ax2.set_ylim(0, 0.04)
     ax2.set_ylabel("Uncertainty/Absolute Error")
     if show_title:
         ax2.set_title("Comparison of Predictive Uncertainty Over Time")

codes/surrogates/AbstractSurrogate/abstract_surrogate.py

@@ -449,8 +449,10 @@ def denormalize(
         Returns:
             Tensor | np.ndarray: The denormalized data.
         """
+        data_type = None
         if self.normalisation is not None:
             if not leave_norm:
+                data_type = data.dtype
                 if self.normalisation["mode"] == "disabled":
                     ...
                 elif self.normalisation["mode"] == "minmax":
@@ -475,6 +477,13 @@ def denormalize(
             if self.normalisation["log10_transform"] and not leave_log:
                 data = 10**data
 
+            # Conserve dtype
+            if data_type is not None:
+                if isinstance(data, Tensor):
+                    return data.to(dtype=data_type)
+                if isinstance(data, np.ndarray):
+                    return data.astype(data_type)
+
         return data
 
     def denormalize_old(self, data: Tensor) -> Tensor: