From 4b457c7d13517fa4a9a1a41f4529eb5a9f7900fa Mon Sep 17 00:00:00 2001 From: Jason Chow Date: Tue, 10 Jun 2025 07:34:33 -0700 Subject: [PATCH 1/2] Remove botorch pin to use latest botorch version (#808) Summary: Pull Request resolved: https://github.com/facebookresearch/aepsych/pull/808 Removed botorch version pin for pypi. CI should cover any change in compatibility. Differential Revision: D76321622 --- setup.py | 2 +- tests/test_mean_covar_factories.py | 11 ++++++++--- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/setup.py b/setup.py index eeb860a1e..f8bceadc9 100644 --- a/setup.py +++ b/setup.py @@ -19,7 +19,7 @@ "pandas", "aepsych_client", "statsmodels", - "botorch==0.13.0", + "botorch", ] BENCHMARK_REQUIRES = ["tqdm", "pathos", "multiprocess", "torcheval"] diff --git a/tests/test_mean_covar_factories.py b/tests/test_mean_covar_factories.py index 9546ad790..02a966fed 100644 --- a/tests/test_mean_covar_factories.py +++ b/tests/test_mean_covar_factories.py @@ -749,7 +749,7 @@ def f_1d(x): [init_strat] generator = SobolGenerator - min_asks = 200 + min_asks = 250 [opt_strat] generator = MixedOptimizeAcqfGenerator @@ -780,5 +780,10 @@ def f_1d(x): y_0, _ = strat.predict(x_0, probability_space=True) y_1, _ = strat.predict(x_1, probability_space=True) - self.assertTrue(torch.allclose(x[torch.argmax(y_0)], torch.tensor([0.0]))) - self.assertTrue(torch.allclose(x[torch.argmax(y_1)], torch.tensor([0.4]))) + # Loose test + self.assertLessEqual( + (torch.abs(x[torch.argmax(y_0)] - torch.tensor([0.0]))).item(), 0.05 + ) + self.assertLessEqual( + (torch.abs(x[torch.argmax(y_1)] - torch.tensor([0.4]))).item(), 0.05 + ) From 99e37186bbe366e3f0e3ad7f27ea6b85b976a666 Mon Sep 17 00:00:00 2001 From: Sebastian Ament Date: Tue, 10 Jun 2025 07:40:31 -0700 Subject: [PATCH 2/2] `average_over_ensemble_models` decorator for acquisition functions" Summary: This commit adds the recently introduced `average_over_ensemble_models` decorator to the AEPsych acquisition functions to ensure continued compatibility with ensemble models. NOTE: Tests will only pass once the [BoTorch commit](https://github.com/pytorch/botorch/pull/2873) has landed and the latest version is pulled to AEPsych. Differential Revision: D76275283 --- aepsych/acquisition/lookahead.py | 7 ++++++- aepsych/acquisition/mc_posterior_variance.py | 6 +++++- aepsych/acquisition/mutual_information.py | 6 +++++- 3 files changed, 16 insertions(+), 3 deletions(-) diff --git a/aepsych/acquisition/lookahead.py b/aepsych/acquisition/lookahead.py index 5374889d4..ac62c1bf6 100644 --- a/aepsych/acquisition/lookahead.py +++ b/aepsych/acquisition/lookahead.py @@ -14,7 +14,10 @@ from botorch.acquisition.input_constructors import acqf_input_constructor from botorch.acquisition.objective import PosteriorTransform from botorch.models.gpytorch import GPyTorchModel -from botorch.utils.transforms import t_batch_mode_transform +from botorch.utils.transforms import ( + average_over_ensemble_models, + t_batch_mode_transform, +) from scipy.stats import norm from torch import Tensor @@ -171,6 +174,7 @@ def __init__( self.posterior_transform = posterior_transform @t_batch_mode_transform(expected_q=1) + @average_over_ensemble_models def forward(self, X: torch.Tensor) -> torch.Tensor: """ Evaluate acquisition function at X. @@ -291,6 +295,7 @@ def __init__( self.register_buffer("Xq", Xq) @t_batch_mode_transform(expected_q=1) + @average_over_ensemble_models def forward(self, X: torch.Tensor) -> torch.Tensor: """ Evaluate acquisition function at X. diff --git a/aepsych/acquisition/mc_posterior_variance.py b/aepsych/acquisition/mc_posterior_variance.py index e8abdd111..aac895726 100644 --- a/aepsych/acquisition/mc_posterior_variance.py +++ b/aepsych/acquisition/mc_posterior_variance.py @@ -15,7 +15,10 @@ from botorch.models.model import Model from botorch.sampling.base import MCSampler from botorch.sampling.normal import SobolQMCNormalSampler -from botorch.utils.transforms import t_batch_mode_transform +from botorch.utils.transforms import ( + average_over_ensemble_models, + t_batch_mode_transform, +) def balv_acq(obj_samps: torch.Tensor) -> torch.Tensor: @@ -61,6 +64,7 @@ def __init__( self.objective = objective @t_batch_mode_transform() + @average_over_ensemble_models def forward(self, X: torch.Tensor) -> torch.Tensor: r"""Evaluate MCPosteriorVariance on the candidate set `X`. diff --git a/aepsych/acquisition/mutual_information.py b/aepsych/acquisition/mutual_information.py index ef498ff54..434fcd415 100644 --- a/aepsych/acquisition/mutual_information.py +++ b/aepsych/acquisition/mutual_information.py @@ -18,7 +18,10 @@ from botorch.models.model import Model from botorch.sampling.base import MCSampler from botorch.sampling.normal import SobolQMCNormalSampler -from botorch.utils.transforms import t_batch_mode_transform +from botorch.utils.transforms import ( + average_over_ensemble_models, + t_batch_mode_transform, +) from torch import Tensor from torch.distributions.bernoulli import Bernoulli @@ -82,6 +85,7 @@ def __init__( ) @t_batch_mode_transform() + @average_over_ensemble_models def forward(self, X: Tensor) -> Tensor: r"""Evaluate mutual information on the candidate set `X`.