Adress review comments

tests/integration/test_mixed_space_bayesian_optimization.py

@@ -39,7 +39,6 @@
 from trieste.bayesian_optimizer import BayesianOptimizer
 from trieste.models import TrainableProbabilisticModel
 from trieste.models.gpflow import GaussianProcessRegression, build_gpr
-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
@@ -49,7 +48,6 @@
     CategoricalSearchSpace,
     DiscreteSearchSpace,
     TaggedProductSearchSpace,
-    one_hot_encoded_space,
     one_hot_encoder,
 )
 from trieste.types import TensorType
@@ -260,9 +258,8 @@ def test_optimizer_finds_minima_of_the_categorical_scaled_branin_function(
 
     # model uses one-hot encoding for the categorical inputs
     encoder = one_hot_encoder(problem.search_space)
-    encoded_space = one_hot_encoded_space(problem.search_space)
     model = GaussianProcessRegression(
-        build_gpr(encode_dataset(initial_data, encoder), encoded_space, likelihood_variance=1e-8),
+        build_gpr(initial_data, problem.search_space, likelihood_variance=1e-8),
         encoder=encoder,
     )
 

tests/unit/test_space.py

@@ -1760,19 +1760,19 @@ def test_categorical_search_space__to_tags_raises_for_non_integers() -> None:
             tf.constant([[1], [1]], dtype=tf.float64),
         ),
         (
-            CategoricalSearchSpace(["R", "G", "B"]),
-            tf.constant([[0], [2], [1]], dtype=tf.float64),
-            tf.constant([[1, 0, 0], [0, 0, 1], [0, 1, 0]], dtype=tf.float64),
+            CategoricalSearchSpace(["R", "G", "B"], dtype=tf.float32),
+            tf.constant([[0], [2], [1]], dtype=tf.float32),
+            tf.constant([[1, 0, 0], [0, 0, 1], [0, 1, 0]], dtype=tf.float32),
         ),
         (
             CategoricalSearchSpace(["R", "G", "B"]),
             tf.constant([[[[[0]]]]], dtype=tf.float64),
             tf.constant([[[[[1, 0, 0]]]]], dtype=tf.float64),
         ),
         (
-            CategoricalSearchSpace(["R", "G", "B", "A"]),
-            tf.constant([[0], [2], [2]], dtype=tf.float64),
-            tf.constant([[1, 0, 0, 0], [0, 0, 1, 0], [0, 0, 1, 0]], dtype=tf.float64),
+            CategoricalSearchSpace(["R", "G", "B", "A"], dtype=tf.float32),
+            tf.constant([[0], [2], [2]], dtype=tf.float32),
+            tf.constant([[1, 0, 0, 0], [0, 0, 1, 0], [0, 0, 1, 0]], dtype=tf.float32),
         ),
         (
             CategoricalSearchSpace([["R", "G", "B"], ["Y", "N"]]),

trieste/models/gpflow/builders.py

@@ -21,7 +21,7 @@
 from __future__ import annotations
 
 import math
-from typing import Optional, Sequence, Type
+from typing import Callable, Optional, Sequence, Type
 
 import gpflow
 import tensorflow as tf
@@ -30,9 +30,10 @@
 from gpflow.models import GPR, SGPR, SVGP, VGP, GPModel
 
 from ...data import Dataset, split_dataset_by_fidelity
-from ...space import Box, SearchSpace
+from ...space import Box, EncoderFunction, SearchSpace, one_hot_encoded_space, one_hot_encoder
 from ...types import TensorType
 from ..gpflow.models import GaussianProcessRegression
+from ..interfaces import encode_dataset
 
 # NOTE: As a static non-Tensor, this should really be a tf.constant (like the other constants).
 # However, changing it breaks serialisation during the expected_improvement.pct.py notebook.
@@ -88,6 +89,8 @@ def build_gpr(
     likelihood_variance: Optional[float] = None,
     trainable_likelihood: bool = False,
     kernel: Optional[gpflow.kernels.Kernel] = None,
+    encoder: EncoderFunction | None = None,
+    space_encoder: Callable[[SearchSpace], SearchSpace] | None = None,
 ) -> GPR:
     """
     Build a :class:`~gpflow.models.GPR` model with sensible initial parameters and
@@ -118,8 +121,20 @@ def build_gpr(
         non-trainable. By default set to `False`.
     :param kernel: The kernel to use in the model, defaults to letting the function set up a
         :class:`~gpflow.kernels.Matern52` kernel.
+    :param encoder: Encoder with which to transform the dataset before training. Defaults to
+        one_hot_encoder if the search_space is specified.
+    :param space_encoder: Encoder with which to transform search_space before generating a kernel.
+        Defaults to one_hot_encoded_space.
     :return: A :class:`~gpflow.models.GPR` model.
     """
+    if search_space is not None:
+        encoder = one_hot_encoder(search_space) if encoder is None else encoder
+        space_encoder = one_hot_encoded_space if space_encoder is None else space_encoder
+        search_space = space_encoder(search_space)
+
+    if encoder is not None:
+        data = encode_dataset(data, encoder)
+
     empirical_mean, empirical_variance, _ = _get_data_stats(data)
 
     if kernel is None: