diff --git a/optimas/core/trial.py b/optimas/core/trial.py index d3815d6b..6879cb32 100644 --- a/optimas/core/trial.py +++ b/optimas/core/trial.py @@ -66,6 +66,8 @@ def __init__( self._custom_parameters = ( [] if custom_parameters is None else custom_parameters ) + self._ignored = False + self._ignored_reason = None # Add custom parameters as trial attributes. for param in self._custom_parameters: @@ -127,6 +129,16 @@ def index(self) -> int: def index(self, value): self._index = value + @property + def ignored(self) -> bool: + """Get whether the trial is ignored by the generator.""" + return self._ignored + + @property + def ignored_reason(self) -> str: + """Get the reason why the trial is ignored by the generator.""" + return self._ignored_reason + @property def custom_parameters(self) -> List[TrialParameter]: """Get the list of custom trial parameters.""" @@ -152,6 +164,26 @@ def evaluated(self) -> bool: """Determine whether the trial has been evaluated.""" return self.completed or self.failed + def ignore(self, reason: str): + """Set trial as ignored. + + Parameters + ---------- + reason : str + The reason why the trial is ignored. + """ + # An alternative implementation of this would have been to add a new + # `IGNORED` trial status. However, this would have an issue: + # when adding old trials to an exploration, the original trial status + # could be overwritten by `IGNORED`, and this information would be lost + # for future explorations where this data is reused (for example, + # when using the `resume` option). + # With the current implementation, the value of `ignored` is controlled + # by (and only relevant for) the current exploration. It won't have + # any impact if the data is attached to a future exploration. + self._ignored = True + self._ignored_reason = reason + def mark_as(self, status) -> None: """Set trial status. diff --git a/optimas/diagnostics/exploration_diagnostics.py b/optimas/diagnostics/exploration_diagnostics.py index 696cf6ee..69148e4e 100644 --- a/optimas/diagnostics/exploration_diagnostics.py +++ b/optimas/diagnostics/exploration_diagnostics.py @@ -946,7 +946,10 @@ def print_best_evaluations( print(best_evals[objective_names + varpar_names + anapar_names]) def build_gp_model( - self, parameter: str, minimize: Optional[bool] = None + self, + parameter: str, + minimize: Optional[bool] = None, + fit_out_of_design: Optional[bool] = False, ) -> AxModelManager: """Build a GP model of the specified parameter. @@ -960,6 +963,11 @@ def build_gp_model( Use it to indicate whether lower or higher values of the parameter are better. This is relevant, e.g. to determine the best point of the model. + fit_out_of_design : bool, optional + Whether to fit the surrogate model taking into account evaluations + outside of the range of the varying parameters. This can be useful + if the range of parameter has been reduced during the optimization. + By default, False. Returns ------- @@ -993,4 +1001,5 @@ def build_gp_model( source=self.history, varying_parameters=self.varying_parameters, objectives=[objective], + fit_out_of_design=fit_out_of_design, ) diff --git a/optimas/explorations/base.py b/optimas/explorations/base.py index 24250cba..2a5732b2 100644 --- a/optimas/explorations/base.py +++ b/optimas/explorations/base.py @@ -136,6 +136,9 @@ def history(self) -> pd.DataFrame: """Get the exploration history.""" history = convert_to_dataframe(self._libe_history.H) ordered_columns = ["trial_index", "trial_status"] + # For backward compatibility, add trial_ignored only if available. + if "trial_ignored" in history.columns.values.tolist(): + ordered_columns += ["trial_ignored"] ordered_columns += [p.name for p in self.generator.varying_parameters] ordered_columns += [p.name for p in self.generator.objectives] ordered_columns += [p.name for p in self.generator.analyzed_parameters] diff --git a/optimas/generators/ax/service/base.py b/optimas/generators/ax/service/base.py index 3b700245..3d2bf42b 100644 --- a/optimas/generators/ax/service/base.py +++ b/optimas/generators/ax/service/base.py @@ -15,7 +15,8 @@ GenerationStep, GenerationStrategy, ) -from ax.exceptions.core import DataRequiredError +from ax.modelbridge.transition_criterion import MaxTrials, MinTrials +from ax import Arm from optimas.core import ( Objective, @@ -165,23 +166,50 @@ def _tell(self, trials: List[Trial]) -> None: trial.varying_parameters, trial.parameter_values ): params[var.name] = value - _, trial_id = self._ax_client.attach_trial(params) + try: + _, trial_id = self._ax_client.attach_trial(params) + except ValueError as error: + # Bypass checks from AxClient and manually add a trial + # outside of the search space. + # https://github.com/facebook/Ax/issues/768#issuecomment-1036515242 + if "not a valid value" in str(error): + if self._fit_out_of_design: + ax_trial = self._ax_client.experiment.new_trial() + ax_trial.add_arm(Arm(parameters=params)) + ax_trial.mark_running(no_runner_required=True) + trial_id = ax_trial.index + else: + ignore_reason = ( + f"The parameters {params} are outside of the " + "range of the varying parameters. " + "Set `fit_out_of_design=True` if you want " + "the model to use these data." + ) + trial.ignore(reason=ignore_reason) + continue + else: + raise error ax_trial = self._ax_client.get_trial(trial_id) # Since data was given externally, reduce number of # initialization trials, but only if they have not failed. - if ( - trial.status != TrialStatus.FAILED - and not self._enforce_n_init - ): + if trial.completed and not self._enforce_n_init: generation_strategy = self._ax_client.generation_strategy current_step = generation_strategy.current_step # Reduce only if there are still Sobol trials left. if current_step.model == Models.SOBOL: - current_step.transition_criteria[0].threshold -= 1 + for tc in current_step.transition_criteria: + # Looping over all criterial makes sure we reduce + # the transition thresholds due to `_n_init` + # (i.e., max trials) and `min_trials_observed=1` ( + # i.e., min trials). + if isinstance(tc, (MinTrials, MaxTrials)): + tc.threshold -= 1 generation_strategy._maybe_move_to_next_step() finally: - if trial.completed: + if trial.ignored: + continue + elif trial.completed: outcome_evals = {} # Add objective evaluations. for ev in trial.objective_evaluations: diff --git a/optimas/generators/base.py b/optimas/generators/base.py index a54b6f48..fa7cba2b 100644 --- a/optimas/generators/base.py +++ b/optimas/generators/base.py @@ -250,12 +250,19 @@ def tell( incorporating the evaluated trials. By default ``True``. """ + self._tell(trials) for trial in trials: if trial not in self._given_trials: self._add_external_evaluated_trial(trial) - self._tell(trials) - for trial in trials: - if not trial.failed: + + if trial.ignored: + # Check first if ignored, because it can be ignored + # and completed at the same time. + log_msg = ( + f"Trial {trial.index} ignored. " + f"Reason: {trial.ignored_reason}" + ) + elif trial.completed: log_msg = "Completed trial {} with objective(s) {}".format( trial.index, trial.objectives_as_dict() ) @@ -263,7 +270,7 @@ def tell( log_msg += " and analyzed parameter(s) {}".format( trial.analyzed_parameters_as_dict() ) - else: + elif trial.failed: log_msg = f"Failed to evaluate trial {trial.index}." logger.info(log_msg) if allow_saving_model and self._save_model: @@ -279,11 +286,15 @@ def incorporate_history(self, history: np.ndarray) -> None: """ # Keep only evaluations where the simulation finished successfully. - history = history[history["sim_ended"]] + history_ended = history[history["sim_ended"]] trials = self._create_trials_from_external_data( - history, ignore_unrecognized_parameters=True + history_ended, ignore_unrecognized_parameters=True ) self.tell(trials, allow_saving_model=False) + # Communicate to history array whether the trial has been ignored. + for trial in trials: + i = np.where(history["trial_index"] == trial.index)[0][0] + history["trial_ignored"][i] = trial.ignored def attach_trials( self, @@ -539,6 +550,7 @@ def get_gen_specs( [(var.name, var.dtype) for var in self._varying_parameters] + [("num_procs", int), ("num_gpus", int)] + [("trial_index", int)] + + [("trial_ignored", bool)] + [ (par.save_name, par.dtype) for par in self._custom_trial_parameters diff --git a/optimas/utils/ax/ax_model_manager.py b/optimas/utils/ax/ax_model_manager.py index 64118033..2e014ea1 100644 --- a/optimas/utils/ax/ax_model_manager.py +++ b/optimas/utils/ax/ax_model_manager.py @@ -24,6 +24,7 @@ convert_optimas_to_ax_parameters, convert_optimas_to_ax_objectives, ) + from ax import Arm ax_installed = True except ImportError: @@ -53,6 +54,11 @@ class AxModelManager: parameters that were varied to scan the value of the objectives. The names and data of these parameters must be contained in the source ``DataFrame``. + fit_out_of_design : bool, optional + Whether to fit the surrogate model taking into account evaluations + outside of the range of the varying parameters. This can be useful + if the range of parameter has been reduced during the optimization. + By default, False. """ def __init__( @@ -60,6 +66,7 @@ def __init__( source: Union[AxClient, str, pd.DataFrame], varying_parameters: Optional[List[VaryingParameter]] = None, objectives: Optional[List[Objective]] = None, + fit_out_of_design: Optional[bool] = False, ) -> None: if not ax_installed: raise ImportError( @@ -72,7 +79,7 @@ def __init__( self.ax_client = AxClient.load_from_json_file(filepath=source) elif isinstance(source, pd.DataFrame): self.ax_client = self._build_ax_client_from_dataframe( - source, varying_parameters, objectives + source, varying_parameters, objectives, fit_out_of_design ) else: raise ValueError( @@ -93,6 +100,7 @@ def _build_ax_client_from_dataframe( df: pd.DataFrame, varying_parameters: List[VaryingParameter], objectives: List[Objective], + fit_out_of_design: Optional[bool] = False, ) -> AxClient: """Initialize the AxClient and the model using the given data. @@ -105,6 +113,9 @@ def _build_ax_client_from_dataframe( varying_parameters : list of `VaryingParameter`. List of parameters that were varied to scan the value of the objectives. + fit_out_of_design : bool, optional + Whether to fit the surrogate model taking into account evaluations + outside of the range of the varying parameters. """ # Define parameters for AxClient axparameters = convert_optimas_to_ax_parameters(varying_parameters) @@ -128,7 +139,29 @@ def _build_ax_client_from_dataframe( # Add trials from DataFrame for _, row in df.iterrows(): params = {vp.name: row[vp.name] for vp in varying_parameters} - _, trial_id = ax_client.attach_trial(params) + try: + _, trial_id = ax_client.attach_trial(params) + except ValueError as error: + # Bypass checks from AxClient and manually add a trial + # outside of the search space. + # https://github.com/facebook/Ax/issues/768#issuecomment-1036515242 + if "not a valid value" in str(error): + if fit_out_of_design: + ax_trial = ax_client.experiment.new_trial() + ax_trial.add_arm(Arm(parameters=params)) + ax_trial.mark_running(no_runner_required=True) + trial_id = ax_trial.index + else: + ignore_reason = ( + f"The parameters {params} are outside of the " + "range of the varying parameters. " + "Set `fit_out_of_design=True` if you want " + "the model to use these data." + ) + print(ignore_reason) + continue + else: + raise error data = {obj.name: (row[obj.name], np.nan) for obj in objectives} ax_client.complete_trial(trial_id, raw_data=data) return ax_client diff --git a/tests/test_ax_generators.py b/tests/test_ax_generators.py index 729a4314..0f91f771 100644 --- a/tests/test_ax_generators.py +++ b/tests/test_ax_generators.py @@ -428,6 +428,73 @@ def test_ax_single_fidelity_updated_params(): make_plots(gen) +def test_ax_single_fidelity_resume(): + """ + Test that an exploration with an AxService generator can resume + with an updated range of the varying parameters, even if some + old trials are out of the updated range. + """ + global trial_count + global trials_to_fail + trial_count = 0 + trials_to_fail = [] + + fit_out_of_design_vals = [False, True] + + for fit_out_of_design in fit_out_of_design_vals: + var1 = VaryingParameter("x0", 5.1, 6.0) + var2 = VaryingParameter("x1", -5.0, 15.0) + obj = Objective("f", minimize=False) + p1 = Parameter("p1") + + gen = AxSingleFidelityGenerator( + varying_parameters=[var1, var2], + objectives=[obj], + analyzed_parameters=[p1], + parameter_constraints=["x0 + x1 <= 10"], + outcome_constraints=["p1 <= 30"], + fit_out_of_design=fit_out_of_design, + ) + ev = FunctionEvaluator(function=eval_func_sf) + exploration = Exploration( + generator=gen, + evaluator=ev, + max_evals=20, + sim_workers=2, + exploration_dir_path="./tests_output/test_ax_single_fidelity", + libe_comms="local_threading", + resume=True, + ) + + # Get reference to original AxClient. + ax_client = gen._ax_client + + # Run exploration. + exploration.run(n_evals=1) + + if not fit_out_of_design: + # Check that no old evaluations were added + assert len(exploration.history) == 11 + assert all(exploration.history.trial_ignored.to_numpy()[:-1]) + # Check that the sobol step has not been skipped. + df = ax_client.get_trials_data_frame() + assert len(df) == 1 + assert df["generation_method"].to_numpy()[0] == "Sobol" + + else: + # Check that the old evaluations were added + assert len(exploration.history) == 12 + assert not all(exploration.history.trial_ignored.to_numpy()) + # Check that the sobol step has been skipped. + df = ax_client.get_trials_data_frame() + assert len(df) == 12 + assert df["generation_method"].to_numpy()[-1] == "GPEI" + + check_run_ax_service( + ax_client, gen, exploration, n_failed_expected=2 + ) + + def test_ax_multi_fidelity(): """Test that an exploration with a multifidelity generator runs""" @@ -808,6 +875,7 @@ def test_ax_service_init(): if __name__ == "__main__": test_ax_single_fidelity() + test_ax_single_fidelity_resume() test_ax_single_fidelity_int() test_ax_single_fidelity_moo() test_ax_single_fidelity_fb()