We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Active learning acquisition functions are missing integration tests


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `docs/notebooks/active_learning.pct.py`
Content:
```
1 # %% [markdown]
2 # # Active Learning
3 
4 # %% [markdown]
5 # Sometimes, we may just want to learn a black-box function, rather than optimizing it. This goal is known as active learning and corresponds to choosing query points that reduce our model uncertainty. This notebook demonstrates how to perform Bayesian active learning using Trieste.
6 
7 # %%
8 # %matplotlib inline
9 import numpy as np
10 import tensorflow as tf
11 
12 np.random.seed(1793)
13 tf.random.set_seed(1793)
14 
15 # %% [markdown]
16 # ## Describe the problem
17 #
18 # In this example, we will perform active learning for the scaled Branin function.
19 
20 
21 # %%
22 from trieste.objectives import scaled_branin
23 from util.plotting_plotly import plot_function_plotly
24 from trieste.space import Box
25 
26 search_space = Box([0, 0], [1, 1])
27 
28 fig = plot_function_plotly(
29     scaled_branin, search_space.lower, search_space.upper, grid_density=20
30 )
31 fig.update_layout(height=400, width=400)
32 fig.show()
33 
34 # %% [markdown]
35 # We begin our Bayesian active learning from a two-point initial design built from a space-filling Halton sequence.
36 
37 # %%
38 import trieste
39 
40 observer = trieste.objectives.utils.mk_observer(scaled_branin)
41 
42 num_initial_points = 4
43 initial_query_points = search_space.sample_halton(num_initial_points)
44 initial_data = observer(initial_query_points)
45 
46 
47 # %% [markdown]
48 # ## Surrogate model
49 #
50 # 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.
51 
52 # %%
53 import gpflow
54 from trieste.models.gpflow.models import GaussianProcessRegression
55 
56 
57 def build_model(data):
58     variance = tf.math.reduce_variance(data.observations)
59     kernel = gpflow.kernels.RBF(variance=variance, lengthscales=[2, 2])
60     gpr = gpflow.models.GPR(data.astuple(), kernel, noise_variance=1e-5)
61     gpflow.set_trainable(gpr.likelihood, False)
62 
63     return GaussianProcessRegression(gpr)
64 
65 
66 model = build_model(initial_data)
67 
68 # %% [markdown]
69 # ## Active learning using predictive variance
70 #
71 # For our first active learning example, we will use a simple acquisition function known as `PredictiveVariance` which chooses points for which we are highly uncertain (i.e. the predictive posterior covariance matrix at these points has large determinant), as discussed in <cite data-cite="MacKay1992"/>. Note that this also implies that our model needs to have `predict_joint` method to be able to return the full covariance, and it's likely to be expensive to compute.
72 #
73 # We will now demonstrate how to choose individual query points using `PredictiveVariance` before moving onto batch active learning. For both cases, we can utilize Trieste's `BayesianOptimizer` to do the active learning steps.
74 #
75 
76 # %%
77 from trieste.acquisition.function import PredictiveVariance
78 from trieste.acquisition.optimizer import generate_continuous_optimizer
79 from trieste.acquisition.rule import EfficientGlobalOptimization
80 
81 acq = PredictiveVariance()
82 rule = EfficientGlobalOptimization(
83     builder=acq, optimizer=generate_continuous_optimizer()
84 )
85 bo = trieste.bayesian_optimizer.BayesianOptimizer(observer, search_space)
86 
87 # %% [markdown]
88 # To plot the contour of variance of our model at each step, we can set the `track_state` parameter to `True` in `bo.optimize()`, this will make Trieste record our model at each iteration.
89 
90 # %%
91 bo_iter = 5
92 result = bo.optimize(bo_iter, initial_data, model, rule, track_state=True)
93 
94 # %% [markdown]
95 # Then we can retrieve our final dataset from the active learning steps.
96 
97 # %%
98 dataset = result.try_get_final_dataset()
99 query_points = dataset.query_points.numpy()
100 observations = dataset.observations.numpy()
101 
102 # %% [markdown]
103 # Finally, we can check the performance of our `PredictiveVariance` active learning acquisition function by plotting the predictive variance landscape of our model. We can see how it samples regions for which our model is highly uncertain.
104 
105 # %%
106 from util.plotting import plot_bo_points, plot_function_2d
107 
108 
109 def plot_active_learning_query(
110     result, bo_iter, num_initial_points, query_points, num_query=1
111 ):
112 
113     for i in range(bo_iter):
114 
115         def pred_var(x):
116             _, var = result.history[i].models["OBJECTIVE"].model.predict_f(x)
117             return var
118 
119         _, ax = plot_function_2d(
120             pred_var,
121             search_space.lower - 0.01,
122             search_space.upper + 0.01,
123             grid_density=100,
124             contour=True,
125             colorbar=True,
126             figsize=(10, 6),
127             title=[
128                 "Variance contour with queried points at iter:" + str(i + 1)
129             ],
130             xlabel="$X_1$",
131             ylabel="$X_2$",
132         )
133 
134         plot_bo_points(
135             query_points[: num_initial_points + (i * num_query)],
136             ax[0, 0],
137             num_initial_points,
138         )
139 
140 
141 plot_active_learning_query(result, bo_iter, num_initial_points, query_points)
142 
143 
144 # %% [markdown]
145 # ## Batch active learning using predictive variance
146 #
147 # 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.
148 
149 # %%
150 bo_iter = 5
151 num_query = 3
152 model = build_model(initial_data)
153 acq = PredictiveVariance()
154 rule = EfficientGlobalOptimization(
155     num_query_points=num_query,
156     builder=acq,
157     optimizer=generate_continuous_optimizer(num_optimization_runs=1),
158 )
159 bo = trieste.bayesian_optimizer.BayesianOptimizer(observer, search_space)
160 
161 result = bo.optimize(bo_iter, initial_data, model, rule, track_state=True)
162 
163 
164 # %% [markdown]
165 # After that, we can retrieve our final dataset.
166 
167 # %%
168 dataset = result.try_get_final_dataset()
169 query_points = dataset.query_points.numpy()
170 observations = dataset.observations.numpy()
171 
172 
173 # %% [markdown]
174 # Now we can visualize the batch predictive variance using our plotting function.
175 
176 # %%
177 from util.plotting import plot_bo_points, plot_function_2d
178 
179 plot_active_learning_query(
180     result, bo_iter, num_initial_points, query_points, num_query
181 )
182 
183 # %% [markdown]
184 # ## LICENSE
185 #
186 # [Apache License 2.0](https://github.com/secondmind-labs/trieste/blob/develop/LICENSE)
187
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
 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 @@
 # %% [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,

{"golden_diff": "diff --git a/docs/notebooks/active_learning.pct.py b/docs/notebooks/active_learning.pct.py\n--- a/docs/notebooks/active_learning.pct.py\n+++ b/docs/notebooks/active_learning.pct.py\n@@ -19,11 +19,11 @@\n \n \n # %%\n-from trieste.objectives import scaled_branin\n+from trieste.objectives import BRANIN_SEARCH_SPACE, scaled_branin\n from util.plotting_plotly import plot_function_plotly\n from trieste.space import Box\n \n-search_space = Box([0, 0], [1, 1])\n+search_space = BRANIN_SEARCH_SPACE\n \n fig = plot_function_plotly(\n     scaled_branin, search_space.lower, search_space.upper, grid_density=20\n@@ -32,7 +32,7 @@\n fig.show()\n \n # %% [markdown]\n-# We begin our Bayesian active learning from a two-point initial design built from a space-filling Halton sequence.\n+# We begin our Bayesian active learning from a small initial design built from a space-filling Halton sequence.\n \n # %%\n import trieste\n@@ -47,16 +47,24 @@\n # %% [markdown]\n # ## Surrogate model\n #\n-# 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.\n+# 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.\n \n # %%\n import gpflow\n from trieste.models.gpflow.models import GaussianProcessRegression\n+import tensorflow_probability as tfp\n \n \n def build_model(data):\n     variance = tf.math.reduce_variance(data.observations)\n-    kernel = gpflow.kernels.RBF(variance=variance, lengthscales=[2, 2])\n+    kernel = gpflow.kernels.Matern52(variance=variance, lengthscales=[0.2, 0.2])\n+    prior_scale = tf.cast(1.0, dtype=tf.float64)\n+    kernel.variance.prior = tfp.distributions.LogNormal(\n+        tf.cast(-2.0, dtype=tf.float64), prior_scale\n+    )\n+    kernel.lengthscales.prior = tfp.distributions.LogNormal(\n+        tf.math.log(kernel.lengthscales), prior_scale\n+    )\n     gpr = gpflow.models.GPR(data.astuple(), kernel, noise_variance=1e-5)\n     gpflow.set_trainable(gpr.likelihood, False)\n \n@@ -79,9 +87,7 @@\n from trieste.acquisition.rule import EfficientGlobalOptimization\n \n acq = PredictiveVariance()\n-rule = EfficientGlobalOptimization(\n-    builder=acq, optimizer=generate_continuous_optimizer()\n-)\n+rule = EfficientGlobalOptimization(builder=acq)  # type: ignore\n bo = trieste.bayesian_optimizer.BayesianOptimizer(observer, search_space)\n \n # %% [markdown]\n@@ -144,12 +150,14 @@\n # %% [markdown]\n # ## Batch active learning using predictive variance\n #\n-# 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.\n+# 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.\n \n # %%\n bo_iter = 5\n num_query = 3\n+\n model = build_model(initial_data)\n+\n acq = PredictiveVariance()\n rule = EfficientGlobalOptimization(\n     num_query_points=num_query,\n", "issue": "Active learning acquisition functions are missing integration tests\n\n", "before_files": [{"content": "# %% [markdown]\n# # Active Learning\n\n# %% [markdown]\n# Sometimes, we may just want to learn a black-box function, rather than optimizing it. This goal is known as active learning and corresponds to choosing query points that reduce our model uncertainty. This notebook demonstrates how to perform Bayesian active learning using Trieste.\n\n# %%\n# %matplotlib inline\nimport numpy as np\nimport tensorflow as tf\n\nnp.random.seed(1793)\ntf.random.set_seed(1793)\n\n# %% [markdown]\n# ## Describe the problem\n#\n# In this example, we will perform active learning for the scaled Branin function.\n\n\n# %%\nfrom trieste.objectives import scaled_branin\nfrom util.plotting_plotly import plot_function_plotly\nfrom trieste.space import Box\n\nsearch_space = Box([0, 0], [1, 1])\n\nfig = plot_function_plotly(\n    scaled_branin, search_space.lower, search_space.upper, grid_density=20\n)\nfig.update_layout(height=400, width=400)\nfig.show()\n\n# %% [markdown]\n# We begin our Bayesian active learning from a two-point initial design built from a space-filling Halton sequence.\n\n# %%\nimport trieste\n\nobserver = trieste.objectives.utils.mk_observer(scaled_branin)\n\nnum_initial_points = 4\ninitial_query_points = search_space.sample_halton(num_initial_points)\ninitial_data = observer(initial_query_points)\n\n\n# %% [markdown]\n# ## Surrogate model\n#\n# 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.\n\n# %%\nimport gpflow\nfrom trieste.models.gpflow.models import GaussianProcessRegression\n\n\ndef build_model(data):\n    variance = tf.math.reduce_variance(data.observations)\n    kernel = gpflow.kernels.RBF(variance=variance, lengthscales=[2, 2])\n    gpr = gpflow.models.GPR(data.astuple(), kernel, noise_variance=1e-5)\n    gpflow.set_trainable(gpr.likelihood, False)\n\n    return GaussianProcessRegression(gpr)\n\n\nmodel = build_model(initial_data)\n\n# %% [markdown]\n# ## Active learning using predictive variance\n#\n# For our first active learning example, we will use a simple acquisition function known as `PredictiveVariance` which chooses points for which we are highly uncertain (i.e. the predictive posterior covariance matrix at these points has large determinant), as discussed in <cite data-cite=\"MacKay1992\"/>. Note that this also implies that our model needs to have `predict_joint` method to be able to return the full covariance, and it's likely to be expensive to compute.\n#\n# We will now demonstrate how to choose individual query points using `PredictiveVariance` before moving onto batch active learning. For both cases, we can utilize Trieste's `BayesianOptimizer` to do the active learning steps.\n#\n\n# %%\nfrom trieste.acquisition.function import PredictiveVariance\nfrom trieste.acquisition.optimizer import generate_continuous_optimizer\nfrom trieste.acquisition.rule import EfficientGlobalOptimization\n\nacq = PredictiveVariance()\nrule = EfficientGlobalOptimization(\n    builder=acq, optimizer=generate_continuous_optimizer()\n)\nbo = trieste.bayesian_optimizer.BayesianOptimizer(observer, search_space)\n\n# %% [markdown]\n# To plot the contour of variance of our model at each step, we can set the `track_state` parameter to `True` in `bo.optimize()`, this will make Trieste record our model at each iteration.\n\n# %%\nbo_iter = 5\nresult = bo.optimize(bo_iter, initial_data, model, rule, track_state=True)\n\n# %% [markdown]\n# Then we can retrieve our final dataset from the active learning steps.\n\n# %%\ndataset = result.try_get_final_dataset()\nquery_points = dataset.query_points.numpy()\nobservations = dataset.observations.numpy()\n\n# %% [markdown]\n# Finally, we can check the performance of our `PredictiveVariance` active learning acquisition function by plotting the predictive variance landscape of our model. We can see how it samples regions for which our model is highly uncertain.\n\n# %%\nfrom util.plotting import plot_bo_points, plot_function_2d\n\n\ndef plot_active_learning_query(\n    result, bo_iter, num_initial_points, query_points, num_query=1\n):\n\n    for i in range(bo_iter):\n\n        def pred_var(x):\n            _, var = result.history[i].models[\"OBJECTIVE\"].model.predict_f(x)\n            return var\n\n        _, ax = plot_function_2d(\n            pred_var,\n            search_space.lower - 0.01,\n            search_space.upper + 0.01,\n            grid_density=100,\n            contour=True,\n            colorbar=True,\n            figsize=(10, 6),\n            title=[\n                \"Variance contour with queried points at iter:\" + str(i + 1)\n            ],\n            xlabel=\"$X_1$\",\n            ylabel=\"$X_2$\",\n        )\n\n        plot_bo_points(\n            query_points[: num_initial_points + (i * num_query)],\n            ax[0, 0],\n            num_initial_points,\n        )\n\n\nplot_active_learning_query(result, bo_iter, num_initial_points, query_points)\n\n\n# %% [markdown]\n# ## Batch active learning using predictive variance\n#\n# 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.\n\n# %%\nbo_iter = 5\nnum_query = 3\nmodel = build_model(initial_data)\nacq = PredictiveVariance()\nrule = EfficientGlobalOptimization(\n    num_query_points=num_query,\n    builder=acq,\n    optimizer=generate_continuous_optimizer(num_optimization_runs=1),\n)\nbo = trieste.bayesian_optimizer.BayesianOptimizer(observer, search_space)\n\nresult = bo.optimize(bo_iter, initial_data, model, rule, track_state=True)\n\n\n# %% [markdown]\n# After that, we can retrieve our final dataset.\n\n# %%\ndataset = result.try_get_final_dataset()\nquery_points = dataset.query_points.numpy()\nobservations = dataset.observations.numpy()\n\n\n# %% [markdown]\n# Now we can visualize the batch predictive variance using our plotting function.\n\n# %%\nfrom util.plotting import plot_bo_points, plot_function_2d\n\nplot_active_learning_query(\n    result, bo_iter, num_initial_points, query_points, num_query\n)\n\n# %% [markdown]\n# ## LICENSE\n#\n# [Apache License 2.0](https://github.com/secondmind-labs/trieste/blob/develop/LICENSE)\n", "path": "docs/notebooks/active_learning.pct.py"}], "after_files": [{"content": "# %% [markdown]\n# # Active Learning\n\n# %% [markdown]\n# Sometimes, we may just want to learn a black-box function, rather than optimizing it. This goal is known as active learning and corresponds to choosing query points that reduce our model uncertainty. This notebook demonstrates how to perform Bayesian active learning using Trieste.\n\n# %%\n# %matplotlib inline\nimport numpy as np\nimport tensorflow as tf\n\nnp.random.seed(1793)\ntf.random.set_seed(1793)\n\n# %% [markdown]\n# ## Describe the problem\n#\n# In this example, we will perform active learning for the scaled Branin function.\n\n\n# %%\nfrom trieste.objectives import BRANIN_SEARCH_SPACE, scaled_branin\nfrom util.plotting_plotly import plot_function_plotly\nfrom trieste.space import Box\n\nsearch_space = BRANIN_SEARCH_SPACE\n\nfig = plot_function_plotly(\n    scaled_branin, search_space.lower, search_space.upper, grid_density=20\n)\nfig.update_layout(height=400, width=400)\nfig.show()\n\n# %% [markdown]\n# We begin our Bayesian active learning from a small initial design built from a space-filling Halton sequence.\n\n# %%\nimport trieste\n\nobserver = trieste.objectives.utils.mk_observer(scaled_branin)\n\nnum_initial_points = 4\ninitial_query_points = search_space.sample_halton(num_initial_points)\ninitial_data = observer(initial_query_points)\n\n\n# %% [markdown]\n# ## Surrogate model\n#\n# 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.\n\n# %%\nimport gpflow\nfrom trieste.models.gpflow.models import GaussianProcessRegression\nimport tensorflow_probability as tfp\n\n\ndef build_model(data):\n    variance = tf.math.reduce_variance(data.observations)\n    kernel = gpflow.kernels.Matern52(variance=variance, lengthscales=[0.2, 0.2])\n    prior_scale = tf.cast(1.0, dtype=tf.float64)\n    kernel.variance.prior = tfp.distributions.LogNormal(\n        tf.cast(-2.0, dtype=tf.float64), prior_scale\n    )\n    kernel.lengthscales.prior = tfp.distributions.LogNormal(\n        tf.math.log(kernel.lengthscales), prior_scale\n    )\n    gpr = gpflow.models.GPR(data.astuple(), kernel, noise_variance=1e-5)\n    gpflow.set_trainable(gpr.likelihood, False)\n\n    return GaussianProcessRegression(gpr)\n\n\nmodel = build_model(initial_data)\n\n# %% [markdown]\n# ## Active learning using predictive variance\n#\n# For our first active learning example, we will use a simple acquisition function known as `PredictiveVariance` which chooses points for which we are highly uncertain (i.e. the predictive posterior covariance matrix at these points has large determinant), as discussed in <cite data-cite=\"MacKay1992\"/>. Note that this also implies that our model needs to have `predict_joint` method to be able to return the full covariance, and it's likely to be expensive to compute.\n#\n# We will now demonstrate how to choose individual query points using `PredictiveVariance` before moving onto batch active learning. For both cases, we can utilize Trieste's `BayesianOptimizer` to do the active learning steps.\n#\n\n# %%\nfrom trieste.acquisition.function import PredictiveVariance\nfrom trieste.acquisition.optimizer import generate_continuous_optimizer\nfrom trieste.acquisition.rule import EfficientGlobalOptimization\n\nacq = PredictiveVariance()\nrule = EfficientGlobalOptimization(builder=acq)  # type: ignore\nbo = trieste.bayesian_optimizer.BayesianOptimizer(observer, search_space)\n\n# %% [markdown]\n# To plot the contour of variance of our model at each step, we can set the `track_state` parameter to `True` in `bo.optimize()`, this will make Trieste record our model at each iteration.\n\n# %%\nbo_iter = 5\nresult = bo.optimize(bo_iter, initial_data, model, rule, track_state=True)\n\n# %% [markdown]\n# Then we can retrieve our final dataset from the active learning steps.\n\n# %%\ndataset = result.try_get_final_dataset()\nquery_points = dataset.query_points.numpy()\nobservations = dataset.observations.numpy()\n\n# %% [markdown]\n# Finally, we can check the performance of our `PredictiveVariance` active learning acquisition function by plotting the predictive variance landscape of our model. We can see how it samples regions for which our model is highly uncertain.\n\n# %%\nfrom util.plotting import plot_bo_points, plot_function_2d\n\n\ndef plot_active_learning_query(\n    result, bo_iter, num_initial_points, query_points, num_query=1\n):\n\n    for i in range(bo_iter):\n\n        def pred_var(x):\n            _, var = result.history[i].models[\"OBJECTIVE\"].model.predict_f(x)\n            return var\n\n        _, ax = plot_function_2d(\n            pred_var,\n            search_space.lower - 0.01,\n            search_space.upper + 0.01,\n            grid_density=100,\n            contour=True,\n            colorbar=True,\n            figsize=(10, 6),\n            title=[\n                \"Variance contour with queried points at iter:\" + str(i + 1)\n            ],\n            xlabel=\"$X_1$\",\n            ylabel=\"$X_2$\",\n        )\n\n        plot_bo_points(\n            query_points[: num_initial_points + (i * num_query)],\n            ax[0, 0],\n            num_initial_points,\n        )\n\n\nplot_active_learning_query(result, bo_iter, num_initial_points, query_points)\n\n\n# %% [markdown]\n# ## Batch active learning using predictive variance\n#\n# 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.\n\n# %%\nbo_iter = 5\nnum_query = 3\n\nmodel = build_model(initial_data)\n\nacq = PredictiveVariance()\nrule = EfficientGlobalOptimization(\n    num_query_points=num_query,\n    builder=acq,\n    optimizer=generate_continuous_optimizer(num_optimization_runs=1),\n)\nbo = trieste.bayesian_optimizer.BayesianOptimizer(observer, search_space)\n\nresult = bo.optimize(bo_iter, initial_data, model, rule, track_state=True)\n\n\n# %% [markdown]\n# After that, we can retrieve our final dataset.\n\n# %%\ndataset = result.try_get_final_dataset()\nquery_points = dataset.query_points.numpy()\nobservations = dataset.observations.numpy()\n\n\n# %% [markdown]\n# Now we can visualize the batch predictive variance using our plotting function.\n\n# %%\nfrom util.plotting import plot_bo_points, plot_function_2d\n\nplot_active_learning_query(\n    result, bo_iter, num_initial_points, query_points, num_query\n)\n\n# %% [markdown]\n# ## LICENSE\n#\n# [Apache License 2.0](https://github.com/secondmind-labs/trieste/blob/develop/LICENSE)\n", "path": "docs/notebooks/active_learning.pct.py"}]}