[BREAKING] Make ProbabilisticModel and TrainableProbabilisticModel pu…

…re interfaces (#797)

docs/notebooks/code_overview.pct.py

@@ -198,12 +198,13 @@ class HasGizmoReparamSamplerAndObservationNoise(
     AcquisitionFunction,
     SingleModelAcquisitionBuilder,
 )
+from trieste.models.interfaces import SupportsPredictY
 from trieste.data import Dataset
 
 
-class ProbabilityOfValidity(SingleModelAcquisitionBuilder[ProbabilisticModel]):
+class ProbabilityOfValidity(SingleModelAcquisitionBuilder[SupportsPredictY]):
     def prepare_acquisition_function(
-        self, model: ProbabilisticModel, dataset: Optional[Dataset] = None
+        self, model: SupportsPredictY, dataset: Optional[Dataset] = None
     ) -> AcquisitionFunction:
         def acquisition(at: TensorType) -> TensorType:
             mean, _ = model.predict_y(tf.squeeze(at, -2))
@@ -217,9 +218,9 @@ def acquisition(at: TensorType) -> TensorType:
 
 
 # %%
-class ProbabilityOfValidity2(SingleModelAcquisitionBuilder[ProbabilisticModel]):
+class ProbabilityOfValidity2(SingleModelAcquisitionBuilder[SupportsPredictY]):
     def prepare_acquisition_function(
-        self, model: ProbabilisticModel, dataset: Optional[Dataset] = None
+        self, model: SupportsPredictY, dataset: Optional[Dataset] = None
     ) -> AcquisitionFunction:
         @tf.function
         def acquisition(at: TensorType) -> TensorType:
@@ -231,7 +232,7 @@ def acquisition(at: TensorType) -> TensorType:
     def update_acquisition_function(
         self,
         function: AcquisitionFunction,
-        model: ProbabilisticModel,
+        model: SupportsPredictY,
         dataset: Optional[Dataset] = None,
     ) -> AcquisitionFunction:
         return function  # no need to update anything

docs/notebooks/expected_improvement.pct.py

@@ -299,7 +299,9 @@ def build_model(data):
 
 # %%
 # save the model to a given path, exporting just the predict method
-module = result.try_get_final_model().get_module_with_variables()
+from trieste.models.utils import get_module_with_variables
+
+module = get_module_with_variables(result.try_get_final_model())
 module.predict = tf.function(
     model.predict,
     input_signature=[tf.TensorSpec(shape=[None, 2], dtype=tf.float64)],

tests/integration/models/multifidelity/test_multifidelity_models.py

@@ -11,6 +11,7 @@
     check_and_extract_fidelity_query_points,
     split_dataset_by_fidelity,
 )
+from trieste.models import TrainableProbabilisticModel
 from trieste.models.gpflow import GaussianProcessRegression
 from trieste.models.gpflow.builders import (
     build_gpr,
@@ -119,7 +120,7 @@ def test_multifidelity_nonlinear_autoregressive_results_better_than_linear() ->
     observer = mk_observer(noisy_nonlinear_multifidelity)
     initial_data = observer(initial_sample)
 
-    nonlinear_model = MultifidelityNonlinearAutoregressive(
+    nonlinear_model: TrainableProbabilisticModel = MultifidelityNonlinearAutoregressive(
         build_multifidelity_nonlinear_autoregressive_models(
             initial_data, n_fidelities, input_search_space
         )

tests/unit/acquisition/function/test_entropy.py

@@ -36,6 +36,7 @@
     MinValueEntropySearch,
     MUMBOModelType,
     SupportsCovarianceObservationNoiseTrajectory,
+    SupportsCovarianceWithTopFidelityPredictY,
     gibbon_quality_term,
     gibbon_repulsion_term,
     min_value_entropy_search,
@@ -48,7 +49,6 @@
     ThompsonSamplerFromTrajectory,
 )
 from trieste.data import Dataset, add_fidelity_column
-from trieste.models import SupportsCovarianceWithTopFidelity
 from trieste.objectives import Branin
 from trieste.space import Box
 from trieste.types import TensorType
@@ -612,7 +612,7 @@ def test_mumbo_raises_when_use_trajectory_sampler_and_model_without_trajectories
 )
 def test_mumbo_builder_builds_min_value_samples(
     mocked_mves: MagicMock,
-    min_value_sampler: ThompsonSampler[SupportsCovarianceWithTopFidelity],
+    min_value_sampler: ThompsonSampler[SupportsCovarianceWithTopFidelityPredictY],
 ) -> None:
     dataset = Dataset(tf.zeros([3, 2], dtype=tf.float64), tf.ones([3, 2], dtype=tf.float64))
     search_space = Box([0, 0], [1, 1])
@@ -638,7 +638,7 @@ def test_mumbo_builder_builds_min_value_samples(
     [ExactThompsonSampler(sample_min_value=True), GumbelSampler(sample_min_value=True)],
 )
 def test_mumbo_builder_updates_acquisition_function(
-    min_value_sampler: ThompsonSampler[SupportsCovarianceWithTopFidelity],
+    min_value_sampler: ThompsonSampler[SupportsCovarianceWithTopFidelityPredictY],
 ) -> None:
     search_space = Box([0.0, 0.0], [1.0, 1.0])
     model = MultiFidelityQuadraticMeanAndRBFKernel(

tests/unit/models/gpflow/test_interface.py

@@ -14,7 +14,7 @@
 
 from __future__ import annotations
 
-from typing import Any
+from typing import Any, Optional
 
 import gpflow
 import numpy.testing as npt
@@ -37,6 +37,9 @@ def optimize(self, dataset: Dataset) -> None:
     def update(self, dataset: Dataset) -> None:
         return
 
+    def log(self, dataset: Optional[Dataset] = None) -> None:
+        return
+
 
 class _QuadraticGPModel(GPModel):
     def __init__(self) -> None:

tests/unit/models/gpflow/test_models.py

@@ -84,6 +84,7 @@
     RandomFourierFeatureTrajectorySampler,
 )
 from trieste.models.optimizer import BatchOptimizer, DatasetTransformer, Optimizer
+from trieste.models.utils import get_last_optimization_result, optimize_model_and_save_result
 from trieste.space import Box
 from trieste.types import TensorType
 from trieste.utils import DEFAULTS
@@ -150,13 +151,14 @@ def test_gpflow_wrappers_default_optimize(
         args = {}
 
     loss = internal_model.training_loss(**args)
-    model.optimize_and_save_result(Dataset(*data))
+    optimize_model_and_save_result(model, Dataset(*data))
 
     new_loss = internal_model.training_loss(**args)
     assert new_loss < loss
     if not isinstance(internal_model, SVGP):
-        assert model.last_optimization_result is not None
-        npt.assert_allclose(new_loss, model.last_optimization_result.fun)
+        optimization_result = get_last_optimization_result(model)
+        assert optimization_result is not None
+        npt.assert_allclose(new_loss, optimization_result.fun)
 
 
 def test_gpflow_wrappers_ref_optimize(gpflow_interface_factory: ModelFactoryType) -> None:

tests/unit/models/gpflow/test_utils.py

@@ -39,6 +39,7 @@
 )
 from trieste.models.interfaces import HasTrajectorySampler
 from trieste.models.optimizer import BatchOptimizer, Optimizer
+from trieste.models.utils import get_module_with_variables
 from trieste.types import TensorType
 
 
@@ -92,7 +93,7 @@ def test_gaussian_process_tf_saved_model(gpflow_interface_factory: ModelFactoryT
         trajectory = trajectory_sampler.get_trajectory()
 
         # generate client model with predict and sample methods
-        module = model.get_module_with_variables(trajectory_sampler, trajectory)
+        module = get_module_with_variables(model, trajectory_sampler, trajectory)
         module.predict = tf.function(
             model.predict, input_signature=[tf.TensorSpec(shape=[None, 1], dtype=tf.float64)]
         )

tests/unit/models/gpflux/test_interface.py

@@ -14,6 +14,8 @@
 
 from __future__ import annotations
 
+from typing import Optional
+
 import gpflow
 import numpy.testing as npt
 import pytest
@@ -70,6 +72,9 @@ def sample(self, query_points: TensorType, num_samples: int) -> TensorType:
     def update(self, dataset: Dataset) -> None:
         return
 
+    def log(self, dataset: Optional[Dataset] = None) -> None:
+        return
+
 
 class _QuadraticGPModel(DeepGP):
     def __init__(

tests/unit/models/gpflux/test_models.py

@@ -50,6 +50,11 @@
 from trieste.models.gpflux import DeepGaussianProcess
 from trieste.models.interfaces import HasTrajectorySampler
 from trieste.models.optimizer import KerasOptimizer
+from trieste.models.utils import (
+    get_last_optimization_result,
+    get_module_with_variables,
+    optimize_model_and_save_result,
+)
 from trieste.types import TensorType
 
 
@@ -297,10 +302,11 @@ def test_deep_gaussian_process_with_lr_scheduler(
     optimizer = KerasOptimizer(tf.optimizers.Adam(lr_schedule), fit_args)
     model = DeepGaussianProcess(two_layer_model(x), optimizer)
 
-    model.optimize_and_save_result(Dataset(x, y))
+    optimize_model_and_save_result(model, Dataset(x, y))
 
-    assert model.last_optimization_result is not None
-    assert len(model.last_optimization_result.history["loss"]) == epochs
+    optimization_result = get_last_optimization_result(model)
+    assert optimization_result is not None
+    assert len(optimization_result.history["loss"]) == epochs
 
 
 def test_deep_gaussian_process_default_optimizer_is_correct(
@@ -395,7 +401,7 @@ def test_deepgp_tf_saved_model() -> None:
         trajectory = trajectory_sampler.get_trajectory()
 
         # generate client model with predict and sample methods
-        module = model.get_module_with_variables(trajectory_sampler, trajectory)
+        module = get_module_with_variables(model, trajectory_sampler, trajectory)
         module.predict = tf.function(
             model.predict, input_signature=[tf.TensorSpec(shape=[None, 1], dtype=tf.float64)]
         )

tests/unit/models/keras/test_models.py

@@ -42,6 +42,11 @@
     sample_with_replacement,
 )
 from trieste.models.optimizer import KerasOptimizer, TrainingData
+from trieste.models.utils import (
+    get_last_optimization_result,
+    get_module_with_variables,
+    optimize_model_and_save_result,
+)
 from trieste.types import TensorType
 
 _ENSEMBLE_SIZE = 3
@@ -215,10 +220,11 @@ def scheduler(epoch: int, lr: float) -> float:
 
     npt.assert_allclose(model.model.optimizer.lr.numpy(), init_lr, rtol=1e-6)
 
-    model.optimize_and_save_result(example_data)
+    optimize_model_and_save_result(model, example_data)
 
-    assert model.last_optimization_result is not None
-    npt.assert_allclose(model.last_optimization_result.history["lr"], [0.5, 0.25])
+    optimization_result = get_last_optimization_result(model)
+    assert optimization_result is not None
+    npt.assert_allclose(optimization_result.history["lr"], [0.5, 0.25])
     npt.assert_allclose(model.model.optimizer.lr.numpy(), init_lr, rtol=1e-6)
 
 
@@ -568,7 +574,7 @@ def test_deep_ensemble_tf_saved_model() -> None:
         trajectory = trajectory_sampler.get_trajectory()
 
         # generate client model with predict and sample methods
-        module = model.get_module_with_variables(trajectory_sampler, trajectory)
+        module = get_module_with_variables(model, trajectory_sampler, trajectory)
         module.predict = tf.function(
             model.predict, input_signature=[tf.TensorSpec(shape=[None, 3], dtype=tf.float64)]
         )

tests/unit/models/test_interfaces.py

@@ -36,9 +36,11 @@
 from trieste.models import TrainableModelStack, TrainableProbabilisticModel
 from trieste.models.interfaces import (
     TrainablePredictJointReparamModelStack,
+    TrainablePredictYModelStack,
     TrainableSupportsPredictJoint,
     TrainableSupportsPredictJointHasReparamSampler,
 )
+from trieste.models.utils import get_last_optimization_result, optimize_model_and_save_result
 from trieste.types import TensorType
 
 
@@ -114,28 +116,11 @@ def test_model_stack_predict_joint() -> None:
     npt.assert_allclose(cov[..., 3:, :, :], cov3)
 
 
-def test_model_missing_predict_y() -> None:
-    model = _QuadraticModel([-1.0], [0.1])
-    x_predict = tf.constant([[0]], gpflow.default_float())
-    with pytest.raises(NotImplementedError):
-        model.predict_y(x_predict)
-
-
-def test_model_stack_missing_predict_y() -> None:
-    x = tf.constant(np.arange(5).reshape(-1, 1), dtype=gpflow.default_float())
-    model1 = gpr_model(x, fnc_3x_plus_10(x))
-    model2 = _QuadraticModel([1.0], [2.0])
-    stack = TrainableModelStack((model1, 1), (model2, 1))
-    x_predict = tf.constant([[0]], gpflow.default_float())
-    with pytest.raises(NotImplementedError):
-        stack.predict_y(x_predict)
-
-
 def test_model_stack_predict_y() -> None:
     x = tf.constant(np.arange(5).reshape(-1, 1), dtype=gpflow.default_float())
     model1 = gpr_model(x, fnc_3x_plus_10(x))
     model2 = sgpr_model(x, fnc_2sin_x_over_3(x))
-    stack = TrainableModelStack((model1, 1), (model2, 1))
+    stack = TrainablePredictYModelStack((model1, 1), (model2, 1))
     mean, variance = stack.predict_y(x)
     npt.assert_allclose(mean[:, 0:1], model1.predict_y(x)[0])
     npt.assert_allclose(mean[:, 1:2], model2.predict_y(x)[0])
@@ -193,8 +178,8 @@ def _assert_data(self, dataset: Dataset) -> None:
     stack = TrainableModelStack((model01, 2), (model2, 1), (model3, 1))
     data = Dataset(tf.random.uniform([5, 7, 3]), tf.random.uniform([5, 7, 4]))
     stack.update(data)
-    stack.optimize_and_save_result(data)
-    assert stack.last_optimization_result == [None] * 3
+    optimize_model_and_save_result(stack, data)
+    assert get_last_optimization_result(stack) == [None] * 3
 
 
 def test_model_stack_reparam_sampler_raises_for_submodels_without_reparam_sampler() -> None:

tests/unit/test_bayesian_optimizer.py

@@ -491,6 +491,9 @@ def update(self, dataset: Dataset) -> NoReturn:
         def optimize(self, dataset: Dataset) -> NoReturn:
             assert False
 
+        def log(self, dataset: Optional[Dataset] = None) -> None:
+            return
+
     class _UnusableRule(AcquisitionRule[NoReturn, Box, ProbabilisticModel]):
         def acquire(
             self,

tests/util/models/gpflow/models.py

@@ -47,6 +47,7 @@
     SupportsGetKernel,
     SupportsGetObservationNoise,
     SupportsPredictJoint,
+    SupportsPredictY,
 )
 from trieste.models.optimizer import Optimizer
 from trieste.types import TensorType
@@ -120,6 +121,9 @@ def covariance_between_points(
         ]
         return tf.concat(covs, axis=-3)
 
+    def log(self, dataset: Optional[Dataset] = None) -> None:
+        return
+
 
 class GaussianProcessWithoutNoise(GaussianMarginal, SupportsPredictJoint, HasReparamSampler):
     """A (static) Gaussian process over a vector random variable with independent reparam sampler
@@ -160,6 +164,9 @@ def reparam_sampler(
     ) -> ReparametrizationSampler[GaussianProcessWithoutNoise]:
         return IndependentReparametrizationSampler(num_samples, self)
 
+    def log(self, dataset: Optional[Dataset] = None) -> None:
+        return
+
 
 class GaussianProcessWithSamplers(GaussianProcess, HasReparamSampler):
     """A (static) Gaussian process over a vector random variable with independent reparam sampler"""
@@ -259,7 +266,7 @@ def optimize(self, dataset: Dataset) -> None:
 
 
 class MultiFidelityQuadraticMeanAndRBFKernel(
-    QuadraticMeanAndRBFKernel, SupportsCovarianceWithTopFidelity
+    QuadraticMeanAndRBFKernel, SupportsPredictY, SupportsCovarianceWithTopFidelity
 ):
     r"""
     A Gaussian process with scalar quadratic mean, an RBF kernel and
@@ -293,7 +300,7 @@ def predict_y(self, query_points: TensorType) -> tuple[TensorType, TensorType]:
 
 
 class MultiFidelityQuadraticMeanAndRBFKernelWithSamplers(
-    QuadraticMeanAndRBFKernelWithSamplers, SupportsCovarianceWithTopFidelity
+    QuadraticMeanAndRBFKernelWithSamplers, SupportsPredictY, SupportsCovarianceWithTopFidelity
 ):
     r"""
     A Gaussian process with scalar quadratic mean, an RBF kernel and
@@ -323,6 +330,12 @@ def covariance_with_top_fidelity(self, x: TensorType) -> TensorType:
         mean, _ = self.predict(x)
         return tf.ones_like(mean, dtype=mean.dtype)  # dummy covariances of correct shape
 
+    @inherit_check_shapes
+    def predict_y(self, query_points: TensorType) -> tuple[TensorType, TensorType]:
+        fmean, fvar = self.predict(query_points)
+        yvar = fvar + tf.constant(1.0, dtype=fmean.dtype)  # dummy noise variance
+        return fmean, yvar
+
 
 class QuadraticMeanAndRBFKernelWithBatchSamplers(
     QuadraticMeanAndRBFKernel, HasTrajectorySampler, HasReparamSampler