Encoded models (#864)

tests/integration/test_mixed_space_bayesian_optimization.py

@@ -39,10 +39,17 @@
 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.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 +197,85 @@ 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)
+    model = GaussianProcessRegression(
+        build_gpr(initial_data, problem.search_space, likelihood_variance=1e-8),
+        encoder=encoder,
+    )
+
+    dataset = (
+        BayesianOptimizer(observer, problem.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)

tests/integration/test_multifidelity_bayesian_optimization.py

@@ -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)

tests/unit/models/conftest.py

@@ -43,6 +43,7 @@
     VariationalGaussianProcess,
 )
 from trieste.models.optimizer import DatasetTransformer, Optimizer
+from trieste.space import EncoderFunction
 from trieste.types import TensorType
 
 
@@ -58,12 +59,15 @@
 )
 def _gpflow_interface_factory(request: Any) -> ModelFactoryType:
     def model_interface_factory(
-        x: TensorType, y: TensorType, optimizer: Optimizer | None = None
+        x: TensorType,
+        y: TensorType,
+        optimizer: Optimizer | None = None,
+        encoder: EncoderFunction | None = None,
     ) -> tuple[GPflowPredictor, Callable[[TensorType, TensorType], GPModel]]:
         model_interface: Callable[..., GPflowPredictor] = request.param[0]
         base_model: GaussianProcessRegression = request.param[1](x, y)
         reference_model: Callable[[TensorType, TensorType], GPModel] = request.param[1]
-        return model_interface(base_model, optimizer=optimizer), reference_model
+        return model_interface(base_model, optimizer=optimizer, encoder=encoder), reference_model
 
     return model_interface_factory
 

tests/unit/models/gpflow/test_interface.py

@@ -24,17 +24,18 @@
 from tests.util.misc import random_seed
 from trieste.data import Dataset
 from trieste.models.gpflow import BatchReparametrizationSampler, GPflowPredictor
+from trieste.space import CategoricalSearchSpace, one_hot_encoder
 
 
 class _QuadraticPredictor(GPflowPredictor):
     @property
     def model(self) -> GPModel:
         return _QuadraticGPModel()
 
-    def optimize(self, dataset: Dataset) -> None:
+    def optimize_encoded(self, dataset: Dataset) -> None:
         self.optimizer.optimize(self.model, dataset)
 
-    def update(self, dataset: Dataset) -> None:
+    def update_encoded(self, dataset: Dataset) -> None:
         return
 
     def log(self, dataset: Optional[Dataset] = None) -> None:
@@ -112,3 +113,14 @@ def test_gpflow_reparam_sampler_returns_reparam_sampler_with_correct_samples() -
     linear_error = 1 / tf.sqrt(tf.cast(num_samples, tf.float32))
     npt.assert_allclose(sample_mean, [[6.25]], rtol=linear_error)
     npt.assert_allclose(sample_variance, 1.0, rtol=2 * linear_error)
+
+
+def test_gpflow_categorical_predict() -> None:
+    search_space = CategoricalSearchSpace(["Red", "Green", "Blue"])
+    query_points = search_space.sample(10)
+    model = _QuadraticPredictor(encoder=one_hot_encoder(search_space))
+    mean, variance = model.predict(query_points)
+    assert mean.shape == [10, 1]
+    assert variance.shape == [10, 1]
+    npt.assert_allclose(mean, [[1.0]] * 10, rtol=0.01)
+    npt.assert_allclose(variance, [[1.0]] * 10, rtol=0.01)

tests/unit/models/test_interfaces.py

@@ -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))

tests/unit/objectives/test_multi_objectives.py

@@ -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,