From d4949ecc3be5cdc63260867ab8475e0c552909e1 Mon Sep 17 00:00:00 2001 From: Uri Granta Date: Tue, 20 Aug 2024 08:28:28 +0100 Subject: [PATCH] Adress review comments --- .../test_mixed_space_bayesian_optimization.py | 5 +---- tests/unit/test_space.py | 12 ++++++------ trieste/models/gpflow/builders.py | 19 +++++++++++++++++-- 3 files changed, 24 insertions(+), 12 deletions(-) diff --git a/tests/integration/test_mixed_space_bayesian_optimization.py b/tests/integration/test_mixed_space_bayesian_optimization.py index 16a9eb78f9..aa73043485 100644 --- a/tests/integration/test_mixed_space_bayesian_optimization.py +++ b/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, ) diff --git a/tests/unit/test_space.py b/tests/unit/test_space.py index 06b44f0486..6fb473d130 100644 --- a/tests/unit/test_space.py +++ b/tests/unit/test_space.py @@ -1760,9 +1760,9 @@ 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"]), @@ -1770,9 +1770,9 @@ def test_categorical_search_space__to_tags_raises_for_non_integers() -> None: 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"]]), diff --git a/trieste/models/gpflow/builders.py b/trieste/models/gpflow/builders.py index 638ae55b9d..813cef6bdf 100644 --- a/trieste/models/gpflow/builders.py +++ b/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: