def test_optimization_result_try_get_final_models_for_multiple_models() -> None: data = {"foo": empty_dataset([1], [1]), "bar": empty_dataset([2], [2])} models = {"foo": _PseudoTrainableQuadratic(), "bar": _PseudoTrainableQuadratic()} result: OptimizationResult[None] = OptimizationResult(Ok(Record(data, models, None)), []) assert result.try_get_final_models() is models with pytest.raises(ValueError): result.try_get_final_model()
def test_single_model_acquisition_builder_using_passes_on_correct_dataset_and_model() -> None: class Builder(SingleModelAcquisitionBuilder): def prepare_acquisition_function( self, dataset: Dataset, model: ProbabilisticModel ) -> AcquisitionFunction: assert dataset is data["foo"] assert model is models["foo"] return raise_exc data = {"foo": empty_dataset([1], [1]), "bar": empty_dataset([1], [1])} models = {"foo": QuadraticMeanAndRBFKernel(), "bar": QuadraticMeanAndRBFKernel()} Builder().using("foo").prepare_acquisition_function(data, models)
def test_optimization_result_try_get_final_datasets_for_successful_optimization( ) -> None: data = {"foo": empty_dataset([1], [1])} result: OptimizationResult[None] = OptimizationResult( Ok(Record(data, {"foo": _PseudoTrainableQuadratic()}, None)), []) assert result.try_get_final_datasets() is data assert result.try_get_final_dataset() is data["foo"]
def test_efficient_global_optimization(optimizer: AcquisitionOptimizer[Box]) -> None: class NegQuadratic(SingleModelAcquisitionBuilder): def __init__(self) -> None: self._updated = False def prepare_acquisition_function( self, model: ProbabilisticModel, dataset: Optional[Dataset] = None, ) -> AcquisitionFunction: return lambda x: -quadratic(tf.squeeze(x, -2) - 1) def update_acquisition_function( self, function: AcquisitionFunction, model: ProbabilisticModel, dataset: Optional[Dataset] = None, ) -> AcquisitionFunction: self._updated = True return function function = NegQuadratic() search_space = Box([-10], [10]) ego = EfficientGlobalOptimization(function, optimizer) data, model = empty_dataset([1], [1]), QuadraticMeanAndRBFKernel(x_shift=1) query_point = ego.acquire_single(search_space, model, dataset=data) npt.assert_allclose(query_point, [[1]], rtol=1e-4) assert not function._updated query_point = ego.acquire(search_space, {OBJECTIVE: model}) npt.assert_allclose(query_point, [[1]], rtol=1e-4) assert function._updated
def test_sum_and_product_for_single_builder(reducer_class: type[Sum | Product]) -> None: data, models = {"": empty_dataset([1], [1])}, {"": QuadraticMeanAndRBFKernel()} acq = reducer_class(_Static(lambda x: x ** 2)).prepare_acquisition_function( models, datasets=data ) xs = tf.random.uniform([3, 5, 1], minval=-1.0) npt.assert_allclose(acq(xs), xs ** 2)
def test_optimization_result_try_get_final_models_for_successful_optimization( ) -> None: models = {"foo": _PseudoTrainableQuadratic()} result: OptimizationResult[None] = OptimizationResult( Ok(Record({"foo": empty_dataset([1], [1])}, models, None)), []) assert result.try_get_final_models() is models assert result.try_get_final_model() is models["foo"]
def test_single_model_acquisition_builder_raises_immediately_for_wrong_key() -> None: builder = _ArbitrarySingleBuilder().using("foo") with pytest.raises(KeyError): builder.prepare_acquisition_function( {"bar": empty_dataset([1], [1])}, {"bar": QuadraticMeanAndRBFKernel()} )
def test_ehvi_builder_raises_for_empty_data() -> None: num_obj = 3 dataset = empty_dataset([2], [num_obj]) model = QuadraticMeanAndRBFKernel() with pytest.raises(TF_DEBUGGING_ERROR_TYPES): ExpectedHypervolumeImprovement().prepare_acquisition_function(dataset, model)
def test_probability_of_feasibility_builder_builds_pof(threshold: float, at: tf.Tensor) -> None: builder = ProbabilityOfFeasibility(threshold) acq = builder.prepare_acquisition_function(empty_dataset([1], [1]), QuadraticMeanAndRBFKernel()) expected = probability_of_feasibility(QuadraticMeanAndRBFKernel(), threshold, at) npt.assert_allclose(acq(at), expected)
def test_single_builder_raises_immediately_for_wrong_key( single_builder: SingleModelAcquisitionBuilder | SingleModelBatchAcquisitionBuilder, ) -> None: builder = single_builder.using("foo") with pytest.raises(KeyError): builder.prepare_acquisition_function( {"bar": empty_dataset([1], [1])}, {"bar": QuadraticMeanAndRBFKernel()})
def test_reducer__reduce() -> None: class Mean(Reducer): def _reduce(self, inputs: Sequence[tf.Tensor]) -> tf.Tensor: return tf.reduce_mean(inputs, axis=0) mean = Mean(_Static(lambda x: -2.0 * x), _Static(lambda x: 3.0 * x)) data, models = {"": empty_dataset([1], [1])}, {"": QuadraticMeanAndRBFKernel()} acq = mean.prepare_acquisition_function(models, datasets=data) xs = tf.random.uniform([3, 5, 1], minval=-1.0) npt.assert_allclose(acq(xs), 0.5 * xs)
def test_product() -> None: prod = Product(_Static(lambda x: x + 1), _Static(lambda x: x + 2)) data, models = { "": empty_dataset([1], [1]) }, { "": QuadraticMeanAndRBFKernel() } acq = prod.prepare_acquisition_function(data, models) xs = tf.random.uniform([3, 5, 1], minval=-1.0, dtype=tf.float64) npt.assert_allclose(acq(xs), (xs + 1) * (xs + 2))
def test_bayesian_optimizer_optimize_raises_for_invalid_rule_keys_and_default_acquisition( ) -> None: optimizer = BayesianOptimizer(lambda x: x[:1], Box([-1], [1])) data, models = { "foo": empty_dataset([1], [1]) }, { "foo": _PseudoTrainableQuadratic() } with pytest.raises(ValueError): optimizer.optimize(3, data, models)
def test_sum() -> None: sum_ = Sum(_Static(lambda x: x), _Static(lambda x: x**2), _Static(lambda x: x**3)) data, models = { "": empty_dataset([1], [1]) }, { "": QuadraticMeanAndRBFKernel() } acq = sum_.prepare_acquisition_function(data, models) xs = tf.random.uniform([3, 5, 1], minval=-1.0) npt.assert_allclose(acq(xs), xs + xs**2 + xs**3)
def test_bayesian_optimizer_optimize_raises_for_negative_steps( num_steps: int) -> None: optimizer = BayesianOptimizer(_quadratic_observer, Box([-1], [1])) data, models = { "": empty_dataset([1], [1]) }, { "": _PseudoTrainableQuadratic() } with pytest.raises(ValueError, match="num_steps"): optimizer.optimize(num_steps, data, models)
def test_bayesian_optimizer_optimize_doesnt_track_state_if_told_not_to() -> None: class _UncopyableModel(_PseudoTrainableQuadratic): def __deepcopy__(self, memo: dict[int, object]) -> NoReturn: assert False data, models = {OBJECTIVE: empty_dataset([1], [1])}, {OBJECTIVE: _UncopyableModel()} history = ( BayesianOptimizer(_quadratic_observer, Box([-1], [1])) .optimize(5, data, models, track_state=False) .history ) assert len(history) == 0
def test_efficient_global_optimization( optimizer: AcquisitionOptimizer[Box]) -> None: class NegQuadratic(AcquisitionFunctionBuilder): def prepare_acquisition_function( self, datasets: Mapping[str, Dataset], models: Mapping[str, ProbabilisticModel]) -> AcquisitionFunction: return lambda x: -quadratic(tf.squeeze(x, -2) - 1) search_space = Box([-10], [10]) ego = EfficientGlobalOptimization(NegQuadratic(), optimizer) data, model = empty_dataset([1], [1]), QuadraticMeanAndRBFKernel(x_shift=1) query_point, _ = ego.acquire(search_space, {"": data}, {"": model}) npt.assert_allclose(query_point, [[1]], rtol=1e-4)
) .astuple() ) final_model = final_opt_state.unwrap().model if fit_initial_model: # optimized at start and end of first BO step assert final_model._optimize_count == 2 # type: ignore else: # optimized just at end of first BO step assert final_model._optimize_count == 1 # type: ignore @pytest.mark.parametrize( "datasets, models", [ ({}, {}), ({"foo": empty_dataset([1], [1])}, {}), ({"foo": empty_dataset([1], [1])}, {"bar": _PseudoTrainableQuadratic()}), ( {"foo": empty_dataset([1], [1])}, {"foo": _PseudoTrainableQuadratic(), "bar": _PseudoTrainableQuadratic()}, ), ], ) def test_bayesian_optimizer_optimize_raises_for_invalid_keys( datasets: dict[str, Dataset], models: dict[str, TrainableProbabilisticModel] ) -> None: search_space = Box([-1], [1]) optimizer = BayesianOptimizer(lambda x: {"foo": Dataset(x, x)}, search_space) rule = FixedAcquisitionRule([[0.0]]) with pytest.raises(ValueError): optimizer.optimize(10, datasets, models, rule)
@pytest.mark.parametrize( "models", [ {}, { "foo": QuadraticMeanAndRBFKernel() }, { "foo": QuadraticMeanAndRBFKernel(), OBJECTIVE: QuadraticMeanAndRBFKernel() }, ], ) @pytest.mark.parametrize("datasets", [{}, { OBJECTIVE: empty_dataset([1], [1]) }]) def test_thompson_sampling_raises_for_invalid_models_keys( datasets: dict[str, Dataset], models: dict[str, ProbabilisticModel]) -> None: search_space = Box([-1], [1]) rule = ThompsonSampling(100, 10) with pytest.raises(ValueError): rule.acquire(search_space, datasets, models) def test_efficient_global_optimization_raises_for_no_query_points() -> None: with pytest.raises(ValueError): EfficientGlobalOptimization(num_query_points=0)
observer = _CountingObserver() optimizer = BayesianOptimizer(observer, Box([-1], [1])) data = mk_dataset([[0.5]], [[0.25]]) optimizer.optimize(steps, data, _PseudoTrainableQuadratic()).final_result.unwrap() assert observer.call_count == steps @pytest.mark.parametrize( "datasets, models", [ ({}, {}), ({ "foo": empty_dataset([1], [1]) }, {}), ({ "foo": empty_dataset([1], [1]) }, { "bar": _PseudoTrainableQuadratic() }), ( { "foo": empty_dataset([1], [1]) }, { "foo": _PseudoTrainableQuadratic(), "bar": _PseudoTrainableQuadratic() }, ),