Remove model protocol

aepsych/acquisition/lookahead_utils.py

@@ -41,6 +41,7 @@ def posterior_at_xstar_xq(
            - Sigma_sq: (b x m) covariance between Xstar and each point in Xq.
     """
     # Evaluate posterior and extract needed components
+    Xq = Xq.to(Xstar)
     Xext = torch.cat((Xstar, Xq), dim=-2)
     posterior = model.posterior(Xext, posterior_transform=posterior_transform)
     mu = posterior.mean[..., :, 0]

aepsych/benchmark/problem.py

@@ -9,7 +9,7 @@
 
 import numpy as np
 import torch
-from aepsych.models.model_protocol import ModelProtocol
+from aepsych.models.base import AEPsychModelMixin
 from aepsych.models.utils import p_below_threshold
 from aepsych.strategy import SequentialStrategy
 from aepsych.utils import make_scaled_sobol
@@ -78,11 +78,11 @@ def sample_y(
         """
         return bernoulli.rvs(self.p(x))
 
-    def f_hat(self, model: ModelProtocol) -> torch.Tensor:
+    def f_hat(self, model: AEPsychModelMixin) -> torch.Tensor:
         """Generate mean predictions from the model over the evaluation grid.
 
         Args:
-            model (TensoModelProtocolr): Model to evaluate.
+            model (AEPsychModelMixin): Model to evaluate.
 
         Returns:
             torch.Tensor: Posterior mean from underlying model over the evaluation grid.
@@ -109,11 +109,11 @@ def p_true(self) -> torch.Tensor:
         normal_dist = torch.distributions.Normal(0, 1)
         return normal_dist.cdf(self.f_true)
 
-    def p_hat(self, model: ModelProtocol) -> torch.Tensor:
+    def p_hat(self, model: AEPsychModelMixin) -> torch.Tensor:
         """Generate mean predictions from the model over the evaluation grid.
 
         Args:
-            model (TensoModelProtocolr): Model to evaluate.
+            model (AEPsychModelMixin): Model to evaluate.
 
         Returns:
             torch.Tensor: Posterior mean from underlying model over the evaluation grid.
@@ -171,9 +171,9 @@ def evaluate(
         # eval in samp-based expectation over posterior instead of just mean
         fsamps = model.sample(self.eval_grid, num_samples=1000)
         try:
-            psamps = (
-                model.sample(self.eval_grid, num_samples=1000, probability_space=True)  # type: ignore
-            )
+            psamps = model.sample(
+                self.eval_grid, num_samples=1000, probability_space=True
+            )  # type: ignore
         except (
             TypeError
         ):  # vanilla models don't have proba_space samps, TODO maybe we should add them

aepsych/generators/acqf_grid_search_generator.py

@@ -10,12 +10,11 @@
 
 import numpy as np
 import torch
-from aepsych.models.model_protocol import ModelProtocol
+from aepsych.generators.grid_eval_acqf_generator import GridEvalAcqfGenerator
+from aepsych.models.base import AEPsychModelMixin
 from aepsych.utils_logging import getLogger
 from numpy.random import choice
 
-from .grid_eval_acqf_generator import GridEvalAcqfGenerator
-
 logger = getLogger()
 
 
@@ -25,7 +24,7 @@ class AcqfGridSearchGenerator(GridEvalAcqfGenerator):
     def _gen(
         self,
         num_points: int,
-        model: ModelProtocol,
+        model: AEPsychModelMixin,
         fixed_features: Optional[Dict[int, float]] = None,
         **gen_options,
     ) -> torch.Tensor:
@@ -34,7 +33,7 @@ def _gen(
 
         Args:
             num_points (int): The number of points to query.
-            model (ModelProtocol): The fitted model used to evaluate the acquisition function.
+            model (AEPsychModelMixin): The fitted model used to evaluate the acquisition function.
             fixed_features: (Dict[int, float], optional): Parameters that are fixed to specific values.
             gen_options (dict): Additional options for generating points, including:
                 - "seed": Random seed for reproducibility.

aepsych/generators/acqf_thompson_sampler_generator.py

@@ -10,7 +10,7 @@
 
 import numpy as np
 import torch
-from aepsych.models.model_protocol import ModelProtocol
+from aepsych.models.base import AEPsychModelMixin
 from aepsych.utils_logging import getLogger
 from numpy.random import choice
 
@@ -25,7 +25,7 @@ class AcqfThompsonSamplerGenerator(GridEvalAcqfGenerator):
     def _gen(
         self,
         num_points: int,
-        model: ModelProtocol,
+        model: AEPsychModelMixin,
         fixed_features: Optional[Dict[int, float]] = None,
         **gen_options,
     ) -> torch.Tensor:
@@ -34,7 +34,7 @@ def _gen(
 
         Args:
             num_points (int): The number of points to query.
-            model (ModelProtocol): The fitted model used to evaluate the acquisition function.
+            model (AEPsychModelMixin): The fitted model used to evaluate the acquisition function.
             fixed_features: (Dict[int, float], optional): Parameters that are fixed to specific values.
             gen_options (dict): Additional options for generating points, including:
                 - "seed": Random seed for reproducibility.

aepsych/generators/base.py

@@ -11,7 +11,7 @@
 
 import torch
 from aepsych.config import Config, ConfigurableMixin
-from aepsych.models.base import AEPsychMixin
+from aepsych.models.base import AEPsychModelMixin
 from botorch.acquisition import (
     AcquisitionFunction,
     LogNoisyExpectedImprovement,
@@ -21,9 +21,7 @@
 )
 from botorch.acquisition.preference import AnalyticExpectedUtilityOfBestOption
 
-from ..models.model_protocol import ModelProtocol
-
-AEPsychModelType = TypeVar("AEPsychModelType", bound=AEPsychMixin)
+AEPsychModelType = TypeVar("AEPsychModelType", bound=AEPsychModelMixin)
 
 
 @runtime_checkable
@@ -166,12 +164,14 @@ def _get_acqf_options(
 
         return extra_acqf_args
 
-    def _instantiate_acquisition_fn(self, model: ModelProtocol) -> AcquisitionFunction:
+    def _instantiate_acquisition_fn(
+        self, model: AEPsychModelMixin
+    ) -> AcquisitionFunction:
         """
         Instantiates the acquisition function with the specified model and additional arguments.
 
         Args:
-            model (ModelProtocol): The model to use with the acquisition function.
+            model (AEPsychModelMixin): The model to use with the acquisition function.
 
         Returns:
             AcquisitionFunction: Configured acquisition function.
@@ -193,6 +193,8 @@ def _instantiate_acquisition_fn(self, model: ModelProtocol) -> AcquisitionFuncti
                 self.acqf_kwargs["ub"] = self.acqf_kwargs["ub"].to(model.device)
 
         if self.acqf in self.baseline_requiring_acqfs:
+            if model.train_inputs is None:
+                raise ValueError(f"model needs data as a baseline for {self.acqf}")
             return self.acqf(model, model.train_inputs[0], **self.acqf_kwargs)
         else:
             return self.acqf(model=model, **self.acqf_kwargs)

aepsych/generators/epsilon_greedy_generator.py

@@ -9,8 +9,8 @@
 
 import numpy as np
 import torch
+from aepsych.models.base import AEPsychModelMixin
 
-from ..models.model_protocol import ModelProtocol
 from .base import AEPsychGenerator
 from .optimize_acqf_generator import OptimizeAcqfGenerator
 
@@ -65,15 +65,15 @@ def get_config_options(
     def gen(
         self,
         num_points: int,
-        model: ModelProtocol,
+        model: AEPsychModelMixin,
         fixed_features: Optional[Dict[int, float]] = None,
         **kwargs,
     ) -> torch.Tensor:
         """Query next point(s) to run by sampling from the subgenerator with probability 1-epsilon, and randomly otherwise.
 
         Args:
             num_points (int): Number of points to query.
-            model (ModelProtocol): Model to use for generating points.
+            model (AEPsychModelMixin): Model to use for generating points.
             fixed_features: (Dict[int, float], optional): Parameters that are fixed to specific values.
             **kwargs: Passed to subgenerator if not exploring
         """

aepsych/generators/grid_eval_acqf_generator.py

@@ -11,7 +11,7 @@
 from aepsych.config import Config
 from aepsych.generators.base import AcqfGenerator, AEPsychGenerator
 from aepsych.generators.sobol_generator import SobolGenerator
-from aepsych.models.model_protocol import ModelProtocol
+from aepsych.models.base import AEPsychModelMixin
 from aepsych.utils_logging import getLogger
 from botorch.acquisition import AcquisitionFunction
 
@@ -53,14 +53,14 @@ def __init__(
     def gen(
         self,
         num_points: int,
-        model: ModelProtocol,
+        model: AEPsychModelMixin,
         fixed_features: Optional[Dict[int, float]] = None,
         **gen_options,
     ) -> torch.Tensor:
         """Query next point(s) to run by optimizing the acquisition function.
         Args:
             num_points (int): Number of points to query.
-            model (ModelProtocol): Fitted model of the data.
+            model (AEPsychModelMixin): Fitted model of the data.
         Returns:
             torch.Tensor: Next set of point(s) to evaluate, [num_points x dim].
         """
@@ -89,7 +89,7 @@ def gen(
     def _gen(
         self,
         num_points: int,
-        model: ModelProtocol,
+        model: AEPsychModelMixin,
         fixed_features: Optional[Dict[int, float]] = None,
         **gen_options,
     ) -> torch.Tensor:
@@ -98,7 +98,7 @@ def _gen(
     def _eval_acqf(
         self,
         num_points: int,
-        model: ModelProtocol,
+        model: AEPsychModelMixin,
         fixed_features: Optional[Dict[int, float]] = None,
         **gen_options,
     ) -> Tuple[torch.Tensor, torch.Tensor]:

aepsych/generators/manual_generator.py

@@ -11,7 +11,7 @@
 import torch
 from aepsych.config import Config
 from aepsych.generators.base import AEPsychGenerator
-from aepsych.models.base import AEPsychMixin
+from aepsych.models.base import AEPsychModelMixin
 from aepsych.utils import _process_bounds
 from torch.quasirandom import SobolEngine
 
@@ -53,14 +53,14 @@ def __init__(
     def gen(
         self,
         num_points: int = 1,
-        model: Optional[AEPsychMixin] = None,  # included for API compatibility
+        model: Optional[AEPsychModelMixin] = None,  # included for API compatibility
         fixed_features: Optional[Dict[int, float]] = None,
         **kwargs,  # Ignored
     ) -> torch.Tensor:
         """Query next point(s) to run by quasi-randomly sampling the parameter space.
         Args:
             num_points (int): Number of points to query. Defaults to 1.
-            model (AEPsychMixin, optional): Model to use for generating points. Not used in this generator. Defaults to None.
+            model (AEPsychModelMixin, optional): Model to use for generating points. Not used in this generator. Defaults to None.
             fixed_features (Dict[int, float], optional): Ignored, kept for consistent
                 API.
             **kwargs: Ignored, API compatibility

aepsych/generators/optimize_acqf_generator.py

@@ -14,7 +14,7 @@
 from aepsych.acquisition.lookahead import LookaheadAcquisitionFunction
 from aepsych.config import Config
 from aepsych.generators.base import AcqfGenerator
-from aepsych.models.model_protocol import ModelProtocol
+from aepsych.models.base import AEPsychModelMixin
 from aepsych.utils_logging import getLogger
 from botorch.acquisition import AcquisitionFunction
 from botorch.optim import optimize_acqf
@@ -60,14 +60,14 @@ def __init__(
     def gen(
         self,
         num_points: int,
-        model: ModelProtocol,
+        model: AEPsychModelMixin,
         fixed_features: Optional[Dict[int, float]] = None,
         **gen_options,
     ) -> torch.Tensor:
         """Query next point(s) to run by optimizing the acquisition function.
         Args:
             num_points (int): Number of points to query.
-            model (ModelProtocol): Fitted model of the data.
+            model (AEPsychModelMixin): Fitted model of the data.
             fixed_features (Dict[int, float], optional): The values where the specified
                 parameters should be at when generating. Should be a dictionary where
                 the keys are the indices of the parameters to fix and the values are the
@@ -116,7 +116,7 @@ def gen(
     def _gen(
         self,
         num_points: int,
-        model: ModelProtocol,
+        model: AEPsychModelMixin,
         acqf: AcquisitionFunction,
         fixed_features: Optional[Dict[int, float]] = None,
         **gen_options: Dict[str, Any],
@@ -126,7 +126,7 @@ def _gen(
 
         Args:
             num_points (int): Number of points to query.
-            model (ModelProtocol): Fitted model of the data.
+            model (AEPsychModelMixin): Fitted model of the data.
             acqf (AcquisitionFunction): Acquisition function.
             fixed_features (Dict[int, float], optional): The values where the specified
                 parameters should be at when generating. Should be a dictionary where

aepsych/generators/random_generator.py

@@ -10,7 +10,7 @@
 import torch
 from aepsych.config import Config
 from aepsych.generators.base import AEPsychGenerator
-from aepsych.models.base import AEPsychMixin
+from aepsych.models.base import AEPsychModelMixin
 from aepsych.utils import _process_bounds
 
 
@@ -38,14 +38,14 @@ def __init__(
     def gen(
         self,
         num_points: int = 1,
-        model: Optional[AEPsychMixin] = None,  # included for API compatibility.
+        model: Optional[AEPsychModelMixin] = None,  # included for API compatibility.
         fixed_features: Optional[Dict[int, float]] = None,
         **kwargs,
     ) -> torch.Tensor:
         """Query next point(s) to run by randomly sampling the parameter space.
         Args:
             num_points (int): Number of points to query. Currently, only 1 point can be queried at a time.
-            model (AEPsychMixin, optional): Model to use for generating points. Not used in this generator.
+            model (AEPsychModelMixin, optional): Model to use for generating points. Not used in this generator.
             fixed_features: (Dict[int, float], optional): Parameters that are fixed to specific values.
             **kwargs: Ignored, API compatibility
 

aepsych/generators/sobol_generator.py

@@ -11,7 +11,7 @@
 import torch
 from aepsych.config import Config
 from aepsych.generators.base import AEPsychGenerator
-from aepsych.models.base import AEPsychMixin
+from aepsych.models.base import AEPsychModelMixin
 from aepsych.utils import _process_bounds
 from torch.quasirandom import SobolEngine
 
@@ -49,14 +49,14 @@ def __init__(
     def gen(
         self,
         num_points: int = 1,
-        model: Optional[AEPsychMixin] = None,  # included for API compatibility
+        model: Optional[AEPsychModelMixin] = None,  # included for API compatibility
         fixed_features: Optional[Dict[int, float]] = None,
         **kwargs,
     ) -> torch.Tensor:
         """Query next point(s) to run by quasi-randomly sampling the parameter space.
         Args:
             num_points (int): Number of points to query. Defaults to 1.
-            moodel (AEPsychMixin, optional): Model to use for generating points. Not used in this generator. Defaults to None.
+            moodel (AEPsychModelMixin, optional): Model to use for generating points. Not used in this generator. Defaults to None.
             fixed_features: (Dict[int, float], optional): Parameters that are fixed to specific values.
             **kwargs: Ignored, API compatibility
         Returns: