Merge branch 'develop' into patch-1

.github/workflows/quality-checks.yaml

@@ -54,7 +54,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        part: [ "1", "2", "3" ]
+        part: [ "1", "2", "3", "4" ]
     name: tests (part${{ matrix.part }})
     steps:
     - uses: actions/checkout@v2
@@ -68,7 +68,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        part: [ "1", "2", "3" ]
+        part: [ "1", "2", "3", "4" ]
     name: tests_old (part${{ matrix.part }})
     steps:
     - uses: actions/checkout@v2

tests/unit/models/gpflow/test_interface.py

@@ -31,6 +31,9 @@ class _QuadraticPredictor(GPflowPredictor):
     def model(self) -> GPModel:
         return _QuadraticGPModel()
 
+    def optimize(self, dataset: Dataset) -> None:
+        self.optimizer.optimize(self.model, dataset)
+
     def update(self, dataset: Dataset) -> None:
         return
 

tests/unit/models/gpflow/test_models.py

@@ -150,9 +150,13 @@ def test_gpflow_wrappers_default_optimize(
         args = {}
 
     loss = internal_model.training_loss(**args)
-    model.optimize(Dataset(*data))
+    model.optimize_and_save_result(Dataset(*data))
 
-    assert internal_model.training_loss(**args) < loss
+    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)
 
 
 def test_gpflow_wrappers_ref_optimize(gpflow_interface_factory: ModelFactoryType) -> None:

tests/unit/models/gpflux/test_models.py

@@ -297,9 +297,10 @@ 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(Dataset(x, y))
+    model.optimize_and_save_result(Dataset(x, y))
 
-    assert len(model.model_keras.history.history["loss"]) == epochs
+    assert model.last_optimization_result is not None
+    assert len(model.last_optimization_result.history["loss"]) == epochs
 
 
 def test_deep_gaussian_process_default_optimizer_is_correct(

tests/unit/models/keras/test_models.py

@@ -215,9 +215,10 @@ def scheduler(epoch: int, lr: float) -> float:
 
     npt.assert_allclose(model.model.optimizer.lr.numpy(), init_lr, rtol=1e-6)
 
-    model.optimize(example_data)
+    model.optimize_and_save_result(example_data)
 
-    npt.assert_allclose(model.model.history.history["lr"], [0.5, 0.25])
+    assert model.last_optimization_result is not None
+    npt.assert_allclose(model.last_optimization_result.history["lr"], [0.5, 0.25])
     npt.assert_allclose(model.model.optimizer.lr.numpy(), init_lr, rtol=1e-6)
 
 

tests/unit/models/keras/test_sampler.py

@@ -407,7 +407,7 @@ def test_ensemble_trajectory_sampler_trajectory_on_subsets_same_as_set(diversify
     eval_2 = trajectory(test_data[100:200, :])
     eval_3 = trajectory(test_data[200:300, :])
 
-    npt.assert_allclose(eval_all, tf.concat([eval_1, eval_2, eval_3], axis=0), rtol=5e-6)
+    npt.assert_allclose(eval_all, tf.concat([eval_1, eval_2, eval_3], axis=0), rtol=1e-5)
 
 
 @random_seed

tests/unit/models/test_interfaces.py

@@ -193,7 +193,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(data)
+    stack.optimize_and_save_result(data)
+    assert stack.last_optimization_result == [None] * 3
 
 
 def test_model_stack_reparam_sampler_raises_for_submodels_without_reparam_sampler() -> None:

tests/util/models/gpflow/models.py

@@ -211,7 +211,7 @@ def mock_data() -> tuple[tf.Tensor, tf.Tensor]:
 
 
 class QuadraticMeanAndRBFKernelWithSamplers(
-    QuadraticMeanAndRBFKernel, HasTrajectorySampler, HasReparamSampler
+    QuadraticMeanAndRBFKernel, HasTrajectorySampler, HasReparamSampler, TrainableProbabilisticModel
 ):
     r"""
     A Gaussian process with scalar quadratic mean, an RBF kernel and

tox.ini

@@ -36,6 +36,7 @@ description =
     tests_1: Run tests part 1
     tests_2: Run tests part 2
     tests_3: Run tests part 3
+    tests_4: Run tests part 4
     coverage: Generate code coverage
     docs: Generate documentation
     quickdocs: Generate dummy documentation
@@ -58,10 +59,11 @@ commands =
     types: mypy {posargs}
     # run tests (run with "-- --runslow yes" to run all tests)
     # note that we currently split this into three separate runs due to memory leak issues
-    tests,tests_1,tests_2,tests_3: pip install .[qhsri] -r tests/latest/requirements.txt -c tests/latest/constraints.txt
+    tests,tests_1,tests_2,tests_3,tests_4: pip install .[qhsri] -r tests/latest/requirements.txt -c tests/latest/constraints.txt
     tests,tests_1: pytest --ignore trieste/experimental --qhsri yes {posargs} trieste tests/unit/models/gpflux
     tests,tests_2: pytest --ignore tests/unit/models/gpflux --qhsri yes {posargs} tests/unit
-    tests,tests_3: pytest --qhsri yes {posargs} tests/util tests/integration
+    tests,tests_3: pytest --qhsri yes {posargs} tests/util tests/integration/test_ask_tell_optimization.py tests/integration/test_bayesian_optimization.py
+    tests,tests_4: pytest --qhsri yes --ignore tests/integration/test_ask_tell_optimization.py --ignore tests/integration/test_bayesian_optimization.py {posargs} tests/integration
     # code coverage
     coverage: pip install .[qhsri] -r tests/latest/requirements.txt -c tests/latest/constraints.txt
     coverage: pip install -r common_build/taskipy/requirements.txt -c common_build/taskipy/constraints.txt
@@ -82,21 +84,23 @@ commands =
 
 # additional tests using python 3.7 and older versions of tensorflow
 
-[testenv:{types_old,tests_old,tests_old_1,tests_old_2,tests_old_3}]
+[testenv:{types_old,tests_old,tests_old_1,tests_old_2,tests_old_3,tests_old_4}]
 basepython = python3.7
 description =
     types_old: Check types [Python 3.7]
     tests_old: Run tests [Python 3.7, no qhsri]
     tests_old_1: Run old tests part 1
     tests_old_2: Run old tests part 2
     tests_old_3: Run old tests part 3
+    tests_old_4: Run old tests part 4
 commands =
     types_old: pip install -r common_build/types/requirements.txt -c common_build/types/constraints.txt
     types_old: pip install . -r tests/old/requirements.txt -c tests/old/constraints.txt
     types_old: mypy {posargs}
     # unlike tests this doesn't include the optional qhsri support
     # note that we currently split this into three separate runs due to memory leak issues
-    tests_old,tests_old_1,tests_old_2,tests_old_3: pip install . -r tests/old/requirements.txt -c tests/old/constraints.txt
+    tests_old,tests_old_1,tests_old_2,tests_old_3,tests_old_4: pip install . -r tests/old/requirements.txt -c tests/old/constraints.txt
     tests_old,tests_old_1: pytest --ignore trieste/experimental {posargs} trieste tests/unit/models/gpflux
     tests_old,tests_old_2: pytest --ignore tests/unit/models/gpflux {posargs} tests/unit
-    tests_old,tests_old_3: pytest {posargs} tests/util tests/integration
+    tests_old,tests_old_3: pytest {posargs} tests/util tests/integration/test_ask_tell_optimization.py tests/integration/test_bayesian_optimization.py
+    tests_old,tests_old_4: pytest --ignore tests/integration/test_ask_tell_optimization.py --ignore tests/integration/test_bayesian_optimization.py {posargs} tests/integration

trieste/ask_tell_optimization.py

@@ -233,7 +233,7 @@ def __init__(
                 for tag, model in self._models.items():
                     dataset = datasets[tag]
                     model.update(dataset)
-                    model.optimize(dataset)
+                    model.optimize_and_save_result(dataset)
 
             summary_writer = logging.get_tensorboard_writer()
             if summary_writer:
@@ -434,7 +434,7 @@ def tell(self, new_data: Mapping[Tag, Dataset] | Dataset) -> None:
             for tag, model in self._models.items():
                 dataset = self._datasets[tag]
                 model.update(dataset)
-                model.optimize(dataset)
+                model.optimize_and_save_result(dataset)
 
         summary_writer = logging.get_tensorboard_writer()
         if summary_writer:

trieste/bayesian_optimizer.py

@@ -719,7 +719,7 @@ def optimize(
                         for tag, model in models.items():
                             dataset = datasets[tag]
                             model.update(dataset)
-                            model.optimize(dataset)
+                            model.optimize_and_save_result(dataset)
                     if summary_writer:
                         logging.set_step_number(0)
                         with summary_writer.as_default(step=0):
@@ -752,7 +752,7 @@ def optimize(
                             for tag, model in models.items():
                                 dataset = datasets[tag]
                                 model.update(dataset)
-                                model.optimize(dataset)
+                                model.optimize_and_save_result(dataset)
 
                 if summary_writer:
                     with summary_writer.as_default(step=step):

trieste/models/gpflow/interface.py

@@ -33,6 +33,7 @@
     SupportsGetKernel,
     SupportsGetObservationNoise,
     SupportsPredictJoint,
+    TrainableProbabilisticModel,
 )
 from ..optimizer import Optimizer
 from ..utils import (
@@ -44,7 +45,12 @@
 
 
 class GPflowPredictor(
-    SupportsPredictJoint, SupportsGetKernel, SupportsGetObservationNoise, HasReparamSampler, ABC
+    SupportsPredictJoint,
+    SupportsGetKernel,
+    SupportsGetObservationNoise,
+    HasReparamSampler,
+    TrainableProbabilisticModel,
+    ABC,
 ):
     """A trainable wrapper for a GPflow Gaussian process model."""
 
@@ -150,14 +156,6 @@ def get_observation_noise(self) -> TensorType:
 
         return noise_variance
 
-    def optimize(self, dataset: Dataset) -> None:
-        """
-        Optimize the model with the specified `dataset`.
-
-        :param dataset: The data with which to optimize the `model`.
-        """
-        self.optimizer.optimize(self.model, dataset)
-
     def log(self, dataset: Optional[Dataset] = None) -> None:
         """
         Log model training information at a given optimization step to the Tensorboard.

trieste/models/gpflow/models.py

@@ -49,7 +49,7 @@
     TrainableProbabilisticModel,
     TrajectorySampler,
 )
-from ..optimizer import BatchOptimizer, Optimizer
+from ..optimizer import BatchOptimizer, Optimizer, OptimizeResult
 from .inducing_point_selectors import InducingPointSelector
 from .interface import GPflowPredictor, SupportsCovarianceBetweenPoints
 from .sampler import DecoupledTrajectorySampler, RandomFourierFeatureTrajectorySampler
@@ -65,7 +65,6 @@
 
 class GaussianProcessRegression(
     GPflowPredictor,
-    TrainableProbabilisticModel,
     FastUpdateModel,
     SupportsCovarianceBetweenPoints,
     SupportsGetInternalData,
@@ -249,7 +248,7 @@ def covariance_between_points(
 
         return cov
 
-    def optimize(self, dataset: Dataset) -> None:
+    def optimize(self, dataset: Dataset) -> OptimizeResult:
         """
         Optimize the model with the specified `dataset`.
 
@@ -284,8 +283,9 @@ def optimize(self, dataset: Dataset) -> None:
             self.find_best_model_initialization(
                 self._num_kernel_samples * num_trainable_params_with_priors_or_constraints
             )
-        self.optimizer.optimize(self.model, dataset)
+        result = self.optimizer.optimize(self.model, dataset)
         self.update_posterior_cache()
+        return result
 
     def find_best_model_initialization(self, num_kernel_samples: int) -> None:
         """
@@ -523,7 +523,6 @@ def conditional_predict_y(
 
 class SparseGaussianProcessRegression(
     GPflowPredictor,
-    TrainableProbabilisticModel,
     SupportsCovarianceBetweenPoints,
     SupportsGetInducingVariables,
     SupportsGetInternalData,
@@ -637,14 +636,15 @@ def _ensure_variable_model_data(self) -> None:
         if not is_variable(self._model.num_data):
             self._model.num_data = tf.Variable(self._model.num_data, trainable=False)
 
-    def optimize(self, dataset: Dataset) -> None:
+    def optimize(self, dataset: Dataset) -> OptimizeResult:
         """
         Optimize the model with the specified `dataset`.
 
         :param dataset: The data with which to optimize the `model`.
         """
-        self.optimizer.optimize(self.model, dataset)
+        result = self.optimizer.optimize(self.model, dataset)
         self.update_posterior_cache()
+        return result
 
     def update(self, dataset: Dataset) -> None:
         self._ensure_variable_model_data()
@@ -836,7 +836,6 @@ def get_internal_data(self) -> Dataset:
 
 class SparseVariational(
     GPflowPredictor,
-    TrainableProbabilisticModel,
     SupportsCovarianceBetweenPoints,
     SupportsGetInducingVariables,
     HasTrajectorySampler,
@@ -979,14 +978,15 @@ def update(self, dataset: Dataset) -> None:
                 self._update_inducing_variables(new_inducing_points)
                 self.update_posterior_cache()
 
-    def optimize(self, dataset: Dataset) -> None:
+    def optimize(self, dataset: Dataset) -> OptimizeResult:
         """
         Optimize the model with the specified `dataset`.
 
         :param dataset: The data with which to optimize the `model`.
         """
-        self.optimizer.optimize(self.model, dataset)
+        result = self.optimizer.optimize(self.model, dataset)
         self.update_posterior_cache()
+        return result
 
     def _update_inducing_variables(self, new_inducing_points: TensorType) -> None:
         """
@@ -1098,7 +1098,6 @@ def trajectory_sampler(self) -> TrajectorySampler[SparseVariational]:
 
 class VariationalGaussianProcess(
     GPflowPredictor,
-    TrainableProbabilisticModel,
     SupportsCovarianceBetweenPoints,
     SupportsGetInducingVariables,
     HasTrajectorySampler,
@@ -1262,7 +1261,7 @@ def update(self, dataset: Dataset, *, jitter: float = DEFAULTS.JITTER) -> None:
         update_vgp_data(self.model, (dataset.query_points, dataset.observations))
         self.update_posterior_cache()
 
-    def optimize(self, dataset: Dataset) -> None:
+    def optimize(self, dataset: Dataset) -> Optional[OptimizeResult]:
         """
         :class:`VariationalGaussianProcess` has a custom `optimize` method that (optionally) permits
         alternating between standard optimization steps (for kernel parameters) and natural gradient
@@ -1297,13 +1296,14 @@ def perform_optimization_step() -> None:  # alternate with natgrad optimizations
             for _ in range(base_optimizer.max_iter):  # type: ignore
                 perform_optimization_step()
 
-            gpflow.set_trainable(model.q_mu, True)  # revert varitional params to trainable
+            gpflow.set_trainable(model.q_mu, True)  # revert variational params to trainable
             gpflow.set_trainable(model.q_sqrt, True)
-
+            result = None  # TODO: find something useful to return
         else:
-            self.optimizer.optimize(model, dataset)
+            result = self.optimizer.optimize(model, dataset)
 
         self.update_posterior_cache()
+        return result
 
     def get_inducing_variables(self) -> Tuple[TensorType, TensorType, TensorType, bool]:
         """

trieste/models/gpflux/models.py

@@ -18,6 +18,7 @@
 
 import dill
 import gpflow
+import keras.callbacks
 import tensorflow as tf
 from check_shapes import inherit_check_shapes
 from gpflow.inducing_variables import InducingPoints
@@ -212,6 +213,10 @@ def __getstate__(self) -> dict[str, Any]:
             finally:
                 self._model_keras.history.model = history_model
 
+        # don't try to serialize any other copies of the history callback
+        if isinstance(state.get("_last_optimization_result"), keras.callbacks.History):
+            state["_last_optimization_result"] = ...
+
         return state
 
     def __setstate__(self, state: dict[str, Any]) -> None:
@@ -266,6 +271,10 @@ def __setstate__(self, state: dict[str, Any]) -> None:
                 model.set_weights(weights)
                 self._model_keras.history.set_model(model)
 
+        # recover optimization result if necessary (and possible)
+        if state.get("_last_optimization_result") is ...:
+            self._last_optimization_result = getattr(self._model_keras, "history")
+
     def __repr__(self) -> str:
         """"""
         return f"DeepGaussianProcess({self.model_gpflux!r}, {self.optimizer.optimizer!r})"
@@ -340,7 +349,7 @@ def update(self, dataset: Dataset) -> None:
 
             inputs = layer(inputs)
 
-    def optimize(self, dataset: Dataset) -> None:
+    def optimize(self, dataset: Dataset) -> keras.callbacks.History:
         """
         Optimize the model with the specified `dataset`.
         :param dataset: The data with which to optimize the `model`.
@@ -372,6 +381,8 @@ def optimize(self, dataset: Dataset) -> None:
         ):
             self.optimizer.optimizer.lr.assign(self.original_lr)
 
+        return hist
+
     def log(self, dataset: Optional[Dataset] = None) -> None:
         """
         Log model training information at a given optimization step to the Tensorboard.