Repository: secondmind-labs/trieste

Issue:
Acquisition function based on the predictive variance
This acquisition function calculates the determinant of the predictive covariance matrix for each set of batch points. For a batch of size 1 this is the same as returning the predictive variance.

This acquisition function is useful in the context of Bayesian active learning.

Notebook and tests in preparation.


Discussion:
Am trying to add some more code to help this PR to be more completed. Though simple and intuitive, am not totally sure which paper of this exactly is linked with. 

@henrymoss Was thinking if we take a log of the det, could we somehow just interpret as the entropy of multivariate normal, or DPP-MAX, or even GIBBON? Is this helps improving this PR?

Also, am not totally sure what kind of test/notebooks we need here, please point out if there is a suggestion :pray:

Resolving patch:
diff --git a/docs/notebooks/active_learning.pct.py b/docs/notebooks/active_learning.pct.py
new file mode 100644
--- /dev/null
+++ b/docs/notebooks/active_learning.pct.py
@@ -0,0 +1,170 @@
+# %% [markdown]
+# # Active Learning
+
+# %% [markdown]
+# 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`.
+
+# %%
+# %matplotlib inline
+import numpy as np
+import tensorflow as tf
+import pandas as pd
+
+np.random.seed(1793)
+tf.random.set_seed(1793)
+
+# %% [markdown]
+# ## Describe the problem
+#
+# In this example, we will perform active learning for the scaled Branin function.
+
+
+# %%
+from trieste.objectives import scaled_branin
+from util.plotting_plotly import plot_function_plotly
+from trieste.space import Box
+
+search_space = Box([0, 0], [1, 1])
+
+fig = plot_function_plotly(scaled_branin, search_space.lower, search_space.upper, grid_density=20)
+fig.update_layout(height=400, width=400)
+fig.show()
+
+# %% [markdown]
+# We begin our Bayesian active learning from a two-point initial design built from a space-filling Halton sequence.
+
+# %%
+import trieste
+
+observer = trieste.objectives.utils.mk_observer(scaled_branin)
+
+num_initial_points = 4
+initial_query_points = search_space.sample_halton(num_initial_points)
+initial_data = observer(initial_query_points)
+
+
+# %% [markdown]
+# ## Surrogate model
+#
+# Just like in sequential optimization, we fit a surrogate Gaussian process model to the initial data.
+
+# %%
+import gpflow
+from trieste.models.gpflow.models import GaussianProcessRegression
+
+
+def build_model(data):
+    variance = tf.math.reduce_variance(data.observations)
+    kernel = gpflow.kernels.RBF(variance=variance, lengthscales=[2, 2])
+    gpr = gpflow.models.GPR(data.astuple(), kernel, noise_variance=1e-5)
+    gpflow.set_trainable(gpr.likelihood, False)
+
+    return GaussianProcessRegression(gpr)
+
+
+model = build_model(initial_data)
+
+# %% [markdown]
+# ## Active learning using predictive variance
+#
+# 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.
+#
+# 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.
+#
+
+# %%
+from trieste.acquisition.optimizer import generate_continuous_optimizer
+from trieste.acquisition.rule import EfficientGlobalOptimization
+from trieste.acquisition.function import PredictiveVariance
+
+acq = PredictiveVariance()
+rule = EfficientGlobalOptimization(
+    builder=acq, optimizer=generate_continuous_optimizer(sigmoid=False)
+)
+bo = trieste.bayesian_optimizer.BayesianOptimizer(observer, search_space)
+
+# %% [markdown]
+# 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.
+
+# %%
+bo_iter = 5
+result = bo.optimize(bo_iter, initial_data, model, rule, track_state=True)
+
+# %% [markdown]
+# Then we can retrieve our final dataset from the active learning steps.
+
+# %%
+dataset = result.try_get_final_dataset()
+query_points = dataset.query_points.numpy()
+observations = dataset.observations.numpy()
+
+# %% [markdown]
+# 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.
+
+# %%
+from util.plotting import plot_bo_points, plot_function_2d
+
+
+def plot_active_learning_query(result, bo_iter, num_initial_points, query_points, num_query=1):
+
+    for i in range(bo_iter):
+
+        def pred_var(x):
+            _, var = result.history[i].models["OBJECTIVE"].model.predict_f(x)
+            return var
+
+        _, ax = plot_function_2d(
+            pred_var,
+            search_space.lower - 0.01,
+            search_space.upper + 0.01,
+            grid_density=100,
+            contour=True,
+            colorbar=True,
+            figsize=(10, 6),
+            title=["Variance contour with queried points at iter:" + str(i + 1)],
+            xlabel="$X_1$",
+            ylabel="$X_2$",
+        )
+
+        plot_bo_points(
+            query_points[: num_initial_points + (i * num_query)], ax[0, 0], num_initial_points
+        )
+
+
+plot_active_learning_query(result, bo_iter, num_initial_points, query_points)
+
+
+# %% [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.
+
+# %%
+bo_iter = 5
+num_query = 3
+model = build_model(initial_data)
+acq = PredictiveVariance()
+rule = EfficientGlobalOptimization(
+    num_query_points=num_query, builder=acq, optimizer=generate_continuous_optimizer(sigmoid=False)
+)
+bo = trieste.bayesian_optimizer.BayesianOptimizer(observer, search_space)
+
+result = bo.optimize(bo_iter, initial_data, model, rule, track_state=True)
+
+
+# %% [markdown]
+# After that, we can retrieve our final dataset.
+
+# %%
+dataset = result.try_get_final_dataset()
+query_points = dataset.query_points.numpy()
+observations = dataset.observations.numpy()
+
+
+# %% [markdown]
+# Now we can visualize the batch predictive variance using our plotting function.
+
+# %%
+from util.plotting import plot_bo_points, plot_function_2d
+
+plot_active_learning_query(result, bo_iter, num_initial_points, query_points, num_query)
diff --git a/trieste/acquisition/function.py b/trieste/acquisition/function.py
--- a/trieste/acquisition/function.py
+++ b/trieste/acquisition/function.py
@@ -2097,3 +2097,71 @@ def _get_min_value_samples(
         sampler = ExactThompsonSampler(num_samples, model, sample_min_value=True)
 
     return sampler.sample(sampled_points)
+
+
+class PredictiveVariance(SingleModelAcquisitionBuilder):
+    """
+    Builder for the determinant of the predictive covariance matrix over the batch points.
+    For a batch of size 1 it is the same as maximizing the predictive variance.
+    """
+
+    def __init__(self, jitter: float = DEFAULTS.JITTER) -> None:
+        """
+        :param jitter: The size of the jitter to use when stabilising the Cholesky decomposition of
+            the covariance matrix.
+        """
+        self._jitter = jitter
+
+    def __repr__(self) -> str:
+        """"""
+        return f"PredictiveVariance(jitter={self._jitter!r})"
+
+    def prepare_acquisition_function(
+        self, dataset: Dataset, model: ProbabilisticModel
+    ) -> AcquisitionFunction:
+        """
+        :param dataset: Unused.
+        :param model: The model over the specified ``dataset``.
+
+        :return: The determinant of the predictive function.
+        """
+
+        return predictive_variance(model, self._jitter)
+
+    def update_acquisition_function(
+        self, function: AcquisitionFunction, dataset: Dataset, model: ProbabilisticModel
+    ) -> AcquisitionFunction:
+        """
+        :param function: The acquisition function to update.
+        :param dataset: Unused.
+        :param model: The model over the specified ``dataset``.
+        """
+        return function  # no need to update anything
+
+
+def predictive_variance(model: ProbabilisticModel, jitter: float) -> TensorType:
+    """
+    The predictive variance acquisition function for active learning, based on
+    the determinant of the covariance (see :cite:`MacKay1992` for details).
+    Note that the model needs to supply covariance of the joint marginal distribution,
+    which can be expensive to compute.
+
+    :param model: The model of the objective function.
+    :param jitter: The size of the jitter to use when stabilising the Cholesky decomposition of
+            the covariance matrix.
+    """
+
+    @tf.function
+    def acquisition(x: TensorType) -> TensorType:
+
+        try:
+            _, covariance = model.predict_joint(x)
+        except NotImplementedError:
+            raise ValueError(
+                """
+                PredictiveVariance only supports models with a predict_joint method.
+                """
+            )
+        return tf.exp(tf.linalg.logdet(covariance + jitter))
+
+    return acquisition


Test patch:
diff --git a/tests/unit/acquisition/test_function.py b/tests/unit/acquisition/test_function.py
--- a/tests/unit/acquisition/test_function.py
+++ b/tests/unit/acquisition/test_function.py
@@ -38,6 +38,7 @@
     various_shapes,
 )
 from tests.util.models.gpflow.models import GaussianProcess, QuadraticMeanAndRBFKernel, rbf
+from tests.util.models.gpflux.models import trieste_deep_gaussian_process
 from trieste.acquisition.function import (
     GIBBON,
     AcquisitionFunction,
@@ -54,6 +55,7 @@
     NegativeLowerConfidenceBound,
     NegativePredictiveMean,
     PenalizationFunction,
+    PredictiveVariance,
     ProbabilityOfFeasibility,
     SingleModelAcquisitionBuilder,
     SingleModelGreedyAcquisitionBuilder,
@@ -67,6 +69,7 @@
     hard_local_penalizer,
     lower_confidence_bound,
     min_value_entropy_search,
+    predictive_variance,
     probability_of_feasibility,
     soft_local_penalizer,
 )
@@ -1289,6 +1292,7 @@ def test_batch_monte_carlo_expected_improvement_updates_without_retracing() -> N
             BatchMonteCarloExpectedImprovement(10_000),
             f"BatchMonteCarloExpectedImprovement(10000, jitter={DEFAULTS.JITTER})",
         ),
+        (PredictiveVariance(), f"PredictiveVariance(jitter={DEFAULTS.JITTER})"),
     ],
 )
 def test_single_model_acquisition_function_builder_reprs(
@@ -1756,3 +1760,96 @@ def acquisition(x: TensorType) -> TensorType:
 
     xs = tf.linspace([[-10.0]], [[10.0]], 100)
     npt.assert_allclose(echvi(xs), constraint_fn(xs))
+
+
+def test_predictive_variance_builder_builds_predictive_variance() -> None:
+    model = QuadraticMeanAndRBFKernel()
+    acq_fn = PredictiveVariance().prepare_acquisition_function(
+        Dataset(tf.zeros([0, 1]), tf.zeros([0, 1])), model
+    )
+    query_at = tf.linspace([[-10]], [[10]], 100)
+    _, covariance = model.predict_joint(query_at)
+    expected = tf.linalg.det(covariance)
+    npt.assert_array_almost_equal(acq_fn(query_at), expected)
+
+
+@pytest.mark.parametrize(
+    "at, acquisition_shape",
+    [
+        (tf.constant([[[1.0]]]), tf.constant([1, 1])),
+        (tf.linspace([[-10.0]], [[10.0]], 5), tf.constant([5, 1])),
+        (tf.constant([[[1.0, 1.0]]]), tf.constant([1, 1])),
+        (tf.linspace([[-10.0, -10.0]], [[10.0, 10.0]], 5), tf.constant([5, 1])),
+    ],
+)
+def test_predictive_variance_returns_correct_shape(
+    at: TensorType, acquisition_shape: TensorType
+) -> None:
+    model = QuadraticMeanAndRBFKernel()
+    acq_fn = PredictiveVariance().prepare_acquisition_function(
+        Dataset(tf.zeros([0, 1]), tf.zeros([0, 1])), model
+    )
+    npt.assert_array_equal(acq_fn(at).shape, acquisition_shape)
+
+
+@random_seed
+@pytest.mark.parametrize(
+    "variance_scale, num_samples_per_point, rtol, atol",
+    [
+        (0.1, 10_000, 0.05, 1e-6),
+        (1.0, 50_000, 0.05, 1e-3),
+        (10.0, 100_000, 0.05, 1e-2),
+        (100.0, 150_000, 0.05, 1e-1),
+    ],
+)
+def test_predictive_variance(
+    variance_scale: float,
+    num_samples_per_point: int,
+    rtol: float,
+    atol: float,
+) -> None:
+    variance_scale = tf.constant(variance_scale, tf.float64)
+
+    x_range = tf.linspace(0.0, 1.0, 11)
+    x_range = tf.cast(x_range, dtype=tf.float64)
+    xs = tf.reshape(tf.stack(tf.meshgrid(x_range, x_range, indexing="ij"), axis=-1), (-1, 2))
+
+    kernel = tfp.math.psd_kernels.MaternFiveHalves(variance_scale, length_scale=0.25)
+    model = GaussianProcess([branin], [kernel])
+
+    mean, variance = model.predict(xs)
+    samples = tfp.distributions.Normal(mean, tf.sqrt(variance)).sample(num_samples_per_point)
+    predvar_approx = tf.math.reduce_variance(samples, axis=0)
+
+    detcov = predictive_variance(model, DEFAULTS.JITTER)
+    predvar = detcov(xs[..., None, :])
+
+    npt.assert_allclose(predvar, predvar_approx, rtol=rtol, atol=atol)
+
+
+def test_predictive_variance_builder_updates_without_retracing() -> None:
+    model = QuadraticMeanAndRBFKernel()
+    builder = PredictiveVariance()
+    acq_fn = builder.prepare_acquisition_function(
+        Dataset(tf.zeros([0, 1]), tf.zeros([0, 1])), model
+    )
+    assert acq_fn._get_tracing_count() == 0  # type: ignore
+    query_at = tf.linspace([[-10]], [[10]], 100)
+    expected = predictive_variance(model, DEFAULTS.JITTER)(query_at)
+    npt.assert_array_almost_equal(acq_fn(query_at), expected)
+    assert acq_fn._get_tracing_count() == 1  # type: ignore
+
+    up_acq_fn = builder.update_acquisition_function(
+        acq_fn, Dataset(tf.ones([0, 1]), tf.ones([0, 1])), model
+    )
+    assert up_acq_fn == acq_fn
+    npt.assert_array_almost_equal(acq_fn(query_at), expected)
+    assert acq_fn._get_tracing_count() == 1  # type: ignore
+
+
+def test_predictive_variance_raises_for_void_predict_joint() -> None:
+    model, _ = trieste_deep_gaussian_process(tf.zeros([0, 1]), 2, 20, 0.01, 100, 100)
+    acq_fn = predictive_variance(model, DEFAULTS.JITTER)
+
+    with pytest.raises(ValueError):
+        acq_fn(tf.zeros([0, 1]))