diff --git a/src/quvac/cluster/optimization.py b/src/quvac/cluster/optimization.py index 08b5631..2038d31 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,45 +763,11 @@ 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 + ax_client = choose_generation_strategy( + ax_client, optimization_params, num_parameters ) - 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) - max_parallel_jobs = cluster_params.get("max_parallel_jobs", 3) executor = setup_job_executor_from_params(cluster_params, save_path, max_parallel_jobs) @@ -850,14 +858,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 +879,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 +904,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 +915,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 +930,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,11 +946,93 @@ 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 +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.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) + + 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, + **gs_kwargs + ) + + 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 + ax_client = _create_new_ax_client( + experiment_name, + 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): + 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: metric_outcome}, + ) + ) + return trials + + def attach_data(self, trials): + for parameters, raw_data in trials: + # 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.