Unify AEPsych Model Mixin (#627)

Summary:
Pull Request resolved: #627

Everything should use the same model mixin now.

Reviewed By: crasanders

Differential Revision: D69194015

fbshipit-source-id: b5a8a1260f4e785760fd605deb853a64864434ff

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,
@@ -23,7 +23,7 @@
 
 from ..models.model_protocol import ModelProtocol
 
-AEPsychModelType = TypeVar("AEPsychModelType", bound=AEPsychMixin)
+AEPsychModelType = TypeVar("AEPsychModelType", bound=AEPsychModelMixin)
 
 
 @runtime_checkable

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/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: