Repository: secondmind-labs/trieste

Issue:
Active learning acquisition functions are missing integration tests



Discussion:

Resolving patch:
diff --git a/docs/notebooks/active_learning.pct.py b/docs/notebooks/active_learning.pct.py
--- a/docs/notebooks/active_learning.pct.py
+++ b/docs/notebooks/active_learning.pct.py
@@ -19,11 +19,11 @@
 
 
 # %%
-from trieste.objectives import scaled_branin
+from trieste.objectives import BRANIN_SEARCH_SPACE, scaled_branin
 from util.plotting_plotly import plot_function_plotly
 from trieste.space import Box
 
-search_space = Box([0, 0], [1, 1])
+search_space = BRANIN_SEARCH_SPACE
 
 fig = plot_function_plotly(
     scaled_branin, search_space.lower, search_space.upper, grid_density=20
@@ -32,7 +32,7 @@
 fig.show()
 
 # %% [markdown]
-# We begin our Bayesian active learning from a two-point initial design built from a space-filling Halton sequence.
+# We begin our Bayesian active learning from a small initial design built from a space-filling Halton sequence.
 
 # %%
 import trieste
@@ -47,16 +47,24 @@
 # %% [markdown]
 # ## Surrogate model
 #
-# Just like in sequential optimization, we fit a surrogate Gaussian process model as implemented in GPflow to the initial data. The GPflow models cannot be used directly in our Bayesian optimization routines, so we build a GPflow's `GPR` model and pass it to the `GaussianProcessRegression` wrapper.
+# Just like in sequential optimization, we fit a surrogate Gaussian process model as implemented in GPflow to the initial data. The GPflow models cannot be used directly in our Bayesian optimization routines, so we build a GPflow's `GPR` model and pass it to the `GaussianProcessRegression` wrapper. As a good practice, we use priors for the kernel hyperparameters.
 
 # %%
 import gpflow
 from trieste.models.gpflow.models import GaussianProcessRegression
+import tensorflow_probability as tfp
 
 
 def build_model(data):
     variance = tf.math.reduce_variance(data.observations)
-    kernel = gpflow.kernels.RBF(variance=variance, lengthscales=[2, 2])
+    kernel = gpflow.kernels.Matern52(variance=variance, lengthscales=[0.2, 0.2])
+    prior_scale = tf.cast(1.0, dtype=tf.float64)
+    kernel.variance.prior = tfp.distributions.LogNormal(
+        tf.cast(-2.0, dtype=tf.float64), prior_scale
+    )
+    kernel.lengthscales.prior = tfp.distributions.LogNormal(
+        tf.math.log(kernel.lengthscales), prior_scale
+    )
     gpr = gpflow.models.GPR(data.astuple(), kernel, noise_variance=1e-5)
     gpflow.set_trainable(gpr.likelihood, False)
 
@@ -79,9 +87,7 @@ def build_model(data):
 from trieste.acquisition.rule import EfficientGlobalOptimization
 
 acq = PredictiveVariance()
-rule = EfficientGlobalOptimization(
-    builder=acq, optimizer=generate_continuous_optimizer()
-)
+rule = EfficientGlobalOptimization(builder=acq)  # type: ignore
 bo = trieste.bayesian_optimizer.BayesianOptimizer(observer, search_space)
 
 # %% [markdown]
@@ -144,12 +150,14 @@ def pred_var(x):
 # %% [markdown]
 # ## Batch active learning using predictive variance
 #
-# For some cases, query several points at a time can be convenient by doing batch active learning. For this case, we must pass a num_query_points input to our `EfficientGlobalOptimization` rule. The drawback of the batch predictive variance is, it tends to query in high variance area less accurately, compared to the sequentially drawing one point at a time.
+# In cases when we can evaluate the black-box function in parallel, it would be useful to produce a batch of points rather than a single point. `PredictiveVariance` acquisition function can also perform batch active learning. We must pass a `num_query_points` input to our `EfficientGlobalOptimization` rule. The drawback of the batch predictive variance is that it tends to query in high variance area less accurately, compared to sequentially drawing one point at a time.
 
 # %%
 bo_iter = 5
 num_query = 3
+
 model = build_model(initial_data)
+
 acq = PredictiveVariance()
 rule = EfficientGlobalOptimization(
     num_query_points=num_query,


Test patch:
diff --git a/tests/integration/test_active_learning.py b/tests/integration/test_active_learning.py
new file mode 100644
--- /dev/null
+++ b/tests/integration/test_active_learning.py
@@ -0,0 +1,251 @@
+# Copyright 2021 The Trieste Contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Integration tests for various forms of active learning implemented in Trieste.
+"""
+
+from __future__ import annotations
+
+import gpflow
+import pytest
+import tensorflow as tf
+import tensorflow_probability as tfp
+
+from tests.util.misc import random_seed
+from trieste.acquisition.function import ExpectedFeasibility, PredictiveVariance
+from trieste.acquisition.rule import AcquisitionRule, EfficientGlobalOptimization
+from trieste.bayesian_optimizer import BayesianOptimizer
+from trieste.data import Dataset
+from trieste.models import TrainableProbabilisticModel
+from trieste.models.gpflow import GaussianProcessRegression
+from trieste.objectives import BRANIN_SEARCH_SPACE, branin, scaled_branin
+from trieste.objectives.utils import mk_observer
+from trieste.observer import Observer
+from trieste.space import Box, SearchSpace
+from trieste.types import TensorType
+
+
+@random_seed
+@pytest.mark.parametrize(
+    "num_steps, acquisition_rule",
+    [
+        (50, EfficientGlobalOptimization(PredictiveVariance())),
+    ],
+)
+def test_optimizer_learns_scaled_branin_function(
+    num_steps: int, acquisition_rule: AcquisitionRule[TensorType, SearchSpace]
+) -> None:
+    """
+    Ensure that the objective function is effectively learned, such that the final model
+    fits well and predictions are close to actual objective values.
+    """
+
+    search_space = BRANIN_SEARCH_SPACE
+
+    def build_model(data: Dataset) -> TrainableProbabilisticModel:
+        variance = tf.math.reduce_variance(data.observations)
+        kernel = gpflow.kernels.Matern52(variance=variance, lengthscales=[0.2, 0.2])
+        prior_scale = tf.cast(1.0, dtype=tf.float64)
+        kernel.variance.prior = tfp.distributions.LogNormal(
+            tf.cast(-2.0, dtype=tf.float64), prior_scale
+        )
+        kernel.lengthscales.prior = tfp.distributions.LogNormal(
+            tf.math.log(kernel.lengthscales), prior_scale
+        )
+        gpr = gpflow.models.GPR(data.astuple(), kernel, noise_variance=1e-5)
+        gpflow.set_trainable(gpr.likelihood, False)
+
+        return GaussianProcessRegression(gpr)
+
+    num_initial_points = 6
+    initial_query_points = search_space.sample_halton(num_initial_points)
+    observer = mk_observer(scaled_branin)
+    initial_data = observer(initial_query_points)
+
+    # we set a performance criterion at 1% of the range
+    # max absolute error needs to be bettter than this criterion
+    test_query_points = search_space.sample_sobol(10000 * search_space.dimension)
+    test_data = observer(test_query_points)
+    test_range = tf.reduce_max(test_data.observations) - tf.reduce_min(test_data.observations)
+    criterion = 0.01 * test_range
+
+    # we expect a model with initial data to fail the criterion
+    initial_model = build_model(initial_data)
+    initial_model.optimize(initial_data)
+    initial_predicted_means, _ = initial_model.model.predict_f(test_query_points)  # type: ignore
+    initial_accuracy = tf.reduce_max(tf.abs(initial_predicted_means - test_data.observations))
+
+    assert not initial_accuracy < criterion
+
+    # after active learning the model should be much more accurate
+    model = build_model(initial_data)
+    final_model = (
+        BayesianOptimizer(observer, search_space)
+        .optimize(num_steps, initial_data, model, acquisition_rule)
+        .try_get_final_model()
+    )
+    final_predicted_means, _ = final_model.model.predict_f(test_query_points)  # type: ignore
+    final_accuracy = tf.reduce_max(tf.abs(final_predicted_means - test_data.observations))
+
+    assert initial_accuracy > final_accuracy
+    assert final_accuracy < criterion
+
+
+@random_seed
+@pytest.mark.slow
+@pytest.mark.parametrize(
+    "num_steps, acquisition_rule, threshold",
+    [
+        (50, EfficientGlobalOptimization(ExpectedFeasibility(80, delta=1)), 80),
+        (50, EfficientGlobalOptimization(ExpectedFeasibility(80, delta=2)), 80),
+        (70, EfficientGlobalOptimization(ExpectedFeasibility(20, delta=1)), 20),
+    ],
+)
+def test_optimizer_learns_feasibility_set_of_thresholded_branin_function(
+    num_steps: int, acquisition_rule: AcquisitionRule[TensorType, SearchSpace], threshold: int
+) -> None:
+    """
+    Ensure that the feasible set is sufficiently well learned, such that the final model
+    classifies with great degree of certainty whether points in the search space are in
+    in the feasible set or not.
+    """
+
+    search_space = BRANIN_SEARCH_SPACE
+
+    def build_model(data: Dataset) -> TrainableProbabilisticModel:
+        variance = tf.math.reduce_variance(data.observations)
+        kernel = gpflow.kernels.Matern52(variance=variance, lengthscales=[0.2, 0.2])
+        prior_scale = tf.cast(1.0, dtype=tf.float64)
+        kernel.variance.prior = tfp.distributions.LogNormal(
+            tf.cast(-2.0, dtype=tf.float64), prior_scale
+        )
+        kernel.lengthscales.prior = tfp.distributions.LogNormal(
+            tf.math.log(kernel.lengthscales), prior_scale
+        )
+        gpr = gpflow.models.GPR(data.astuple(), kernel, noise_variance=1e-5)
+        gpflow.set_trainable(gpr.likelihood, False)
+
+        return GaussianProcessRegression(gpr)
+
+    num_initial_points = 6
+    initial_query_points = search_space.sample_halton(num_initial_points)
+    observer = mk_observer(branin)
+    initial_data = observer(initial_query_points)
+
+    # we set a performance criterion at 0.001 probability of required precision per point
+    # for global points and 0.01 close to the boundary
+    n_global = 10000 * search_space.dimension
+    n_boundary = 2000 * search_space.dimension
+    global_test, boundary_test = _get_feasible_set_test_data(
+        search_space, observer, n_global, n_boundary, threshold
+    )
+    global_criterion = 0.001 * (1 - 0.001) * tf.cast(n_global, tf.float64)
+    boundary_criterion = 0.01 * (1 - 0.01) * tf.cast(n_boundary, tf.float64)
+
+    # we expect a model with initial data to fail the criteria
+    initial_model = build_model(initial_data)
+    initial_model.optimize(initial_data)
+    initial_accuracy_global = _get_excursion_accuracy(global_test, initial_model, threshold)
+    initial_accuracy_boundary = _get_excursion_accuracy(boundary_test, initial_model, threshold)
+
+    assert not initial_accuracy_global < global_criterion
+    assert not initial_accuracy_boundary < boundary_criterion
+
+    # after active learning the model should be much more accurate
+    model = build_model(initial_data)
+    final_model = (
+        BayesianOptimizer(observer, search_space)
+        .optimize(num_steps, initial_data, model, acquisition_rule)
+        .try_get_final_model()
+    )
+    final_accuracy_global = _get_excursion_accuracy(global_test, final_model, threshold)
+    final_accuracy_boundary = _get_excursion_accuracy(boundary_test, final_model, threshold)
+
+    assert initial_accuracy_global > final_accuracy_global
+    assert initial_accuracy_boundary > final_accuracy_boundary
+    assert final_accuracy_global < global_criterion
+    assert final_accuracy_boundary < boundary_criterion
+
+
+def _excursion_probability(
+    x: TensorType, model: TrainableProbabilisticModel, threshold: int
+) -> tfp.distributions.Distribution:
+    mean, variance = model.model.predict_f(x)  # type: ignore
+    normal = tfp.distributions.Normal(tf.cast(0, x.dtype), tf.cast(1, x.dtype))
+    t = (mean - threshold) / tf.sqrt(variance)
+    return normal.cdf(t)
+
+
+def _get_excursion_accuracy(
+    x: TensorType, model: TrainableProbabilisticModel, threshold: int
+) -> float:
+    prob = _excursion_probability(x, model, threshold)
+    accuracy = tf.reduce_sum(prob * (1 - prob))
+
+    return accuracy
+
+
+def _get_feasible_set_test_data(
+    search_space: Box,
+    observer: Observer,
+    n_global: int,
+    n_boundary: int,
+    threshold: float,
+    range_pct: float = 0.01,
+) -> tuple[TensorType, TensorType]:
+
+    boundary_done = False
+    global_done = False
+    boundary_points = tf.constant(0, dtype=tf.float64, shape=(0, search_space.dimension))
+    global_points = tf.constant(0, dtype=tf.float64, shape=(0, search_space.dimension))
+
+    while not boundary_done and not global_done:
+        test_query_points = search_space.sample(100000)
+        test_data = observer(test_query_points)
+        threshold_deviation = range_pct * (
+            tf.reduce_max(test_data.observations)  # type: ignore
+            - tf.reduce_min(test_data.observations)  # type: ignore
+        )
+
+        mask = tf.reduce_all(
+            tf.concat(
+                [
+                    test_data.observations > threshold - threshold_deviation,  # type: ignore
+                    test_data.observations < threshold + threshold_deviation,  # type: ignore
+                ],
+                axis=1,
+            ),
+            axis=1,
+        )
+        boundary_points = tf.concat(
+            [boundary_points, tf.boolean_mask(test_query_points, mask)], axis=0
+        )
+        global_points = tf.concat(
+            [global_points, tf.boolean_mask(test_query_points, tf.logical_not(mask))], axis=0
+        )
+
+        if boundary_points.shape[0] > n_boundary:
+            boundary_done = True
+        if global_points.shape[0] > n_global:
+            global_done = True
+
+    return (
+        global_points[
+            :n_global,
+        ],
+        boundary_points[
+            :n_boundary,
+        ],
+    )