From fa0c6baeca65205bcd67c1e0f73f6d9599828926 Mon Sep 17 00:00:00 2001 From: maxbalrog Date: Fri, 5 Jun 2026 10:32:52 +0200 Subject: [PATCH 1/6] add surrogate model from npz trial data --- src/quvac/cluster/optimization.py | 59 +++++++++++++++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/src/quvac/cluster/optimization.py b/src/quvac/cluster/optimization.py index 08b5631..0347990 100755 --- a/src/quvac/cluster/optimization.py +++ b/src/quvac/cluster/optimization.py @@ -934,6 +934,65 @@ def predict_2d(self, param_names, param_values, fixed_params=None): return mean, covariance +class SurrogateModelFromNPZ(SurrogateModel): + """ + Main info is taken from this guide + https://ax.dev/docs/recipes/existing-data + """ + def __init__(self, data, optimization_params, metric="N_total"): + self.metric = metric + + self.parameter_space = optimization_params["parameters"] + self.ax_client = self.create_ax_client(optimization_params) + + trials = self.prepare_data(data) + self.attach_data(trials) + + self.experiment_data = self.ax_client._experiment.fetch_data() + # refit the model + self.model = Generators.BOTORCH_MODULAR( + experiment=self.ax_client._experiment, + data=self.experiment_data, + ) + + def create_ax_client(self, optimization_params): + params_for_ax = prepare_params_for_ax(self.parameter_space) + experiment_name = optimization_params.get("experiment_name", "imported_data") + parameter_constraints = optimization_params.get("parameter_constraints", None) + # Set up optimization client + ax_client = _create_new_ax_client( + experiment_name, + params_for_ax, + parameter_constraints, + ) + return ax_client + + def prepare_data(self, data): + metric = self.metric + n_trials = len(data[metric]) + trials = [] + for i in range(n_trials): + trials.append( + ( + {key: data[key][i] for key in self.parameter_space.keys()}, + {metric: data[metric][i]}, + ) + ) + return trials + + def attach_data(self, trials): + for parameters, raw_data in trials.items(): + # First attach the trial and note the trial index + trial_index = self.ax_client.attach_trial( + parameters=parameters, + ) + + # Then complete the trial with the existing data + self.ax_client.complete_trial( + trial_index=trial_index, raw_data=raw_data, + ) + + def main_optimization(): """ Main function to run the optimization. From 78f48aa1d4f3ea495a8a0500cd37139a7d2b1329 Mon Sep 17 00:00:00 2001 From: maxbalrog Date: Fri, 5 Jun 2026 10:42:33 +0200 Subject: [PATCH 2/6] add option for predicting the confidence interval instead of covariance --- src/quvac/cluster/optimization.py | 25 +++++++++++++++++-------- 1 file changed, 17 insertions(+), 8 deletions(-) diff --git a/src/quvac/cluster/optimization.py b/src/quvac/cluster/optimization.py index 0347990..bd0da8d 100755 --- a/src/quvac/cluster/optimization.py +++ b/src/quvac/cluster/optimization.py @@ -850,14 +850,16 @@ def _update_input_parameters(self, model_input, fixed_params): model_input = [ObservationFeatures(parameters=pt) for pt in model_input] return model_input - def _make_model_prediction(self, model_input, fixed_params): + def _make_model_prediction(self, model_input, fixed_params, ci=False): model_input = self._update_input_parameters(model_input, fixed_params) mean, covariance = self.model.predict(model_input) mean = np.array(mean[self.metric]) covariance = np.array(covariance[self.metric][self.metric]) + if ci: + covariance = 1.96*np.sqrt(covariance) return mean, covariance - def predict_at_point(self, param_name, param_value, fixed_params=None): + def predict_at_point(self, param_name, param_value, fixed_params=None, ci=False): """ Make a prediction for a particular parameter value. @@ -869,17 +871,20 @@ def predict_at_point(self, param_name, param_value, fixed_params=None): Parameter value (prediction point). fixed_params: dict of {param_name: param_value}, optional Fixed values of other parameters (useful for high-dimentional optimization). + ci: bool + Return 95% confidence interval instead of covariance. Returns ------- (np.ndarray, np.ndarray) - Mean and covariance values at a given parameter point. + Mean and covariance (or confidence interval) values at a given parameter + point. """ model_input = [{param_name: param_value}] - mean, covariance = self._make_model_prediction(model_input, fixed_params) + mean, covariance = self._make_model_prediction(model_input, fixed_params, ci=ci) return mean, covariance - def predict_1d(self, param_name, param_values, fixed_params=None): + def predict_1d(self, param_name, param_values, fixed_params=None, ci=False): """ Make a prediction for a particular parameter value. @@ -891,6 +896,8 @@ def predict_1d(self, param_name, param_values, fixed_params=None): Array of parameter values. fixed_params: dict of {param_name: param_value}, optional Fixed values of other parameters (useful for high-dimentional optimization). + ci: bool + Return 95% confidence interval instead of covariance. Returns ------- @@ -900,10 +907,10 @@ def predict_1d(self, param_name, param_values, fixed_params=None): err_msg = "Method `predict_1d` works only with sequences" assert isinstance(param_values, Iterable), err_msg model_input = [{param_name: value} for value in param_values] - mean, covariance = self._make_model_prediction(model_input, fixed_params) + mean, covariance = self._make_model_prediction(model_input, fixed_params, ci=ci) return mean, covariance - def predict_2d(self, param_names, param_values, fixed_params=None): + def predict_2d(self, param_names, param_values, fixed_params=None, ci=False): """ Make a prediction for a particular parameter value. @@ -915,6 +922,8 @@ def predict_2d(self, param_names, param_values, fixed_params=None): Two arrays of parameter values. fixed_params: dict of {param_name: param_value}, optional Fixed values of other parameters (useful for high-dimentional optimization). + ci: bool + Return 95% confidence interval instead of covariance. Returns ------- @@ -929,7 +938,7 @@ def predict_2d(self, param_names, param_values, fixed_params=None): {param_name_1: value1, param_name_2: value2} for value1,value2 in itertools.product(*param_values) ] - mean, covariance = self._make_model_prediction(model_input, fixed_params) + mean, covariance = self._make_model_prediction(model_input, fixed_params, ci=ci) mean, covariance = mean.reshape((n1,n2)), covariance.reshape((n1,n2)) return mean, covariance From 3255570bf4c6703060aa7bca78baea5b1a9a2a98 Mon Sep 17 00:00:00 2001 From: maxbalrog Date: Fri, 5 Jun 2026 15:12:09 +0200 Subject: [PATCH 3/6] change attribute for parameter names --- src/quvac/cluster/optimization.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/quvac/cluster/optimization.py b/src/quvac/cluster/optimization.py index bd0da8d..07e0329 100755 --- a/src/quvac/cluster/optimization.py +++ b/src/quvac/cluster/optimization.py @@ -966,6 +966,7 @@ def __init__(self, data, optimization_params, metric="N_total"): def create_ax_client(self, optimization_params): params_for_ax = prepare_params_for_ax(self.parameter_space) + self.parameter_names = [param.name for param in params_for_ax] experiment_name = optimization_params.get("experiment_name", "imported_data") parameter_constraints = optimization_params.get("parameter_constraints", None) # Set up optimization client @@ -983,7 +984,7 @@ def prepare_data(self, data): for i in range(n_trials): trials.append( ( - {key: data[key][i] for key in self.parameter_space.keys()}, + {key: data[key][i] for key in self.parameter_names}, {metric: data[metric][i]}, ) ) From 268b1097587224db7edbb1cb4a21d522fc74e1c6 Mon Sep 17 00:00:00 2001 From: maxbalrog Date: Fri, 5 Jun 2026 15:14:40 +0200 Subject: [PATCH 4/6] fix issue: trials is a list --- src/quvac/cluster/optimization.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/quvac/cluster/optimization.py b/src/quvac/cluster/optimization.py index 07e0329..264394a 100755 --- a/src/quvac/cluster/optimization.py +++ b/src/quvac/cluster/optimization.py @@ -991,7 +991,7 @@ def prepare_data(self, data): return trials def attach_data(self, trials): - for parameters, raw_data in trials.items(): + for parameters, raw_data in trials: # First attach the trial and note the trial index trial_index = self.ax_client.attach_trial( parameters=parameters, From f1d70c3192114692cc32c2ffd690e0d5ad94bf05 Mon Sep 17 00:00:00 2001 From: maxbalrog Date: Mon, 15 Jun 2026 11:16:51 +0200 Subject: [PATCH 5/6] complete SurrogateModelFromNPZ class --- src/quvac/cluster/optimization.py | 114 ++++++++++++++++++------------ 1 file changed, 70 insertions(+), 44 deletions(-) diff --git a/src/quvac/cluster/optimization.py b/src/quvac/cluster/optimization.py index 264394a..6808051 100755 --- a/src/quvac/cluster/optimization.py +++ b/src/quvac/cluster/optimization.py @@ -651,6 +651,51 @@ def setup_ax_client( return ax_client +def choose_generation_strategy(ax_client, optimization_params, num_parameters): + number_of_ini_trials = optimization_params.get("number_of_ini_trials", + max(5, 2 * num_parameters)) + gs_initialization_random_seed = optimization_params.get( + "gs_initialization_random_seed", None + ) + if gs_initialization_random_seed is None: + gs_initialization_random_seed = int(np.random.randint(low=0, high=1000000)) + _logger.info( + f"Initial random seed for generation strategy: {gs_initialization_random_seed}" + ) + + generation_strategy_type = optimization_params.get("generation_strategy_type", + "fast") + + if generation_strategy_type in ["fast", "quality"]: + ax_client.configure_generation_strategy( + method=generation_strategy_type, + initialization_budget=number_of_ini_trials, + ) + _logger.info( + f"Using default generation strategy with {generation_strategy_type} option." + ) + elif generation_strategy_type == "custom": + generation_strategy = construct_generation_strategy( + num_parameters, + node_name="Custom-GP", + num_random_trials=number_of_ini_trials, + gs_random_seed=gs_initialization_random_seed, + ) + ax_client.set_generation_strategy( + generation_strategy=generation_strategy, + ) + _logger.info( + "Using custom generation strategy with Matern(2.5) kernel." + ) + else: + err_msg = ( + "Only generation strategy types ['fast', 'quality', 'custom'] are supported" + f"but you passed {generation_strategy_type}." + ) + raise NotImplementedError(err_msg) + return ax_client + + def cluster_optimization(ini_file, save_path=None, wisdom_file=None): """ Launch optimization of quvac simulation for a given initial configuration file. @@ -698,9 +743,6 @@ def cluster_optimization(ini_file, save_path=None, wisdom_file=None): params_for_ax = prepare_params_for_ax(optimization_params["parameters"]) num_parameters = len(params_for_ax) - # custom generation strategy - number_of_ini_trials = optimization_params.get("number_of_ini_trials", - max(5, 2 * num_parameters)) experiment_name = optimization_params.get("experiment_name", "test_optimization") # Set up optimization client @@ -721,44 +763,8 @@ def cluster_optimization(ini_file, save_path=None, wisdom_file=None): if metrics_to_track: ax_client.configure_tracking_metrics(metrics_to_track) - gs_initialization_random_seed = optimization_params.get( - "gs_initialization_random_seed", None - ) - if gs_initialization_random_seed is None: - gs_initialization_random_seed = int(np.random.randint(low=0, high=1000000)) - _logger.info( - f"Initial random seed for generation strategy: {gs_initialization_random_seed}" - ) - - generation_strategy_type = optimization_params.get("generation_strategy_type", - "fast") - if generation_strategy_type in ["fast", "quality"]: - ax_client.configure_generation_strategy( - method=generation_strategy_type, - initialization_budget=number_of_ini_trials, - ) - _logger.info( - f"Using default generation strategy with {generation_strategy_type} option." - ) - elif generation_strategy_type == "custom": - generation_strategy = construct_generation_strategy( - num_parameters, - node_name="Custom-GP", - num_random_trials=number_of_ini_trials, - gs_random_seed=gs_initialization_random_seed, - ) - ax_client.set_generation_strategy( - generation_strategy=generation_strategy, - ) - _logger.info( - "Using custom generation strategy with Matern(2.5) kernel." - ) - else: - err_msg = ( - "Only generation strategy types ['fast', 'quality', 'custom'] are supported" - f"but you passed {generation_strategy_type}." - ) - raise NotImplementedError(err_msg) + # generation strategy + ax_client = choose_generation_strategy(ax_client, optimization_params, num_parameters) max_parallel_jobs = cluster_params.get("max_parallel_jobs", 3) executor = setup_job_executor_from_params(cluster_params, save_path, @@ -952,7 +958,17 @@ def __init__(self, data, optimization_params, metric="N_total"): self.metric = metric self.parameter_space = optimization_params["parameters"] - self.ax_client = self.create_ax_client(optimization_params) + self.noiseless_observations = optimization_params.get("noiseless_observations", False) + ax_client = self.create_ax_client(optimization_params) + + # define generation strategy to have the ability to choose a different kernel + self.ax_client = choose_generation_strategy( + ax_client, optimization_params, len(self.parameter_names) + ) + gs = ax_client._generation_strategy._nodes[-1].generator_spec + gs_kwargs = gs.generator_kwargs + if gs_kwargs is None: + gs_kwargs = {} trials = self.prepare_data(data) self.attach_data(trials) @@ -962,6 +978,7 @@ def __init__(self, data, optimization_params, metric="N_total"): self.model = Generators.BOTORCH_MODULAR( experiment=self.ax_client._experiment, data=self.experiment_data, + **gs_kwargs ) def create_ax_client(self, optimization_params): @@ -975,17 +992,26 @@ def create_ax_client(self, optimization_params): params_for_ax, parameter_constraints, ) + objective = optimization_params.get("parameter_constraints", "N_disc") + ax_client.configure_optimization( + objective=f"{objective}", + outcome_constraints=optimization_params.get("outcome_constraints", None), + ) return ax_client def prepare_data(self, data): metric = self.metric n_trials = len(data[metric]) trials = [] - for i in range(n_trials): + for i in range(n_trials): + if self.noiseless_observations: + metric_outcome = (data[metric][i],0.0) + else: + metric_outcome = data[metric][i] trials.append( ( {key: data[key][i] for key in self.parameter_names}, - {metric: data[metric][i]}, + {metric: metric_outcome}, ) ) return trials From e2745c6fc6fdb55c6a59e99497d3de2a9764ffc6 Mon Sep 17 00:00:00 2001 From: maxbalrog Date: Mon, 15 Jun 2026 12:58:00 +0200 Subject: [PATCH 6/6] fix ruff style issues --- src/quvac/cluster/optimization.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/quvac/cluster/optimization.py b/src/quvac/cluster/optimization.py index 6808051..2038d31 100755 --- a/src/quvac/cluster/optimization.py +++ b/src/quvac/cluster/optimization.py @@ -764,7 +764,9 @@ def cluster_optimization(ini_file, save_path=None, wisdom_file=None): ax_client.configure_tracking_metrics(metrics_to_track) # generation strategy - ax_client = choose_generation_strategy(ax_client, optimization_params, num_parameters) + ax_client = choose_generation_strategy( + ax_client, optimization_params, num_parameters + ) max_parallel_jobs = cluster_params.get("max_parallel_jobs", 3) executor = setup_job_executor_from_params(cluster_params, save_path, @@ -958,7 +960,9 @@ def __init__(self, data, optimization_params, metric="N_total"): self.metric = metric self.parameter_space = optimization_params["parameters"] - self.noiseless_observations = optimization_params.get("noiseless_observations", False) + self.noiseless_observations = optimization_params.get( + "noiseless_observations", False + ) ax_client = self.create_ax_client(optimization_params) # define generation strategy to have the ability to choose a different kernel