From 4043d86fc6efe608c2fc6bdc6b26d6b173002750 Mon Sep 17 00:00:00 2001 From: nabenabe0928 Date: Fri, 9 Aug 2024 08:53:26 +0200 Subject: [PATCH 1/2] Refactor botorch --- optuna_integration/botorch/botorch.py | 54 +++++++++++++-------------- 1 file changed, 25 insertions(+), 29 deletions(-) diff --git a/optuna_integration/botorch/botorch.py b/optuna_integration/botorch/botorch.py index d18d4cc2..44886fd5 100644 --- a/optuna_integration/botorch/botorch.py +++ b/optuna_integration/botorch/botorch.py @@ -81,6 +81,19 @@ def _get_sobol_qmc_normal_sampler(num_samples: int) -> SobolQMCNormalSampler: ) +def _validate_botorch_version_for_constrained_opt(func_name: str) -> None: + if version.parse(botorch.version.version) < version.parse("0.9.0"): + raise ImportError( + f"{func_name} requires botorch>=0.9.0 for constrained problems, but got " + f"botorch={botorch.version.version}.\n" + "Please run ``pip install botorch --upgrade``." + ) + + +def _get_constraint_funcs(n_constraints: int) -> list[Callable[["torch.Tensor"], "torch.Tensor"]]: + return [lambda Z: Z[..., -n_constraints + i] for i in range(n_constraints)] + + @experimental_func("3.3.0") def logei_candidates_func( train_x: "torch.Tensor", @@ -235,11 +248,7 @@ def qei_candidates_func( if train_obj.size(-1) != 1: raise ValueError("Objective may only contain single values with qEI.") if train_con is not None: - if version.parse(botorch.version.version) < version.parse("0.9.0"): - raise ImportError( - "qei_candidates_func requires botorch >=0.9.0. for constrained problems." - "Please upgrade botorch" - ) + _validate_botorch_version_for_constrained_opt("qei_candidates_func") train_y = torch.cat([train_obj, train_con], dim=-1) is_feas = (train_con <= 0).all(dim=-1) @@ -257,7 +266,7 @@ def qei_candidates_func( n_constraints = train_con.size(1) additonal_qei_kwargs = { "objective": GenericMCObjective(lambda Z, X: Z[..., 0]), - "constraints": [lambda Z: Z[..., -n_constraints + i] for i in range(n_constraints)], + "constraints": _get_constraint_funcs(n_constraints), } else: train_y = train_obj @@ -320,17 +329,13 @@ def qnei_candidates_func( if train_obj.size(-1) != 1: raise ValueError("Objective may only contain single values with qNEI.") if train_con is not None: - if version.parse(botorch.version.version) < version.parse("0.9.0"): - raise ImportError( - "qnei_candidates_func requires botorch >=0.9.0. for constrained problems." - "Please upgrade botorch" - ) + _validate_botorch_version_for_constrained_opt("qnei_candidates_func") train_y = torch.cat([train_obj, train_con], dim=-1) n_constraints = train_con.size(1) additional_qnei_kwargs = { "objective": GenericMCObjective(lambda Z, X: Z[..., 0]), - "constraints": [lambda Z: Z[..., -n_constraints + i] for i in range(n_constraints)], + "constraints": _get_constraint_funcs(n_constraints), } else: train_y = train_obj @@ -400,7 +405,7 @@ def qehvi_candidates_func( n_constraints = train_con.size(1) additional_qehvi_kwargs = { "objective": IdentityMCMultiOutputObjective(outcomes=list(range(n_objectives))), - "constraints": [lambda Z: Z[..., -n_constraints + i] for i in range(n_constraints)], + "constraints": _get_constraint_funcs(n_constraints), } else: train_y = train_obj @@ -548,9 +553,7 @@ def qnehvi_candidates_func( n_constraints = train_con.size(1) additional_qnehvi_kwargs = { "objective": IdentityMCMultiOutputObjective(outcomes=list(range(n_objectives))), - "constraints": [ - (lambda Z, i=i: Z[..., -n_constraints + i]) for i in range(n_constraints) - ], + "constraints": _get_constraint_funcs(n_constraints), } else: train_y = train_obj @@ -631,23 +634,18 @@ def qparego_candidates_func( scalarization = get_chebyshev_scalarization(weights=weights, Y=train_obj) if train_con is not None: - if version.parse(botorch.version.version) < version.parse("0.9.0"): - raise ImportError( - "qparego_candidates_func requires botorch >=0.9.0. for constrained problems." - "Please upgrade botorch" - ) - + _validate_botorch_version_for_constrained_opt("qparego_candidates_func") train_y = torch.cat([train_obj, train_con], dim=-1) n_constraints = train_con.size(1) objective = GenericMCObjective(lambda Z, X: scalarization(Z[..., :n_objectives])) - additional_kwargs = { - "constraints": [lambda Z: Z[..., -n_constraints + i] for i in range(n_constraints)], + additional_qei_kwargs = { + "constraints": _get_constraint_funcs(n_constraints), } else: train_y = train_obj objective = GenericMCObjective(scalarization) - additional_kwargs = {} + additional_qei_kwargs = {} train_x = normalize(train_x, bounds=bounds) if pending_x is not None: @@ -663,7 +661,7 @@ def qparego_candidates_func( sampler=_get_sobol_qmc_normal_sampler(256), objective=objective, X_pending=pending_x, - **additional_kwargs, + **additional_qei_kwargs, ) standard_bounds = torch.zeros_like(bounds) @@ -711,9 +709,7 @@ def qkg_candidates_func( n_constraints = train_con.size(1) objective = ConstrainedMCObjective( objective=lambda Z, X: Z[..., 0], - constraints=[ - (lambda Z, i=i: Z[..., -n_constraints + i]) for i in range(n_constraints) - ], + constraints=_get_constraint_funcs(n_constraints), ) else: train_y = train_obj From 0e815a38c22e94e653b30957d0960ba9da0a2059 Mon Sep 17 00:00:00 2001 From: nabenabe0928 Date: Wed, 14 Aug 2024 07:18:38 +0200 Subject: [PATCH 2/2] Add future annotations --- optuna_integration/botorch/botorch.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/optuna_integration/botorch/botorch.py b/optuna_integration/botorch/botorch.py index 44886fd5..3249659a 100644 --- a/optuna_integration/botorch/botorch.py +++ b/optuna_integration/botorch/botorch.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from typing import Any from typing import Callable from typing import Dict