From 46b0999b396734bff71e08215cc65afe278a9be4 Mon Sep 17 00:00:00 2001 From: Robin Date: Thu, 24 Jul 2025 16:17:20 +0200 Subject: [PATCH 01/13] First draft of iterative error eval --- codes/benchmark/__init__.py | 3 +- codes/benchmark/bench_fcts.py | 115 ++++++++++++++++++++++++++ codes/benchmark/bench_plots.py | 92 +++++++++++++++++++++ codes/surrogates/DeepONet/deeponet.py | 22 ++--- codes/surrogates/FCNN/fcnn.py | 16 ++-- 5 files changed, 227 insertions(+), 21 deletions(-) diff --git a/codes/benchmark/__init__.py b/codes/benchmark/__init__.py index 5a2c7804..64d2e72a 100644 --- a/codes/benchmark/__init__.py +++ b/codes/benchmark/__init__.py @@ -35,6 +35,7 @@ plot_error_correlation_heatmap, plot_error_distribution_comparative, plot_error_distribution_per_quantity, + plot_example_iterative_predictions, plot_example_mode_predictions, plot_example_predictions_with_uncertainty, plot_generalization_error_comparison, @@ -69,11 +70,11 @@ get_surrogate, load_model, make_comparison_csv, + measure_inference_time, measure_memory_footprint, read_yaml_config, save_table_csv, write_metrics_to_yaml, - measure_inference_time, ) __all__ = [ diff --git a/codes/benchmark/bench_fcts.py b/codes/benchmark/bench_fcts.py index b73cfbca..a10ada29 100644 --- a/codes/benchmark/bench_fcts.py +++ b/codes/benchmark/bench_fcts.py @@ -20,6 +20,7 @@ plot_error_correlation_heatmap, plot_error_distribution_comparative, plot_error_distribution_per_quantity, + plot_example_iterative_predictions, plot_example_mode_predictions, plot_example_predictions_with_uncertainty, plot_generalization_error_comparison, @@ -137,6 +138,13 @@ def run_benchmark(surr_name: str, surrogate_class, conf: dict) -> dict[str, Any] model, surr_name, timesteps, val_loader, conf, labels ) + if conf["iterative"]: + # Iterative training benchmark + print("Running iterative training benchmark...") + metrics["iterative"] = iterative_training_benchmark( + model, surr_name, timesteps, val_loader, conf, labels + ) + # Gradients benchmark if conf["gradients"]: print("Running gradients benchmark...") @@ -281,6 +289,113 @@ def evaluate_accuracy( return accuracy_metrics +def iterative_training_benchmark( + model, + surr_name: str, + timesteps: np.ndarray, + val_loader: DataLoader, + conf: dict, + labels: list | None = None, +) -> dict[str, Any]: + """ + Benchmark error accumulation when running the model iteratively in chunks. + + Returns the same set of error metrics as evaluate_accuracy, but over the + full trajectory built by re-feeding the last prediction as the next initial state. + """ + # load trained model + training_id = conf["training_id"] + model.load(training_id, surr_name, model_identifier=f"{surr_name.lower()}_main") + + # get full ground truth (targets) and ignore one-shot preds + _, targets = model.predict(data_loader=val_loader) + n_samples, n_timesteps, n_quantities = targets.shape + + # how many timesteps per chunk + iter_interval = 10 # conf["iterative"]["interval"] + # batch size same as in run_benchmark + surr_idx = conf["surrogates"].index(surr_name) + if isinstance(conf["batch_size"], list): + batch_size = conf["batch_size"][surr_idx] + else: + batch_size = conf["batch_size"] + + # container for the piecewise predictions + preds_all = np.zeros_like(targets) + + # number of chunks + n_chunks = (n_timesteps + iter_interval - 1) // iter_interval + + for i in range(n_chunks): + start = i * iter_interval + end = min(start + iter_interval, n_timesteps) + + # choose initial state + if i == 0: + init_state = targets[:, 0, :] + else: + init_state = preds_all[:, start - 1, :] + + # build dummy dataset: only first slice matters for prepare_data + ds = np.zeros((n_samples, iter_interval, n_quantities)) + ds[:, 0, :] = init_state + + # only need the "train" loader for prediction + dt = timesteps[:iter_interval] + train_loader, _, _ = model.prepare_data( + dataset_train=ds, + dataset_test=None, + dataset_val=None, + timesteps=dt, + batch_size=batch_size, + shuffle=False, + dataset_train_params=None, + dataset_test_params=None, + dataset_val_params=None, + dummy_timesteps=True, + ) + + # predict this chunk and insert into the global array + preds_chunk, _ = model.predict(data_loader=train_loader) + preds_all[:, start:end, :] = preds_chunk[:, : end - start, :] + + # compute error metrics + errors = preds_all - targets + abs_errors = np.abs(errors) + mse = float(np.mean(errors**2)) + mae = float(np.mean(abs_errors)) + + thresh = float(conf.get("relative_error_threshold", 0.0)) + rel_errors = abs_errors / np.maximum(np.abs(targets), thresh) + + errors = np.mean(np.abs(preds_all - targets), axis=(1, 2)) + example_idx = int(np.argsort(np.abs(errors - np.median(errors)))[0]) + + plot_example_iterative_predictions( + surr_name, + conf, + preds_all, + targets, + timesteps, + conf["iterative"]["interval"], + example_idx=example_idx, + labels=labels, + save=True, + show_title=TITLE, + ) + + return { + "mean_squared_error": mse, + "mean_absolute_error": mae, + "mean_relative_error": float(np.mean(rel_errors)), + "median_relative_error": float(np.median(rel_errors)), + "max_relative_error": float(np.max(rel_errors)), + "min_relative_error": float(np.min(rel_errors)), + "absolute_errors": abs_errors, + "relative_errors": rel_errors, + } + + def evaluate_dynamic_accuracy( model, surr_name: str, diff --git a/codes/benchmark/bench_plots.py b/codes/benchmark/bench_plots.py index 371e6ff4..b21376d2 100644 --- a/codes/benchmark/bench_plots.py +++ b/codes/benchmark/bench_plots.py @@ -503,6 +503,98 @@ def plot_example_mode_predictions( plt.close() +def plot_example_iterative_predictions( + surr_name: str, + conf: dict, + preds: np.ndarray, + targets: np.ndarray, + timesteps: np.ndarray, + iter_interval: int, + example_idx: int | None = None, + num_quantities: int = 100, + labels: list[str] | None = None, + save: bool = False, + show_title: bool = True, +) -> None: + """ + Plot one sample's full iterative trajectory: + ground truth vs. chained predictions, with retrigger lines. + """ + # choose example if not given + if example_idx is None: + errors = np.mean(np.abs(preds - targets), axis=(1, 2)) + example_idx = int(np.argsort(np.abs(errors - np.median(errors)))[0]) + + n_q = min(preds.shape[2], num_quantities) + per_plot = 10 + n_plots = int(np.ceil(n_q / per_plot)) + colors = plt.cm.viridis(np.linspace(0, 1, n_q)) + + fig = plt.figure(figsize=(6, 4 * n_plots)) + gs = GridSpec(n_plots, 1, figure=fig) + + for pi in range(n_plots): + ax = fig.add_subplot(gs[pi]) + start, end = pi * per_plot, min((pi + 1) * per_plot, n_q) + for qi in range(start, end): + c = colors[qi] + gt = targets[example_idx, :, qi] + pr = preds[example_idx, :, qi] + ax.plot(timesteps, gt, "--", color=c) + ax.plot(timesteps, pr, "-", color=c) + # retrigger lines + for t in timesteps[::iter_interval]: + ax.axvline(x=t, linestyle=":", linewidth=0.8, alpha=0.7) + if conf.get("dataset", {}).get("log10_transform", False): + ax.set_yscale("log") + ax.set_xlim(timesteps.min(), timesteps.max()) + if conf["dataset"].get("log_timesteps", False): + ax.set_xscale("log") + ax.set_ylabel("Abundance") + if labels is not None: + legend_lines = [ + plt.Line2D([0], [0], color=colors[i]) for i in range(start, end) + ] + ax.legend( + legend_lines, + labels[start:end], + loc="center left", + bbox_to_anchor=(1, 0.5), + ) + + fig.text(0.5, 0.04, "Time", ha="center", va="center", fontsize=12) + + handles = [ + plt.Line2D([0], [0], color="black", linestyle="--", label="Ground Truth"), + plt.Line2D([0], [0], color="black", linestyle="-", label="Prediction"), + ] + pos = 0.95 - (0.06 / n_plots) + fig.legend( + handles, + ["Ground Truth", "Prediction"], + loc="upper center", + bbox_to_anchor=(0.5, pos), + ncol=2, + fontsize="small", + ) + fig.align_ylabels() + + if show_title: + title = ( + f"Iterative Prediction Example ({surr_name})\n" + f"Sample {example_idx}, Interval {iter_interval}" + ) + plt.suptitle(title, y=0.97) + + plt.tight_layout(rect=[0.05, 0.03, 0.95, 0.92]) + + if save and conf: + fname = f"iterative_example_{surr_name}.png" + save_plot(plt, fname, conf, surr_name) + + plt.close() + + def plot_example_predictions_with_uncertainty( surr_name: str, conf: dict, diff --git a/codes/surrogates/DeepONet/deeponet.py b/codes/surrogates/DeepONet/deeponet.py index 3cc21d72..9c633915 100644 --- a/codes/surrogates/DeepONet/deeponet.py +++ b/codes/surrogates/DeepONet/deeponet.py @@ -419,21 +419,23 @@ def prepare_data( dataset_train, timesteps, batch_size, - True, + shuffle=shuffle, dataset_params=dataset_train_params, params_in_branch=self.config.params_branch, num_workers=nw, ) - test_loader = self.create_dataloader( - dataset_test, - timesteps, - batch_size, - False, - dataset_params=dataset_test_params, - params_in_branch=self.config.params_branch, - num_workers=nw, - ) + test_loader = None + if dataset_test is not None: + test_loader = self.create_dataloader( + dataset_test, + timesteps, + batch_size, + False, + dataset_params=dataset_test_params, + params_in_branch=self.config.params_branch, + num_workers=nw, + ) val_loader = None if dataset_val is not None: diff --git a/codes/surrogates/FCNN/fcnn.py b/codes/surrogates/FCNN/fcnn.py index 5d7630e6..41fcb612 100644 --- a/codes/surrogates/FCNN/fcnn.py +++ b/codes/surrogates/FCNN/fcnn.py @@ -186,8 +186,9 @@ def prepare_data( num_workers=nw, ) - test_loader = ( - self.create_dataloader( + test_loader = None + if dataset_test is not None: + test_loader = self.create_dataloader( dataset_test, timesteps, batch_size, @@ -195,12 +196,10 @@ def prepare_data( dataset_params=dataset_test_params, num_workers=nw, ) - if dataset_test is not None - else None - ) - val_loader = ( - self.create_dataloader( + val_loader = None + if dataset_val is not None: + val_loader = self.create_dataloader( dataset_val, timesteps, batch_size, @@ -208,9 +207,6 @@ def prepare_data( dataset_params=dataset_val_params, num_workers=nw, ) - if dataset_val is not None - else None - ) return train_loader, test_loader, val_loader From be60abd0cd28955e1de71a8f8a432e205f806668 Mon Sep 17 00:00:00 2001 From: Robin Date: Tue, 29 Jul 2025 15:02:05 +0200 Subject: [PATCH 02/13] update config files to run older models --- datasets/osu2008/surrogates_config.py | 1 + datasets/simple_ode/surrogates_config.py | 1 + datasets/simple_reaction/surrogates_config.py | 1 + 3 files changed, 3 insertions(+) diff --git a/datasets/osu2008/surrogates_config.py b/datasets/osu2008/surrogates_config.py index 0287dfe2..888067b1 100644 --- a/datasets/osu2008/surrogates_config.py +++ b/datasets/osu2008/surrogates_config.py @@ -48,3 +48,4 @@ class LatentPolyConfig: learning_rate: float = 0.0004 # 0.001 layers_factor: int = 84 activation: nn.Module = nn.ReLU() + model_version: str = "v1" diff --git a/datasets/simple_ode/surrogates_config.py b/datasets/simple_ode/surrogates_config.py index cdfc783a..d773a4b1 100644 --- a/datasets/simple_ode/surrogates_config.py +++ b/datasets/simple_ode/surrogates_config.py @@ -48,3 +48,4 @@ class LatentPolyConfig: learning_rate: float = 0.002 layers_factor: int = 64 activation: nn.Module = nn.LeakyReLU() + model_version: str = "v1" diff --git a/datasets/simple_reaction/surrogates_config.py b/datasets/simple_reaction/surrogates_config.py index 289f1ebf..417abc2f 100644 --- a/datasets/simple_reaction/surrogates_config.py +++ b/datasets/simple_reaction/surrogates_config.py @@ -48,3 +48,4 @@ class LatentPolyConfig: learning_rate: float = 0.0005 layers_factor: int = 159 activation: nn.Module = nn.LeakyReLU() + model_version: str = "v1" From 0d814bf099ce01b8c5fff606de8d7bad1c3141b0 Mon Sep 17 00:00:00 2001 From: Robin Date: Tue, 29 Jul 2025 15:03:10 +0200 Subject: [PATCH 03/13] ensure correct timesteps shape in prepare_data --- codes/surrogates/DeepONet/deeponet.py | 4 ++ codes/surrogates/FCNN/fcnn.py | 4 ++ .../LatentNeuralODE/latent_neural_ode.py | 53 ++++++++++++++----- .../LatentPolynomial/latent_poly.py | 4 ++ 4 files changed, 51 insertions(+), 14 deletions(-) diff --git a/codes/surrogates/DeepONet/deeponet.py b/codes/surrogates/DeepONet/deeponet.py index 9c633915..835b024f 100644 --- a/codes/surrogates/DeepONet/deeponet.py +++ b/codes/surrogates/DeepONet/deeponet.py @@ -413,6 +413,10 @@ def prepare_data( if dummy_timesteps: timesteps = np.linspace(0, 1, dataset_train.shape[1]) + assert ( + timesteps.shape[0] == dataset_train.shape[1] + ), "Number of timesteps in timesteps array and dataset must match." + nw = getattr(self.config, "num_workers", 0) train_loader = self.create_dataloader( diff --git a/codes/surrogates/FCNN/fcnn.py b/codes/surrogates/FCNN/fcnn.py index 41fcb612..7c239291 100644 --- a/codes/surrogates/FCNN/fcnn.py +++ b/codes/surrogates/FCNN/fcnn.py @@ -175,6 +175,10 @@ def prepare_data( if dummy_timesteps: timesteps = np.linspace(0, 1, dataset_train.shape[1]) + assert ( + timesteps.shape[0] == dataset_train.shape[1] + ), "Number of timesteps in timesteps array and dataset must match." + nw = getattr(self.config, "num_workers", 0) train_loader = self.create_dataloader( diff --git a/codes/surrogates/LatentNeuralODE/latent_neural_ode.py b/codes/surrogates/LatentNeuralODE/latent_neural_ode.py index 83746c72..c426de91 100644 --- a/codes/surrogates/LatentNeuralODE/latent_neural_ode.py +++ b/codes/surrogates/LatentNeuralODE/latent_neural_ode.py @@ -87,6 +87,10 @@ def prepare_data( if dummy_timesteps: timesteps = np.linspace(0, 1, dataset_train.shape[1]) + assert ( + timesteps.shape[0] == dataset_train.shape[1] + ), "Number of timesteps in timesteps array and dataset must match." + nw = getattr(self.config, "num_workers", 0) train_loader = self.create_dataloader( @@ -266,13 +270,26 @@ def __init__(self, config, n_quantities: int, n_parameters: int = 0): enc_in = n_quantities + n_parameters else: enc_in = n_quantities - self.encoder = Encoder( - in_features=enc_in, - latent_features=latent_dim, - coder_layers=config.coder_layers, - coder_width=config.coder_width, - activation=config.activation, - ) + if self.config.model_version == "v1": + self.encoder = OldEncoder( + in_features=enc_in, + latent_features=latent_dim, + layers_factor=self.config.layers_factor, + activation=self.config.activation, + ) + elif self.config.model_version == "v2": + self.encoder = Encoder( + in_features=enc_in, + latent_features=latent_dim, + coder_layers=config.coder_layers, + coder_width=config.coder_width, + activation=config.activation, + ) + else: + raise ValueError( + f"Unknown model version {self.config.model_version}. " + "Supported versions: 'v1', 'v2'." + ) # --- Build ODE --- if n_parameters == 0 or config.encode_params: @@ -305,13 +322,21 @@ def __init__(self, config, n_quantities: int, n_parameters: int = 0): self.solver = to.AutoDiffAdjoint(step, ctrl) # --- Build decoder --- - self.decoder = Decoder( - out_features=n_quantities, - latent_features=latent_dim, - coder_layers=config.coder_layers, - coder_width=config.coder_width, - activation=config.activation, - ) + if self.config.model_version == "v1": + self.decoder = OldDecoder( + out_features=n_quantities, + latent_features=latent_dim, + layers_factor=self.config.layers_factor, + activation=self.config.activation, + ) + elif self.config.model_version == "v2": + self.decoder = Decoder( + out_features=n_quantities, + latent_features=latent_dim, + coder_layers=config.coder_layers, + coder_width=config.coder_width, + activation=config.activation, + ) def forward(self, x0: Tensor, t_range: Tensor, params: Tensor = None): # encode initial state diff --git a/codes/surrogates/LatentPolynomial/latent_poly.py b/codes/surrogates/LatentPolynomial/latent_poly.py index 56216803..534c6baf 100644 --- a/codes/surrogates/LatentPolynomial/latent_poly.py +++ b/codes/surrogates/LatentPolynomial/latent_poly.py @@ -126,6 +126,10 @@ def prepare_data( if dummy_timesteps: timesteps = np.linspace(0, 1, dataset_train.shape[1]) + assert ( + timesteps.shape[0] == dataset_train.shape[1] + ), "Number of timesteps in timesteps array and dataset must match." + nw = getattr(self.config, "num_workers", 0) train_loader = self.create_dataloader( From 56cb6d9d0ef31a756e9219cece023a713a2a0f66 Mon Sep 17 00:00:00 2001 From: Robin Date: Tue, 29 Jul 2025 15:04:29 +0200 Subject: [PATCH 04/13] unrelated tuning improvements/results --- codes/tune/optuna_config.yaml | 100 +++++++++------------- codes/tune/optuna_fcts.py | 124 ++++++++++++++-------------- datasets/cloud/surrogates_config.py | 95 ++++++++++++++------- 3 files changed, 168 insertions(+), 151 deletions(-) diff --git a/codes/tune/optuna_config.yaml b/codes/tune/optuna_config.yaml index 42c6d07f..264404b6 100644 --- a/codes/tune/optuna_config.yaml +++ b/codes/tune/optuna_config.yaml @@ -1,20 +1,20 @@ -tuning_id: cloud_tuning_timetest +tuning_id: cloud_tuning_fine seed: 42 dataset: name: cloud log10_transform: True normalise: minmax - subset_factor: 4 + subset_factor: 1 tolerance: 1e-25 normalise_per_species: True log_timesteps: True -devices: ["cuda:1", "cuda:2", "cuda:3", "cuda:4", "cuda:5", "cuda:6", "cuda:8", "cuda:9"] +devices: ["cuda:2", "cuda:3", "cuda:4", "cuda:5", "cuda:6", "cuda:7", "cuda:8", "cuda:9"] optuna_logs: False prune: True verbose: False multi_objective: True -population_size: 100 +population_size: 30 target_percentile: 0.99 postgres_config: mode: "local" # "local" or "remote" @@ -22,7 +22,7 @@ postgres_config: user: "optuna_user" host: "localhost" # "localhost" for local use database_folder: "/export/home/rjanssen/postgres/" # only for local use - db_name: "optuna_cloud" # remote mode: single DB for all runs + db_name: "optuna_cloud_2" # remote mode: single DB for all runs sslmode: "require" # if needed password: "" # optional; prefer env PGPASSWORD @@ -35,88 +35,67 @@ global_optuna_params: regularization_factor: type: float low: 1.0e-6 - high: 1.0 - log: true - optimizer: - type: categorical - choices: ["AdamW", "SGD"] - momentum: - type: float - low: 0.0 - high: 0.99 - step: 0.01 - scheduler: - type: categorical - choices: ["cosine", "poly", "schedulefree"] - poly_power: - type: float - low: 0.5 - high: 2.0 - step: 0.1 - eta_min: - type: float - low: 1.0e-3 - high: 1.0 + high: 1.0e-2 log: true activation: type: categorical choices: ["ReLU", "LeakyReLU", "PReLU", "Tanh", "GELU", "Mish", "SiLU", "ELU"] - loss_function: - type: categorical - choices: ["mse", "smoothl1"] - beta: - type: float - low: 0.1 - high: 10.0 - log: true surrogates: - name: MultiONet batch_size: 65536 - epochs: 200 - trials: 200 + epochs: 4096 + trials: 100 optuna_params: branch_hidden_layers: type: int low: 1 - high: 10 + high: 6 hidden_size: type: int low: 10 - high: 1000 + high: 500 step: 10 output_factor: type: int - low: 1 - high: 200 + low: 30 + high: 300 + step: 10 trunk_hidden_layers: type: int low: 1 - high: 10 - # params_branch: - # type: categorical - # choices: ["True", "False"] + high: 7 + beta: + type: float + low: 0.1 + high: 10.0 + log: true - name: FullyConnected batch_size: 65536 - epochs: 30 - trials: 200 + epochs: 4096 + trials: 100 optuna_params: hidden_size: type: int low: 10 - high: 1000 + high: 500 step: 10 num_hidden_layers: type: int low: 1 - high: 10 + high: 6 + beta: + type: float + low: 0.1 + high: 10.0 + log: true - name: LatentPoly batch_size: 512 - epochs: 200 - trials: 200 + epochs: 4096 + trials: 100 optuna_params: degree: type: int @@ -129,19 +108,19 @@ surrogates: coder_layers: type: int low: 1 - high: 10 + high: 6 coder_width: type: int low: 10 - high: 1000 - step: 10 + high: 600 + step: 100 # coeff_network: # type: categorical # choices: ["True", "False"] # coeff_width: # type: int # low: 10 - # high: 1000 + # high: 700 # step: 10 # coeff_layers: # type: int @@ -150,13 +129,13 @@ surrogates: - name: LatentNeuralODE batch_size: 1024 - epochs: 200 + epochs: 4096 trials: 200 optuna_params: latent_features: type: int low: 1 - high: 10 + high: 15 coder_layers: type: int low: 1 @@ -164,15 +143,12 @@ surrogates: coder_width: type: int low: 10 - high: 1000 + high: 400 step: 10 - ode_tanh_reg: - type: categorical - choices: ["True", "False"] ode_width: type: int low: 10 - high: 1000 + high: 400 step: 10 ode_layers: type: int diff --git a/codes/tune/optuna_fcts.py b/codes/tune/optuna_fcts.py index c8fc65df..d345b19c 100644 --- a/codes/tune/optuna_fcts.py +++ b/codes/tune/optuna_fcts.py @@ -2,7 +2,6 @@ import os import queue from datetime import datetime -from distutils.util import strtobool import numpy as np import optuna @@ -65,83 +64,88 @@ def _suggest_param(trial: optuna.Trial, name: str, opts: dict): # categorical or bool raw = trial.suggest_categorical(name, opts.get("choices", [])) if isinstance(raw, str) and raw.lower() in ("true", "false"): - return bool(strtobool(raw)) + return raw == "true" return raw def make_optuna_params(trial: optuna.Trial, optuna_params: dict) -> dict: suggested: dict[str, any] = {} - # Sample switch parameters + # Sample any switches that exist for switch in ("scheduler", "optimizer", "coeff_network", "loss_function"): - if switch not in optuna_params: - continue - suggested[switch] = _suggest_param(trial, switch, optuna_params[switch]) - - # Sample conditional parameters - # scheduler - poly_power or eta_min - sched = suggested.get("scheduler") - if sched == "poly": - suggested["poly_power"] = _suggest_param( - trial, "poly_power", optuna_params["poly_power"] - ) - elif sched == "cosine": - suggested["eta_min"] = _suggest_param( - trial, "eta_min", optuna_params["eta_min"] - ) - - # optimizer - momentum if SGD - optm = suggested.get("optimizer") - if isinstance(optm, str) and optm.lower() == "sgd": - suggested["momentum"] = _suggest_param( - trial, "momentum", optuna_params["momentum"] - ) - - # coeff_network - coeff_width, coeff_layers - coeff_net = suggested.get("coeff_network") - if bool(coeff_net): - suggested["coeff_width"] = _suggest_param( - trial, "coeff_width", optuna_params["coeff_width"] - ) - suggested["coeff_layers"] = _suggest_param( - trial, "coeff_layers", optuna_params["coeff_layers"] - ) - - # loss_function - beta if smoothl1 - lf = suggested.get("loss_function") - if isinstance(lf, str) and lf.lower() == "smoothl1": - suggested["beta"] = _suggest_param(trial, "beta", optuna_params["beta"]) - - # Sample independent parameters - excluded = { - "scheduler", - "poly_power", - "eta_min", - "optimizer", - "momentum", - "coeff_network", - "coeff_width", - "coeff_layers", - "loss_function", - "beta", + if switch in optuna_params: + suggested[switch] = _suggest_param(trial, switch, optuna_params[switch]) + + # Child‐sampling rules + # mapping: switch to its conditional children + mapping = { + "scheduler": ("poly_power", "eta_min"), + "optimizer": ("momentum",), + "coeff_network": ("coeff_width", "coeff_layers"), + "loss_function": ("beta",), } + for switch, children in mapping.items(): + if switch in optuna_params: + # switch _was_ sampled: sample each child only if the switch value demands it + val = suggested.get(switch) + if switch == "scheduler": + if val == "poly" and "poly_power" in optuna_params: + suggested["poly_power"] = _suggest_param( + trial, "poly_power", optuna_params["poly_power"] + ) + elif val == "cosine" and "eta_min" in optuna_params: + suggested["eta_min"] = _suggest_param( + trial, "eta_min", optuna_params["eta_min"] + ) + elif switch == "optimizer": + if ( + isinstance(val, str) + and val.lower() == "sgd" + and "momentum" in optuna_params + ): + suggested["momentum"] = _suggest_param( + trial, "momentum", optuna_params["momentum"] + ) + elif switch == "coeff_network": + if bool(val): + for child in ("coeff_width", "coeff_layers"): + if child in optuna_params: + suggested[child] = _suggest_param( + trial, child, optuna_params[child] + ) + elif switch == "loss_function": + if ( + isinstance(val, str) + and val.lower() == "smoothl1" + and "beta" in optuna_params + ): + suggested["beta"] = _suggest_param( + trial, "beta", optuna_params["beta"] + ) + + else: + # switch _not_ in config, but user still might want to tune child directly + for child in children: + if child in optuna_params: + suggested[child] = _suggest_param( + trial, child, optuna_params[child] + ) + + # Sample everything else exactly once + excluded = set(mapping.keys()) | {c for kids in mapping.values() for c in kids} for name, opts in optuna_params.items(): if name in excluded: continue suggested[name] = _suggest_param(trial, name, opts) - # map activation and loss_function to actual callables + # Map activation & loss_function strings to actual nn classes/instances for name, val in list(suggested.items()): if "activation" in name.lower(): - cls = MODULE_REGISTRY.get(val.lower()) - if cls is None: - raise ValueError(f"Unknown activation: {val}") + cls = MODULE_REGISTRY[val.lower()] suggested[name] = cls() elif name == "loss_function": - cls = MODULE_REGISTRY.get(val.lower()) - if cls is None: - raise ValueError(f"Unknown loss function: {val}") + cls = MODULE_REGISTRY[val.lower()] suggested[name] = cls() return suggested diff --git a/datasets/cloud/surrogates_config.py b/datasets/cloud/surrogates_config.py index c3ce8ed4..d3046f87 100644 --- a/datasets/cloud/surrogates_config.py +++ b/datasets/cloud/surrogates_config.py @@ -7,50 +7,87 @@ class MultiONetConfig: """Model config for MultiONet for the simple_ode dataset""" - # cloud, trial 69 - branch_hidden_layers: int = 1 - trunk_hidden_layers: int = 9 - hidden_size: int = 225 - output_factor: int = 63 - learning_rate: float = 4e-5 # optimal for ~4000 epochs - activation: nn.Module = nn.Tanh() + loss_function: nn.Module = nn.SmoothL1Loss() + optimizer: str = "adamw" + scheduler: str = "schedulefree" @dataclass class LatentNeuralODEConfig: """Model config for LatentNeuralODE for the simple_ode dataset""" - # cloud, trial 63 - latent_features: int = 6 - coder_layers: int = 2 - coder_width: int = 103 - learning_rate: float = 3e-4 - ode_layers: int = 4 - ode_width: int = 197 - ode_tanh_reg: bool = True - activation: nn.Module = nn.SiLU() - model_version: str = "v2" + loss_function: nn.Module = nn.MSELoss() + optimizer: str = "adamw" + scheduler: str = "schedulefree" + ode_tanh_reg: bool = False @dataclass class FullyConnectedConfig: """Model config for FullyConnected for the simple_ode dataset""" - # cloud, trial 44 - hidden_size: int = 261 - num_hidden_layers: int = 1 - learning_rate: float = 1e-4 - activation: nn.Module = nn.LeakyReLU() + loss_function: nn.Module = nn.SmoothL1Loss() + optimizer: str = "adamw" + scheduler: str = "schedulefree" @dataclass class LatentPolyConfig: """Model config for LatentPoly for the simple_ode dataset""" - # cloud, trial 92 - latent_features: int = 9 - degree: int = 5 - learning_rate: float = 3e-4 - coder_layers: int = 1 - coder_width: int = 86 - activation: nn.Module = nn.Mish() + loss_function: nn.Module = nn.MSELoss() + optimizer: str = "adamw" + scheduler: str = "schedulefree" + + +# @dataclass +# class MultiONetConfig: +# """Model config for MultiONet for the simple_ode dataset""" + +# # cloud, trial 69 +# branch_hidden_layers: int = 1 +# trunk_hidden_layers: int = 9 +# hidden_size: int = 225 +# output_factor: int = 63 +# learning_rate: float = 4e-5 # optimal for ~4000 epochs +# activation: nn.Module = nn.Tanh() + + +# @dataclass +# class LatentNeuralODEConfig: +# """Model config for LatentNeuralODE for the simple_ode dataset""" + +# # cloud, trial 63 +# latent_features: int = 6 +# coder_layers: int = 2 +# coder_width: int = 103 +# learning_rate: float = 3e-4 +# ode_layers: int = 4 +# ode_width: int = 197 +# ode_tanh_reg: bool = True +# activation: nn.Module = nn.SiLU() +# model_version: str = "v2" + + +# @dataclass +# class FullyConnectedConfig: +# """Model config for FullyConnected for the simple_ode dataset""" + +# # cloud, trial 44 +# hidden_size: int = 261 +# num_hidden_layers: int = 1 +# learning_rate: float = 1e-4 +# activation: nn.Module = nn.LeakyReLU() + + +# @dataclass +# class LatentPolyConfig: +# """Model config for LatentPoly for the simple_ode dataset""" + +# # cloud, trial 92 +# latent_features: int = 9 +# degree: int = 5 +# learning_rate: float = 3e-4 +# coder_layers: int = 1 +# coder_width: int = 86 +# activation: nn.Module = nn.Mish() From 2a2f678af356e8d060832e1251038b415539f32f Mon Sep 17 00:00:00 2001 From: Robin Date: Tue, 29 Jul 2025 15:05:24 +0200 Subject: [PATCH 05/13] implement iterative prediction evaluation --- codes/benchmark/__init__.py | 1 + codes/benchmark/bench_fcts.py | 72 ++++++++++++------- codes/benchmark/bench_plots.py | 33 +++++---- .../AbstractSurrogate/abstract_surrogate.py | 42 ++++++----- config.yaml | 42 +++++------ 5 files changed, 114 insertions(+), 76 deletions(-) diff --git a/codes/benchmark/__init__.py b/codes/benchmark/__init__.py index 64d2e72a..e1f4be85 100644 --- a/codes/benchmark/__init__.py +++ b/codes/benchmark/__init__.py @@ -107,6 +107,7 @@ "plot_average_errors_over_time", "plot_example_predictions_with_uncertainty", "plot_example_mode_predictions", + "plot_example_iterative_predictions", "plot_average_uncertainty_over_time", "plot_uncertainty_vs_errors", "plot_uncertainty_confidence", diff --git a/codes/benchmark/bench_fcts.py b/codes/benchmark/bench_fcts.py index a10ada29..daccfbab 100644 --- a/codes/benchmark/bench_fcts.py +++ b/codes/benchmark/bench_fcts.py @@ -141,7 +141,7 @@ def run_benchmark(surr_name: str, surrogate_class, conf: dict) -> dict[str, Any] if conf["iterative"]: # Iterative training benchmark print("Running iterative training benchmark...") - metrics["iterative"] = iterative_training_benchmark( + metrics["iterative"] = evaluate_iterative_predictions( model, surr_name, timesteps, val_loader, conf, labels ) @@ -215,18 +215,18 @@ def evaluate_accuracy( labels: list | None = None, ) -> dict[str, Any]: """ - Evaluate the accuracy of the surrogate model. - quantitiesquantities - Args: - model: Instance of the surrogate model class. - surr_name (str): The name of the surrogate model. - timesteps (np.ndarray): The timesteps array. - test_loader (DataLoader): The DataLoader object containing the test data. - conf (dict): The configuration dictionary. - labels (list, optional): The labels for the quantities. - - Returns: - dict: A dictionary containing accuracy metrics. + Evaluate the accuracy of the surrogate model. + + Args: + model: Instance of the surrogate model class. + surr_name (str): The name of the surrogate model. + timesteps (np.ndarray): The timesteps array. + test_loader (DataLoader): The DataLoader object containing the test data. + conf (dict): The configuration dictionary. + labels (list, optional): The labels for the quantities. + + Returns: + dict: A dictionary containing accuracy metrics. """ training_id = conf["training_id"] @@ -289,7 +289,7 @@ def evaluate_accuracy( return accuracy_metrics -def iterative_training_benchmark( +def evaluate_iterative_predictions( model, surr_name: str, timesteps: np.ndarray, @@ -298,7 +298,7 @@ def iterative_training_benchmark( labels: list | None = None, ) -> dict[str, Any]: """ - Benchmark error accumulation when running the model iteratively in chunks. + Evaluate the iterative predictions of the surrogate model. Returns the same set of error metrics as evaluate_accuracy, but over the full trajectory built by re-feeding the last prediction as the next initial state. @@ -308,9 +308,14 @@ def iterative_training_benchmark( model.load(training_id, surr_name, model_identifier=f"{surr_name.lower()}_main") # get full ground truth (targets) and ignore one-shot preds - _, targets = model.predict(data_loader=val_loader) + full_preds, targets = model.predict( + data_loader=val_loader, leave_log=True, leave_norm=True + ) + targets = targets.detach().cpu().numpy() n_samples, n_timesteps, n_quantities = targets.shape + original_n_timesteps = model.n_timesteps + # how many timesteps per chunk iter_interval = 10 # conf["iterative"]["interval"] # batch size same as in run_benchmark @@ -321,7 +326,7 @@ def iterative_training_benchmark( batch_size = conf["batch_size"] # container for the piecewise predictions - preds_all = np.zeros_like(targets) + iterative_preds = np.zeros_like(targets) # number of chunks n_chunks = (n_timesteps + iter_interval - 1) // iter_interval @@ -329,19 +334,22 @@ def iterative_training_benchmark( for i in range(n_chunks): start = i * iter_interval end = min(start + iter_interval, n_timesteps) + model.n_timesteps = ( + end - start + 1 + ) # set the number of timesteps for this chunk # choose initial state if i == 0: init_state = targets[:, 0, :] else: - init_state = preds_all[:, start - 1, :] + init_state = iterative_preds[:, start - 1, :] # build dummy dataset: only first slice matters for prepare_data - ds = np.zeros((n_samples, iter_interval, n_quantities)) + ds = np.zeros((n_samples, model.n_timesteps, n_quantities)) ds[:, 0, :] = init_state # only need the "train" loader for prediction - dt = timesteps[:iter_interval] + dt = timesteps[: model.n_timesteps] train_loader, _, _ = model.prepare_data( dataset_train=ds, dataset_test=None, @@ -356,11 +364,19 @@ def iterative_training_benchmark( ) # predict this chunk and insert into the global array - preds_chunk, _ = model.predict(data_loader=train_loader) - preds_all[:, start:end, :] = preds_chunk[:, : end - start, :] + preds_chunk, _ = model.predict( + data_loader=train_loader, leave_log=True, leave_norm=True + ) + iterative_preds[:, start:end, :] = ( + preds_chunk[:, 1 : model.n_timesteps, :].detach().cpu().numpy() + ) + + iterative_preds = model.denormalize(iterative_preds) + full_preds = model.denormalize(full_preds.detach().cpu().numpy()) + targets = model.denormalize(targets) # compute error metrics - errors = preds_all - targets + errors = iterative_preds - targets abs_errors = np.abs(errors) mse = float(np.mean(errors**2)) mae = float(np.mean(abs_errors)) @@ -368,16 +384,20 @@ def iterative_training_benchmark( thresh = float(conf.get("relative_error_threshold", 0.0)) rel_errors = abs_errors / np.maximum(np.abs(targets), thresh) - errors = np.mean(np.abs(preds_all - targets), axis=(1, 2)) + errors = np.mean(np.abs(iterative_preds - targets), axis=(1, 2)) example_idx = int(np.argsort(np.abs(errors - np.median(errors)))[0]) + # Restore original number of timesteps + model.n_timesteps = original_n_timesteps + plot_example_iterative_predictions( surr_name, conf, - preds_all, + iterative_preds, + full_preds, targets, timesteps, - conf["iterative"]["interval"], + iter_interval=iter_interval, example_idx=example_idx, labels=labels, save=True, diff --git a/codes/benchmark/bench_plots.py b/codes/benchmark/bench_plots.py index b21376d2..f0b42831 100644 --- a/codes/benchmark/bench_plots.py +++ b/codes/benchmark/bench_plots.py @@ -394,7 +394,7 @@ def plot_example_mode_predictions( num_quantities = preds.shape[2] # Define the color palette for plotting quantities - colors = plt.cm.viridis(np.linspace(0, 1, num_quantities)) + colors = plt.cm.viridis(np.linspace(0, 0.9, num_quantities)) # Create the overall figure and subplots fig = plt.figure(figsize=(6, 4 * num_plots)) @@ -506,7 +506,8 @@ def plot_example_mode_predictions( def plot_example_iterative_predictions( surr_name: str, conf: dict, - preds: np.ndarray, + iterative_preds: np.ndarray, + full_preds: np.ndarray, targets: np.ndarray, timesteps: np.ndarray, iter_interval: int, @@ -522,13 +523,13 @@ def plot_example_iterative_predictions( """ # choose example if not given if example_idx is None: - errors = np.mean(np.abs(preds - targets), axis=(1, 2)) + errors = np.mean(np.abs(iterative_preds - targets), axis=(1, 2)) example_idx = int(np.argsort(np.abs(errors - np.median(errors)))[0]) - n_q = min(preds.shape[2], num_quantities) - per_plot = 10 + n_q = min(iterative_preds.shape[2], num_quantities) + per_plot = min(10, n_q) n_plots = int(np.ceil(n_q / per_plot)) - colors = plt.cm.viridis(np.linspace(0, 1, n_q)) + colors = plt.cm.viridis(np.linspace(0, 0.9, per_plot)) fig = plt.figure(figsize=(6, 4 * n_plots)) gs = GridSpec(n_plots, 1, figure=fig) @@ -537,11 +538,13 @@ def plot_example_iterative_predictions( ax = fig.add_subplot(gs[pi]) start, end = pi * per_plot, min((pi + 1) * per_plot, n_q) for qi in range(start, end): - c = colors[qi] + c = colors[qi % per_plot] gt = targets[example_idx, :, qi] - pr = preds[example_idx, :, qi] + pr = iterative_preds[example_idx, :, qi] ax.plot(timesteps, gt, "--", color=c) ax.plot(timesteps, pr, "-", color=c) + init_pr = full_preds[example_idx, :, qi] + ax.plot(timesteps, init_pr, ":", color=c) # retrigger lines for t in timesteps[::iter_interval]: ax.axvline(x=t, linestyle=":", linewidth=0.8, alpha=0.7) @@ -553,7 +556,8 @@ def plot_example_iterative_predictions( ax.set_ylabel("Abundance") if labels is not None: legend_lines = [ - plt.Line2D([0], [0], color=colors[i]) for i in range(start, end) + plt.Line2D([0], [0], color=colors[i % per_plot]) + for i in range(start, end) ] ax.legend( legend_lines, @@ -566,12 +570,15 @@ def plot_example_iterative_predictions( handles = [ plt.Line2D([0], [0], color="black", linestyle="--", label="Ground Truth"), - plt.Line2D([0], [0], color="black", linestyle="-", label="Prediction"), + plt.Line2D( + [0], [0], color="black", linestyle="-", label="Iterative Prediction" + ), + plt.Line2D([0], [0], color="black", linestyle=":", label="Full Prediction"), ] pos = 0.95 - (0.06 / n_plots) fig.legend( handles, - ["Ground Truth", "Prediction"], + ["Ground Truth", "Iterative Prediction", "Full Prediction"], loc="upper center", bbox_to_anchor=(0.5, pos), ncol=2, @@ -589,7 +596,7 @@ def plot_example_iterative_predictions( plt.tight_layout(rect=[0.05, 0.03, 0.95, 0.92]) if save and conf: - fname = f"iterative_example_{surr_name}.png" + fname = "iterative_example_preds.png" save_plot(plt, fname, conf, surr_name) plt.close() @@ -632,7 +639,7 @@ def plot_example_predictions_with_uncertainty( num_plots = int(np.ceil(num_quantities / quantities_per_plot)) # Define the color palette - colors = plt.cm.viridis(np.linspace(0, 1, quantities_per_plot)) + colors = plt.cm.viridis(np.linspace(0, 0.9, quantities_per_plot)) # Create subplots fig = plt.figure(figsize=(8, 4 * num_plots)) diff --git a/codes/surrogates/AbstractSurrogate/abstract_surrogate.py b/codes/surrogates/AbstractSurrogate/abstract_surrogate.py index dcd5aedc..831dc923 100644 --- a/codes/surrogates/AbstractSurrogate/abstract_surrogate.py +++ b/codes/surrogates/AbstractSurrogate/abstract_surrogate.py @@ -215,7 +215,7 @@ def fit( pass def predict( - self, data_loader: DataLoader, leave_log: bool = False + self, data_loader: DataLoader, leave_log: bool = False, leave_norm: bool = False ) -> tuple[Tensor, Tensor]: """ Evaluate the model on the given dataloader. @@ -224,6 +224,7 @@ def predict( data_loader (DataLoader): The DataLoader object containing the data the model is evaluated on. leave_log (bool): If True, do not exponentiate the data even if log10_transform is True. + leave_norm (bool): If True, do not denormalize the data even if normalisation is applied. Returns: tuple[Tensor, Tensor]: The predictions and targets. @@ -265,8 +266,8 @@ def predict( predictions = predictions[:processed_samples, ...] targets = targets[:processed_samples, ...] - predictions = self.denormalize(predictions, leave_log=leave_log) - targets = self.denormalize(targets, leave_log=leave_log) + predictions = self.denormalize(predictions, leave_log, leave_norm) + targets = self.denormalize(targets, leave_log, leave_norm) predictions = predictions.reshape(-1, self.n_timesteps, self.n_quantities) targets = targets.reshape(-1, self.n_timesteps, self.n_quantities) @@ -428,29 +429,36 @@ def setup_progress_bar(self, epochs: int, position: int, description: str): return progress_bar - def denormalize(self, data: Tensor, leave_log: bool = False) -> Tensor: + def denormalize( + self, + data: Tensor | np.ndarray, + leave_log: bool = False, + leave_norm: bool = False, + ) -> Tensor | np.ndarray: """ Denormalize the data. Args: - data (np.ndarray): The data to denormalize. + data (Tensor | np.ndarray): The data to denormalize. leave_log (bool): If True, do not exponentiate the data even if log10_transform is True. + leave_norm (bool): If True, do not denormalize the data even if normalisation is applied. Returns: - np.ndarray: The denormalized data. + Tensor | np.ndarray: The denormalized data. """ if self.normalisation is not None: - if self.normalisation["mode"] == "disabled": - ... - elif self.normalisation["mode"] == "minmax": - dmax = self.normalisation["max"] - dmin = self.normalisation["min"] - data = data.to("cpu") - data = (data + 1) * (dmax - dmin) / 2 + dmin - elif self.normalisation["mode"] == "standardize": - mean = self.normalisation["mean"] - std = self.normalisation["std"] - data = data * std + mean + if not leave_norm: + if self.normalisation["mode"] == "disabled": + ... + elif self.normalisation["mode"] == "minmax": + dmax = self.normalisation["max"] + dmin = self.normalisation["min"] + # data = data.to("cpu") + data = (data + 1) * (dmax - dmin) / 2 + dmin + elif self.normalisation["mode"] == "standardize": + mean = self.normalisation["mean"] + std = self.normalisation["std"] + data = data * std + mean if self.normalisation["log10_transform"] and not leave_log: data = 10**data diff --git a/config.yaml b/config.yaml index 37dab2db..97f40b4e 100644 --- a/config.yaml +++ b/config.yaml @@ -1,41 +1,43 @@ # Global settings for the benchmark -training_id: "cloud_full" -surrogates: ["MultiONet", "FullyConnected", "LatentNeuralODE", "LatentPoly"] -batch_size: [65536, 65536, 512, 512] -epochs: [20000, 20000, 11000, 20000] # [20000, 7500, 20000, 15000] +training_id: "simple_reaction_final" +surrogates: ["MultiONet", "FullyConnected", "LatentPoly", "LatentNeuralODE", ] +batch_size: [4096, 4096, 256, 256] +epochs: [12000, 10000, 10000, 7000] dataset: - name: "cloud" + name: "simple_reaction" + log_timesteps: False log10_transform: True - log10_transform_params: False normalise: "minmax" # "minmax" # "standardise", "minmax", "disable" use_optimal_params: True - tolerance: 1e-25 + tolerance: 1e-30 subset_factor: 1 - log_timesteps: True -devices: ["cuda:2", "cuda:3", "cuda:4", "cuda:5", "cuda:6", "cuda:7", "cuda:8", "cuda:9"] +devices: ["cuda:1"] seed: 42 verbose: False # Models to train interpolation: - enabled: True + enabled: False intervals: [2, 3, 4, 5, 6, 7, 8, 9, 10] extrapolation: - enabled: True + enabled: False cutoffs: [50, 60, 70, 80, 90] sparse: - enabled: True + enabled: False factors: [2, 4, 8, 16, 32] batch_scaling: - enabled: True - sizes: [1/16, 1/8, 1/4, 1/2] + enabled: False + sizes: [64, 512, 1024, 4096] uncertainty: - enabled: True + enabled: False ensemble_size: 5 # Number of models for deep ensemble # Evaluations during benchmark -losses: True -gradients: True -timing: True -compute: True -compare: True # Whether to compare the surrogates +iterative: True +losses: False +gradients: False +timing: False +compute: False +compare: False # Whether to compare the surrogates + + From d6108e0c5d01a864928279e7be855c320e8a4b1f Mon Sep 17 00:00:00 2001 From: Robin Date: Tue, 29 Jul 2025 15:34:22 +0200 Subject: [PATCH 06/13] add iterative eval test --- test/README.md | 206 ---------------------------------- test/test_bench_main.py | 76 ++++++++++++- test/test_eval_pipeline.py | 8 +- test/test_model_comparison.py | 1 - 4 files changed, 78 insertions(+), 213 deletions(-) delete mode 100644 test/README.md diff --git a/test/README.md b/test/README.md deleted file mode 100644 index 174cb9f2..00000000 --- a/test/README.md +++ /dev/null @@ -1,206 +0,0 @@ -# CODES Benchmark Test Suite - -This directory contains comprehensive unit tests for the CODES benchmark framework, focusing on surrogate models and datasets. - -## Test Files - -### Core Test Modules - -- **`test_surrogate_models.py`** - Comprehensive tests for all surrogate model implementations - - Tests AbstractSurrogateModel interface compliance - - Tests model initialization, training, prediction, and save/load functionality - - Tests all registered surrogate model classes (FCNN, DeepONet, LatentNeuralODE, LatentPolynomial) - -- **`test_datasets.py`** - Comprehensive tests for dataset functionality - - Tests dataset loading from local and remote sources - - Tests data normalization and preprocessing - - Tests dataset creation and validation - - Tests data_sources.yaml configuration - -### Legacy Test Files - -- **`test_data.py`** - Legacy data loading tests (kept for compatibility) -- **`test_surrogates.py`** - Legacy surrogate model tests (kept for compatibility) -- **`test_run.py`** - Integration tests for training/evaluation pipelines - -### Configuration - -- **`conftest.py`** - Pytest configuration and shared fixtures -- **`README.md`** - This documentation file - -## Running Tests - -### Run All Tests -```bash -pytest test/ -``` - -### Run Specific Test Modules -```bash -# Test surrogate models only -pytest test/test_surrogate_models.py - -# Test datasets only -pytest test/test_datasets.py - -# Test legacy functionality -pytest test/test_data.py test/test_surrogates.py -``` - -### Run Tests with Specific Markers -```bash -# Skip slow tests -pytest test/ -m "not slow" - -# Skip download tests (tests that require internet) -pytest test/ -m "not download" - -# Run only GPU tests (if GPU available) -pytest test/ -m "gpu" -``` - -### Run Tests with Coverage -```bash -pytest test/ --cov=codes --cov-report=html -``` - -### Run Tests in Parallel (if pytest-xdist installed) -```bash -pytest test/ -n auto -``` - -## Test Structure - -### Surrogate Model Tests - -The surrogate model tests are organized into several test classes: - -1. **TestAbstractSurrogateModelInterface** - Tests the registry system and interface compliance -2. **TestSurrogateModelInitialization** - Tests model initialization and basic attributes -3. **TestDataPreparation** - Tests data loading and dataloader creation -4. **TestForwardPass** - Tests model forward pass functionality -5. **TestTraining** - Tests model training functionality -6. **TestPrediction** - Tests model prediction functionality -7. **TestSaveLoad** - Tests model serialization and deserialization -8. **TestDenormalization** - Tests data denormalization functionality -9. **TestOptimizer** - Tests optimizer and scheduler setup -10. **TestProgressBar** - Tests progress bar functionality - -### Dataset Tests - -The dataset tests are organized into several test classes: - -1. **TestDataSourcesYaml** - Tests data_sources.yaml configuration -2. **TestDownloadData** - Tests dataset downloading functionality -3. **TestLocalDatasets** - Tests loading of locally available datasets -4. **TestCheckAndLoadData** - Tests the main data loading function -5. **TestCreateDataset** - Tests dataset creation functionality -6. **TestNormalizeData** - Tests data normalization functionality -7. **TestDatasetError** - Tests error handling -8. **TestIntegration** - Integration tests combining multiple operations - -## Test Configuration - -### Fixtures - -The test suite uses several shared fixtures defined in `conftest.py`: - -- `device` - Provides the device (CPU/GPU) to use for testing -- `test_constants` - Provides test constants (dimensions, batch sizes, etc.) -- `random_seed` - Ensures reproducible test results -- `temp_dir` - Provides temporary directories for file operations -- `sample_3d_data` - Provides sample training/test/validation data -- `sample_parameters` - Provides sample parameter arrays -- `mock_normalisation` - Provides mock normalization parameters - -### Markers - -The test suite uses custom pytest markers: - -- `@pytest.mark.slow` - For tests that take longer to run -- `@pytest.mark.download` - For tests that require internet access -- `@pytest.mark.gpu` - For tests that require GPU - -### Parameterization - -Many tests are parameterized to run across: -- All registered surrogate model classes -- All available datasets (local and remote) -- Different normalization modes -- Different configuration options - -## Test Coverage - -The test suite aims for comprehensive coverage of: - -### Surrogate Models -- ✅ Model initialization and configuration -- ✅ Data preparation and dataloader creation -- ✅ Forward pass functionality -- ✅ Training loop execution -- ✅ Prediction and evaluation -- ✅ Model serialization (save/load) -- ✅ Data denormalization -- ✅ Progress tracking and optimization -- ✅ Interface compliance with AbstractSurrogateModel - -### Datasets -- ✅ Data loading from HDF5 files -- ✅ Dataset downloading from remote sources -- ✅ Data validation and structure checking -- ✅ Data normalization (minmax, standardization) -- ✅ Dataset creation and export -- ✅ Parameter handling -- ✅ Error handling and edge cases -- ✅ Integration workflows - -## Development Guidelines - -### Adding New Tests - -1. **For new surrogate models**: Add tests to `test_surrogate_models.py` or create model-specific test files -2. **For new dataset functionality**: Add tests to `test_datasets.py` -3. **For integration tests**: Add to existing integration test classes or create new ones - -### Test Naming Convention - -- Test functions should start with `test_` -- Test classes should start with `Test` -- Use descriptive names that clearly indicate what is being tested - -### Assertion Guidelines - -- Use descriptive assertion messages -- Test both positive and negative cases -- Use appropriate pytest features (parametrize, fixtures, markers) -- Keep tests focused and atomic - -### Performance Considerations - -- Use minimal data sizes for unit tests -- Mark slow tests with `@pytest.mark.slow` -- Skip expensive operations when possible (use mocks/stubs) -- Use temporary directories for file operations - -## Troubleshooting - -### Common Issues - -1. **Import errors**: Ensure CODES package is properly installed -2. **Missing datasets**: Some tests require local datasets - download them first -3. **GPU tests failing**: Ensure CUDA is available or skip GPU tests -4. **Network timeouts**: Skip download tests if running without internet - -### Debug Mode - -Run tests with verbose output: -```bash -pytest test/ -v -s -``` - -### Test Isolation - -Each test should be independent. If tests are interfering with each other: -- Check for global state modifications -- Ensure proper cleanup in fixtures -- Use isolated temporary directories diff --git a/test/test_bench_main.py b/test/test_bench_main.py index 3bd20a7f..04b73c9f 100644 --- a/test/test_bench_main.py +++ b/test/test_bench_main.py @@ -1,6 +1,5 @@ -# test_test_bench_functions.py -import pytest import numpy as np +import pytest import torch import codes.benchmark.bench_fcts as bf @@ -164,3 +163,76 @@ def __iter__(self): assert model.load_calls == [("TID", surr, f"{surr.lower()}_main")] assert out["num_trainable_parameters"] == 12345 assert out["memory_footprint"] is fake_mem + + +def test_evaluate_iterative_predictions(simple_conf, simple_loader, monkeypatch): + import numpy as np + import torch + + import codes.benchmark.bench_fcts as bf + + # stub out the final plot_example_iterative_predictions so it doesn't call save_plot + monkeypatch.setattr( + bf, "plot_example_iterative_predictions", lambda *args, **kwargs: None + ) + + # model with T=12 timesteps and Q=2 quantities + T, Q = 12, 2 + timesteps = np.arange(1.0, T + 1.0) + + class FakeIterModel: + def __init__(self, n_timesteps, n_quantities): + self.n_timesteps = n_timesteps + self.n_quantities = n_quantities + self.load_calls = [] + + def load(self, training_id, surr_name, model_identifier): + self.load_calls.append((training_id, surr_name, model_identifier)) + + def predict(self, *, data_loader, leave_log=None, leave_norm=None): + # preds == targets == ones + shape = (1, self.n_timesteps, self.n_quantities) + ones = torch.ones(shape, dtype=torch.float32) + return ones, ones + + def prepare_data(self, **kwargs): + # only the returned loader is used by predict + return "train_loader", None, None + + def denormalize(self, arr): + return arr + + model = FakeIterModel(n_timesteps=T, n_quantities=Q) + + surr = "SurrA" + simple_conf["surrogates"] = [surr] + simple_conf["batch_size"] = 4 + simple_conf["relative_error_threshold"] = 0.0 + simple_conf["training_id"] = "TID" # ensure load uses this + + metrics = bf.evaluate_iterative_predictions( + model=model, + surr_name=surr, + timesteps=timesteps, + val_loader="dummy_val_loader", + conf=simple_conf, + labels=["q1", "q2"], + ) + + # ensure we loaded the main model + assert model.load_calls == [("TID", surr, f"{surr.lower()}_main")] + + # since preds == targets, all errors should be zero + for key in [ + "mean_squared_error", + "mean_absolute_error", + "mean_relative_error", + "median_relative_error", + "max_relative_error", + "min_relative_error", + ]: + assert metrics[key] == pytest.approx(0.0) + + # array shapes should be (1, T, Q) + assert metrics["absolute_errors"].shape == (1, T, Q) + assert metrics["relative_errors"].shape == (1, T, Q) diff --git a/test/test_eval_pipeline.py b/test/test_eval_pipeline.py index f4090456..99e3a55c 100644 --- a/test/test_eval_pipeline.py +++ b/test/test_eval_pipeline.py @@ -1,12 +1,11 @@ -import pytest from types import SimpleNamespace from unittest.mock import patch + import numpy as np +import pytest import run_eval -from codes.benchmark.bench_fcts import ( - run_benchmark, -) +from codes.benchmark.bench_fcts import run_benchmark @pytest.fixture @@ -26,6 +25,7 @@ def minimal_bench_config(): "gradients": False, "timing": False, "compute": False, + "iterative": False, "interpolation": {"enabled": False}, "extrapolation": {"enabled": False}, "sparse": {"enabled": False}, diff --git a/test/test_model_comparison.py b/test/test_model_comparison.py index 063dc3b2..4bffb392 100644 --- a/test/test_model_comparison.py +++ b/test/test_model_comparison.py @@ -1,4 +1,3 @@ -# test/test_compare_models.py import pytest from codes.benchmark import bench_fcts From b769ffd6a6409e6ca6cb5b38607eed47cde6bdf1 Mon Sep 17 00:00:00 2001 From: Robin Date: Fri, 1 Aug 2025 18:23:18 +0200 Subject: [PATCH 07/13] fix iterative eval error, small bugfixes --- codes/benchmark/__init__.py | 4 +- codes/benchmark/bench_fcts.py | 94 ++++++++--- codes/benchmark/bench_plots.py | 113 +++++++------ .../AbstractSurrogate/abstract_surrogate.py | 6 + codes/tune/optuna_config.yaml | 153 ++++++++++-------- codes/tune/optuna_fcts.py | 14 +- codes/utils/data_utils.py | 2 +- config.yaml | 20 ++- datasets/cloud/surrogates_config.py | 32 ++++ .../cloud_parametric/surrogates_config_old.py | 115 +++++++------ datasets/primordial/surrogates_config_old.py | 56 +++++++ 11 files changed, 381 insertions(+), 228 deletions(-) create mode 100644 datasets/primordial/surrogates_config_old.py diff --git a/codes/benchmark/__init__.py b/codes/benchmark/__init__.py index e1f4be85..30a7686d 100644 --- a/codes/benchmark/__init__.py +++ b/codes/benchmark/__init__.py @@ -35,6 +35,7 @@ plot_error_correlation_heatmap, plot_error_distribution_comparative, plot_error_distribution_per_quantity, + plot_error_percentiles_over_time, plot_example_iterative_predictions, plot_example_mode_predictions, plot_example_predictions_with_uncertainty, @@ -45,7 +46,6 @@ plot_losses, plot_MAE_comparison, plot_relative_errors, - plot_relative_errors_over_time, plot_surr_losses, plot_uncertainty_confidence, plot_uncertainty_over_time_comparison, @@ -101,7 +101,6 @@ "tabular_comparison", "save_plot", "save_plot_counter", - "plot_relative_errors_over_time", "plot_dynamic_correlation", "plot_generalization_errors", "plot_average_errors_over_time", @@ -124,6 +123,7 @@ "plot_error_correlation_heatmap", "plot_dynamic_correlation_heatmap", "plot_error_distribution_comparative", + "plot_error_percentiles_over_time", "plot_comparative_error_correlation_heatmaps", "plot_comparative_dynamic_correlation_heatmaps", "get_custom_palette", diff --git a/codes/benchmark/bench_fcts.py b/codes/benchmark/bench_fcts.py index 7c65817b..fed34c5f 100644 --- a/codes/benchmark/bench_fcts.py +++ b/codes/benchmark/bench_fcts.py @@ -20,6 +20,7 @@ plot_error_correlation_heatmap, plot_error_distribution_comparative, plot_error_distribution_per_quantity, + plot_error_percentiles_over_time, plot_example_iterative_predictions, plot_example_mode_predictions, plot_example_predictions_with_uncertainty, @@ -28,7 +29,6 @@ plot_loss_comparison_equal, plot_loss_comparison_train_duration, plot_relative_errors, - plot_relative_errors_over_time, plot_surr_losses, plot_uncertainty_confidence, plot_uncertainty_over_time_comparison, @@ -224,11 +224,13 @@ def evaluate_accuracy( test_loader (DataLoader): The DataLoader object containing the test data. conf (dict): The configuration dictionary. labels (list, optional): The labels for the quantities. + percentile (int, optional): The percentile for error metrics. Returns: dict: A dictionary containing accuracy metrics. """ training_id = conf["training_id"] + percentile = conf.get("error_percentile", 99) # Load the model model.load(training_id, surr_name, model_identifier=f"{surr_name.lower()}_main") @@ -237,27 +239,56 @@ def evaluate_accuracy( model_index = conf["surrogates"].index(surr_name) n_epochs = conf["epochs"][model_index] - # Use the model's predict method - criterion = torch.nn.MSELoss() - preds, targets = model.predict(data_loader=test_loader) - mean_squared_error = criterion(preds, targets).item() # / torch.numel(preds) + # Obtain log-space predictions and targets + preds, targets = model.predict(data_loader=test_loader, leave_log=True) + preds, targets = preds.detach().cpu().numpy(), targets.detach().cpu().numpy() + + # Compute log-space error metrics + absolute_errors_log = np.abs(preds - targets) + root_mean_squared_error_log = np.sqrt(np.mean(absolute_errors_log**2)) + median_absolute_error_log = np.median(absolute_errors_log) + mean_absolute_error_log = np.mean(absolute_errors_log) + percentile_absolute_error_log = np.percentile(absolute_errors_log, percentile) + + # Obtain real-space predictions and targets + preds, targets = model.predict(data_loader=test_loader, leave_log=False) preds, targets = preds.detach().cpu().numpy(), targets.detach().cpu().numpy() - # Calculate relative errors + # Compute real-space error metrics absolute_errors = np.abs(preds - targets) - mean_absolute_error = np.mean(absolute_errors) + root_mean_squared_error_real = np.sqrt(np.mean(absolute_errors**2)) + median_absolute_error_real = np.median(absolute_errors) + mean_absolute_error_real = np.mean(absolute_errors) + percentile_absolute_error_real = np.percentile(absolute_errors, percentile) + + # Additional real-space errors: Relative error relative_error_threshold = float(conf.get("relative_error_threshold", 0.0)) relative_errors = np.abs( absolute_errors / np.maximum(np.abs(targets), relative_error_threshold) ) + median_relative_error = np.median(relative_errors) + mean_relative_error = np.mean(relative_errors) + percentile_relative_error = np.percentile(relative_errors, percentile) - # Plot relative errors over time - plot_relative_errors_over_time( + plot_error_percentiles_over_time( surr_name, conf, relative_errors, timesteps, title=f"Relative Errors over Time for {surr_name}", + mode="relative", + save=True, + show_title=TITLE, + ) + + plot_error_percentiles_over_time( + surr_name, + conf, + absolute_errors_log, + timesteps, + title=r"$\Delta dex$ (Absolute Log-Space) Errors over Time for " + + f"{surr_name}", + mode="deltadex", save=True, show_title=TITLE, ) @@ -274,14 +305,18 @@ def evaluate_accuracy( # Store metrics accuracy_metrics = { - "mean_squared_error": mean_squared_error, - "mean_absolute_error": mean_absolute_error, - "mean_relative_error": np.mean(relative_errors), - "median_relative_error": np.median(relative_errors), - "max_relative_error": np.max(relative_errors), - "min_relative_error": np.min(relative_errors), - "absolute_errors": absolute_errors, - "relative_errors": relative_errors, + "root_mean_squared_error_log": root_mean_squared_error_log, + "median_absolute_error_log": median_absolute_error_log, + "mean_absolute_error_log": mean_absolute_error_log, + "percentile_absolute_error_log": percentile_absolute_error_log, + "root_mean_squared_error_real": root_mean_squared_error_real, + "median_absolute_error_real": median_absolute_error_real, + "mean_absolute_error_real": mean_absolute_error_real, + "percentile_absolute_error_real": percentile_absolute_error_real, + "median_relative_error": median_relative_error, + "mean_relative_error": mean_relative_error, + "percentile_relative_error": percentile_relative_error, + "error_percentile": percentile, "main_model_training_time": train_time, "main_model_epochs": n_epochs, } @@ -331,6 +366,9 @@ def evaluate_iterative_predictions( # number of chunks n_chunks = (n_timesteps + iter_interval - 1) // iter_interval + # create timesteps array for the iterative predictions + timesteps_full = np.linspace(0, 1, original_n_timesteps) + for i in range(n_chunks): start = i * iter_interval end = min(start + iter_interval, n_timesteps) @@ -349,7 +387,7 @@ def evaluate_iterative_predictions( ds[:, 0, :] = init_state # only need the "train" loader for prediction - dt = timesteps[: model.n_timesteps] + dt = timesteps_full[: model.n_timesteps] train_loader, _, _ = model.prepare_data( dataset_train=ds, dataset_test=None, @@ -360,7 +398,7 @@ def evaluate_iterative_predictions( dataset_train_params=None, dataset_test_params=None, dataset_val_params=None, - dummy_timesteps=True, + dummy_timesteps=False, ) # predict this chunk and insert into the global array @@ -371,6 +409,8 @@ def evaluate_iterative_predictions( preds_chunk[:, 1 : model.n_timesteps, :].detach().cpu().numpy() ) + iterative_preds_log = model.denormalize(iterative_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()) targets = model.denormalize(targets) @@ -381,8 +421,12 @@ def evaluate_iterative_predictions( mse = float(np.mean(errors**2)) mae = float(np.mean(abs_errors)) - thresh = float(conf.get("relative_error_threshold", 0.0)) - rel_errors = abs_errors / np.maximum(np.abs(targets), thresh) + # compute log-space errors + abs_errors_log = np.abs(iterative_preds_log - targets_log) + rmse_log = float(np.mean(abs_errors_log**2)) + mae_log = float(np.mean(abs_errors_log)) + percentile = conf.get("error_percentile", 99) + percentile_abs_error_log = float(np.percentile(abs_errors_log, percentile)) errors = np.mean(np.abs(iterative_preds - targets), axis=(1, 2)) example_idx = int(np.argsort(np.abs(errors - np.median(errors)))[0]) @@ -405,14 +449,12 @@ def evaluate_iterative_predictions( ) return { + "root_mean_squared_error_log": rmse_log, + "mean_absolute_error_log": mae_log, + "percentile_absolute_error_log": percentile_abs_error_log, "mean_squared_error": mse, "mean_absolute_error": mae, - "mean_relative_error": float(np.mean(rel_errors)), - "median_relative_error": float(np.median(rel_errors)), - "max_relative_error": float(np.max(rel_errors)), - "min_relative_error": float(np.min(rel_errors)), "absolute_errors": abs_errors, - "relative_errors": rel_errors, } diff --git a/codes/benchmark/bench_plots.py b/codes/benchmark/bench_plots.py index f0b42831..32753dde 100644 --- a/codes/benchmark/bench_plots.py +++ b/codes/benchmark/bench_plots.py @@ -89,85 +89,84 @@ def save_plot_counter( # Per-surrogate model plots -def plot_relative_errors_over_time( +def plot_error_percentiles_over_time( surr_name: str, conf: dict, - relative_errors: np.ndarray, + errors: np.ndarray, timesteps: np.ndarray, title: str, + mode: str = "relative", # "relative" or "deltadex" save: bool = False, show_title: bool = True, ) -> None: """ - Plot the mean and median relative errors over time with shaded regions for - the 50th, 90th, and 99th percentiles. - - Args: - surr_name (str): The name of the surrogate model. - conf (dict): The configuration dictionary. - relative_errors (np.ndarray): The relative errors of the model. - timesteps (np.ndarray): Array of timesteps. - title (str): The title of the plot. - save (bool): Whether to save the plot. - show_title (bool): Whether to show the title on the plot. - """ - # Calculate the mean, median, and percentiles across all samples and quantities - mean_errors = np.mean(relative_errors, axis=(0, 2)) - mean = np.mean(mean_errors) - median_errors = np.median(relative_errors, axis=(0, 2)) - median = np.median(median_errors) - p50_upper = np.percentile(relative_errors, 75, axis=(0, 2)) - p50_lower = np.percentile(relative_errors, 25, axis=(0, 2)) - p90_upper = np.percentile(relative_errors, 95, axis=(0, 2)) - p90_lower = np.percentile(relative_errors, 5, axis=(0, 2)) - p99_upper = np.percentile(relative_errors, 99.5, axis=(0, 2)) - p99_lower = np.percentile(relative_errors, 0.5, axis=(0, 2)) + Plot mean, median, and percentiles (50th, 90th, 99th) over time. + mode="relative" treats `errors` as relative errors; + mode="deltadex" treats them as log-space absolute errors. + """ + # compute statistics across samples and quantities + mean_ts = np.mean(errors, axis=(0, 2)) + median_ts = np.median(errors, axis=(0, 2)) + stats = { + "50": (25, 75), + "90": (5, 95), + "99": (0.5, 99.5), + } + percentiles = {} + for p, (low, high) in stats.items(): + percentiles[p] = ( + np.percentile(errors, low, axis=(0, 2)), + np.percentile(errors, high, axis=(0, 2)), + ) + # overall means for legend + mean_val = mean_ts.mean() + median_val = median_ts.mean() plt.figure(figsize=(6, 4)) - mean_label = f"Mean Error\nMean={mean * 100:.2f}%" - plt.plot(timesteps, mean_errors, label=mean_label, color="blue") - median_label = f"Median Error\nMedian={median * 100:.2f}%" - plt.plot(timesteps, median_errors, label=median_label, color="red") - - # Shading areas - plt.fill_between( - timesteps, - p50_lower, - p50_upper, - color="grey", - alpha=0.45, - label="50th Percentile", - ) - plt.fill_between( + plt.plot( timesteps, - p90_lower, - p90_upper, - color="grey", - alpha=0.4, - label="90th Percentile", + mean_ts, + label=f"Mean Error\nMean={mean_val * 100:.2f}%", + color="blue", ) - plt.fill_between( + plt.plot( timesteps, - p99_lower, - p99_upper, - color="grey", - alpha=0.15, - label="99th Percentile", + median_ts, + label=f"Median Error\nMedian={median_val * 100:.2f}%", + color="red", ) - plt.yscale("log") + for p, (low_ts, high_ts) in percentiles.items(): + alpha = {"50": 0.45, "90": 0.4, "99": 0.15}[p] + plt.fill_between( + timesteps, + low_ts, + high_ts, + color="grey", + alpha=alpha, + label=f"{p}th Percentile", + ) + plt.xlabel("Time") - plt.ylabel("Relative Error") + if mode == "relative": + plt.yscale("log") + plt.ylabel("Relative Error") + filename = "accuracy_rel_errors_time.pdf" + elif mode == "deltadex": + plt.ylabel(r"$\Delta dex$ (log-space absolute error)") + filename = "accuracy_delta_dex_time.pdf" + else: + raise ValueError(f"Unknown mode: {mode}") + plt.xlim(timesteps[0], timesteps[-1]) - plt.ylim(bottom=1e-8) - if conf["dataset"]["log_timesteps"]: + if conf.get("dataset", {}).get("log_timesteps"): plt.xscale("log") if show_title: plt.title(title) plt.legend(loc="center left", bbox_to_anchor=(1, 0.5)) - if save and conf: - save_plot(plt, "accuracy_rel_errors_time.pdf", conf, surr_name) + if save: + save_plot(plt, filename, conf, surr_name) plt.close() diff --git a/codes/surrogates/AbstractSurrogate/abstract_surrogate.py b/codes/surrogates/AbstractSurrogate/abstract_surrogate.py index 831dc923..a33a7016 100644 --- a/codes/surrogates/AbstractSurrogate/abstract_surrogate.py +++ b/codes/surrogates/AbstractSurrogate/abstract_surrogate.py @@ -453,11 +453,17 @@ def denormalize( elif self.normalisation["mode"] == "minmax": dmax = self.normalisation["max"] dmin = self.normalisation["min"] + if isinstance(data, Tensor) and isinstance(dmax, np.ndarray): + dmax = Tensor(dmax).to(data.device) + dmin = Tensor(dmin).to(data.device) # data = data.to("cpu") data = (data + 1) * (dmax - dmin) / 2 + dmin elif self.normalisation["mode"] == "standardize": mean = self.normalisation["mean"] std = self.normalisation["std"] + if isinstance(data, Tensor) and isinstance(mean, np.ndarray): + mean = Tensor(mean).to(data.device) + std = Tensor(std).to(data.device) data = data * std + mean if self.normalisation["log10_transform"] and not leave_log: diff --git a/codes/tune/optuna_config.yaml b/codes/tune/optuna_config.yaml index 264404b6..0099e763 100644 --- a/codes/tune/optuna_config.yaml +++ b/codes/tune/optuna_config.yaml @@ -1,20 +1,21 @@ -tuning_id: cloud_tuning_fine +tuning_id: primordial_tuning seed: 42 dataset: - name: cloud + name: primordial log10_transform: True normalise: minmax + per_species: True subset_factor: 1 - tolerance: 1e-25 + tolerance: 1e-15 normalise_per_species: True log_timesteps: True -devices: ["cuda:2", "cuda:3", "cuda:4", "cuda:5", "cuda:6", "cuda:7", "cuda:8", "cuda:9"] +devices: ["cuda:0", "cuda:2", "cuda:3", "cuda:5", "cuda:6", "cuda:8", "cuda:9"] optuna_logs: False prune: True verbose: False multi_objective: True -population_size: 30 +population_size: 100 target_percentile: 0.99 postgres_config: mode: "local" # "local" or "remote" @@ -22,7 +23,7 @@ postgres_config: user: "optuna_user" host: "localhost" # "localhost" for local use database_folder: "/export/home/rjanssen/postgres/" # only for local use - db_name: "optuna_cloud_2" # remote mode: single DB for all runs + db_name: "optuna_primordial" # remote mode: single DB for all runs sslmode: "require" # if needed password: "" # optional; prefer env PGPASSWORD @@ -35,67 +36,88 @@ global_optuna_params: regularization_factor: type: float low: 1.0e-6 - high: 1.0e-2 + high: 1.0 + log: true + optimizer: + type: categorical + choices: ["AdamW", "SGD"] + momentum: + type: float + low: 0.0 + high: 0.99 + step: 0.01 + scheduler: + type: categorical + choices: ["cosine", "poly", "schedulefree"] + poly_power: + type: float + low: 0.5 + high: 2.0 + step: 0.1 + eta_min: + type: float + low: 1.0e-3 + high: 1.0 log: true activation: type: categorical choices: ["ReLU", "LeakyReLU", "PReLU", "Tanh", "GELU", "Mish", "SiLU", "ELU"] + loss_function: + type: categorical + choices: ["mse", "smoothl1"] + beta: + type: float + low: 0.1 + high: 10.0 + log: true surrogates: - name: MultiONet batch_size: 65536 - epochs: 4096 - trials: 100 + epochs: 2048 + trials: 200 optuna_params: branch_hidden_layers: type: int low: 1 - high: 6 + high: 10 hidden_size: type: int low: 10 - high: 500 + high: 700 step: 10 output_factor: type: int - low: 30 - high: 300 - step: 10 + low: 1 + high: 400 trunk_hidden_layers: type: int low: 1 - high: 7 - beta: - type: float - low: 0.1 - high: 10.0 - log: true + high: 10 + # params_branch: + # type: categorical + # choices: ["True", "False"] - name: FullyConnected batch_size: 65536 - epochs: 4096 - trials: 100 + epochs: 2048 + trials: 200 optuna_params: hidden_size: type: int low: 10 - high: 500 + high: 700 step: 10 num_hidden_layers: type: int low: 1 - high: 6 - beta: - type: float - low: 0.1 - high: 10.0 - log: true + high: 10 - name: LatentPoly batch_size: 512 - epochs: 4096 - trials: 100 + epochs: 2048 + trials: 200 optuna_params: degree: type: int @@ -108,12 +130,12 @@ surrogates: coder_layers: type: int low: 1 - high: 6 + high: 10 coder_width: type: int low: 10 - high: 600 - step: 100 + high: 700 + step: 10 # coeff_network: # type: categorical # choices: ["True", "False"] @@ -127,33 +149,36 @@ surrogates: # low: 1 # high: 5 - - name: LatentNeuralODE - batch_size: 1024 - epochs: 4096 - trials: 200 - optuna_params: - latent_features: - type: int - low: 1 - high: 15 - coder_layers: - type: int - low: 1 - high: 10 - coder_width: - type: int - low: 10 - high: 400 - step: 10 - ode_width: - type: int - low: 10 - high: 400 - step: 10 - ode_layers: - type: int - low: 1 - high: 10 - # encode_params: - # type: categorical - # choices: ["True", "False"] + # - name: LatentNeuralODE + # batch_size: 1024 + # epochs: 2048 + # trials: 200 + # optuna_params: + # latent_features: + # type: int + # low: 1 + # high: 10 + # coder_layers: + # type: int + # low: 1 + # high: 10 + # coder_width: + # type: int + # low: 10 + # high: 700 + # step: 10 + # ode_tanh_reg: + # type: categorical + # choices: ["True", "False"] + # ode_width: + # type: int + # low: 10 + # high: 700 + # step: 10 + # ode_layers: + # type: int + # low: 1 + # high: 10 + # # encode_params: + # # type: categorical + # # choices: ["True", "False"] diff --git a/codes/tune/optuna_fcts.py b/codes/tune/optuna_fcts.py index 5330279b..87970485 100644 --- a/codes/tune/optuna_fcts.py +++ b/codes/tune/optuna_fcts.py @@ -17,7 +17,7 @@ measure_inference_time, ) from codes.utils import check_and_load_data, make_description, set_random_seeds -from codes.utils.data_utils import download_data, get_data_subset +from codes.utils.data_utils import download_data MODULE_REGISTRY: dict[str, type[nn.Module]] = { "relu": nn.ReLU, @@ -237,14 +237,10 @@ def training_run( ) subset_factor = config["dataset"].get("subset_factor", 1) - # Get the appropriate data subset - (train_data, test_data), (train_params, test_params), timesteps = get_data_subset( - (train_data, test_data), - timesteps, - "sparse", - subset_factor, - (train_params, test_params), - ) + # Get the appropriate subset of the training data + # We nevertheless use the full test data to measure performance. + train_data = train_data[::subset_factor] + train_params = train_params[::subset_factor] if train_params is not None else None set_random_seeds(config["seed"], device=device) surr_name = config["surrogate"]["name"] diff --git a/codes/utils/data_utils.py b/codes/utils/data_utils.py index 3ef5e3cf..e868f220 100644 --- a/codes/utils/data_utils.py +++ b/codes/utils/data_utils.py @@ -575,7 +575,7 @@ def get_data_subset( Returns: tuple: (data_subset, params_subset, timesteps_subset) """ - # First, subsample the data based on subset_factor. + # First, subsample the dataset based on subset_factor. data_sub = tuple(d[::subset_factor] for d in data) # Handle params: diff --git a/config.yaml b/config.yaml index 97f40b4e..076aea50 100644 --- a/config.yaml +++ b/config.yaml @@ -1,17 +1,18 @@ # Global settings for the benchmark -training_id: "simple_reaction_final" -surrogates: ["MultiONet", "FullyConnected", "LatentPoly", "LatentNeuralODE", ] -batch_size: [4096, 4096, 256, 256] -epochs: [12000, 10000, 10000, 7000] +training_id: "profile_lnode32_improved" +surrogates: ["LatentNeuralODE"] # ["MultiONet", "FullyConnected", "LatentNeuralODE", "LatentPoly"] +batch_size: [716] # [8192, 8192, 512, 512] +epochs: [50] # [20000, 7500, 20000, 15000] dataset: - name: "simple_reaction" - log_timesteps: False + name: "primordial" log10_transform: True + log10_transform_params: False normalise: "minmax" # "minmax" # "standardise", "minmax", "disable" use_optimal_params: True - tolerance: 1e-30 + tolerance: 1e-20 subset_factor: 1 -devices: ["cuda:1"] + log_timesteps: True +devices: ["cuda:5"] seed: 42 verbose: False @@ -33,11 +34,8 @@ uncertainty: ensemble_size: 5 # Number of models for deep ensemble # Evaluations during benchmark -iterative: True losses: False gradients: False timing: False compute: False compare: False # Whether to compare the surrogates - - diff --git a/datasets/cloud/surrogates_config.py b/datasets/cloud/surrogates_config.py index d3046f87..67c05310 100644 --- a/datasets/cloud/surrogates_config.py +++ b/datasets/cloud/surrogates_config.py @@ -10,6 +10,15 @@ class MultiONetConfig: loss_function: nn.Module = nn.SmoothL1Loss() optimizer: str = "adamw" scheduler: str = "schedulefree" + # other params from cloud_tuning_fine, trial 81: + beta: float = 1.78 + branch_hidden_layers: int = 6 + hidden_size = 170 + output_factor: int = 80 + trunk_hidden_layers: int = 7 + learning_rate: float = 8.3e-4 + regularization_factor: float = 9.6e-06 + activation: nn.Module = nn.Tanh() @dataclass @@ -20,6 +29,15 @@ class LatentNeuralODEConfig: optimizer: str = "adamw" scheduler: str = "schedulefree" ode_tanh_reg: bool = False + # other params from cloud_tuning_fine, trial 35: + latent_features: int = 14 + coder_layers: int = 4 + coder_width: int = 70 + ode_width: int = 30 + ode_layers: int = 2 + learning_rate: float = 7.1e-03 + regularization_factor: float = 3.4e-03 + activation: nn.Module = nn.Tanh() @dataclass @@ -29,6 +47,13 @@ class FullyConnectedConfig: loss_function: nn.Module = nn.SmoothL1Loss() optimizer: str = "adamw" scheduler: str = "schedulefree" + # other params from cloud_tuning_fine, trial 77 + beta = 0.817 + hidden_size: int = 380 + num_hidden_layers: int = 2 + learning_rate: float = 6.8e-03 + regularization_factor: float = 8.7e-06 + activation: nn.Module = nn.ReLU() @dataclass @@ -38,6 +63,13 @@ class LatentPolyConfig: loss_function: nn.Module = nn.MSELoss() optimizer: str = "adamw" scheduler: str = "schedulefree" + # other params from cloud_tuning_fine, trial 89: + degree: int = 3 + latent_features: int = 9 + coder_layers: int = 3 + coder_width: int = 210 + learning_rate: float = 2.2e-5 + regularization_factor: float = 1.7e-03 # @dataclass diff --git a/datasets/cloud_parametric/surrogates_config_old.py b/datasets/cloud_parametric/surrogates_config_old.py index d2c8c134..980e4a9f 100644 --- a/datasets/cloud_parametric/surrogates_config_old.py +++ b/datasets/cloud_parametric/surrogates_config_old.py @@ -2,61 +2,60 @@ from torch import nn - -@dataclass -class MultiONetConfig: - """Model config for MultiONet for the simple_ode dataset""" - - # cloud, trial 69 - branch_hidden_layers: int = 1 - trunk_hidden_layers: int = 9 - hidden_size: int = 225 - output_factor: int = 63 - learning_rate: float = 4e-5 # optimal for ~4000 epochs - activation: nn.Module = nn.Tanh() - - -@dataclass -class LatentNeuralODEConfig: - """Model config for LatentNeuralODE for the simple_ode dataset""" - - # cloudparams, trial 40 - latent_features: int = 3 - coder_layers: int = 3 - coder_width: int = 377 - learning_rate: float = 3e-4 - ode_layers: int = 5 - ode_width: int = 167 - regularization_factor: float = 0.000127 - encode_params: bool = False - optimizer: str = "sgd" - momentum: float = 0.226 - scheduler: str = "cosine" - eta_min: float = 0.0222 - ode_tanh_reg: bool = False - activation: nn.Module = nn.SiLU() - model_version: str = "v2" - - -@dataclass -class FullyConnectedConfig: - """Model config for FullyConnected for the simple_ode dataset""" - - # cloud, trial 44 - hidden_size: int = 261 - num_hidden_layers: int = 1 - learning_rate: float = 1e-4 - activation: nn.Module = nn.LeakyReLU() - - -@dataclass -class LatentPolyConfig: - """Model config for LatentPoly for the simple_ode dataset""" - - # cloud, trial 92 - latent_features: int = 9 - degree: int = 5 - learning_rate: float = 3e-4 - coder_layers: int = 1 - coder_width: int = 86 - activation: nn.Module = nn.Mish() +# @dataclass +# class MultiONetConfig: +# """Model config for MultiONet for the simple_ode dataset""" + +# # cloud, trial 69 +# branch_hidden_layers: int = 1 +# trunk_hidden_layers: int = 9 +# hidden_size: int = 225 +# output_factor: int = 63 +# learning_rate: float = 4e-5 # optimal for ~4000 epochs +# activation: nn.Module = nn.Tanh() + + +# @dataclass +# class LatentNeuralODEConfig: +# """Model config for LatentNeuralODE for the simple_ode dataset""" + +# # cloudparams, trial 40 +# latent_features: int = 3 +# coder_layers: int = 3 +# coder_width: int = 377 +# learning_rate: float = 3e-4 +# ode_layers: int = 5 +# ode_width: int = 167 +# regularization_factor: float = 0.000127 +# encode_params: bool = False +# optimizer: str = "sgd" +# momentum: float = 0.226 +# scheduler: str = "cosine" +# eta_min: float = 0.0222 +# ode_tanh_reg: bool = False +# activation: nn.Module = nn.SiLU() +# model_version: str = "v2" + + +# @dataclass +# class FullyConnectedConfig: +# """Model config for FullyConnected for the simple_ode dataset""" + +# # cloud, trial 44 +# hidden_size: int = 261 +# num_hidden_layers: int = 1 +# learning_rate: float = 1e-4 +# activation: nn.Module = nn.LeakyReLU() + + +# @dataclass +# class LatentPolyConfig: +# """Model config for LatentPoly for the simple_ode dataset""" + +# # cloud, trial 92 +# latent_features: int = 9 +# degree: int = 5 +# learning_rate: float = 3e-4 +# coder_layers: int = 1 +# coder_width: int = 86 +# activation: nn.Module = nn.Mish() diff --git a/datasets/primordial/surrogates_config_old.py b/datasets/primordial/surrogates_config_old.py new file mode 100644 index 00000000..fb8d4b73 --- /dev/null +++ b/datasets/primordial/surrogates_config_old.py @@ -0,0 +1,56 @@ +from dataclasses import dataclass + +from torch import nn + + +@dataclass +class MultiONetConfig: + """Model config for MultiONet for the simple_ode dataset""" + + # primordial3, trial 41 + branch_hidden_layers: int = 8 + trunk_hidden_layers: int = 8 + hidden_size: int = 130 + output_factor: int = 98 + learning_rate: float = 5e-4 + activation: nn.Module = nn.Tanh() + + +@dataclass +class LatentNeuralODEConfig: + """Model config for LatentNeuralODE for the simple_ode dataset""" + + # primordial5, trial 3 + latent_features: int = 10 + coder_layers: int = 4 + coder_width: int = 350 + learning_rate: float = 8e-4 + ode_layers: int = 10 + ode_width: int = 157 + ode_tanh_reg: bool = True + activation: nn.Module = nn.Softplus() + model_version: str = "v2" + + +@dataclass +class FullyConnectedConfig: + """Model config for FullyConnected for the simple_ode dataset""" + + # primordial4, trial 64 + hidden_size: int = 453 + num_hidden_layers: int = 1 + learning_rate: float = 5e-4 + activation: nn.Module = nn.ReLU() + + +@dataclass +class LatentPolyConfig: + """Model config for LatentPoly for the simple_ode dataset""" + + # primordial3, trial 90 + latent_features: int = 9 + degree: int = 1 + learning_rate: float = 8e-4 + coder_layers: int = 2 + coder_width: int = 264 + activation: nn.Module = nn.LeakyReLU() From 4cc43284efd8a45bc33487f6924241cb2f79932f Mon Sep 17 00:00:00 2001 From: Robin Date: Fri, 1 Aug 2025 18:24:06 +0200 Subject: [PATCH 08/13] first improvement: use float64 only in solver --- .../LatentNeuralODE/latent_neural_ode.py | 153 ++++++++++++++++-- datasets/primordial/surrogates_config.py | 56 ------- 2 files changed, 143 insertions(+), 66 deletions(-) delete mode 100644 datasets/primordial/surrogates_config.py diff --git a/codes/surrogates/LatentNeuralODE/latent_neural_ode.py b/codes/surrogates/LatentNeuralODE/latent_neural_ode.py index 566ffc8f..ccdd7c27 100644 --- a/codes/surrogates/LatentNeuralODE/latent_neural_ode.py +++ b/codes/surrogates/LatentNeuralODE/latent_neural_ode.py @@ -35,7 +35,7 @@ def __init__( n_parameters: int = 0, training_id: str | None = None, config: dict | None = None, - dtype: torch.dtype = torch.float64, + dtype: torch.dtype = torch.float32, ): super().__init__( device=device, @@ -48,9 +48,12 @@ def __init__( self.n_parameters = n_parameters # Instantiate the model wrapper with the additional n_parameters. self.model = ModelWrapper( - config=self.config, n_quantities=n_quantities, n_parameters=n_parameters + config=self.config, + n_quantities=n_quantities, + n_parameters=n_parameters, + dtype=dtype, ).to(device) - self.to(dtype=dtype) + self.dtype = dtype def forward(self, inputs): """ @@ -59,7 +62,7 @@ def forward(self, inputs): """ inputs = tuple( ( - x.to(self.device, dtype=torch.float64, non_blocking=True) + x.to(self.device, dtype=self.dtype, non_blocking=True) if isinstance(x, Tensor) else x ) @@ -158,7 +161,7 @@ def create_dataloader( ) @time_execution - def fit( + def fit_normal( self, train_loader: DataLoader, test_loader: DataLoader, @@ -238,6 +241,125 @@ def fit( self.n_epochs = epoch + 1 self.get_checkpoint(test_loader, criterion) + @time_execution + def fit( + self, + train_loader: DataLoader, + test_loader: DataLoader, + epochs: int, + position: int = 0, + description: str = "Training LatentNeuralODE", + multi_objective: bool = False, + ) -> None: + from torch.profiler import ProfilerActivity, profile, record_function + + optimizer, scheduler = self.setup_optimizer_and_scheduler(epochs) + criterion = self.config.loss_function + + loss_length = (epochs + self.update_epochs - 1) // self.update_epochs + self.train_loss, self.test_loss, self.MAE = [ + np.zeros(loss_length) for _ in range(3) + ] + + progress_bar = self.setup_progress_bar(epochs, position, description) + + self.model.train() + # If optimizer.train() is a no-op or not needed, consider removing it; + # leave it only if it's a custom wrapper that actually requires it. + try: + optimizer.train() + except AttributeError: + pass # standard PyTorch optimizers don't have .train() + + self.setup_checkpoint() + + profiled = False # flag to do profiling only once + for epoch in progress_bar: + for i, batch in enumerate(train_loader): + batch = tuple( + ( + x.to(device=self.device, non_blocking=True) + if isinstance(x, Tensor) + else x + ) + for x in batch + ) + x_true, t_range, params = batch + + # Profile only the first batch of the first epoch + if not profiled and epoch == 2 and i == 1: + with profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + record_shapes=True, + with_stack=True, # optional: deeper stack traces + profile_memory=True, # track memory allocs + ) as prof: + with record_function("zero_grad"): + optimizer.zero_grad() + + with record_function("model_forward"): + x_pred, x_true = self((x_true, t_range, params)) + + with record_function("loss_compute"): + loss = self.model.total_loss( + x_true, x_pred, params, criterion + ) + + with record_function("backward"): + loss.backward() + + with record_function("optimizer_step"): + optimizer.step() + + # Advance scheduler if your original logic expects it here + scheduler.step() + + # Output profiling summary + print("=== Profiler summary for epoch 0 batch 0 ===") + print( + prof.key_averages().table( + sort_by="self_cuda_time_total", row_limit=60 + ) + ) + + # Export trace for timeline inspection (e.g., chrome://tracing or TensorBoard) + prof.export_chrome_trace(f"prof_trace_epoch{epoch}_batch{i}.json") + + profiled = True # don't profile again + + else: + # Normal training step + optimizer.zero_grad() + x_pred, x_true = self((x_true, t_range, params)) + loss = self.model.total_loss(x_true, x_pred, params, criterion) + 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() + + self.validate( + epoch=epoch, + train_loader=train_loader, + test_loader=test_loader, + optimizer=optimizer, + progress_bar=progress_bar, + total_epochs=epochs, + multi_objective=multi_objective, + ) + + progress_bar.close() + self.n_epochs = epoch + 1 + self.get_checkpoint(test_loader, criterion) + class ModelWrapper(nn.Module): """ @@ -260,12 +382,19 @@ class ModelWrapper(nn.Module): - Decoder: latent_dim -> output dimensions """ - def __init__(self, config, n_quantities: int, n_parameters: int = 0): + def __init__( + self, + config, + n_quantities: int, + n_parameters: int = 0, + dtype: torch.dtype = torch.float64, + ): super().__init__() self.config = config self.n_parameters = n_parameters latent_dim = config.latent_features self.loss_weights = getattr(config, "loss_weights", [100.0, 1.0, 1.0, 1.0]) + self.dtype = dtype # --- Build encoder --- if n_parameters == 0: @@ -288,6 +417,7 @@ def __init__(self, config, n_quantities: int, n_parameters: int = 0): coder_layers=config.coder_layers, coder_width=config.coder_width, activation=config.activation, + dtype=dtype, ) else: raise ValueError( @@ -305,6 +435,7 @@ def __init__(self, config, n_quantities: int, n_parameters: int = 0): ode_layers=config.ode_layers, ode_width=config.ode_width, tanh_reg=config.ode_tanh_reg, + dtype=dtype, ) ode_module = ode_net else: @@ -316,6 +447,7 @@ def __init__(self, config, n_quantities: int, n_parameters: int = 0): ode_layers=config.ode_layers, ode_width=config.ode_width, tanh_reg=config.ode_tanh_reg, + dtype=dtype, ) ode_module = ODEWithParams(base_ode, n_parameters, latent_dim) @@ -340,6 +472,7 @@ def __init__(self, config, n_quantities: int, n_parameters: int = 0): coder_layers=config.coder_layers, coder_width=config.coder_width, activation=config.activation, + dtype=dtype, ) def forward(self, x0: Tensor, t_range: Tensor, params: Tensor = None): @@ -349,7 +482,7 @@ def forward(self, x0: Tensor, t_range: Tensor, params: Tensor = None): enc_in = torch.cat([x0, params], dim=1) else: enc_in = x0 - z0 = self.encoder(enc_in) + z0 = self.encoder(enc_in) # .to(torch.float64) # if using closure to inject params if self.n_parameters > 0 and not self.config.encode_params: @@ -357,12 +490,12 @@ def forward(self, x0: Tensor, t_range: Tensor, params: Tensor = None): self.ode.set_params(params) # solve dynamics - t_eval = t_range.repeat(x0.size(0), 1) + t_eval = t_range.repeat(x0.size(0), 1) # .to(torch.float64) sol = self.solver.solve(to.InitialValueProblem(y0=z0, t_eval=t_eval)) latent_traj = sol.ys # [timesteps, batch, latent_dim] # decode - return self.decoder(latent_traj) + return self.decoder(latent_traj.to(self.dtype)) def renormalize_loss_weights( self, x_true, x_pred, params, criterion: nn.Module = nn.MSELoss() @@ -522,7 +655,7 @@ def forward(self, t, x): output = self.mlp(x) if self.tanh_reg: return self.reg_factor * torch.tanh(output / self.reg_factor) - return output + return output.to(torch.float64) class ODEWithParams(nn.Module): diff --git a/datasets/primordial/surrogates_config.py b/datasets/primordial/surrogates_config.py deleted file mode 100644 index fb8d4b73..00000000 --- a/datasets/primordial/surrogates_config.py +++ /dev/null @@ -1,56 +0,0 @@ -from dataclasses import dataclass - -from torch import nn - - -@dataclass -class MultiONetConfig: - """Model config for MultiONet for the simple_ode dataset""" - - # primordial3, trial 41 - branch_hidden_layers: int = 8 - trunk_hidden_layers: int = 8 - hidden_size: int = 130 - output_factor: int = 98 - learning_rate: float = 5e-4 - activation: nn.Module = nn.Tanh() - - -@dataclass -class LatentNeuralODEConfig: - """Model config for LatentNeuralODE for the simple_ode dataset""" - - # primordial5, trial 3 - latent_features: int = 10 - coder_layers: int = 4 - coder_width: int = 350 - learning_rate: float = 8e-4 - ode_layers: int = 10 - ode_width: int = 157 - ode_tanh_reg: bool = True - activation: nn.Module = nn.Softplus() - model_version: str = "v2" - - -@dataclass -class FullyConnectedConfig: - """Model config for FullyConnected for the simple_ode dataset""" - - # primordial4, trial 64 - hidden_size: int = 453 - num_hidden_layers: int = 1 - learning_rate: float = 5e-4 - activation: nn.Module = nn.ReLU() - - -@dataclass -class LatentPolyConfig: - """Model config for LatentPoly for the simple_ode dataset""" - - # primordial3, trial 90 - latent_features: int = 9 - degree: int = 1 - learning_rate: float = 8e-4 - coder_layers: int = 2 - coder_width: int = 264 - activation: nn.Module = nn.LeakyReLU() From 0ed0252482302b55ea238ec8d648099039baf6e6 Mon Sep 17 00:00:00 2001 From: Robin Date: Mon, 4 Aug 2025 10:43:34 +0200 Subject: [PATCH 09/13] results from primordial_tuning --- datasets/primordial/surrogates_config.py | 79 ++++++++++++++++++++++++ 1 file changed, 79 insertions(+) create mode 100644 datasets/primordial/surrogates_config.py diff --git a/datasets/primordial/surrogates_config.py b/datasets/primordial/surrogates_config.py new file mode 100644 index 00000000..1b9124a2 --- /dev/null +++ b/datasets/primordial/surrogates_config.py @@ -0,0 +1,79 @@ +from dataclasses import dataclass + +from torch import nn + + +@dataclass +class MultiONetConfig: + """Model config for MultiONet for the simple_ode dataset""" + + # primordial_tuning, trial 18 + scheduler: str = "poly" + optimizer: str = "AdamW" + loss_function: nn.Module = nn.SmoothL1Loss() + poly_power: float = 0.725 + beta: float = 0.568 + branch_hidden_layers: int = 5 + hidden_size: int = 560 + output_factor: int = 293 + trunk_hidden_layers: int = 5 + learning_rate: float = 5.4e-04 + regularization_factor: float = 1.8e-02 + activation: nn.Module = nn.GELU() + + +@dataclass +class LatentNeuralODEConfig: + """Model config for LatentNeuralODE for the simple_ode dataset""" + + # primordial_tuning, trial 186 + scheduler: str = "poly" + optimizer: str = "sgd" + loss_function: nn.Module = nn.SmoothL1Loss() + poly_power: float = 0.948 + momentum: float = 0.702 + beta: float = 0.684 + latent_features: int = 10 + coder_layers: int = 1 + coder_width: int = 490 + ode_layers: int = 3 + ode_width: int = 50 # went up from 20 - very cheap increase in expressivity + learning_rate: float = 1.99e-4 + regularization_factor: float = 7.66e-02 + activation: nn.Module = nn.Mish() + + +@dataclass +class FullyConnectedConfig: + """Model config for FullyConnected for the simple_ode dataset""" + + # primordial_tuning, trial 63 + scheduler: str = "poly" + optimizer: str = "AdamW" + loss_function: nn.Module = nn.SmoothL1Loss() + poly_power: float = 1.691 + beta: float = 2.611 + hidden_size: int = 470 + num_hidden_layers: int = 4 + learning_rate: float = 2.3e-03 + regularization_factor: float = 0.309 + activation: nn.Module = nn.LeakyReLU() + + +@dataclass +class LatentPolyConfig: + """Model config for LatentPoly for the simple_ode dataset""" + + # primordial_tuning, trial 176 + scheduler: str = "poly" + optimizer: str = "AdamW" + loss_function: nn.Module = nn.SmoothL1Loss() + poly_power: float = 0.845 + beta: float = 2.462 + degree: int = 4 + latent_features: int = 8 + coder_layers: int = 2 + coder_width: int = 150 + learning_rate: float = 9.36e-04 + regularization_factor: float = 4.9e-04 + activation: nn.Module = nn.GELU() From e884be0b31502830e63b2e207d4faaffaaf2751c Mon Sep 17 00:00:00 2001 From: Robin Date: Mon, 4 Aug 2025 10:44:34 +0200 Subject: [PATCH 10/13] Further latentneuralode improvements --- .../AbstractSurrogate/abstract_surrogate.py | 8 +- .../LatentNeuralODE/latent_neural_ode.py | 244 +++++++++--------- codes/train/train_fcts.py | 2 +- codes/tune/optuna_config.yaml | 233 ++++++++--------- config.yaml | 14 +- 5 files changed, 242 insertions(+), 259 deletions(-) diff --git a/codes/surrogates/AbstractSurrogate/abstract_surrogate.py b/codes/surrogates/AbstractSurrogate/abstract_surrogate.py index a33a7016..edee3a6e 100644 --- a/codes/surrogates/AbstractSurrogate/abstract_surrogate.py +++ b/codes/surrogates/AbstractSurrogate/abstract_surrogate.py @@ -537,17 +537,11 @@ def time_pruning(self, current_epoch: int, total_epochs: int) -> None: else: threshold = None - # print( - # f"[time_pruning] Epoch: {current_epoch}/{total_epochs} | " - # f"Elapsed: {elapsed:.1f}s | Avg per epoch: {average_epoch_time:.1f}s | " - # f"Projected total: {projected_total_time:.1f}s | Threshold: {threshold:.1f}s" - # ) - if threshold is not 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." + f"[time_pruning] Projected total time {projected_total_time:.1f}s exceeds threshold {threshold:.1f}s. Pruning trial {self.optuna_trial.number}." ) self.optuna_trial.set_user_attr( "prune_reason", diff --git a/codes/surrogates/LatentNeuralODE/latent_neural_ode.py b/codes/surrogates/LatentNeuralODE/latent_neural_ode.py index ccdd7c27..da0adfd9 100644 --- a/codes/surrogates/LatentNeuralODE/latent_neural_ode.py +++ b/codes/surrogates/LatentNeuralODE/latent_neural_ode.py @@ -51,6 +51,7 @@ def __init__( config=self.config, n_quantities=n_quantities, n_parameters=n_parameters, + n_timesteps=self.n_timesteps, dtype=dtype, ).to(device) self.dtype = dtype @@ -161,7 +162,7 @@ def create_dataloader( ) @time_execution - def fit_normal( + def fit( self, train_loader: DataLoader, test_loader: DataLoader, @@ -218,12 +219,12 @@ def fit_normal( 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 - ) + # # 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 + # ) scheduler.step() @@ -242,7 +243,7 @@ def fit_normal( self.get_checkpoint(test_loader, criterion) @time_execution - def fit( + def fit_profile( self, train_loader: DataLoader, test_loader: DataLoader, @@ -335,12 +336,12 @@ def fit( 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 - ) + # # 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 @@ -387,7 +388,10 @@ def __init__( config, n_quantities: int, n_parameters: int = 0, + n_timesteps: int = 101, dtype: torch.dtype = torch.float64, + use_pid: bool = False, # switch from Integral to PID + adjoint_type: str = "autodiff", # "autodiff" | "backsolve" | "joint_backsolve" ): super().__init__() self.config = config @@ -395,6 +399,8 @@ def __init__( latent_dim = config.latent_features self.loss_weights = getattr(config, "loss_weights", [100.0, 1.0, 1.0, 1.0]) self.dtype = dtype + self.n_timesteps = n_timesteps + self.l2 = nn.MSELoss() # --- Build encoder --- if n_parameters == 0: @@ -451,11 +457,36 @@ def __init__( ) ode_module = ODEWithParams(base_ode, n_parameters, latent_dim) + # --- ODE and solver setup --- self.ode = ode_module term = to.ODETerm(self.ode) - step = to.Tsit5(term=term) - ctrl = to.IntegralController(atol=config.atol, rtol=config.rtol, term=term) - self.solver = to.AutoDiffAdjoint(step, ctrl) + step = to.Tsit5(term=term) # or expose choice of step method if desired + + # choose controller + if use_pid: + # tune pcoeff/icoeff/dcoeff as hyperparams if needed + controller = to.PIDController( + atol=config.atol, + rtol=config.rtol, + pcoeff=getattr(config, "pid_pcoeff", 0.2), + icoeff=getattr(config, "pid_icoeff", 0.5), + dcoeff=getattr(config, "pid_dcoeff", 0.0), + term=term, + ) + else: + controller = to.IntegralController( + atol=config.atol, rtol=config.rtol, term=term + ) + + # choose adjoint/backprop method + if adjoint_type == "autodiff": + self.solver = to.AutoDiffAdjoint(step, controller) + elif adjoint_type == "backsolve": + self.solver = to.BacksolveAdjoint(term, step, controller) + elif adjoint_type == "joint_backsolve": + self.solver = to.JointBacksolveAdjoint(term, step, controller) + else: + raise ValueError(f"Unknown adjoint_type {adjoint_type}") # --- Build decoder --- if self.config.model_version == "v1": @@ -489,32 +520,19 @@ def forward(self, x0: Tensor, t_range: Tensor, params: Tensor = None): assert params is not None self.ode.set_params(params) - # solve dynamics - t_eval = t_range.repeat(x0.size(0), 1) # .to(torch.float64) + # use float64 for ODE solver + z0 = z0.to(torch.float64) # solver state must be double + t_eval = t_range.repeat(x0.size(0), 1).to(torch.float64) + assert z0.dtype == torch.float64, f"z0 dtype {z0.dtype}" + assert t_eval.dtype == torch.float64, f"t_eval dtype {t_eval.dtype}" + sol = self.solver.solve(to.InitialValueProblem(y0=z0, t_eval=t_eval)) + latent_traj = sol.ys # [timesteps, batch, latent_dim] # decode return self.decoder(latent_traj.to(self.dtype)) - def renormalize_loss_weights( - self, x_true, x_pred, params, criterion: nn.Module = nn.MSELoss() - ): - """ - Renormalize the loss weights based on the current loss values so that they are accurately - weighted based on the provided weights. To be used once after a short burn in phase. - - Args: - x_true (Tensor): The true trajectory. - x_pred (Tensor): The predicted trajectory - params (Tensor): Fixed parameters (batch, n_parameters). - criterion (nn.Module): Loss function to use for calculating the losses. - """ - self.loss_weights[0] = 1 / criterion(x_pred, x_true).item() * 100 - self.loss_weights[1] = 1 / self.identity_loss(x_true, params).item() - self.loss_weights[2] = 1 / self.deriv_loss(x_true, x_pred).item() - self.loss_weights[3] = 1 / self.deriv2_loss(x_true, x_pred).item() - def total_loss( self, x_true: Tensor, @@ -523,14 +541,59 @@ def total_loss( criterion: nn.Module = nn.MSELoss(), ): """ - Calculate the total loss based on the loss weights, including params for identity. + Total loss: weighted sum of trajectory reconstruction, identity, first derivative, + and second derivative losses. All terms remain in the computation graph. """ - return ( - self.loss_weights[0] * criterion(x_pred, x_true).item() - + self.loss_weights[1] * self.identity_loss(x_true, params) - + self.loss_weights[2] * self.deriv_loss(x_true, x_pred) - + self.loss_weights[3] * self.deriv2_loss(x_true, x_pred) - ) + w0, w1, w2, w3 = ( + self.loss_weights + ) # assume these are set in config and are floats + + # primary trajectory loss + traj_loss = criterion(x_pred, x_true) + + # identity loss (reconstruct x0) + identity = self.identity_loss(x_true, params) + + # derivative losses: compute once + d_pred = self.first_derivative(x_pred) + d_true = self.first_derivative(x_true) + deriv_loss = self.l2(d_pred, d_true) + + d2_pred = self.second_derivative(x_pred) + d2_true = self.second_derivative(x_true) + deriv2_loss = self.l2(d2_pred, d2_true) + + return w0 * traj_loss + w1 * identity + w2 * deriv_loss + w3 * deriv2_loss + + def first_derivative(self, x: Tensor): + # x: [B, T, F] + h = 1.0 / self.n_timesteps + # central differences for interior + d_center = (x[:, 2:, :] - x[:, :-2, :]) / (2 * h) # [B, T-2, F] + # forward/backward for boundaries + d_first = (x[:, 1:2, :] - x[:, :1, :]) / h # [B,1,F] + d_last = (x[:, -1:, :] - x[:, -2:-1, :]) / h # [B,1,F] + derivative = torch.cat([d_first, d_center, d_last], dim=1) # [B,T,F] + return derivative + + def second_derivative(self, x: Tensor): + # x: [B, T, F] + h = 1.0 / self.n_timesteps + # standard second derivative central + d2_center = (x[:, 2:, :] - 2 * x[:, 1:-1, :] + x[:, :-2, :]) / ( + h * h + ) # [B, T-2, F] + # one-sided approximations at ends (second-order): + # at t0: f''(t0) ≈ (2 f0 - 5 f1 + 4 f2 - f3) / h^2 + d2_first = ( + 2 * x[:, :1, :] - 5 * x[:, 1:2, :] + 4 * x[:, 2:3, :] - x[:, 3:4, :] + ) / (h * h) + # at t_{N-1}: symmetric formula + d2_last = ( + 2 * x[:, -1:, :] - 5 * x[:, -2:-1, :] + 4 * x[:, -3:-2, :] - x[:, -4:-3, :] + ) / (h * h) + d2 = torch.cat([d2_first, d2_center, d2_last], dim=1) # [B,T,F] + return d2 def identity_loss(self, x_true: Tensor, params: Tensor = None): """ @@ -552,75 +615,7 @@ def identity_loss(self, x_true: Tensor, params: Tensor = None): # encode-decode z0 = self.encoder(enc_input) x0_hat = self.decoder(z0) - return self.l2_loss(x0, x0_hat) - - @staticmethod - def l2_loss(x_true: Tensor, x_pred: Tensor): - """ - Calculate the L2 loss. - - Args: - x_true (Tensor): The true trajectory. - x_pred (Tensor): The predicted trajectory - - Returns: - Tensor: The L2 loss. - """ - return torch.mean(torch.abs(x_true - x_pred) ** 2) - - @classmethod - def deriv_loss(cls, x_true, x_pred): - """ - Difference between the slopes of the predicted and true trajectories. - - Args: - x_true (Tensor): The true trajectory. - x_pred (Tensor): The predicted trajectory - - Returns: - Tensor: The derivative loss. - """ - return cls.l2_loss(cls.deriv(x_pred), cls.deriv(x_true)) - - @classmethod - def deriv2_loss(cls, x_true, x_pred): - """ - Difference between the curvature of the predicted and true trajectories. - - Args: - x_true (Tensor): The true trajectory. - x_pred (Tensor): The predicted trajectory - - Returns: - Tensor: The second derivative loss. - """ - return cls.l2_loss(cls.deriv2(x_pred), cls.deriv2(x_true)) - - @staticmethod - def deriv(x): - """ - Calculate the numerical derivative. - - Args: - x (Tensor): The input tensor. - - Returns: - Tensor: The numerical derivative. - """ - return torch.gradient(x, dim=1)[0].squeeze(0) - - @classmethod - def deriv2(cls, x): - """ - Calculate the numerical second derivative. - - Args: - x (Tensor): The input tensor. - - Returns: - Tensor: The numerical second derivative. - """ - return cls.deriv(cls.deriv(x)) + return self.l2(x0, x0_hat) class ODE(nn.Module): @@ -642,6 +637,7 @@ def __init__( self.tanh_reg = tanh_reg self.reg_factor = nn.Parameter(torch.tensor(1.0)) self.activation = activation + self.dtype = dtype layers = [] layers.append(nn.Linear(input_shape, ode_width, dtype=dtype)) layers.append(activation) @@ -652,10 +648,26 @@ def __init__( self.mlp = nn.Sequential(*layers) def forward(self, t, x): - output = self.mlp(x) + # Expect solver to always pass float64 state + if x.dtype != torch.float64: + raise RuntimeError( + f"ODE.forward expected float64 input state, got {x.dtype}" + ) + + # Downcast for MLP computation + x32 = x.to(torch.float32) + out32 = self.mlp(x32) # float32 + if self.tanh_reg: - return self.reg_factor * torch.tanh(output / self.reg_factor) - return output.to(torch.float64) + reg32 = self.reg_factor.to(torch.float32) + activated32 = reg32 * torch.tanh(out32 / reg32) + out64 = activated32.to(torch.float64) + if out64.dtype != torch.float64: + raise RuntimeError("Output not float64 after upcast") + return out64 + + out64 = out32.to(torch.float64) + return out64 class ODEWithParams(nn.Module): diff --git a/codes/train/train_fcts.py b/codes/train/train_fcts.py index e85c6096..b81bf091 100644 --- a/codes/train/train_fcts.py +++ b/codes/train/train_fcts.py @@ -102,7 +102,7 @@ def train_and_save_model( training_id=config["training_id"], ) model.normalisation = data_info - model.checkpointing = config.get("checkpointing", False) + model.checkpointing = config.get("checkpoint", False) surr_idx = config["surrogates"].index(surr_name) batch_size = determine_batch_size(config, surr_idx, mode, metric) diff --git a/codes/tune/optuna_config.yaml b/codes/tune/optuna_config.yaml index 0099e763..56fa3633 100644 --- a/codes/tune/optuna_config.yaml +++ b/codes/tune/optuna_config.yaml @@ -1,13 +1,12 @@ -tuning_id: primordial_tuning +tuning_id: cloud_tuning_fine seed: 42 dataset: - name: primordial + name: cloud log10_transform: True normalise: minmax - per_species: True subset_factor: 1 - tolerance: 1e-15 + tolerance: 1e-25 normalise_per_species: True log_timesteps: True devices: ["cuda:0", "cuda:2", "cuda:3", "cuda:5", "cuda:6", "cuda:8", "cuda:9"] @@ -15,7 +14,7 @@ optuna_logs: False prune: True verbose: False multi_objective: True -population_size: 100 +population_size: 30 target_percentile: 0.99 postgres_config: mode: "local" # "local" or "remote" @@ -23,7 +22,7 @@ postgres_config: user: "optuna_user" host: "localhost" # "localhost" for local use database_folder: "/export/home/rjanssen/postgres/" # only for local use - db_name: "optuna_primordial" # remote mode: single DB for all runs + db_name: "optuna_cloud_2" # remote mode: single DB for all runs sslmode: "require" # if needed password: "" # optional; prefer env PGPASSWORD @@ -36,97 +35,107 @@ global_optuna_params: regularization_factor: type: float low: 1.0e-6 - high: 1.0 - log: true - optimizer: - type: categorical - choices: ["AdamW", "SGD"] - momentum: - type: float - low: 0.0 - high: 0.99 - step: 0.01 - scheduler: - type: categorical - choices: ["cosine", "poly", "schedulefree"] - poly_power: - type: float - low: 0.5 - high: 2.0 - step: 0.1 - eta_min: - type: float - low: 1.0e-3 - high: 1.0 + high: 1.0e-2 log: true activation: type: categorical choices: ["ReLU", "LeakyReLU", "PReLU", "Tanh", "GELU", "Mish", "SiLU", "ELU"] - loss_function: - type: categorical - choices: ["mse", "smoothl1"] - beta: - type: float - low: 0.1 - high: 10.0 - log: true surrogates: - - name: MultiONet - batch_size: 65536 - epochs: 2048 - trials: 200 - optuna_params: - branch_hidden_layers: - type: int - low: 1 - high: 10 - hidden_size: - type: int - low: 10 - high: 700 - step: 10 - output_factor: - type: int - low: 1 - high: 400 - trunk_hidden_layers: - type: int - low: 1 - high: 10 - # params_branch: - # type: categorical - # choices: ["True", "False"] + # - name: MultiONet + # batch_size: 65536 + # epochs: 4096 + # trials: 100 + # optuna_params: + # branch_hidden_layers: + # type: int + # low: 1 + # high: 6 + # hidden_size: + # type: int + # low: 10 + # high: 500 + # step: 10 + # output_factor: + # type: int + # low: 30 + # high: 300 + # step: 10 + # trunk_hidden_layers: + # type: int + # low: 1 + # high: 7 + # beta: + # type: float + # low: 0.1 + # high: 10.0 + # log: true - - name: FullyConnected - batch_size: 65536 - epochs: 2048 - trials: 200 - optuna_params: - hidden_size: - type: int - low: 10 - high: 700 - step: 10 - num_hidden_layers: - type: int - low: 1 - high: 10 + # - name: FullyConnected + # batch_size: 65536 + # epochs: 4096 + # trials: 100 + # optuna_params: + # hidden_size: + # type: int + # low: 10 + # high: 500 + # step: 10 + # num_hidden_layers: + # type: int + # low: 1 + # high: 6 + # beta: + # type: float + # low: 0.1 + # high: 10.0 + # log: true - - name: LatentPoly - batch_size: 512 - epochs: 2048 - trials: 200 + # - name: LatentPoly + # batch_size: 512 + # epochs: 4096 + # trials: 100 + # optuna_params: + # degree: + # type: int + # low: 1 + # high: 10 + # latent_features: + # type: int + # low: 1 + # high: 10 + # coder_layers: + # type: int + # low: 1 + # high: 6 + # coder_width: + # type: int + # low: 10 + # high: 600 + # step: 100 + # # coeff_network: + # # type: categorical + # # choices: ["True", "False"] + # # coeff_width: + # # type: int + # # low: 10 + # # high: 700 + # # step: 10 + # # coeff_layers: + # # type: int + # # low: 1 + # # high: 5 + + - name: LatentNeuralODE + batch_size: 1024 + epochs: 4096 + trials: 100 optuna_params: - degree: - type: int - low: 1 - high: 10 latent_features: type: int low: 1 - high: 10 + high: 15 coder_layers: type: int low: 1 @@ -134,51 +143,17 @@ surrogates: coder_width: type: int low: 10 - high: 700 + high: 400 + step: 10 + ode_width: + type: int + low: 10 + high: 400 step: 10 - # coeff_network: + ode_layers: + type: int + low: 1 + high: 10 + # encode_params: # type: categorical # choices: ["True", "False"] - # coeff_width: - # type: int - # low: 10 - # high: 700 - # step: 10 - # coeff_layers: - # type: int - # low: 1 - # high: 5 - - # - name: LatentNeuralODE - # batch_size: 1024 - # epochs: 2048 - # trials: 200 - # optuna_params: - # latent_features: - # type: int - # low: 1 - # high: 10 - # coder_layers: - # type: int - # low: 1 - # high: 10 - # coder_width: - # type: int - # low: 10 - # high: 700 - # step: 10 - # ode_tanh_reg: - # type: categorical - # choices: ["True", "False"] - # ode_width: - # type: int - # low: 10 - # high: 700 - # step: 10 - # ode_layers: - # type: int - # low: 1 - # high: 10 - # # encode_params: - # # type: categorical - # # choices: ["True", "False"] diff --git a/config.yaml b/config.yaml index 076aea50..93d308b8 100644 --- a/config.yaml +++ b/config.yaml @@ -1,20 +1,22 @@ # Global settings for the benchmark -training_id: "profile_lnode32_improved" -surrogates: ["LatentNeuralODE"] # ["MultiONet", "FullyConnected", "LatentNeuralODE", "LatentPoly"] -batch_size: [716] # [8192, 8192, 512, 512] -epochs: [50] # [20000, 7500, 20000, 15000] +training_id: "primordial_new_main" +surrogates: ["MultiONet", "FullyConnected", "LatentNeuralODE", "LatentPoly"] +batch_size: [71650, 71650, 717, 717] # [65536, 65536, 512, 512] +epochs: [10000, 10000, 10000, 10000] # [4096, 4096, 4096, 4096] dataset: name: "primordial" log10_transform: True log10_transform_params: False + per_species: True normalise: "minmax" # "minmax" # "standardise", "minmax", "disable" use_optimal_params: True - tolerance: 1e-20 + tolerance: 1e-15 subset_factor: 1 log_timesteps: True -devices: ["cuda:5"] +devices: ["cuda:0", "cuda:1", "cuda:2", "cuda:3"] seed: 42 verbose: False +checkpoint: True # Models to train interpolation: From 80ec8c9ebd082a583dc64dc93f8ef8e891bbfaf6 Mon Sep 17 00:00:00 2001 From: Robin Date: Mon, 4 Aug 2025 11:20:59 +0200 Subject: [PATCH 11/13] fix unit tests --- codes/benchmark/bench_fcts.py | 2 ++ test/test_bench_main.py | 23 +++++++++++++---------- test/test_tuning_pipeline.py | 22 +++------------------- 3 files changed, 18 insertions(+), 29 deletions(-) diff --git a/codes/benchmark/bench_fcts.py b/codes/benchmark/bench_fcts.py index fed34c5f..bb64fdc4 100644 --- a/codes/benchmark/bench_fcts.py +++ b/codes/benchmark/bench_fcts.py @@ -319,6 +319,8 @@ def evaluate_accuracy( "error_percentile": percentile, "main_model_training_time": train_time, "main_model_epochs": n_epochs, + "absolute_errors": absolute_errors, + "relative_errors": relative_errors, } return accuracy_metrics diff --git a/test/test_bench_main.py b/test/test_bench_main.py index 04b73c9f..77b71175 100644 --- a/test/test_bench_main.py +++ b/test/test_bench_main.py @@ -14,7 +14,7 @@ def __init__(self, *, train_duration=2.5, n_quantities=2): def load(self, training_id, surr_name, model_identifier): self.load_calls.append((training_id, surr_name, model_identifier)) - def predict(self, *, data_loader): + def predict(self, *, data_loader, leave_log=None, leave_norm=None): """ Return preds and targets of shape [1, 3, n_quantities] where preds are always 2x targets. @@ -26,12 +26,16 @@ def predict(self, *, data_loader): preds = (targets * 2).expand(batch, T, Q) # [[0,2,8], …] return preds, targets + def denormalize(self, arr, leave_log=None, leave_norm=None): + # return the array as is, no normalization + return arr + @pytest.fixture(autouse=True) def no_plots(monkeypatch): # patch out all the plotting functions so they don't error or try to open displays for fn in [ - "plot_relative_errors_over_time", + "plot_error_percentiles_over_time", "plot_error_distribution_per_quantity", "plot_dynamic_correlation_heatmap", ]: @@ -85,10 +89,10 @@ def test_evaluate_accuracy(simple_conf, simple_loader): # check load called with main identifier assert model.load_calls == [("TID", surr, f"{surr.lower()}_main")] # mean squared error ≈ 98/3 - assert metrics["mean_squared_error"] == pytest.approx(98 / 3) + assert metrics["root_mean_squared_error_real"] == pytest.approx(np.sqrt(98 / 3)) # mean absolute error ≈ 14/3 - assert metrics["mean_absolute_error"] == pytest.approx(14 / 3) + assert metrics["mean_absolute_error_real"] == pytest.approx(14 / 3) # relative errors: abs(1)/max(abs(0),0.0) -> 1/0 -> inf; but threshold=0 so yields 1.0 assert metrics["mean_relative_error"] == pytest.approx(1.0) assert metrics["main_model_training_time"] == 3.14 @@ -199,7 +203,7 @@ def prepare_data(self, **kwargs): # only the returned loader is used by predict return "train_loader", None, None - def denormalize(self, arr): + def denormalize(self, arr, leave_log=None, leave_norm=None): return arr model = FakeIterModel(n_timesteps=T, n_quantities=Q) @@ -224,15 +228,14 @@ def denormalize(self, arr): # since preds == targets, all errors should be zero for key in [ + "root_mean_squared_error_log", + "mean_absolute_error_log", + "percentile_absolute_error_log", "mean_squared_error", "mean_absolute_error", - "mean_relative_error", - "median_relative_error", - "max_relative_error", - "min_relative_error", + "absolute_errors", ]: assert metrics[key] == pytest.approx(0.0) # array shapes should be (1, T, Q) assert metrics["absolute_errors"].shape == (1, T, Q) - assert metrics["relative_errors"].shape == (1, T, Q) diff --git a/test/test_tuning_pipeline.py b/test/test_tuning_pipeline.py index f7c4baad..351a5c23 100644 --- a/test/test_tuning_pipeline.py +++ b/test/test_tuning_pipeline.py @@ -1,15 +1,15 @@ -import queue import math +import queue from datetime import datetime, timedelta import pytest from optuna.trial import TrialState from codes.tune.optuna_fcts import ( + MODULE_REGISTRY, + create_objective, make_optuna_params, maybe_set_runtime_threshold, - create_objective, - MODULE_REGISTRY, ) @@ -212,14 +212,6 @@ def test_training_run_single_objective(monkeypatch, tmp_path): None, ), ) - monkeypatch.setattr( - "codes.tune.optuna_fcts.get_data_subset", - lambda *args, **kw: ( - (dummy_data, dummy_data), - (dummy_params, dummy_params), - dummy_timesteps, - ), - ) monkeypatch.setattr("codes.tune.optuna_fcts.set_random_seeds", lambda *a, **k: None) monkeypatch.setattr("codes.tune.optuna_fcts.get_surrogate", lambda name: DummyModel) @@ -274,14 +266,6 @@ def test_training_run_multi_objective(monkeypatch, tmp_path): None, ), ) - monkeypatch.setattr( - "codes.tune.optuna_fcts.get_data_subset", - lambda *args, **kw: ( - (dummy_data, dummy_data), - (dummy_params, dummy_params), - dummy_timesteps, - ), - ) monkeypatch.setattr("codes.tune.optuna_fcts.set_random_seeds", lambda *a, **k: None) class DummyModel2(DummyModel): From 5cf4330dbc6268fe14367ff58762b1a882f935f1 Mon Sep 17 00:00:00 2001 From: Robin Date: Mon, 4 Aug 2025 11:34:15 +0200 Subject: [PATCH 12/13] add loss improvements to latentpoly --- .../LatentNeuralODE/latent_neural_ode.py | 7 - .../LatentPolynomial/latent_poly.py | 172 +++++++----------- 2 files changed, 68 insertions(+), 111 deletions(-) diff --git a/codes/surrogates/LatentNeuralODE/latent_neural_ode.py b/codes/surrogates/LatentNeuralODE/latent_neural_ode.py index da0adfd9..ba295204 100644 --- a/codes/surrogates/LatentNeuralODE/latent_neural_ode.py +++ b/codes/surrogates/LatentNeuralODE/latent_neural_ode.py @@ -219,13 +219,6 @@ def fit( 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 - # ) - scheduler.step() self.validate( diff --git a/codes/surrogates/LatentPolynomial/latent_poly.py b/codes/surrogates/LatentPolynomial/latent_poly.py index f75cf8e6..67545842 100644 --- a/codes/surrogates/LatentPolynomial/latent_poly.py +++ b/codes/surrogates/LatentPolynomial/latent_poly.py @@ -381,24 +381,6 @@ def forward(self, x, t_range, params=None): z_pred = poly_out + z0.unsqueeze(1) return self.decoder(z_pred) - def renormalize_loss_weights( - self, x_true, x_pred, params, criterion: nn.Module = nn.MSELoss - ): - """ - Renormalize the loss weights based on the current loss values so that they are accurately - weighted based on the provided weights. To be used once after a short burn in phase. - - Args: - x_true (Tensor): The true trajectory. - x_pred (Tensor): The predicted trajectory - params (Tensor): Fixed parameters (batch, n_parameters). - criterion (nn.Module): Loss function to use for calculating the losses. - """ - self.loss_weights[0] = 1 / criterion(x_pred, x_true).item() * 100 - self.loss_weights[1] = 1 / self.identity_loss(x_true, params).item() - self.loss_weights[2] = 1 / self.deriv_loss(x_true, x_pred).item() - self.loss_weights[3] = 1 / self.deriv2_loss(x_true, x_pred).item() - def total_loss( self, x_true: Tensor, @@ -407,99 +389,81 @@ def total_loss( criterion: nn.Module = nn.MSELoss(), ): """ - Calculate the total loss based on the loss weights, including params for identity. - """ - return ( - self.loss_weights[0] * criterion(x_pred, x_true) - + self.loss_weights[1] * self.identity_loss(x_true, params) - + self.loss_weights[2] * self.deriv_loss(x_true, x_pred) - + self.loss_weights[3] * self.deriv2_loss(x_true, x_pred) - ) - - def identity_loss(self, x_true: Tensor, params: Tensor = None) -> Tensor: - """ - Identity loss on the initial state x0, handling three cases: - 1. No params: params is None → encode x0 only. - 2. coeff_network=True: encode x0 only, ignore params here. - 3. coeff_network=False and params provided: encode [x0, params]. - """ - x0 = x_true[:, 0, :] # [batch, n_quantities] - # decide what to feed into the encoder: - if params is None or self.config.coeff_network: - enc_in = x0 - else: - enc_in = torch.cat([x0, params], dim=1) - z0 = self.encoder(enc_in) - x0_hat = self.decoder(z0) - return self.l2_loss(x0, x0_hat) - - @classmethod - def l2_loss(cls, x_true: Tensor, x_pred: Tensor): + Total loss: weighted sum of trajectory reconstruction, identity, first derivative, + and second derivative losses. All terms remain in the computation graph. """ - Compute the L2 loss. - - Args: - x_true (Tensor): Ground truth. - x_pred (Tensor): Predictions. - - Returns: - Tensor: L2 loss. - """ - return torch.mean(torch.abs(x_true - x_pred) ** 2) - - @classmethod - def deriv_loss(cls, x_true, x_pred): + w0, w1, w2, w3 = ( + self.loss_weights + ) # assume these are set in config and are floats + + # primary trajectory loss + traj_loss = criterion(x_pred, x_true) + + # identity loss (reconstruct x0) + identity = self.identity_loss(x_true, params) + + # derivative losses: compute once + d_pred = self.first_derivative(x_pred) + d_true = self.first_derivative(x_true) + deriv_loss = self.l2(d_pred, d_true) + + d2_pred = self.second_derivative(x_pred) + d2_true = self.second_derivative(x_true) + deriv2_loss = self.l2(d2_pred, d2_true) + + return w0 * traj_loss + w1 * identity + w2 * deriv_loss + w3 * deriv2_loss + + def first_derivative(self, x: Tensor): + # x: [B, T, F] + h = 1.0 / self.n_timesteps + # central differences for interior + d_center = (x[:, 2:, :] - x[:, :-2, :]) / (2 * h) # [B, T-2, F] + # forward/backward for boundaries + d_first = (x[:, 1:2, :] - x[:, :1, :]) / h # [B,1,F] + d_last = (x[:, -1:, :] - x[:, -2:-1, :]) / h # [B,1,F] + derivative = torch.cat([d_first, d_center, d_last], dim=1) # [B,T,F] + return derivative + + def second_derivative(self, x: Tensor): + # x: [B, T, F] + h = 1.0 / self.n_timesteps + # standard second derivative central + d2_center = (x[:, 2:, :] - 2 * x[:, 1:-1, :] + x[:, :-2, :]) / ( + h * h + ) # [B, T-2, F] + # one-sided approximations at ends (second-order): + # at t0: f''(t0) ≈ (2 f0 - 5 f1 + 4 f2 - f3) / h^2 + d2_first = ( + 2 * x[:, :1, :] - 5 * x[:, 1:2, :] + 4 * x[:, 2:3, :] - x[:, 3:4, :] + ) / (h * h) + # at t_{N-1}: symmetric formula + d2_last = ( + 2 * x[:, -1:, :] - 5 * x[:, -2:-1, :] + 4 * x[:, -3:-2, :] - x[:, -4:-3, :] + ) / (h * h) + d2 = torch.cat([d2_first, d2_center, d2_last], dim=1) # [B,T,F] + return d2 + + def identity_loss(self, x_true: Tensor, params: Tensor = None): """ - Compute the loss based on the difference of first derivatives. + Calculate the identity loss (Encoder -> Decoder) on the initial state x0. Args: - x_true (Tensor): Ground truth. - x_pred (Tensor): Predictions. - + x_true (Tensor): The full trajectory (batch, timesteps, features). + params (Tensor | None): Fixed parameters (batch, n_parameters). Returns: - Tensor: Derivative loss. + Tensor: The identity loss on x0. """ - return cls.l2_loss(cls.deriv(x_pred), cls.deriv(x_true)) - - @classmethod - def deriv2_loss(cls, x_true, x_pred): - """ - Compute the loss based on the difference of second derivatives. - - Args: - x_true (Tensor): Ground truth. - x_pred (Tensor): Predictions. - - Returns: - Tensor: Second derivative loss. - """ - return cls.l2_loss(cls.deriv2(x_pred), cls.deriv2(x_true)) - - @classmethod - def deriv(cls, x): - """ - Compute the numerical first derivative. - - Args: - x (Tensor): Input tensor. - - Returns: - Tensor: First derivative. - """ - return torch.gradient(x, dim=1)[0].squeeze(0) - - @classmethod - def deriv2(cls, x): - """ - Compute the numerical second derivative. - - Args: - x (Tensor): Input tensor. + # only reconstruct the initial state + x0 = x_true[:, 0, :] + if self.config.encode_params and params is not None: + enc_input = torch.cat([x0, params], dim=1) + else: + enc_input = x0 - Returns: - Tensor: Second derivative. - """ - return cls.deriv(cls.deriv(x)) + # encode-decode + z0 = self.encoder(enc_input) + x0_hat = self.decoder(z0) + return self.l2(x0, x0_hat) class Polynomial(nn.Module): From 6ade2be9120c1e3d8bc25947cca4461c4935f155 Mon Sep 17 00:00:00 2001 From: Robin Date: Mon, 4 Aug 2025 13:10:56 +0200 Subject: [PATCH 13/13] add required class attributes --- codes/surrogates/LatentPolynomial/latent_poly.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/codes/surrogates/LatentPolynomial/latent_poly.py b/codes/surrogates/LatentPolynomial/latent_poly.py index 67545842..435ab12a 100644 --- a/codes/surrogates/LatentPolynomial/latent_poly.py +++ b/codes/surrogates/LatentPolynomial/latent_poly.py @@ -61,7 +61,10 @@ def __init__( # else: # self.config.in_features = n_quantities + n_parameters self.model = PolynomialModelWrapper( - config=self.config, device=self.device, n_parameters=n_parameters + config=self.config, + device=self.device, + n_parameters=n_parameters, + n_timesteps=n_timesteps, ) def forward(self, inputs) -> tuple[Tensor, Tensor]: @@ -256,12 +259,14 @@ class PolynomialModelWrapper(nn.Module): coefficient_net (Module | None): The coefficient network (if config.coeff_network is True). """ - def __init__(self, config, device, n_parameters: int = 0): + def __init__(self, config, device, n_parameters: int = 0, n_timesteps: int = 101): super().__init__() self.config = config self.loss_weights = getattr(config, "loss_weights", [100.0, 1.0, 1.0, 1.0]) self.device = device latent_dim = self.config.latent_features + self.n_timesteps = n_timesteps + self.l2 = nn.MSELoss() # Use coeff_network as the single switch. if self.config.coeff_network and n_parameters > 0: @@ -455,7 +460,7 @@ def identity_loss(self, x_true: Tensor, params: Tensor = None): """ # only reconstruct the initial state x0 = x_true[:, 0, :] - if self.config.encode_params and params is not None: + if not self.config.coeff_network and params is not None: enc_input = torch.cat([x0, params], dim=1) else: enc_input = x0