Merge pull request #49 from robin-janssen/improve-surrogate-evaluation

Improve surrogate evaluation

Author: robin-janssen

README.md

@@ -1,6 +1,6 @@
 # CODES Benchmark
 
-[![codecov](https://codecov.io/github/robin-janssen/CODES-Benchmark/graph/badge.svg?token=TNF9ISCAJK)](https://codecov.io/github/robin-janssen/CODES-Benchmark)
+[![codecov](https://codecov.io/github/robin-janssen/CODES-Benchmark/branch/main/graph/badge.svg?token=TNF9ISCAJK)](https://codecov.io/github/robin-janssen/CODES-Benchmark)
 ![Static Badge](https://img.shields.io/badge/license-GPLv3-blue)
 ![Static Badge](https://img.shields.io/badge/NeurIPS-2024-green)
 

codes/surrogates/AbstractSurrogate/abstract_surrogate.py

@@ -13,7 +13,7 @@
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
-from codes.utils import create_model_dir
+from codes.utils import create_model_dir, parse_hyperparameters
 
 
 class AbstractSurrogateModel(ABC, nn.Module):
@@ -334,6 +334,9 @@ def save(
             )
         hyperparameters["date"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
 
+        # Recursively parse hyperparameters to make them yaml-serializable
+        hyperparameters = parse_hyperparameters(hyperparameters)
+
         # Reduce the precision of the losses and accuracy
         for attribute in ["train_loss", "test_loss", "MAE"]:
             value = getattr(self, attribute)
@@ -518,10 +521,6 @@ def time_pruning(self, current_epoch: int, total_epochs: int) -> None:
         # Define warmup period based on 10% of total epochs.
         warmup_epochs = max(10, int(total_epochs * 0.02))
         if current_epoch < warmup_epochs:
-            # Do not attempt to prune before the warmup period is complete.
-            # print(
-            #     f"[time_pruning] Warmup period: {current_epoch}/{warmup_epochs} epochs completed. Skipping pruning check."
-            # )
             return
 
         elapsed = time.time() - self._trial_start_time
@@ -541,7 +540,7 @@ def time_pruning(self, current_epoch: int, total_epochs: int) -> None:
             if projected_total_time > threshold:
                 if self.optuna_trial is not None:
                     tqdm.write(
-                        f"[time_pruning] Projected total time {projected_total_time:.1f}s exceeds threshold {threshold:.1f}s. Pruning trial {self.optuna_trial.number}."
+                        f"[Trial {self.optuna_trial.number}] Projected total time {projected_total_time:.1f}s exceeds threshold {threshold:.1f}s. Pruning trial."
                     )
                     self.optuna_trial.set_user_attr(
                         "prune_reason",

codes/surrogates/LatentNeuralODE/latent_neural_ode.py

@@ -329,13 +329,6 @@ def fit_profile(
                     loss.backward()
                     optimizer.step()
 
-                # # renormalize once after 10 epochs
-                # if epoch == 10 and i == 0:
-                #     with torch.no_grad():
-                #         self.model.renormalize_loss_weights(
-                #             x_true, x_pred, params, criterion
-                #         )
-
             if not (profiled and epoch == 0):
                 # Only step here if you didn't already step inside profiled block
                 scheduler.step()

codes/surrogates/LatentPolynomial/latent_poly.py

@@ -218,12 +218,6 @@ def fit(
                 loss.backward()
                 optimizer.step()
 
-                if epoch == 10 and i == 0:
-                    with torch.no_grad():
-                        self.model.renormalize_loss_weights(
-                            x_true, x_pred, params, criterion
-                        )
-
             scheduler.step()
 
             self.validate(
@@ -461,6 +455,7 @@ def identity_loss(self, x_true: Tensor, params: Tensor = None):
         # only reconstruct the initial state
         x0 = x_true[:, 0, :]
         if not self.config.coeff_network and params is not None:
+            params = params.to(self.device)
             enc_input = torch.cat([x0, params], dim=1)
         else:
             enc_input = x0

codes/tune/__init__.py

@@ -5,20 +5,23 @@
     plot_test_losses,
 )
 from .optuna_fcts import (
+    MaxValidTrialsCallback,
+    _count_valid_trials,
+    build_fine_optuna_params,
     create_objective,
     load_yaml_config,
     make_optuna_params,
     maybe_set_runtime_threshold,
     training_run,
 )
 from .postgres_fcts import (
-    _make_db_url,
-    initialize_optuna_database,
     _check_postgres_running_local,
-    _start_postgres_server_local,
     _check_remote_reachable,
     _initialize_postgres_local,
     _initialize_postgres_remote,
+    _make_db_url,
+    _start_postgres_server_local,
+    initialize_optuna_database,
 )
 from .tune_utils import (
     build_study_names,
@@ -30,6 +33,8 @@
 
 __all__ = [
     "create_objective",
+    "MaxValidTrialsCallback",
+    "build_fine_optuna_params",
     "load_yaml_config",
     "make_optuna_params",
     "maybe_set_runtime_threshold",
@@ -50,4 +55,5 @@
     "_check_remote_reachable",
     "_initialize_postgres_local",
     "_initialize_postgres_remote",
+    "_count_valid_trials",
 ]

codes/tune/evaluate_tuning.py

@@ -23,6 +23,58 @@
 from codes.utils import nice_print
 
 
+def pareto_front(points: np.ndarray) -> np.ndarray:
+    # lower-is-better for both objectives
+    is_efficient = np.ones(points.shape[0], dtype=bool)
+    for i, p in enumerate(points):
+        if not is_efficient[i]:
+            continue
+        # any other point strictly better in both dims dominates p
+        better = np.all(points <= p, axis=1) & np.any(points < p, axis=1)
+        dominated = better & (np.arange(points.shape[0]) != i)
+        if np.any(dominated):
+            is_efficient[i] = False
+    return points[is_efficient]
+
+
+def hypervolume_2d(pareto_points: np.ndarray, reference: np.ndarray) -> float:
+    # assumes minimize-minimize; reference worse than all pareto_points
+    if pareto_points.size == 0:
+        return 0.0
+    pts = pareto_points[np.argsort(pareto_points[:, 0])]  # sort by first objective
+    hv = 0.0
+    prev_f2 = reference[1]
+    for f1, f2 in pts:
+        width = reference[0] - f1
+        height = prev_f2 - f2
+        if width > 0 and height > 0:
+            hv += width * height
+        prev_f2 = f2
+    return hv
+
+
+def compute_hypervolume_over_time(study: optuna.Study, ref_slack=1.1):
+    from optuna.trial import TrialState
+
+    completed = [t for t in study.trials if t.state == TrialState.COMPLETE]
+    if not completed:
+        return [], None
+
+    # Order by completion time
+    completed.sort(key=lambda t: t.datetime_complete or t.datetime_start)
+    all_vals = np.array([t.values for t in completed])  # shape (N, 2)
+    reference = all_vals.max(axis=0) * ref_slack  # slightly worse than worst seen
+
+    hypervolumes = []
+    for k in range(1, len(completed) + 1):
+        subset = completed[:k]
+        pts = np.array([t.values for t in subset])
+        pareto = pareto_front(pts)
+        hv = hypervolume_2d(pareto, reference)
+        hypervolumes.append(hv)
+    return hypervolumes, reference
+
+
 def load_loss_history(model_path: str) -> tuple[np.ndarray, np.ndarray, int]:
     """
     Load loss histories from a saved model file (.pth).
@@ -218,6 +270,49 @@ def evaluate_tuning(
             print(f"Could not load study '{full_name}'")
             continue
 
+        # Compute hypervolume over time
+        if len(study.directions) == 2:
+            hvs, reference = compute_hypervolume_over_time(study)
+            if hvs:
+                # Normalize to final hypervolume for relative curve
+                final_hv = hvs[-1]
+                rel_hvs = [hv / final_hv if final_hv > 0 else 0 for hv in hvs]
+
+                # Plot absolute and relative hypervolume
+                plt.figure(figsize=(6, 4))
+                plt.plot(np.arange(1, len(hvs) + 1), hvs, label="Hypervolume")
+                plt.xlabel("Completed Trials")
+                plt.ylabel("Hypervolume")
+                plt.title(f"{suffix} Hypervolume over trials")
+                plt.grid(True)
+                plt.tight_layout()
+                plt.savefig(
+                    os.path.join(save_dir, f"hypervolume_{suffix}.png"), dpi=300
+                )
+                plt.close()
+
+                plt.figure(figsize=(6, 4))
+                plt.plot(
+                    np.arange(1, len(rel_hvs) + 1),
+                    rel_hvs,
+                    label="Relative Hypervolume",
+                )
+                plt.xlabel("Completed Trials")
+                plt.ylabel("Fraction of Final HV")
+                plt.title(f"{suffix} Relative Hypervolume")
+                plt.grid(True)
+                plt.tight_layout()
+                plt.savefig(
+                    os.path.join(save_dir, f"hypervolume_relative_{suffix}.png"),
+                    dpi=300,
+                )
+                plt.close()
+                print(f"Saved hypervolume plots for {suffix} (final HV={final_hv:.3e})")
+            else:
+                print("No hypervolume computed (no complete trials).")
+        else:
+            print("Skipping hypervolume: study is not two-objective.")
+
         best = get_best_trials(study, top_n)
         if not best:
             print(f"No completed trials in {full_name}")
@@ -285,19 +380,19 @@ def parse_args():
     p.add_argument(
         "--study_name",
         type=str,
-        default="cloud_tuning_rough",
+        default="cloud_tuning_fine",
         help="Main study prefix (e.g. lvparams5)",
     )
     p.add_argument(
         "--storage_name",
         type=str,
-        default="optuna_cloud",
+        default="optuna_cloud_2",
         help="Main study prefix (e.g. lvparams5)",
     )
     p.add_argument(
         "--top_n",
         type=int,
-        default=10,
+        default=20,
         help="Number of top trials to plot per surrogate",
     )
     return p.parse_args()