Add TorchAdapter.botorch_model property (#4827)

Summary:
Pull Request resolved: #4827

This is a convenience method that should save us from doing calls like `client._generation_strategy.adapter.generator.surrogate.model` and instead simplify it to `client._generation_strategy.adapter.botorch_model`.

Reviewed By: Balandat

Differential Revision: D91588495

fbshipit-source-id: 2cc66a3dbe847ca37ae2b5ccf50b8eb813bfd54c

Author: saitcakmak

ax/adapter/tests/test_torch_adapter.py

@@ -43,6 +43,7 @@
 from ax.core.search_space import SearchSpace, SearchSpaceDigest
 from ax.core.types import ComparisonOp
 from ax.exceptions.core import DataRequiredError, UnsupportedError, UserInputError
+from ax.exceptions.model import ModelError
 from ax.generators.torch.botorch_modular.generator import BoTorchGenerator
 from ax.generators.torch.botorch_modular.surrogate import Surrogate, SurrogateSpec
 from ax.generators.torch.botorch_modular.utils import ModelConfig
@@ -1391,6 +1392,37 @@ def test_moo_with_derived_parameter(self) -> None:
         )
         self.assertEqual(len(objective_thresholds), len(experiment.metrics))
 
+    def test_botorch_model_property(self) -> None:
+        experiment = get_branin_experiment(with_completed_trial=True)
+        # Case: Invalid generator.
+        adapter = TorchAdapter(
+            generator=TorchGenerator(),
+            experiment=experiment,
+            transforms=Cont_X_trans,
+        )
+        with self.assertRaisesRegex(UnsupportedError, "BoTorchGenerator"):
+            adapter.botorch_model
+
+        # Case: Model not fitted yet.
+        adapter = TorchAdapter(
+            generator=BoTorchGenerator(),
+            experiment=experiment,
+            transforms=Cont_X_trans,
+            fit_on_init=False,
+        )
+        with self.assertRaisesRegex(ModelError, "has not yet been constructed"):
+            adapter.botorch_model
+
+        # Case: Model fitted.
+        generator = BoTorchGenerator()
+        adapter = TorchAdapter(
+            generator=generator,
+            experiment=experiment,
+            transforms=Cont_X_trans,
+        )
+        self.assertIs(adapter.botorch_model, generator.surrogate.model)
+        self.assertIsInstance(adapter.botorch_model, SingleTaskGP)
+
 
 class AdapterWithPLBOTest(TestCase):
     """Test the PLBO (Preference-Learning-guided BO) step in BOPE (Bayesian

ax/adapter/torch.py

@@ -86,6 +86,7 @@
 from ax.generators.types import TConfig
 from ax.utils.common.constants import Keys
 from ax.utils.common.logger import get_logger
+from botorch.models.model import Model
 from botorch.utils.datasets import MultiTaskDataset, SupervisedDataset
 from pyre_extensions import none_throws
 from torch import Tensor
@@ -216,6 +217,16 @@ def feature_importances(self, metric_signature: str) -> dict[str, float]:
         importances_arr = importances_dict[metric_signature].flatten()
         return dict(zip(self.parameters, importances_arr, strict=True))
 
+    @property
+    def botorch_model(self) -> Model:
+        """Returns the underlying BoTorch model for BoTorchGenerator."""
+        if not isinstance(self.generator, BoTorchGenerator):
+            raise UnsupportedError(
+                "Generator must be a BoTorchGenerator to "
+                f"access botorch_model. Found {type(self.generator)}."
+            )
+        return self.generator.surrogate.model
+
     def infer_objective_thresholds(
         self,
         search_space: SearchSpace | None = None,
@@ -254,7 +265,7 @@ def infer_objective_thresholds(
         )
         # Infer objective thresholds.
         if isinstance(self.generator, BoTorchGenerator):
-            model = self.generator.surrogate.model
+            model = self.botorch_model
             Xs = self.generator.surrogate.Xs
         else:
             raise UnsupportedError(

ax/benchmark/benchmark_test_functions/surrogate.py

@@ -104,12 +104,7 @@ def wrap_surrogate_in_deterministic_model(self) -> None:
     def surrogate(self) -> TorchAdapter:
         if self._surrogate is None:
             self._surrogate = none_throws(self.get_surrogate)()
-        if not isinstance(
-            # pyre-ignore[16]: `ax.generators.torch_base.TorchGenerator` has no
-            # attribute `surrogate`.
-            self._surrogate.generator.surrogate.model,
-            DeterministicModel,
-        ):
+        if not isinstance(self._surrogate.botorch_model, DeterministicModel):
             self.wrap_surrogate_in_deterministic_model()
         return none_throws(self._surrogate)
 

ax/benchmark/tests/benchmark_test_functions/test_surrogate_test_function.py

@@ -24,6 +24,7 @@
 )
 from botorch.models.deterministic import PosteriorMeanModel
 from botorch.sampling.pathwise.posterior_samplers import MatheronPathModel
+from pyre_extensions import assert_is_instance
 
 
 class TestSurrogateTestFunction(TestCase):
@@ -237,11 +238,11 @@ def test_ensemble_sampling(self) -> None:
 
         # Access surrogate to trigger wrapping
         surrogate = test_function.surrogate
-        # pyre-ignore[16]: Access base_model through deterministic wrapper
-        wrapped_model = surrogate.generator.surrogate.model
+        # Access base_model through deterministic wrapper
+        wrapped_model = assert_is_instance(surrogate.botorch_model, PosteriorMeanModel)
 
         # Check that exactly one model has weight 1.0 and others have weight 0.0
-        weights = wrapped_model.ensemble_weights
+        weights = assert_is_instance(wrapped_model.ensemble_weights, torch.Tensor)
         self.assertEqual(weights.sum().item(), 1.0)
         self.assertEqual((weights == 1.0).sum().item(), 1)
         self.assertEqual((weights == 0.0).sum().item(), len(weights) - 1)
@@ -262,5 +263,4 @@ def test_ensemble_no_sampling(self) -> None:
 
         # Access surrogate to trigger wrapping
         surrogate = test_function.surrogate
-        # pyre-ignore[16]: Access base_model through deterministic wrapper
-        self.assertIsNone(surrogate.generator.surrogate.model.ensemble_weights)
+        self.assertIsNone(surrogate.botorch_model.ensemble_weights)

ax/utils/sensitivity/tests/test_sensitivity.py

@@ -73,8 +73,8 @@ class SensitivityAnalysisTest(TestCase):
     def setUp(self) -> None:
         super().setUp()
         set_rng_seed(0)
-        self.model = get_adapter().generator.surrogate.model
-        self.saas_model = get_adapter(saasbo=True).generator.surrogate.model
+        self.model = get_adapter().botorch_model
+        self.saas_model = get_adapter(saasbo=True).botorch_model
 
     def test_DgsmGpMean(self) -> None:
         bounds = torch.tensor([(0.0, 1.0) for _ in range(2)]).t()
@@ -418,7 +418,7 @@ def test_SobolGPMean_SAASBO_Ax_utils(self) -> None:
                     **sobol_kwargs,
                 )
                 ind_deriv = compute_derivatives_from_model_list(
-                    model_list=[adapter.generator.surrogate.model],
+                    model_list=[adapter.botorch_model],
                     bounds=torch.tensor(adapter.generator.search_space_digest.bounds).T,
                     discrete_features=discrete_features,
                     fixed_features=None,