Merge branch 'develop' into neil/compile_args

docs/notebooks/ask_tell_optimization.pct.py

@@ -12,7 +12,10 @@
 import matplotlib.pyplot as plt
 import gpflow
 
-from trieste.ask_tell_optimization import AskTellOptimizer
+from trieste.ask_tell_optimization import (
+    AskTellOptimizer,
+    AskTellOptimizerNoTraining,
+)
 from trieste.bayesian_optimizer import Record
 from trieste.data import Dataset
 from trieste.models.gpflow.models import GaussianProcessRegression
@@ -96,7 +99,7 @@ def plot_ask_tell_regret(ask_tell_result):
 # %% [markdown]
 # ## Model selection: using only Ask part
 #
-# We now turn to a slightly more complex use case. Let's suppose we want to switch between two models depending on some criteria dynamically during the optimization loop, e.g. we want to be able to train a model outside of Trieste. In this case we can only use Ask part of the Ask-Tell interface.
+# We now turn to a slightly more complex use case. Let's suppose we want to switch between two models depending on some criteria dynamically during the optimization loop, e.g. we want to be able to train a model outside of Trieste. In this case we can only use Ask part of the Ask-Tell interface. For this it is recommended to use the `AskTellOptimizerNoTraining` class, which performs no training during the Tell stage and can therefore be used with any probabilistic model, including ones which aren't trainable.
 
 # %%
 model1 = build_model(
@@ -118,8 +121,8 @@ def plot_ask_tell_regret(ask_tell_result):
         model = model2
 
     print("Asking for new point to observe")
-    ask_tell = AskTellOptimizer(search_space, dataset, model)
-    new_point = ask_tell.ask()
+    ask_only = AskTellOptimizerNoTraining(search_space, dataset, model)
+    new_point = ask_only.ask()
 
     new_data_point = observer(new_point)
     dataset = dataset + new_data_point
@@ -131,7 +134,7 @@ def plot_ask_tell_regret(ask_tell_result):
     model2.update(dataset)
     model2.optimize(dataset)
 
-plot_ask_tell_regret(ask_tell.to_result())
+plot_ask_tell_regret(ask_only.to_result())
 
 
 # %% [markdown]
@@ -151,7 +154,7 @@ def plot_ask_tell_regret(ask_tell_result):
     new_config = ask_tell.ask()
 
     print("Saving Trieste state to re-use later")
-    state: Record[None] = ask_tell.to_record()
+    state: Record[None, GaussianProcessRegression] = ask_tell.to_record()
     saved_state = pickle.dumps(state)
 
     print(f"In the lab running the experiment #{step}.")

docs/notebooks/expected_improvement.pct.py

@@ -292,7 +292,7 @@ def build_model(data):
 saved_result = trieste.bayesian_optimizer.OptimizationResult.from_path(  # type: ignore
     "results_path"
 )
-saved_result.try_get_final_model().model  # type: ignore
+saved_result.try_get_final_model().model
 
 # %% [markdown]
 # The second approach is to save the model using the tensorflow SavedModel format. This requires explicitly exporting the methods to be saved and results in a portable model than can be safely loaded and evaluated, but which can no longer be used in subsequent BO steps.

docs/notebooks/failure_ego.pct.py

@@ -238,7 +238,7 @@ def acquisition(at):
 
 # %%
 fig, ax = plot_gp_2d(
-    result.models[FAILURE].model,  # type: ignore
+    result.models[FAILURE].model,
     search_space.lower,
     search_space.upper,
     grid_density=20,

tests/integration/test_active_learning.py

@@ -107,7 +107,7 @@ def test_optimizer_learns_scaled_branin_function(
         .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_predicted_means, _ = final_model.model.predict_f(test_query_points)
     final_accuracy = tf.reduce_max(tf.abs(final_predicted_means - test_data.observations))
 
     assert initial_accuracy > final_accuracy
@@ -357,7 +357,7 @@ def ilink(f: TensorType) -> TensorType:
         .optimize(num_steps, initial_data, model, rule)
         .try_get_final_model()
     )
-    final_predicted_means, _ = ilink(final_model.model.predict_f(test_query_points))  # type: ignore
+    final_predicted_means, _ = ilink(final_model.model.predict_f(test_query_points))
     final_error = tf.reduce_mean(tf.abs(final_predicted_means - test_data.observations))
 
     assert initial_error > final_error

tests/integration/test_ask_tell_optimization.py

@@ -207,7 +207,8 @@ def _test_ask_tell_optimization_finds_minima(
 
                 if reload_state:
                     state: Record[
-                        None | State[TensorType, AsynchronousRuleState | BatchTrustRegionBox.State]
+                        None | State[TensorType, AsynchronousRuleState | BatchTrustRegionBox.State],
+                        GaussianProcessRegression,
                     ] = ask_tell.to_record()
                     written_state = pickle.dumps(state)
 
@@ -228,7 +229,8 @@ def _test_ask_tell_optimization_finds_minima(
                 ask_tell.tell(new_data_point)
 
     result: OptimizationResult[
-        None | State[TensorType, AsynchronousRuleState | BatchTrustRegionBox.State]
+        None | State[TensorType, AsynchronousRuleState | BatchTrustRegionBox.State],
+        GaussianProcessRegression,
     ] = ask_tell.to_result()
     dataset = result.try_get_final_dataset()
 

tests/integration/test_bayesian_optimization.py

@@ -701,9 +701,9 @@ def _test_optimizer_finds_minimum(
             # check history saved ok
             assert len(result.history) <= (num_steps or 2)
             assert len(result.loaded_history) == len(result.history)
-            loaded_result: OptimizationResult[None] = OptimizationResult.from_path(
-                Path(tmpdirname) / "history"
-            )
+            loaded_result: OptimizationResult[
+                None, TrainableProbabilisticModel
+            ] = OptimizationResult.from_path(Path(tmpdirname) / "history")
             assert loaded_result.final_result.is_ok
             assert len(loaded_result.history) == len(result.history)