Fixed optimizer tests for stochasticity in ensemble predictions

obsidian/tests/test_optimizer_MOO.py

@@ -7,6 +7,8 @@
 from obsidian.optimizer import BayesianOptimizer
 from obsidian.experiment.benchmark import two_leaves
 from obsidian.tests.utils import equal_state_dicts
+
+from botorch.models.ensemble import EnsembleModel
 
 import pandas as pd
 import numpy as np
@@ -51,10 +53,17 @@ def test_optimizer_fit(X_space, surrogate, Z0, serial_test=True):
         obj_dict2 = optimizer_2.save_state()
         assert equal_state_dicts(obj_dict, obj_dict2)
         optimizer_2.__repr__()
+
         y_pred = optimizer.predict(optimizer.X_train)
         y_pred_2 = optimizer_2.predict(optimizer.X_train)
-        y_error = ((y_pred_2-y_pred)/y_pred.max(axis=0)).values
-        assert abs(y_error).max() < tol, 'Prediction error in loading parameters of saved optimizer'
+        y_error = ((y_pred_2-y_pred)/y_pred.max(axis=0))
+
+        # Only check the mean for ensemble models, as the lb/ub are not deterministic
+        if isinstance(optimizer.surrogate[0].torch_model, EnsembleModel):
+            check_cols = [col for col in y_pred if '(pred)' in col]
+            y_error = y_error[check_cols]
+
+        assert abs(y_error.values).max() < tol, 'Prediction error in loading parameters of saved optimizer'
 
 
 # Generate a baseline optimizer to use for future tests

obsidian/tests/test_optimizer_SOO.py

@@ -8,6 +8,8 @@
 from obsidian.experiment.benchmark import shifted_parab
 from obsidian.tests.utils import approx_equal, equal_state_dicts
 
+from botorch.models.ensemble import EnsembleModel
+
 import pandas as pd
 import numpy as np
 import pytest
@@ -80,8 +82,14 @@ def test_optimizer_fit(X_space, surrogate, Z0, serial_test=True):
         optimizer_2.__repr__()
         y_pred = optimizer.predict(optimizer.X_train)
         y_pred_2 = optimizer_2.predict(optimizer.X_train)
-        y_error = ((y_pred_2-y_pred)/y_pred.max(axis=0)).values
-        assert abs(y_error).max() < tol, 'Prediction error in loading parameters of saved optimizer'
+        y_error = ((y_pred_2-y_pred)/y_pred.max(axis=0))
+
+        # Only check the mean for ensemble models, as the lb/ub are not deterministic
+        if isinstance(optimizer.surrogate[0].torch_model, EnsembleModel):
+            check_cols = [col for col in y_pred if '(pred)' in col]
+            y_error = y_error[check_cols]
+
+        assert abs(y_error.values).max() < tol, 'Prediction error in loading parameters of saved optimizer'
 
 
 # Generate a baseline optimizer to use for future tests