Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Categorical trust regions #865

Merged
merged 68 commits into from
Aug 27, 2024
Merged
Changes from 10 commits
Commits
Show all changes
68 commits
Select commit Hold shift + click to select a range
c82a549
CategoricalSearchSpace
Jul 26, 2024
260fff7
Fix AutoGraph error
Jul 26, 2024
976953f
EncoderFunction
Jul 29, 2024
2a863b5
Move one hot encoder to space.py
Jul 29, 2024
5e3cae0
DiscreteSearchSpaceABC
Jul 29, 2024
33aa239
Support one-hot encoding mixed search spaces
Jul 29, 2024
5e135eb
Not yet using latest gpflow
Jul 29, 2024
8a541a1
mypy
Jul 29, 2024
6bc5d4d
Refactor to allow categorical TR spaces
Jul 29, 2024
74122e5
Add more tests
Jul 30, 2024
d0f9376
More tests
Jul 30, 2024
6d2a4e9
Test to_tags
Jul 30, 2024
c0cfd42
has_bounds property
Jul 30, 2024
9c588a9
encode_query_points decorator
Jul 30, 2024
e6a0692
Encode some more query points
Jul 30, 2024
003d8a7
Categorical Trust Regions
Jul 31, 2024
25d20ca
Tweaks
Jul 31, 2024
69e7c0b
Categorical search spaces
Jul 31, 2024
c071112
Remove superfluous encodings
Jul 31, 2024
563e765
Migrate to encode method
Aug 1, 2024
6acd6b7
Experiment with encoded model approaches
Aug 2, 2024
d864468
Merge branch 'uri/categorical_search_spaces' into uri/experiment_with…
Aug 2, 2024
d8be1b7
Fix typing
Aug 2, 2024
cd9b89e
Better name
Aug 2, 2024
7a0dde0
Make encoded methods final
Aug 2, 2024
5867830
Docstrings
Aug 2, 2024
110fc28
EncodedFastUpdateModel
Aug 2, 2024
b05d413
Missed finals
Aug 2, 2024
485e284
inherit_check_shapes
Aug 3, 2024
1209ff4
Review comments
Aug 7, 2024
a6a9545
Merge branch 'uri/categorical_search_spaces' into uri/experiment_with…
Aug 7, 2024
ef78736
Remove trust region rule changes
Aug 7, 2024
f97d062
Revert "Remove trust region rule changes"
Aug 7, 2024
3612083
Remove internal check shapes (external ones are still there)
Aug 7, 2024
14e2710
Merge branch 'uri/experiment_with_encoded_models' into uri/categorica…
Aug 7, 2024
eee39cc
Optimize GeneralDiscreteSearchSpaces
Aug 7, 2024
ee04910
Merge branch 'uri/experiment_with_encoded_models' into uri/categorica…
Aug 7, 2024
b13c147
Use float indices to support product search spaces
Aug 7, 2024
63b62b1
Test non-integer indices
Aug 7, 2024
d162eff
Merge branch 'uri/categorical_search_spaces' into uri/experiment_with…
Aug 7, 2024
ad3e5ca
Merge branch 'uri/categorical_search_spaces' into uri/categorical_tru…
Aug 7, 2024
fc8b5df
Docstring example
Aug 7, 2024
ed45943
Merge branch 'uri/categorical_search_spaces' into uri/experiment_with…
Aug 8, 2024
df56c13
Merge branch 'uri/experiment_with_encoded_models' into uri/categorica…
Aug 8, 2024
3d2b286
Merge remote-tracking branch 'origin/develop' into uri/experiment_wit…
Aug 8, 2024
78ca4ed
Add a few unit tests
Aug 8, 2024
d3254f4
mypy
Aug 8, 2024
355d9e1
Check we can use Embedding layer as an encoder
Aug 12, 2024
ab37c52
Start writing integration test (and fix one_hot_encoder dtype issue)
Aug 14, 2024
8a8dcd8
Encode initial model data too
Aug 14, 2024
1ea01df
Custom gpr kernel
Aug 14, 2024
073c95e
Merge remote-tracking branch 'origin/develop' into uri/experiment_wit…
Aug 14, 2024
09c727e
Consistent dtype in encoder unit test
Aug 14, 2024
1546cce
Merge branch 'uri/experiment_with_encoded_models' into uri/categorica…
Aug 15, 2024
cc018c8
Eps
Aug 15, 2024
d39679b
one_hot_encoded_space
Aug 16, 2024
b682395
Couple of unit tests
Aug 19, 2024
d4949ec
Adress review comments
Aug 20, 2024
4833999
Merge branch 'uri/experiment_with_encoded_models' into uri/categorica…
Aug 20, 2024
bd6ed69
Unit test
Aug 20, 2024
bba0f1b
Fix typo and hidden optimizer issue
Aug 20, 2024
4b77fd8
Merge branch 'uri/experiment_with_encoded_models' into uri/categorica…
Aug 20, 2024
d1cb88f
Integration test and fix thompson sampling
Aug 20, 2024
e91ce75
Merge remote-tracking branch 'origin/develop' into uri/categorical_tr…
Aug 21, 2024
c48a40a
See whether increasing steps fixes test_old
Aug 21, 2024
31fb555
Review comments
Aug 22, 2024
e2e1a06
Switch num_steps
Aug 23, 2024
2b64c8f
Revert to 8
Aug 25, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
98 changes: 96 additions & 2 deletions tests/integration/test_mixed_space_bayesian_optimization.py
Original file line number Diff line number Diff line change
@@ -15,6 +15,7 @@

from typing import cast

import gpflow
import numpy as np
import numpy.testing as npt
import pytest
@@ -39,10 +40,18 @@
from trieste.bayesian_optimizer import BayesianOptimizer
from trieste.models import TrainableProbabilisticModel
from trieste.models.gpflow import GaussianProcessRegression, build_gpr
from trieste.objectives import ScaledBranin
from trieste.models.interfaces import encode_dataset
from trieste.objectives import ScaledBranin, SingleObjectiveTestProblem
from trieste.objectives.single_objectives import scaled_branin
from trieste.objectives.utils import mk_observer
from trieste.observer import OBJECTIVE
from trieste.space import Box, DiscreteSearchSpace, TaggedProductSearchSpace
from trieste.space import (
Box,
CategoricalSearchSpace,
DiscreteSearchSpace,
TaggedProductSearchSpace,
one_hot_encoder,
)
from trieste.types import TensorType


@@ -190,3 +199,88 @@ def test_optimizer_finds_minima_of_the_scaled_branin_function(
acquisition_function = acquisition_rule._acquisition_function
if isinstance(acquisition_function, AcquisitionFunctionClass):
assert acquisition_function.__call__._get_tracing_count() <= 4 # type: ignore


def categorical_scaled_branin(
categories_to_points: TensorType,
) -> SingleObjectiveTestProblem[TaggedProductSearchSpace]:
"""
Generate a Scaled Branin test problem defined on the product of a categorical space and a
continuous space, with categories mapped to points using the given 1D tensor.
"""
categorical_space = CategoricalSearchSpace([str(float(v)) for v in categories_to_points])
continuous_space = Box([0], [1])
search_space = TaggedProductSearchSpace(
spaces=[categorical_space, continuous_space],
tags=["discrete", "continuous"],
)

def objective(x: TensorType) -> TensorType:
points = tf.gather(categories_to_points, tf.cast(x[..., 0], tf.int32))
x_mapped = tf.concat([tf.expand_dims(points, -1), x[..., 1:]], axis=-1)
return scaled_branin(x_mapped)

minimizer_indices = []
for minimizer0 in ScaledBranin.minimizers[..., 0]:
indices = tf.where(tf.equal(categories_to_points, minimizer0))
minimizer_indices.append(indices[0][0])
category_indices = tf.expand_dims(tf.convert_to_tensor(minimizer_indices, dtype=tf.float64), -1)
minimizers = tf.concat([category_indices, ScaledBranin.minimizers[..., 1:]], axis=-1)

return SingleObjectiveTestProblem(
name="Categorical scaled Branin",
objective=objective,
search_space=search_space,
minimizers=minimizers,
minimum=ScaledBranin.minimum,
)


@random_seed
@pytest.mark.parametrize(
"num_steps, acquisition_rule",
[
pytest.param(25, EfficientGlobalOptimization(), id="EfficientGlobalOptimization"),
],
)
def test_optimizer_finds_minima_of_the_categorical_scaled_branin_function(
num_steps: int,
acquisition_rule: AcquisitionRule[
TensorType, TaggedProductSearchSpace, TrainableProbabilisticModel
],
) -> None:
# 6 categories mapping to 3 random points plus the 3 minimizer points
points = tf.concat(
[tf.random.uniform([3], dtype=tf.float64), ScaledBranin.minimizers[..., 0]], 0
)
problem = categorical_scaled_branin(tf.random.shuffle(points))
initial_query_points = problem.search_space.sample(5)
observer = mk_observer(problem.objective)
initial_data = observer(initial_query_points)

# model uses one-hot encoding for the categorical inputs
encoder = one_hot_encoder(problem.search_space)
kernel = gpflow.kernels.Matern52(
variance=tf.math.reduce_variance(initial_data.observations), lengthscales=0.1
)
model = GaussianProcessRegression(
build_gpr(encode_dataset(initial_data, encoder), kernel=kernel, likelihood_variance=1e-8),
encoder=encoder,
)

dataset = (
BayesianOptimizer(observer, mixed_search_space)
.optimize(num_steps, initial_data, model, acquisition_rule)
.try_get_final_dataset()
)

arg_min_idx = tf.squeeze(tf.argmin(dataset.observations, axis=0))

best_y = dataset.observations[arg_min_idx]
best_x = dataset.query_points[arg_min_idx]

relative_minimizer_err = tf.abs((best_x - problem.minimizers) / problem.minimizers)
assert tf.reduce_any(
tf.reduce_all(relative_minimizer_err < 0.1, axis=-1), axis=0
), relative_minimizer_err
npt.assert_allclose(best_y, problem.minimum, rtol=0.005)
10 changes: 6 additions & 4 deletions tests/integration/test_multifidelity_bayesian_optimization.py
Original file line number Diff line number Diff line change
@@ -38,11 +38,13 @@
)
from trieste.objectives.utils import mk_observer
from trieste.observer import SingleObserver
from trieste.space import TaggedProductSearchSpace
from trieste.space import SearchSpaceType, TaggedProductSearchSpace
from trieste.types import TensorType


def _build_observer(problem: SingleObjectiveMultifidelityTestProblem) -> SingleObserver:
def _build_observer(
problem: SingleObjectiveMultifidelityTestProblem[SearchSpaceType],
) -> SingleObserver:
objective_function = problem.objective

def noisy_objective(x: TensorType) -> TensorType:
@@ -57,7 +59,7 @@ def noisy_objective(x: TensorType) -> TensorType:


def _build_nested_multifidelity_dataset(
problem: SingleObjectiveMultifidelityTestProblem, observer: SingleObserver
problem: SingleObjectiveMultifidelityTestProblem[SearchSpaceType], observer: SingleObserver
) -> Dataset:
num_fidelities = problem.num_fidelities
initial_sample_sizes = [10 + 2 * (num_fidelities - i) for i in range(num_fidelities)]
@@ -83,7 +85,7 @@ def _build_nested_multifidelity_dataset(
@random_seed
@pytest.mark.parametrize("problem", ((Linear2Fidelity), (Linear3Fidelity), (Linear5Fidelity)))
def test_multifidelity_bo_finds_minima_of_linear_problem(
problem: SingleObjectiveMultifidelityTestProblem,
problem: SingleObjectiveMultifidelityTestProblem[SearchSpaceType],
) -> None:
observer = _build_observer(problem)
initial_data = _build_nested_multifidelity_dataset(problem, observer)
96 changes: 96 additions & 0 deletions tests/unit/models/test_interfaces.py
Original file line number Diff line number Diff line change
@@ -15,6 +15,7 @@
from __future__ import annotations

from collections.abc import Callable, Sequence
from typing import Optional

import gpflow
import numpy as np
@@ -35,12 +36,17 @@
from trieste.data import Dataset
from trieste.models import TrainableModelStack, TrainableProbabilisticModel
from trieste.models.interfaces import (
EncodedProbabilisticModel,
EncodedSupportsPredictJoint,
EncodedSupportsPredictY,
EncodedTrainableProbabilisticModel,
TrainablePredictJointReparamModelStack,
TrainablePredictYModelStack,
TrainableSupportsPredictJoint,
TrainableSupportsPredictJointHasReparamSampler,
)
from trieste.models.utils import get_last_optimization_result, optimize_model_and_save_result
from trieste.space import EncoderFunction
from trieste.types import TensorType


@@ -216,3 +222,93 @@ def test_model_stack_reparam_sampler() -> None:
npt.assert_allclose(var[..., :2], var01, rtol=0.04)
npt.assert_allclose(var[..., 2:3], var2, rtol=0.04)
npt.assert_allclose(var[..., 3:], var3, rtol=0.04)


class _EncodedModel(
EncodedTrainableProbabilisticModel,
EncodedSupportsPredictJoint,
EncodedSupportsPredictY,
EncodedProbabilisticModel,
):
def __init__(self, encoder: EncoderFunction | None = None) -> None:
self.dataset: Dataset | None = None
self._encoder = (lambda x: x + 1) if encoder is None else encoder

@property
def encoder(self) -> EncoderFunction | None:
return self._encoder

def predict_encoded(self, query_points: TensorType) -> tuple[TensorType, TensorType]:
return query_points, query_points

def sample_encoded(self, query_points: TensorType, num_samples: int) -> TensorType:
return tf.tile(tf.expand_dims(query_points, 0), [num_samples, 1, 1])

def log(self, dataset: Optional[Dataset] = None) -> None:
pass

def update_encoded(self, dataset: Dataset) -> None:
self.dataset = dataset

def optimize_encoded(self, dataset: Dataset) -> None:
self.dataset = dataset

def predict_joint_encoded(self, query_points: TensorType) -> tuple[TensorType, TensorType]:
b, d = query_points.shape
return query_points, tf.zeros([d, b, b])

def predict_y_encoded(self, query_points: TensorType) -> tuple[TensorType, TensorType]:
return self.predict_encoded(query_points)


def test_encoded_probabilistic_model() -> None:
model = _EncodedModel()
query_points = tf.random.uniform([3, 5])
mean, var = model.predict(query_points)
npt.assert_allclose(mean, query_points + 1)
npt.assert_allclose(var, query_points + 1)
samples = model.sample(query_points, 10)
assert len(samples) == 10
for i in range(10):
npt.assert_allclose(samples[i], query_points + 1)


def test_encoded_trainable_probabilistic_model() -> None:
model = _EncodedModel()
assert model.dataset is None
for method in model.update, model.optimize:
query_points = tf.random.uniform([3, 5])
observations = tf.random.uniform([3, 1])
dataset = Dataset(query_points, observations)
method(dataset)
assert model.dataset is not None
# no idea why mypy thinks model.dataset couldn't have changed here
npt.assert_allclose( # type: ignore[unreachable]
model.dataset.query_points, query_points + 1
)
npt.assert_allclose(model.dataset.observations, observations)


def test_encoded_supports_predict_joint() -> None:
model = _EncodedModel()
query_points = tf.random.uniform([3, 5])
mean, var = model.predict_joint(query_points)
npt.assert_allclose(mean, query_points + 1)
npt.assert_allclose(var, tf.zeros([5, 3, 3]))


def test_encoded_supports_predict_y() -> None:
model = _EncodedModel()
query_points = tf.random.uniform([3, 5])
mean, var = model.predict_y(query_points)
npt.assert_allclose(mean, query_points + 1)
npt.assert_allclose(var, query_points + 1)


def test_encoded_probabilistic_model_keras_embedding() -> None:
encoder = tf.keras.layers.Embedding(3, 2)
model = _EncodedModel(encoder=encoder)
query_points = tf.random.uniform([3, 5], minval=0, maxval=3, dtype=tf.int32)
mean, var = model.predict(query_points)
assert mean.shape == (3, 5, 2)
npt.assert_allclose(mean, encoder(query_points))
7 changes: 4 additions & 3 deletions tests/unit/objectives/test_multi_objectives.py
Original file line number Diff line number Diff line change
@@ -19,6 +19,7 @@
from check_shapes.exceptions import ShapeMismatchError

from trieste.objectives.multi_objectives import DTLZ1, DTLZ2, VLMOP2, MultiObjectiveTestProblem
from trieste.space import SearchSpaceType
from trieste.types import TensorType


@@ -117,7 +118,7 @@ def test_dtlz2_has_expected_output(
],
)
def test_gen_pareto_front_is_equal_to_math_defined(
obj_type: Callable[[int, int], MultiObjectiveTestProblem],
obj_type: Callable[[int, int], MultiObjectiveTestProblem[SearchSpaceType]],
input_dim: int,
num_obj: int,
gen_pf_num: int,
@@ -140,7 +141,7 @@ def test_gen_pareto_front_is_equal_to_math_defined(
],
)
def test_func_raises_specified_input_dim_not_align_with_actual_input_dim(
obj_inst: MultiObjectiveTestProblem, actual_x: TensorType
obj_inst: MultiObjectiveTestProblem[SearchSpaceType], actual_x: TensorType
) -> None:
with pytest.raises(ShapeMismatchError):
obj_inst.objective(actual_x)
@@ -160,7 +161,7 @@ def test_func_raises_specified_input_dim_not_align_with_actual_input_dim(
@pytest.mark.parametrize("num_obs", [1, 5, 10])
@pytest.mark.parametrize("dtype", [tf.float32, tf.float64])
def test_objective_has_correct_shape_and_dtype(
problem: MultiObjectiveTestProblem,
problem: MultiObjectiveTestProblem[SearchSpaceType],
input_dim: int,
num_obj: int,
num_obs: int,
11 changes: 6 additions & 5 deletions tests/unit/objectives/test_single_objectives.py
Original file line number Diff line number Diff line change
@@ -36,6 +36,7 @@
SingleObjectiveTestProblem,
Trid10,
)
from trieste.space import Box, SearchSpaceType


@pytest.fixture(
@@ -58,12 +59,12 @@
Levy8,
],
)
def _problem_fixture(request: Any) -> Tuple[SingleObjectiveTestProblem, int]:
def _problem_fixture(request: Any) -> Tuple[SingleObjectiveTestProblem[SearchSpaceType], int]:
return request.param


def test_objective_maps_minimizers_to_minimum(
problem: SingleObjectiveTestProblem,
problem: SingleObjectiveTestProblem[SearchSpaceType],
) -> None:
objective = problem.objective
minimizers = problem.minimizers
@@ -74,7 +75,7 @@ def test_objective_maps_minimizers_to_minimum(


def test_no_function_values_are_less_than_global_minimum(
problem: SingleObjectiveTestProblem,
problem: SingleObjectiveTestProblem[Box],
) -> None:
objective = problem.objective
space = problem.search_space
@@ -86,7 +87,7 @@ def test_no_function_values_are_less_than_global_minimum(
@pytest.mark.parametrize("num_obs", [5, 1])
@pytest.mark.parametrize("dtype", [tf.float32, tf.float64])
def test_objective_has_correct_shape_and_dtype(
problem: SingleObjectiveTestProblem,
problem: SingleObjectiveTestProblem[SearchSpaceType],
num_obs: int,
dtype: tf.DType,
) -> None:
@@ -120,7 +121,7 @@ def test_objective_has_correct_shape_and_dtype(
)
@pytest.mark.parametrize("num_obs", [5, 1])
def test_search_space_has_correct_shape_and_default_dtype(
problem: SingleObjectiveTestProblem,
problem: SingleObjectiveTestProblem[SearchSpaceType],
input_dim: int,
num_obs: int,
) -> None:
Loading