WIP fix bug where variational GPs wouldn't use correct likelihoods; u…

…se botorch inducing point selection; use botorch default mean/covar; update botorch/ax versions (#323)

Summary:

This ensures variational GPs will use correct likelihoods, and it moves some logic out of AEPsych and into botorch. Anecdotally, the botorch priors seem to produce better models, and since we are so resource-strapped on the AEPsych side, we should rely on botorch logic as much as possible.

Differential Revision: D48891019

aepsych/config.py

@@ -372,4 +372,5 @@ def from_config(cls, config: Config, name: Optional[str] = None):
 Config.register_module(gpytorch.likelihoods)
 Config.register_module(gpytorch.kernels)
 Config.register_module(botorch.acquisition)
+Config.register_module(botorch.acquisition.multi_objective)
 Config.registered_names["None"] = None

aepsych/generators/__init__.py

@@ -12,7 +12,6 @@
 from .manual_generator import ManualGenerator
 from .monotonic_rejection_generator import MonotonicRejectionGenerator
 from .monotonic_thompson_sampler_generator import MonotonicThompsonSamplerGenerator
-from .multi_outcome_generator import MultiOutcomeOptimizationGenerator
 from .optimize_acqf_generator import AxOptimizeAcqfGenerator, OptimizeAcqfGenerator
 from .pairwise_optimize_acqf_generator import PairwiseOptimizeAcqfGenerator
 from .pairwise_sobol_generator import PairwiseSobolGenerator
@@ -33,7 +32,6 @@
     "AxOptimizeAcqfGenerator",
     "AxSobolGenerator",
     "IntensityAwareSemiPGenerator",
-    "MultiOutcomeOptimizationGenerator",
     "AxRandomGenerator",
 ]
 

aepsych/generators/optimize_acqf_generator.py

@@ -15,7 +15,7 @@
 from aepsych.config import Config, ConfigurableMixin
 from aepsych.generators.base import AEPsychGenerationStep, AEPsychGenerator
 from aepsych.models.base import ModelProtocol
-from aepsych.models.surrogate import AEPsychSurrogate
+from ax.models.torch.botorch_modular.surrogate import Surrogate
 from aepsych.utils_logging import getLogger
 from ax.modelbridge import Models
 from ax.modelbridge.registry import Cont_X_trans
@@ -152,14 +152,11 @@ def get_config_options(cls, config: Config, name: str) -> Dict:
         acqf_options = cls._get_acqf_options(acqf_cls, config)
         gen_options = cls._get_gen_options(config)
 
-        max_fit_time = model_options["max_fit_time"]
-
         model_kwargs = {
-            "surrogate": AEPsychSurrogate(
+            "surrogate": Surrogate(
                 botorch_model_class=model_class,
                 mll_class=model_class.get_mll_class(),
                 model_options=model_options,
-                max_fit_time=max_fit_time,
             ),
             "acquisition_class": AEPsychAcquisition,
             "botorch_acqf_class": acqf_cls,

aepsych/models/__init__.py

@@ -11,6 +11,7 @@
 from .exact_gp import ContinuousRegressionGP, ExactGP
 from .gp_classification import GPBetaRegressionModel, GPClassificationModel
 from .gp_regression import GPRegressionModel
+from .model_list import AEPsychModelListGP
 from .monotonic_projection_gp import MonotonicProjectionGP
 from .monotonic_rejection_gp import MonotonicRejectionGP
 from .multitask_regression import IndependentMultitaskGPRModel, MultitaskGPRModel
@@ -21,8 +22,12 @@
     semi_p_posterior_transform,
     SemiParametricGPModel,
 )
-from .variational_gp import BetaRegressionGP, BinaryClassificationGP, OrdinalGP, VariationalGP
-
+from .variational_gp import (
+    BetaRegressionGP,
+    BinaryClassificationGP,
+    OrdinalGP,
+    VariationalGP,
+)
 
 
 __all__ = [
@@ -44,6 +49,7 @@
     "semi_p_posterior_transform",
     "OrdinalGP",
     "GPBetaRegressionModel",
+    "AEPsychModelListGP",
 ]
 
 Config.register_module(sys.modules[__name__])