diff --git a/aepsych/acquisition/lookahead.py b/aepsych/acquisition/lookahead.py index aa906a264..c81c0418c 100644 --- a/aepsych/acquisition/lookahead.py +++ b/aepsych/acquisition/lookahead.py @@ -16,6 +16,7 @@ from botorch.models.gpytorch import GPyTorchModel from botorch.utils.transforms import t_batch_mode_transform from scipy.stats import norm +from torch import Tensor from .lookahead_utils import ( approximate_lookahead_levelset_at_xstar, @@ -245,6 +246,8 @@ def construct_inputs_local_lookahead( class GlobalLookaheadAcquisitionFunction(LookaheadAcquisitionFunction): def __init__( self, + lb: Tensor, + ub: Tensor, model: GPyTorchModel, lookahead_type: Literal["levelset", "posterior"] = "levelset", target: Optional[float] = None, @@ -256,14 +259,16 @@ def __init__( A global look-ahead acquisition function. Args: - model (GPyTorchModel): The gpytorch model to use. + lb (Tensor): Lower bounds of the input space, used to generate the query set (Xq). + ub (Tensor): Upper bounds of the input space, used to generate the query set (Xq). + model (GPyTorchModel): The gpytorch model. lookahead_type (Literal["levelset", "posterior"]): The type of look-ahead to perform (default is "levelset"). - If the lookahead_type is "levelset", the acqf will consider the posterior probability that a point is above or below the target level set. - If the lookahead_type is "posterior", the acqf will consider the posterior probability that a point will be detected or not. target (float, optional): Threshold value to target in p-space. - posterior_transform (PosteriorTransform, optional): Optional transformation to apply to the posterior. - query_set_size (int, optional): Number of points in the query set. - Xq (torch.Tensor, optional): (m x d) global reference set. + posterior_transform (PosteriorTransform, optional): Posterior transform to use. Defaults to None. + query_set_size (int, optional): Size of the query set. Defaults to 256. + Xq (Tensor, optional): (m x d) global reference set. Defaults to None. """ super().__init__(model=model, target=target, lookahead_type=lookahead_type) self.posterior_transform = posterior_transform @@ -282,7 +287,7 @@ def __init__( assert int(query_set_size) == query_set_size # make sure casting is safe # if the asserts above pass and Xq is None, query_set_size is not None so this is safe query_set_size = int(query_set_size) # cast - Xq = make_scaled_sobol(model.lb, model.ub, query_set_size) + Xq = make_scaled_sobol(lb, ub, query_set_size) self.register_buffer("Xq", Xq) @t_batch_mode_transform(expected_q=1) @@ -335,8 +340,10 @@ def _compute_acqf( class ApproxGlobalSUR(GlobalSUR): def __init__( self, + lb: Tensor, + ub: Tensor, model: GPyTorchModel, - lookahead_type="levelset", + lookahead_type: Literal["levelset", "poserior"] = "levelset", target: Optional[float] = None, query_set_size: Optional[int] = 256, Xq: Optional[torch.Tensor] = None, @@ -346,7 +353,9 @@ def __init__( Args: model (GPyTorchModel): The gpytorch model to use. - lookahed_type (str): The type of look-ahead to perform (default is "levelset"). + lookahead_type (Literal["levelset", "posterior"]): The type of look-ahead to perform (default is "levelset"). + - If the lookahead_type is "levelset", the acqf will consider the posterior probability that a point is above or below the target level set. + - If the lookahead_type is "posterior", the acqf will consider the posterior probability that a point will be detected or not. target (float, optional): Threshold value to target in p-space. query_set_size (int, optional): Number of points in the query set. Xq (torch.Tensor, optional): (m x d) global reference set. @@ -355,6 +364,8 @@ def __init__( lookahead_type == "levelset" ), f"ApproxGlobalSUR only supports lookahead on level set, got {lookahead_type}!" super().__init__( + lb=lb, + ub=ub, model=model, target=target, lookahead_type=lookahead_type, @@ -431,8 +442,10 @@ class SMOCU(GlobalLookaheadAcquisitionFunction): def __init__( self, + lb: Tensor, + ub: Tensor, model: GPyTorchModel, - lookahead_type="posterior", + lookahead_type: Literal["levelset", "posterior"] = "posterior", target: Optional[float] = None, query_set_size: Optional[int] = 256, Xq: Optional[torch.Tensor] = None, @@ -440,7 +453,9 @@ def __init__( ) -> None: """ model (GPyTorchModel): The gpytorch model to use. - lookahead_type (str): The type of look-ahead to perform (default is "posterior"). + lookahead_type (Literal["levelset", "posterior"]): The type of look-ahead to perform (default is "posterior"). + - If the lookahead_type is "levelset", the acqf will consider the posterior probability that a point is above or below the target level set. + - If the lookahead_type is "posterior", the acqf will consider the posterior probability that a point will be detected or not. target (float, optional): Threshold value to target in p-space. Default is None. query_set_size (int, optional): Number of points in the query set. Default is 256. Xq (torch.Tensor, optional): (m x d) global reference set. Default is None. @@ -448,6 +463,8 @@ def __init__( """ super().__init__( + lb=lb, + ub=ub, model=model, target=target, lookahead_type=lookahead_type, @@ -530,7 +547,7 @@ def _compute_acqf( def construct_inputs_global_lookahead( model: GPyTorchModel, training_data: None, - lookahead_type="levelset", + lookahead_type: Literal["levelset", "posterior"] = "levelset", target: Optional[float] = None, posterior_transform: Optional[PosteriorTransform] = None, query_set_size: Optional[int] = 256, @@ -543,7 +560,9 @@ def construct_inputs_global_lookahead( Args: model (GPyTorchModel): The gpytorch model to use. training_data (None): Placeholder for compatibility; not used in this function. - lookahead_type (str): Type of look-ahead to perform. Default is "levelset". + lookahead_type (Literal["levelset", "posterior"]): The type of look-ahead to perform (default is "levelset"). + - If the lookahead_type is "levelset", the acqf will consider the posterior probability that a point is above or below the target level set. + - If the lookahead_type is "posterior", the acqf will consider the posterior probability that a point will be detected or not. target (float, optional): Target threshold value in probability space. Default is None. posterior_transform (PosteriorTransform, optional): Optional transformation to apply to the posterior. Default is None. query_set_size (int, optional): Number of points in the query set. Default is 256. diff --git a/aepsych/benchmark/benchmark.py b/aepsych/benchmark/benchmark.py index 917690b7e..a2fcd8804 100644 --- a/aepsych/benchmark/benchmark.py +++ b/aepsych/benchmark/benchmark.py @@ -155,6 +155,7 @@ def run_experiment( np.random.seed(seed) config_dict["common"]["lb"] = str(problem.lb.tolist()) config_dict["common"]["ub"] = str(problem.ub.tolist()) + config_dict["common"]["dim"] = str(problem.lb.shape[0]) config_dict["common"]["parnames"] = str( [f"par{i}" for i in range(len(problem.ub.tolist()))] ) diff --git a/aepsych/benchmark/example_problems.py b/aepsych/benchmark/example_problems.py index 01500f56b..8f4b8e9e3 100644 --- a/aepsych/benchmark/example_problems.py +++ b/aepsych/benchmark/example_problems.py @@ -1,6 +1,6 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved. import os -from typing import List, Optional, Union +from typing import List, Union import numpy as np import torch @@ -11,7 +11,7 @@ novel_discrimination_testfun, ) from aepsych.models import GPClassificationModel -from aepsych.models.inducing_point_allocators import KMeansAllocator +from aepsych.models.inducing_points import KMeansAllocator """The DiscrimLowDim, DiscrimHighDim, ContrastSensitivity6d, and Hartmann6Binary classes are copied from bernoulli_lse github repository (https://github.com/facebookresearch/bernoulli_lse) @@ -104,13 +104,13 @@ def __init__( ) y = torch.LongTensor(self.data[:, 0]) x = torch.Tensor(self.data[:, 1:]) + inducing_size = 100 # Fit a model, with a large number of inducing points self.m = GPClassificationModel( - lb=self.bounds[0], - ub=self.bounds[1], - inducing_size=100, - inducing_point_method=KMeansAllocator(bounds=self.bounds), + dim=6, + inducing_size=inducing_size, + inducing_point_method=KMeansAllocator(dim=6), ) self.m.fit( diff --git a/aepsych/generators/acqf_thompson_sampler_generator.py b/aepsych/generators/acqf_thompson_sampler_generator.py index 02df99c29..cb14e0b8b 100644 --- a/aepsych/generators/acqf_thompson_sampler_generator.py +++ b/aepsych/generators/acqf_thompson_sampler_generator.py @@ -41,6 +41,8 @@ class AcqfThompsonSamplerGenerator(AEPsychGenerator): def __init__( self, + lb: torch.Tensor, + ub: torch.Tensor, acqf: AcquisitionFunction, acqf_kwargs: Optional[Dict[str, Any]] = None, samps: int = 1000, @@ -48,6 +50,8 @@ def __init__( ) -> None: """Initialize OptimizeAcqfGenerator. Args: + lb (torch.Tensor): Lower bounds for the optimization. + ub (torch.Tensor): Upper bounds for the optimization. acqf (AcquisitionFunction): Acquisition function to use. acqf_kwargs (Dict[str, object], optional): Extra arguments to pass to acquisition function. Defaults to no arguments. @@ -61,6 +65,8 @@ def __init__( self.acqf_kwargs = acqf_kwargs self.samps = samps self.stimuli_per_trial = stimuli_per_trial + self.lb = lb + self.ub = ub def _instantiate_acquisition_fn(self, model: ModelProtocol) -> AcquisitionFunction: """Instantiate the acquisition function with the model and any extra arguments. @@ -124,7 +130,7 @@ def _gen( starttime = time.time() seed = gen_options.get("seed") - bounds = torch.tensor(np.c_[model.lb, model.ub]).T.cpu() + bounds = torch.tensor(np.c_[self.lb, self.ub]).T.cpu() bounds_cpu = bounds.cpu() effective_dim = bounds.shape[-1] * num_points if effective_dim <= SobolEngine.MAXDIM: @@ -160,12 +166,16 @@ def from_config(cls, config: Config) -> AcqfThompsonSamplerGenerator: AcqfThompsonSamplerGenerator: The initialized generator. """ classname = cls.__name__ + lb = config.gettensor(classname, "lb") + ub = config.gettensor(classname, "ub") acqf = config.getobj(classname, "acqf", fallback=None) extra_acqf_args = cls._get_acqf_options(acqf, config) stimuli_per_trial = config.getint(classname, "stimuli_per_trial") samps = config.getint(classname, "samps", fallback=1000) return cls( + lb=lb, + ub=ub, acqf=acqf, acqf_kwargs=extra_acqf_args, samps=samps, diff --git a/aepsych/generators/epsilon_greedy_generator.py b/aepsych/generators/epsilon_greedy_generator.py index 7986de316..9608511cc 100644 --- a/aepsych/generators/epsilon_greedy_generator.py +++ b/aepsych/generators/epsilon_greedy_generator.py @@ -15,15 +15,25 @@ class EpsilonGreedyGenerator(AEPsychGenerator): - def __init__(self, subgenerator: AEPsychGenerator, epsilon: float = 0.1) -> None: + def __init__( + self, + lb: torch.Tensor, + ub: torch.Tensor, + subgenerator: AEPsychGenerator, + epsilon: float = 0.1, + ) -> None: """Initialize EpsilonGreedyGenerator. Args: + lb (torch.Tensor): Lower bounds for the optimization. + ub (torch.Tensor): Upper bounds for the optimization. subgenerator (AEPsychGenerator): The generator to use when not exploiting. epsilon (float): The probability of exploration. Defaults to 0.1. """ self.subgenerator = subgenerator self.epsilon = epsilon + self.lb = lb + self.ub = ub @classmethod def from_config(cls, config: Config) -> "EpsilonGreedyGenerator": @@ -36,12 +46,14 @@ def from_config(cls, config: Config) -> "EpsilonGreedyGenerator": EpsilonGreedyGenerator: The generator. """ classname = cls.__name__ + lb = torch.tensor(config.getlist(classname, "lb")) + ub = torch.tensor(config.getlist(classname, "ub")) subgen_cls = config.getobj( classname, "subgenerator", fallback=OptimizeAcqfGenerator ) subgen = subgen_cls.from_config(config) epsilon = config.getfloat(classname, "epsilon", fallback=0.1) - return cls(subgenerator=subgen, epsilon=epsilon) + return cls(lb=lb, ub=ub, subgenerator=subgen, epsilon=epsilon) def gen(self, num_points: int, model: ModelProtocol) -> torch.Tensor: """Query next point(s) to run by sampling from the subgenerator with probability 1-epsilon, and randomly otherwise. @@ -53,7 +65,7 @@ def gen(self, num_points: int, model: ModelProtocol) -> torch.Tensor: if num_points > 1: raise NotImplementedError("Epsilon-greedy batched gen is not implemented!") if np.random.uniform() < self.epsilon: - sample = np.random.uniform(low=model.lb, high=model.ub) + sample = np.random.uniform(low=self.lb, high=self.ub) return torch.tensor(sample).reshape(1, -1) else: return self.subgenerator.gen(num_points, model) diff --git a/aepsych/generators/monotonic_rejection_generator.py b/aepsych/generators/monotonic_rejection_generator.py index 3aca51e89..c655c6e13 100644 --- a/aepsych/generators/monotonic_rejection_generator.py +++ b/aepsych/generators/monotonic_rejection_generator.py @@ -12,6 +12,7 @@ from aepsych.config import Config from aepsych.generators.base import AEPsychGenerator from aepsych.models.monotonic_rejection_gp import MonotonicRejectionGP +from aepsych.utils import _process_bounds from botorch.acquisition import AcquisitionFunction from botorch.logging import logger from botorch.optim.initializers import gen_batch_initial_conditions @@ -43,13 +44,17 @@ class MonotonicRejectionGenerator(AEPsychGenerator[MonotonicRejectionGP]): def __init__( self, acqf: MonotonicMCAcquisition, + lb: torch.Tensor, + ub: torch.Tensor, acqf_kwargs: Optional[Dict[str, Any]] = None, model_gen_options: Optional[Dict[str, Any]] = None, explore_features: Optional[Sequence[int]] = None, ) -> None: """Initialize MonotonicRejectionGenerator. Args: - acqf (MonotonicMCAcquisition): Acquisition function to use. + acqf (AcquisitionFunction): Acquisition function to use. + lb (torch.Tensor): Lower bounds for the optimization. + ub (torch.Tensor): Upper bounds for the optimization. acqf_kwargs (Dict[str, object], optional): Extra arguments to pass to acquisition function. Defaults to None. model_gen_options (Dict[str, Any], optional): Dictionary with options for generating candidate, such as @@ -63,6 +68,8 @@ def __init__( self.acqf_kwargs = acqf_kwargs self.model_gen_options = model_gen_options self.explore_features = explore_features + self.lb, self.ub, _ = _process_bounds(lb, ub, None) + self.bounds = torch.stack((self.lb, self.ub)) def _instantiate_acquisition_fn( self, model: MonotonicRejectionGP @@ -110,7 +117,7 @@ def gen( ) # Augment bounds with deriv indicator - bounds = torch.cat((model.bounds_, torch.zeros(2, 1)), dim=1) + bounds = torch.cat((self.bounds, torch.zeros(2, 1)), dim=1) # Fix deriv indicator to 0 during optimization fixed_features = {(bounds.shape[1] - 1): 0.0} # Fix explore features to random values @@ -192,6 +199,8 @@ def from_config(cls, config: Config) -> "MonotonicRejectionGenerator": classname = cls.__name__ acqf = config.getobj("common", "acqf", fallback=None) extra_acqf_args = cls._get_acqf_options(acqf, config) + lb = torch.tensor(config.getlist(classname, "lb")) + ub = torch.tensor(config.getlist(classname, "ub")) options = {} options["num_restarts"] = config.getint(classname, "restarts", fallback=10) @@ -217,6 +226,8 @@ def from_config(cls, config: Config) -> "MonotonicRejectionGenerator": return cls( acqf=acqf, + lb=lb, + ub=ub, acqf_kwargs=extra_acqf_args, model_gen_options=options, explore_features=explore_features, diff --git a/aepsych/generators/optimize_acqf_generator.py b/aepsych/generators/optimize_acqf_generator.py index d5d746159..82630dfa7 100644 --- a/aepsych/generators/optimize_acqf_generator.py +++ b/aepsych/generators/optimize_acqf_generator.py @@ -5,10 +5,10 @@ # LICENSE file in the root directory of this source tree. from __future__ import annotations +import inspect import time from typing import Any, Dict, Optional -import numpy as np import torch from aepsych.config import Config from aepsych.generators.base import AEPsychGenerator @@ -39,6 +39,8 @@ class OptimizeAcqfGenerator(AEPsychGenerator): def __init__( self, + lb: torch.Tensor, + ub: torch.Tensor, acqf: AcquisitionFunction, acqf_kwargs: Optional[Dict[str, Any]] = None, restarts: int = 10, @@ -48,6 +50,8 @@ def __init__( ) -> None: """Initialize OptimizeAcqfGenerator. Args: + lb (torch.Tensor): Lower bounds for the optimization. + ub (torch.Tensor): Upper bounds for the optimization. acqf (AcquisitionFunction): Acquisition function to use. acqf_kwargs (Dict[str, object], optional): Extra arguments to pass to acquisition function. Defaults to no arguments. @@ -65,6 +69,8 @@ def __init__( self.samps = samps self.max_gen_time = max_gen_time self.stimuli_per_trial = stimuli_per_trial + self.lb = lb + self.ub = ub def _instantiate_acquisition_fn(self, model: ModelProtocol) -> AcquisitionFunction: """ @@ -76,14 +82,33 @@ def _instantiate_acquisition_fn(self, model: ModelProtocol) -> AcquisitionFuncti Returns: AcquisitionFunction: Configured acquisition function. """ + if ( + "lb" in inspect.signature(self.acqf).parameters + and "ub" in inspect.signature(self.acqf).parameters + ): + if self.acqf == AnalyticExpectedUtilityOfBestOption: + return self.acqf(pref_model=model, lb=self.lb, ub=self.ub) + + self.lb = self.lb.to(model.device) + self.ub = self.ub.to(model.device) + if self.acqf in self.baseline_requiring_acqfs: + return self.acqf( + model, + model.train_inputs[0], + lb=self.lb, + ub=self.ub, + **self.acqf_kwargs, + ) + + return self.acqf(model=model, lb=self.lb, ub=self.ub, **self.acqf_kwargs) if self.acqf == AnalyticExpectedUtilityOfBestOption: return self.acqf(pref_model=model) if self.acqf in self.baseline_requiring_acqfs: return self.acqf(model, model.train_inputs[0], **self.acqf_kwargs) - else: - return self.acqf(model=model, **self.acqf_kwargs) + + return self.acqf(model=model, **self.acqf_kwargs) def gen(self, num_points: int, model: ModelProtocol, **gen_options) -> torch.Tensor: """Query next point(s) to run by optimizing the acquisition function. @@ -124,12 +149,16 @@ def _gen( model.eval() # type: ignore acqf = self._instantiate_acquisition_fn(model) + if hasattr(model, "device"): + self.lb = self.lb.to(model.device) + self.ub = self.ub.to(model.device) + logger.info("Starting gen...") starttime = time.time() new_candidate, _ = optimize_acqf( acq_function=acqf, - bounds=torch.stack([model.lb, model.ub]), + bounds=torch.stack([self.lb, self.ub]), q=num_points, num_restarts=self.restarts, raw_samples=self.samps, @@ -153,6 +182,8 @@ def from_config(cls, config: Config) -> "OptimizeAcqfGenerator": restart and sample parameters, maximum generation time, and stimuli per trial. """ classname = cls.__name__ + lb = config.gettensor(classname, "lb") + ub = config.gettensor(classname, "ub") acqf = config.getobj(classname, "acqf", fallback=None) extra_acqf_args = cls._get_acqf_options(acqf, config) stimuli_per_trial = config.getint(classname, "stimuli_per_trial") @@ -161,6 +192,8 @@ def from_config(cls, config: Config) -> "OptimizeAcqfGenerator": max_gen_time = config.getfloat(classname, "max_gen_time", fallback=None) return cls( + lb=lb, + ub=ub, acqf=acqf, acqf_kwargs=extra_acqf_args, restarts=restarts, diff --git a/aepsych/generators/sobol_generator.py b/aepsych/generators/sobol_generator.py index 22e8b9a77..93bd4b783 100644 --- a/aepsych/generators/sobol_generator.py +++ b/aepsych/generators/sobol_generator.py @@ -8,7 +8,6 @@ from typing import Optional -import numpy as np import torch from aepsych.config import Config from aepsych.generators.base import AEPsychGenerator diff --git a/aepsych/models/__init__.py b/aepsych/models/__init__.py index 463ab7f3e..a3acf1294 100644 --- a/aepsych/models/__init__.py +++ b/aepsych/models/__init__.py @@ -10,14 +10,6 @@ from ..config import Config from .gp_classification import GPBetaRegressionModel, GPClassificationModel from .gp_regression import GPRegressionModel -from .inducing_point_allocators import ( - AutoAllocator, - DummyAllocator, - FixedAllocator, - GreedyVarianceReduction, - KMeansAllocator, - SobolAllocator, -) from .monotonic_projection_gp import MonotonicProjectionGP from .monotonic_rejection_gp import MonotonicRejectionGP from .multitask_regression import IndependentMultitaskGPRModel, MultitaskGPRModel @@ -42,12 +34,6 @@ "semi_p_posterior_transform", "GPBetaRegressionModel", "PairwiseProbitModel", - "AutoAllocator", - "KMeansAllocator", - "SobolAllocator", - "DummyAllocator", - "FixedAllocator", - "GreedyVarianceReduction", ] Config.register_module(sys.modules[__name__]) diff --git a/aepsych/models/base.py b/aepsych/models/base.py index 431aa9643..87636eedf 100644 --- a/aepsych/models/base.py +++ b/aepsych/models/base.py @@ -6,25 +6,19 @@ # LICENSE file in the root directory of this source tree. from __future__ import annotations -import abc import time from collections.abc import Iterable from copy import deepcopy -from typing import Any, Callable, Dict, List, Mapping, Optional, Protocol, Tuple, Union +from typing import Any, Callable, Dict, List, Mapping, Optional, Protocol, Tuple import gpytorch -import numpy as np import torch -from aepsych.config import Config, ConfigurableMixin -from aepsych.models.utils import get_extremum, inv_query -from aepsych.utils import dim_grid, get_jnd_multid, make_scaled_sobol, promote_0d from aepsych.utils_logging import getLogger from botorch.fit import fit_gpytorch_mll, fit_gpytorch_mll_scipy from botorch.models.gpytorch import GPyTorchModel from botorch.posteriors import GPyTorchPosterior from gpytorch.likelihoods import Likelihood from gpytorch.mlls import MarginalLogLikelihood -from scipy.stats import norm logger = getLogger() @@ -119,218 +113,6 @@ class AEPsychMixin(GPyTorchModel): train_inputs: Optional[Tuple[torch.Tensor]] train_targets: Optional[torch.Tensor] - @property - def bounds(self) -> torch.Tensor: - return torch.stack((self.lb, self.ub)) - - def get_max( - self: ModelProtocol, - locked_dims: Optional[Mapping[int, float]] = None, - probability_space: bool = False, - n_samples: int = 1000, - max_time: Optional[float] = None, - ) -> Tuple[float, torch.Tensor]: - """Return the maximum of the modeled function, subject to constraints - - Args: - locked_dims (Mapping[int, List[float]], optional): Dimensions to fix, so that the - max is along a slice of the full surface. Defaults to None. - probability_space (bool): Is y (and therefore the returned nearest_y) in - probability space instead of latent function space? Defaults to False. - n_samples (int): number of coarse grid points to sample for optimization estimate. - max_time (float, optional): Maximum time to spend optimizing. Defaults to None. - - Returns: - Tuple[float, torch.Tensor]: Tuple containing the max and its location (argmax). - """ - locked_dims = locked_dims or {} - _, _arg = get_extremum( - self, "max", self.bounds, locked_dims, n_samples, max_time=max_time - ) - arg = torch.tensor(_arg.reshape(1, self.dim)) - if probability_space: - val, _ = self.predict_probability(arg) - else: - val, _ = self.predict(arg) - return float(val.item()), arg - - def get_min( - self: ModelProtocol, - locked_dims: Optional[Mapping[int, float]] = None, - probability_space: bool = False, - n_samples: int = 1000, - max_time: Optional[float] = None, - ) -> Tuple[float, torch.Tensor]: - """Return the minimum of the modeled function, subject to constraints - Args: - locked_dims (Mapping[int, List[float]], optional): Dimensions to fix, so that the - min is along a slice of the full surface. - probability_space (bool): Is y (and therefore the returned nearest_y) in - probability space instead of latent function space? Defaults to False. - n_samples (int): number of coarse grid points to sample for optimization estimate. - max_time (float, optional): Maximum time to spend optimizing. Defaults to None. - - Returns: - Tuple[float, torch.Tensor]: Tuple containing the min and its location (argmin). - """ - locked_dims = locked_dims or {} - _, _arg = get_extremum( - self, "min", self.bounds, locked_dims, n_samples, max_time=max_time - ) - arg = torch.tensor(_arg.reshape(1, self.dim)) - if probability_space: - val, _ = self.predict_probability(arg) - else: - val, _ = self.predict(arg) - return float(val.item()), arg - - def inv_query( - self, - y: float, - locked_dims: Optional[Mapping[int, float]] = None, - probability_space: bool = False, - n_samples: int = 1000, - max_time: Optional[float] = None, - weights: Optional[torch.Tensor] = None, - ) -> Tuple[float, torch.Tensor]: - """Query the model inverse. - Return nearest x such that f(x) = queried y, and also return the - value of f at that point. - - Args: - y (float): Points at which to find the inverse. - locked_dims (Mapping[int, float], optional): Dimensions to fix, so that the - inverse is along a slice of the full surface. - probability_space (bool): Is y (and therefore the returned nearest_y) in - probability space instead of latent function space? Defaults to False. - n_samples (int): number of coarse grid points to sample for optimization estimate. Defaults to 1000. - max_time (float, optional): Maximum time to spend optimizing. Defaults to None. - weights (torch.Tensor, optional): Weights for the optimization. Defaults to None. - - Returns: - Tuple[float, torch.Tensor]: Tuple containing the value of f - nearest to queried y and the x position of this value. - """ - _, _arg = inv_query( - self, - y=y, - bounds=self.bounds, - locked_dims=locked_dims, - probability_space=probability_space, - n_samples=n_samples, - max_time=max_time, - weights=weights, - ) - arg = torch.tensor(_arg.reshape(1, self.dim)) - if probability_space: - val, _ = self.predict_probability(arg.reshape(1, self.dim)) - else: - val, _ = self.predict(arg.reshape(1, self.dim)) - return float(val.item()), arg - - def get_jnd( - self: ModelProtocol, - grid: Optional[torch.Tensor] = None, - cred_level: Optional[float] = None, - intensity_dim: int = -1, - confsamps: int = 500, - method: str = "step", - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: - """Calculate the JND. - - Note that JND can have multiple plausible definitions - outside of the linear case, so we provide options for how to compute it. - For method="step", we report how far one needs to go over in stimulus - space to move 1 unit up in latent space (this is a lot of people's - conventional understanding of the JND). - For method="taylor", we report the local derivative, which also maps to a - 1st-order Taylor expansion of the latent function. This is a formal - generalization of JND as defined in Weber's law. - Both definitions are equivalent for linear psychometric functions. - - Args: - grid (torch.Tensor, optional): Mesh grid over which to find the JND. - Defaults to a square grid of size as determined by aepsych.utils.dim_grid. - cred_level (float, optional): Credible level for computing an interval. - Defaults to None, computing no interval. - intensity_dim (int): Dimension over which to compute the JND. - Defaults to -1. - confsamps (int): Number of posterior samples to use for - computing the credible interval. Defaults to 500. - method (str): "taylor" or "step" method (see docstring). - Defaults to "step". - - Returns: - Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: either the - mean JND, or a median, lower, upper tuple of the JND posterior. - """ - if grid is None: - grid = self.dim_grid() - elif isinstance(grid, np.ndarray): - grid = torch.tensor(grid) - - # this is super awkward, back into intensity dim grid assuming a square grid - gridsize = int(grid.shape[0] ** (1 / grid.shape[1])) - coords = torch.linspace( - self.lb[intensity_dim].item(), self.ub[intensity_dim].item(), gridsize - ) - - if cred_level is None: - fmean, _ = self.predict(grid) - fmean = fmean.reshape(*[gridsize for i in range(self.dim)]) - - if method == "taylor": - return torch.tensor(1 / np.gradient(fmean, coords, axis=intensity_dim)) - elif method == "step": - return torch.clip( - get_jnd_multid( - fmean, - coords, - mono_dim=intensity_dim, - ), - 0, - np.inf, - ) - - alpha = 1 - cred_level # type: ignore - qlower = alpha / 2 - qupper = 1 - alpha / 2 - - fsamps = self.sample(grid, confsamps) - if method == "taylor": - jnds = torch.tensor( - 1 - / np.gradient( - fsamps.reshape(confsamps, *[gridsize for i in range(self.dim)]), - coords, - axis=intensity_dim, - ) - ) - elif method == "step": - samps = [s.reshape((gridsize,) * self.dim) for s in fsamps] - jnds = torch.stack( - [get_jnd_multid(s, coords, mono_dim=intensity_dim) for s in samps] - ) - else: - raise RuntimeError(f"Unknown method {method}!") - upper = torch.clip(torch.quantile(jnds, qupper, axis=0), 0, np.inf) # type: ignore - lower = torch.clip(torch.quantile(jnds, qlower, axis=0), 0, np.inf) # type: ignore - median = torch.clip(torch.quantile(jnds, 0.5, axis=0), 0, np.inf) # type: ignore - return median, lower, upper - - def dim_grid( - self: ModelProtocol, - gridsize: int = 30, - slice_dims: Optional[Mapping[int, float]] = None, - ) -> torch.Tensor: - """Generate a grid based on lower, upper, and dim. - - Args: - gridsize (int): Number of points in each dimension. Defaults to 30. - slice_dims (Mapping[int, float], optional): Dimensions to fix at a certain value. Defaults to None. - """ - return dim_grid(self.lb, self.ub, gridsize, slice_dims) - def set_train_data( self, inputs: Optional[torch.Tensor] = None, diff --git a/aepsych/models/gp_classification.py b/aepsych/models/gp_classification.py index e1c666439..6e07f864f 100644 --- a/aepsych/models/gp_classification.py +++ b/aepsych/models/gp_classification.py @@ -6,7 +6,6 @@ # LICENSE file in the root directory of this source tree. from __future__ import annotations -import warnings from copy import deepcopy from typing import Any, Dict, Optional, Tuple @@ -16,16 +15,10 @@ from aepsych.config import Config from aepsych.factory.default import default_mean_covar_factory from aepsych.models.base import AEPsychModelDeviceMixin -from aepsych.models.inducing_point_allocators import ( - AutoAllocator, - DummyAllocator, - KMeansAllocator, - SobolAllocator, -) -from aepsych.models.utils import select_inducing_points -from aepsych.utils import _process_bounds, get_optimizer_options, promote_0d +from aepsych.models.inducing_points import GreedyVarianceReduction +from aepsych.models.inducing_points.base import InducingPointAllocator +from aepsych.utils import get_dims, get_optimizer_options, promote_0d from aepsych.utils_logging import getLogger -from botorch.models.utils.inducing_point_allocators import InducingPointAllocator from gpytorch.likelihoods import BernoulliLikelihood, BetaLikelihood, Likelihood from gpytorch.models import ApproximateGP from gpytorch.variational import CholeskyVariationalDistribution, VariationalStrategy @@ -55,40 +48,35 @@ class GPClassificationModel(AEPsychModelDeviceMixin, ApproximateGP): def __init__( self, - lb: torch.Tensor, - ub: torch.Tensor, - inducing_point_method: InducingPointAllocator, - dim: Optional[int] = None, + dim: int, mean_module: Optional[gpytorch.means.Mean] = None, covar_module: Optional[gpytorch.kernels.Kernel] = None, likelihood: Optional[Likelihood] = None, - inducing_size: Optional[int] = None, + inducing_point_method: Optional[InducingPointAllocator] = None, + inducing_size: int = 100, max_fit_time: Optional[float] = None, optimizer_options: Optional[Dict[str, Any]] = None, - inducing_points: Optional[torch.Tensor] = None, ) -> None: """Initialize the GP Classification model Args: - lb (torch.Tensor): Lower bounds of the parameters. - ub (torch.Tensor): Upper bounds of the parameters. - inducing_point_method (InducingPointAllocator): The method to use for selecting inducing points. - dim (int, optional): The number of dimensions in the parameter space. If None, it is inferred from the size - of lb and ub. + dim (int): The number of dimensions in the parameter space. mean_module (gpytorch.means.Mean, optional): GP mean class. Defaults to a constant with a normal prior. covar_module (gpytorch.kernels.Kernel, optional): GP covariance kernel class. Defaults to scaled RBF with a gamma prior. likelihood (gpytorch.likelihood.Likelihood, optional): The likelihood function to use. If None defaults to Bernouli likelihood. - inducing_size (int, optional): Number of inducing points. Defaults to 99. + inducing_point_method (InducingPointAllocator, optional): The method to use for selecting inducing points. + If not set, a GreedyVarianceReduction is made. + inducing_size (int): Number of inducing points. Defaults to 100. max_fit_time (float, optional): The maximum amount of time, in seconds, to spend fitting the model. If None, there is no limit to the fitting time. optimizer_options (Dict[str, Any], optional): Optimizer options to pass to the SciPy optimizer during fitting. Assumes we are using L-BFGS-B. """ - lb, ub, self.dim = _process_bounds(lb, ub, dim) + self.dim = dim self.max_fit_time = max_fit_time - self.inducing_size = inducing_size or 99 + self.inducing_size = inducing_size self.optimizer_options = ( {"options": optimizer_options} if optimizer_options else {"options": {}} @@ -111,15 +99,14 @@ def __init__( dim=self.dim, stimuli_per_trial=self.stimuli_per_trial ) - self.inducing_point_method = inducing_point_method - inducing_points = select_inducing_points( - allocator=self.inducing_point_method, - inducing_size=self.inducing_size, + self.inducing_point_method = inducing_point_method or GreedyVarianceReduction( + dim=self.dim + ) + inducing_points = self.inducing_point_method.allocate_inducing_points( + num_inducing=self.inducing_size, covar_module=covar_module or default_covar, ) - self.last_inducing_points_method = self.inducing_point_method.allocator_used - self.inducing_points = inducing_points variational_distribution = CholeskyVariationalDistribution( inducing_points.size(0), batch_shape=torch.Size([self._batch_size]) ).to(inducing_points) @@ -132,9 +119,6 @@ def __init__( ) super().__init__(variational_strategy) - # Tensors need to be directly registered, Modules themselves can be assigned as attr - self.register_buffer("lb", lb) - self.register_buffer("ub", ub) self.likelihood = likelihood self.mean_module = mean_module or default_mean self.covar_module = covar_module or default_covar @@ -157,11 +141,11 @@ def from_config(cls, config: Config) -> GPClassificationModel: """ classname = cls.__name__ - inducing_size = config.getint(classname, "inducing_size", fallback=None) + inducing_size = config.getint(classname, "inducing_size", fallback=100) - lb = config.gettensor(classname, "lb") - ub = config.gettensor(classname, "ub") dim = config.getint(classname, "dim", fallback=None) + if dim is None: + dim = get_dims(config) mean_covar_factory = config.getobj( classname, "mean_covar_factory", fallback=default_mean_covar_factory @@ -171,7 +155,7 @@ def from_config(cls, config: Config) -> GPClassificationModel: max_fit_time = config.getfloat(classname, "max_fit_time", fallback=None) inducing_point_method_class = config.getobj( - classname, "inducing_point_method", fallback=AutoAllocator + classname, "inducing_point_method", fallback=GreedyVarianceReduction ) # Check if allocator class has a `from_config` method if hasattr(inducing_point_method_class, "from_config"): @@ -192,8 +176,6 @@ def from_config(cls, config: Config) -> GPClassificationModel: optimizer_options = get_optimizer_options(config, classname) return cls( - lb=lb, - ub=ub, dim=dim, inducing_size=inducing_size, mean_module=mean, @@ -220,14 +202,11 @@ def _reset_variational_strategy(self) -> None: if self.train_inputs is not None: # remember original device device = self.device - inducing_points = select_inducing_points( - allocator=self.inducing_point_method, - inducing_size=self.inducing_size, + inducing_points = self.inducing_point_method.allocate_inducing_points( + num_inducing=self.inducing_size, covar_module=self.covar_module, - X=self.train_inputs[0], - bounds=self.bounds, + inputs=self.train_inputs[0], ).to(device) - self.last_inducing_points_method = self.inducing_point_method.allocator_used variational_distribution = CholeskyVariationalDistribution( inducing_points.size(0), batch_shape=torch.Size([self._batch_size]) ).to(device) @@ -264,8 +243,7 @@ def fit( self._reset_hyperparameters() if not warmstart_induc or ( - self.last_inducing_points_method == "DummyAllocator" - and self.inducing_point_method.__class__.__name__ != "DummyAllocator" + self.inducing_point_method.last_allocator_used is None ): self._reset_variational_strategy() @@ -363,40 +341,34 @@ class GPBetaRegressionModel(GPClassificationModel): def __init__( self, - lb: torch.Tensor, - ub: torch.Tensor, - inducing_point_method: InducingPointAllocator, - dim: Optional[int] = None, + dim: int, mean_module: Optional[gpytorch.means.Mean] = None, covar_module: Optional[gpytorch.kernels.Kernel] = None, likelihood: Optional[Likelihood] = None, - inducing_size: Optional[int] = None, + inducing_point_method: Optional[InducingPointAllocator] = None, + inducing_size: int = 100, max_fit_time: Optional[float] = None, optimizer_options: Optional[Dict[str, Any]] = None, ) -> None: """Initialize the GP Beta Regression model Args: - lb (torch.Tensor): Lower bounds of the parameters. - ub (torch.Tensor): Upper bounds of the parameters. - inducing_point_method (InducingPointAllocator): The method to use to select the inducing points. - dim (int, optional): The number of dimensions in the parameter space. If None, it is inferred from the size - of lb and ub. Defaults to None. + dim (int): The number of dimensions in the parameter space. mean_module (gpytorch.means.Mean, optional): GP mean class. Defaults to a constant with a normal prior. Defaults to None. covar_module (gpytorch.kernels.Kernel, optional): GP covariance kernel class. Defaults to scaled RBF with a gamma prior. likelihood (gpytorch.likelihood.Likelihood, optional): The likelihood function to use. If None defaults to Beta likelihood. - inducing_size (int, optional): Number of inducing points. Defaults to 100. + inducing_point_method (InducingPointAllocator, optional): The method to use for selecting inducing points. + If not set, a GreedyVarianceReduction is made. + inducing_size (int): Number of inducing points. Defaults to 100. max_fit_time (float, optional): The maximum amount of time, in seconds, to spend fitting the model. If None, there is no limit to the fitting time. Defaults to None. """ if likelihood is None: likelihood = BetaLikelihood() - self.inducing_point_method = inducing_point_method + super().__init__( - lb=lb, - ub=ub, dim=dim, mean_module=mean_module, covar_module=covar_module, diff --git a/aepsych/models/gp_regression.py b/aepsych/models/gp_regression.py index bc418f2dc..81c8c95e1 100644 --- a/aepsych/models/gp_regression.py +++ b/aepsych/models/gp_regression.py @@ -10,12 +10,11 @@ from typing import Any, Dict, Optional, Tuple import gpytorch -import numpy as np import torch from aepsych.config import Config from aepsych.factory.default import default_mean_covar_factory from aepsych.models.base import AEPsychModelDeviceMixin -from aepsych.utils import _process_bounds, get_optimizer_options, promote_0d +from aepsych.utils import get_dims, get_optimizer_options, promote_0d from aepsych.utils_logging import getLogger from gpytorch.likelihoods import GaussianLikelihood, Likelihood from gpytorch.models import ExactGP @@ -33,9 +32,7 @@ class GPRegressionModel(AEPsychModelDeviceMixin, ExactGP): def __init__( self, - lb: torch.Tensor, - ub: torch.Tensor, - dim: Optional[int] = None, + dim: int, mean_module: Optional[gpytorch.means.Mean] = None, covar_module: Optional[gpytorch.kernels.Kernel] = None, likelihood: Optional[Likelihood] = None, @@ -45,10 +42,7 @@ def __init__( """Initialize the GP regression model Args: - lb (torch.Tensor): Lower bounds of the parameters. - ub (torch.Tensor): Upper bounds of the parameters. - dim (int, optional): The number of dimensions in the parameter space. If None, it is inferred from the size - of lb and ub. + dim (int): The number of dimensions in the parameter space. mean_module (gpytorch.means.Mean, optional): GP mean class. Defaults to a constant with a normal prior. covar_module (gpytorch.kernels.Kernel, optional): GP covariance kernel class. Defaults to scaled RBF with a gamma prior. @@ -59,12 +53,13 @@ def __init__( optimizer_options (Dict[str, Any], optional): Optimizer options to pass to the SciPy optimizer during fitting. Assumes we are using L-BFGS-B. """ + self.dim = dim + if likelihood is None: likelihood = GaussianLikelihood() super().__init__(None, None, likelihood) - lb, ub, self.dim = _process_bounds(lb, ub, dim) self.max_fit_time = max_fit_time self.optimizer_options = ( @@ -73,12 +68,10 @@ def __init__( if mean_module is None or covar_module is None: default_mean, default_covar = default_mean_covar_factory( - dim=self.dim, stimuli_per_trial=self.stimuli_per_trial + dim=self.dim, + stimuli_per_trial=self.stimuli_per_trial, ) - # Tensors need to be directly registered, Modules themselves can be assigned as attr - self.register_buffer("lb", lb) - self.register_buffer("ub", ub) self.likelihood = likelihood self.mean_module = mean_module or default_mean self.covar_module = covar_module or default_covar @@ -98,9 +91,9 @@ def construct_inputs(cls, config: Config) -> Dict: """ classname = cls.__name__ - lb = config.gettensor(classname, "lb") - ub = config.gettensor(classname, "ub") dim = config.getint(classname, "dim", fallback=None) + if dim is None: + dim = get_dims(config) mean_covar_factory = config.getobj( classname, "mean_covar_factory", fallback=default_mean_covar_factory @@ -123,8 +116,6 @@ def construct_inputs(cls, config: Config) -> Dict: optimizer_options = get_optimizer_options(config, classname) return { - "lb": lb, - "ub": ub, "dim": dim, "mean_module": mean, "covar_module": covar, diff --git a/aepsych/models/inducing_point_allocators.py b/aepsych/models/inducing_point_allocators.py deleted file mode 100644 index 53fc80ba0..000000000 --- a/aepsych/models/inducing_point_allocators.py +++ /dev/null @@ -1,604 +0,0 @@ -#!/usr/bin/env python3 -# Copyright (c) Facebook, Inc. and its affiliates. -# All rights reserved. - -# This source code is licensed under the license found in the -# LICENSE file in the root directory of this source tree. - -from abc import ABC, abstractmethod -from typing import Any, Dict, Optional, Union - -import numpy as np -import torch - -from aepsych.config import Config, ConfigurableMixin -from aepsych.utils import get_bounds -from botorch.models.utils.inducing_point_allocators import ( - GreedyVarianceReduction as BaseGreedyVarianceReduction, - InducingPointAllocator, -) -from botorch.utils.sampling import draw_sobol_samples -from scipy.cluster.vq import kmeans2 - - -class BaseAllocator(InducingPointAllocator, ConfigurableMixin): - """Base class for inducing point allocators.""" - - def __init__(self, bounds: Optional[torch.Tensor] = None) -> None: - """ - Initialize the allocator with optional bounds. - - Args: - bounds (torch.Tensor, optional): Bounds for allocating points. Should be of shape (2, d). - """ - self.bounds = bounds - self.dim = self._initialize_dim() - - def _initialize_dim(self) -> Optional[int]: - """ - Initialize the dimension `dim` based on the bounds, if available. - - Returns: - int: The dimension `d` if bounds are provided, or None otherwise. - """ - if self.bounds is not None: - # Validate bounds and extract dimension - assert self.bounds.shape[0] == 2, "Bounds must have shape (2, d)!" - lb, ub = self.bounds[0], self.bounds[1] - for i, (l, u) in enumerate(zip(lb, ub)): - assert ( - l <= u - ), f"Lower bound {l} is not less than or equal to upper bound {u} on dimension {i}!" - return self.bounds.shape[1] # Number of dimensions (d) - return None - - def _determine_dim_from_inputs(self, inputs: torch.Tensor) -> int: - """ - Determine dimension `dim` from the inputs tensor. - - Args: - inputs (torch.Tensor): Input tensor of shape (..., d). - - Returns: - int: The inferred dimension `d`. - """ - return inputs.shape[-1] - - @abstractmethod - def allocate_inducing_points( - self, - inputs: Optional[torch.Tensor], - covar_module: Optional[torch.nn.Module], - num_inducing: int, - input_batch_shape: torch.Size, - ) -> torch.Tensor: - """ - Abstract method for allocating inducing points. - - Args: - inputs (torch.Tensor, optional): Input tensor, implementation-specific. - covar_module (torch.nn.Module, optional): Kernel covariance module. - num_inducing (int): Number of inducing points to allocate. - input_batch_shape (torch.Size): Shape of the input batch. - - Returns: - torch.Tensor: Allocated inducing points. - """ - if self.dim is None and inputs is not None: - self.dim = self._determine_dim_from_inputs(inputs) - - raise NotImplementedError("This method should be implemented by subclasses.") - - @abstractmethod - def _get_quality_function(self) -> Optional[Any]: - """ - Abstract method for returning a quality function if required. - - Returns: - None or Callable: Quality function if needed. - """ - raise NotImplementedError("This method should be implemented by subclasses.") - - -class SobolAllocator(BaseAllocator): - """An inducing point allocator that uses Sobol sequences to allocate inducing points.""" - - def __init__(self, bounds: torch.Tensor) -> None: - """Initialize the SobolAllocator with bounds.""" - self.bounds: torch.Tensor = bounds - super().__init__(bounds=bounds) - - def _get_quality_function(self) -> None: - """Sobol sampling does not require a quality function, so this returns None.""" - return None - - def allocate_inducing_points( - self, - inputs: Optional[torch.Tensor] = None, - covar_module: Optional[torch.nn.Module] = None, - num_inducing: int = 10, - input_batch_shape: torch.Size = torch.Size([]), - ) -> torch.Tensor: - """ - Generates `num_inducing` inducing points within the specified bounds using Sobol sampling. - - Args: - inputs (torch.Tensor): Input tensor, not required for Sobol sampling. - covar_module (torch.nn.Module, optional): Kernel covariance module; included for API compatibility, but not used here. - num_inducing (int, optional): The number of inducing points to generate. Defaults to 10. - input_batch_shape (torch.Size, optional): Batch shape, defaults to an empty size; included for API compatibility, but not used here. - - - Returns: - torch.Tensor: A (num_inducing, d)-dimensional tensor of inducing points within the specified bounds. - - Raises: - ValueError: If `bounds` is not provided. - """ - - # Validate bounds shape - assert ( - self.bounds.shape[0] == 2 - ), "Bounds must have shape (2, d) for Sobol sampling." - # if bounds are long, make them float - if self.bounds.dtype == torch.long: - self.bounds = self.bounds.float() - # Generate Sobol samples within the unit cube [0,1]^d and rescale to [bounds[0], bounds[1]] - inducing_points = draw_sobol_samples( - bounds=self.bounds, n=num_inducing, q=1 - ).squeeze() - - # Ensure correct shape in case Sobol sampling returns a 1D tensor - if inducing_points.ndim == 1: - inducing_points = inducing_points.view(-1, 1) - self.allocator_used = "SobolAllocator" - return inducing_points - - @classmethod - def get_config_options( - cls, - config: Config, - name: Optional[str] = None, - options: Optional[Dict[str, Any]] = None, - ) -> Dict[str, Any]: - """Get configuration options for the SobolAllocator. - - Args: - config (Config): Configuration object. - name (str, optional): Name of the allocator, defaults to None. - options (Dict[str, Any], optional): Additional options, defaults to None. - - Returns: - Dict[str, Any]: Configuration options for the SobolAllocator. - """ - if name is None: - name = cls.__name__ - lb = config.gettensor("common", "lb") - ub = config.gettensor("common", "ub") - bounds = torch.stack((lb, ub)) - return {"bounds": bounds} - - -class KMeansAllocator(BaseAllocator): - """An inducing point allocator that uses k-means++ to allocate inducing points.""" - - def __init__(self, bounds: Optional[torch.Tensor] = None) -> None: - """Initialize the KMeansAllocator.""" - super().__init__(bounds=bounds) - if bounds is not None: - self.bounds = bounds - self.dummy_allocator = DummyAllocator(bounds) - - def _get_quality_function(self) -> None: - """K-means++ does not require a quality function, so this returns None.""" - return None - - def allocate_inducing_points( - self, - inputs: Optional[torch.Tensor] = None, - covar_module: Optional[torch.nn.Module] = None, - num_inducing: int = 10, - input_batch_shape: torch.Size = torch.Size([]), - ) -> torch.Tensor: - """ - Generates `num_inducing` inducing points using k-means++ initialization on the input data. - - Args: - inputs (torch.Tensor): A tensor of shape (n, d) containing the input data. - covar_module (torch.nn.Module, optional): Kernel covariance module; included for API compatibility, but not used here. - num_inducing (int, optional): The number of inducing points to generate. Defaults to 10. - input_batch_shape (torch.Size, optional): Batch shape, defaults to an empty size; included for API compatibility, but not used here. - - Returns: - torch.Tensor: A (num_inducing, d)-dimensional tensor of inducing points selected via k-means++. - """ - if inputs is None and self.bounds is not None: - self.allocator_used = self.dummy_allocator.__class__.__name__ - return self.dummy_allocator.allocate_inducing_points( - inputs=inputs, - covar_module=covar_module, - num_inducing=num_inducing, - input_batch_shape=input_batch_shape, - ) - elif inputs is None and self.bounds is None: - raise ValueError("Either inputs or bounds must be provided.") - # Ensure inputs are unique to avoid duplication issues with k-means++ - unique_inputs = torch.unique(inputs, dim=0) - - # If unique inputs are less than or equal to the required inducing points, return them directly - if unique_inputs.shape[0] <= num_inducing: - self.allocator_used = self.__class__.__name__ - return unique_inputs - - # Run k-means++ on the unique inputs to select inducing points - inducing_points = torch.tensor( - kmeans2(unique_inputs.cpu().numpy(), num_inducing, minit="++")[0] - ) - self.allocator_used = self.__class__.__name__ - return inducing_points - - @classmethod - def get_config_options( - cls, - config: Config, - name: Optional[str] = None, - options: Optional[Dict[str, Any]] = None, - ) -> Dict[str, Any]: - """Get configuration options for the KMeansAllocator. - - Args: - config (Config): Configuration object. - name (str, optional): Name of the allocator, defaults to None. - options (Dict[str, Any], optional): Additional options, defaults to None. - - Returns: - Dict[str, Any]: Configuration options for the KMeansAllocator. - """ - if name is None: - name = cls.__name__ - lb = config.gettensor("common", "lb") - ub = config.gettensor("common", "ub") - bounds = torch.stack((lb, ub)) - return {"bounds": bounds} - - -class DummyAllocator(BaseAllocator): - def __init__(self, bounds: torch.Tensor) -> None: - """Initialize the DummyAllocator with bounds. - - Args: - bounds (torch.Tensor): Bounds for allocating points. Should be of shape (2, d). - """ - super().__init__(bounds=bounds) - self.bounds: torch.Tensor = bounds - - def _get_quality_function(self) -> None: - """DummyAllocator does not require a quality function, so this returns None.""" - return None - - def allocate_inducing_points( - self, - inputs: Optional[torch.Tensor] = None, - covar_module: Optional[torch.nn.Module] = None, - num_inducing: int = 10, - input_batch_shape: torch.Size = torch.Size([]), - ) -> torch.Tensor: - """Allocate inducing points by returning zeros of the appropriate shape. - - Args: - inputs (torch.Tensor): Input tensor, not required for DummyAllocator. - covar_module (torch.nn.Module, optional): Kernel covariance module; included for API compatibility, but not used here. - num_inducing (int, optional): The number of inducing points to generate. Defaults to 10. - input_batch_shape (torch.Size, optional): Batch shape, defaults to an empty size; included for API compatibility, but not used here. - - Returns: - torch.Tensor: A (num_inducing, d)-dimensional tensor of zeros. - """ - self.allocator_used = self.__class__.__name__ - return torch.zeros(num_inducing, self.bounds.shape[-1]) - - @classmethod - def get_config_options( - cls, - config: Config, - name: Optional[str] = None, - options: Optional[Dict[str, Any]] = None, - ) -> Dict[str, Any]: - """Get configuration options for the DummyAllocator. - - Args: - config (Config): Configuration object. - name (str, optional): Name of the allocator, defaults to None. - options (Dict[str, Any], optional): Additional options, defaults to None. - - Returns: - Dict[str, Any]: Configuration options for the DummyAllocator. - """ - if name is None: - name = cls.__name__ - lb = config.gettensor("common", "lb") - ub = config.gettensor("common", "ub") - bounds = torch.stack((lb, ub)) - return {"bounds": bounds} - - -class AutoAllocator(BaseAllocator): - """An inducing point allocator that dynamically chooses an allocation strategy - based on the number of unique data points available.""" - - def __init__( - self, - bounds: Optional[torch.Tensor] = None, - fallback_allocator: InducingPointAllocator = KMeansAllocator(), - ) -> None: - """ - Initialize the AutoAllocator with a fallback allocator. - - Args: - fallback_allocator (InducingPointAllocator, optional): Allocator to use if there are - more unique points than required. - """ - super().__init__(bounds=bounds) - self.fallback_allocator = fallback_allocator - if bounds is not None: - self.bounds = bounds - self.dummy_allocator = DummyAllocator(bounds=bounds) - - def _get_quality_function(self) -> None: - """AutoAllocator does not require a quality function, so this returns None.""" - return None - - def allocate_inducing_points( - self, - inputs: Optional[torch.Tensor], - covar_module: Optional[torch.nn.Module] = None, - num_inducing: int = 10, - input_batch_shape: torch.Size = torch.Size([]), - ) -> torch.Tensor: - """ - Allocate inducing points by either using the unique input data directly - or falling back to another allocation method if there are too many unique points. - - Args: - inputs (torch.Tensor): A tensor of shape (n, d) containing the input data. - covar_module (torch.nn.Module, optional): Kernel covariance module; included for API compatibility, but not used here. - num_inducing (int, optional): The number of inducing points to generate. - input_batch_shape (torch.Size, optional): Batch shape, defaults to an empty size; included for API compatibility, but not used here. - - Returns: - torch.Tensor: A (num_inducing, d)-dimensional tensor of inducing points. - """ - # Ensure inputs are not None - if inputs is None and self.bounds is not None: - self.allocator_used = self.dummy_allocator.__class__.__name__ - return self.dummy_allocator.allocate_inducing_points( - inputs=inputs, - covar_module=covar_module, - num_inducing=num_inducing, - input_batch_shape=input_batch_shape, - ) - elif inputs is None and self.bounds is None: - raise ValueError(f"Either inputs or bounds must be provided.{self.bounds}") - - assert ( - inputs is not None - ), "inputs should not be None here" # to make mypy happy - - unique_inputs = torch.unique(inputs, dim=0) - - # If there are fewer unique points than required, return unique inputs directly - if unique_inputs.shape[0] <= num_inducing: - self.allocator_used = self.__class__.__name__ - return unique_inputs - - # Otherwise, fall back to the provided allocator (e.g., KMeansAllocator) - if inputs.shape[0] <= num_inducing: - self.allocator_used = self.__class__.__name__ - return inputs - else: - self.allocator_used = self.fallback_allocator.__class__.__name__ - return self.fallback_allocator.allocate_inducing_points( - inputs=inputs, - covar_module=covar_module, - num_inducing=num_inducing, - input_batch_shape=input_batch_shape, - ) - - @classmethod - def get_config_options( - cls, - config: Config, - name: Optional[str] = None, - options: Optional[Dict[str, Any]] = None, - ) -> Dict[str, Any]: - """Get configuration options for the AutoAllocator. - - Args: - config (Config): Configuration object. - name (str, optional): Name of the allocator, defaults to None. - options (Dict[str, Any], optional): Additional options, defaults to None. - - Returns: - Dict[str, Any]: Configuration options for the AutoAllocator. - """ - if name is None: - name = cls.__name__ - lb = config.gettensor("common", "lb") - ub = config.gettensor("common", "ub") - bounds = torch.stack((lb, ub)) - fallback_allocator_cls = config.getobj( - name, "fallback_allocator", fallback=KMeansAllocator - ) - fallback_allocator = ( - fallback_allocator_cls.from_config(config) - if hasattr(fallback_allocator_cls, "from_config") - else fallback_allocator_cls() - ) - - return {"fallback_allocator": fallback_allocator, "bounds": bounds} - - -class FixedAllocator(BaseAllocator): - def __init__( - self, points: torch.Tensor, bounds: Optional[torch.Tensor] = None - ) -> None: - """Initialize the FixedAllocator with inducing points and bounds. - - Args: - points (torch.Tensor): Inducing points to use. - bounds (torch.Tensor, optional): Bounds for allocating points. Should be of shape (2, d). - """ - super().__init__(bounds=bounds) - self.points = points - - def _get_quality_function(self) -> None: - """FixedAllocator does not require a quality function, so this returns None.""" - return None - - def allocate_inducing_points( - self, - inputs: Optional[torch.Tensor] = None, - covar_module: Optional[torch.nn.Module] = None, - num_inducing: int = 10, - input_batch_shape: torch.Size = torch.Size([]), - ) -> torch.Tensor: - """Allocate inducing points by returning the fixed inducing points. - - Args: - inputs (torch.Tensor): Input tensor, not required for FixedAllocator. - covar_module (torch.nn.Module, optional): Kernel covariance module; included for API compatibility, but not used here. - num_inducing (int, optional): The number of inducing points to generate. Defaults to 10. - input_batch_shape (torch.Size, optional): Batch shape, defaults to an empty size; included for API compatibility, but not used here. - - Returns: - torch.Tensor: The fixed inducing points. - """ - self.allocator_used = self.__class__.__name__ - return self.points - - @classmethod - def get_config_options( - cls, - config: Config, - name: Optional[str] = None, - options: Optional[Dict[str, Any]] = None, - ) -> Dict[str, Any]: - """Get configuration options for the FixedAllocator. - - Args: - config (Config): Configuration object. - name (str, optional): Name of the allocator, defaults to None. - options (Dict[str, Any], optional): Additional options, defaults to None. - - Returns: - Dict[str, Any]: Configuration options for the FixedAllocator. - """ - if name is None: - name = cls.__name__ - lb = config.gettensor("common", "lb") - ub = config.gettensor("common", "ub") - bounds = torch.stack((lb, ub)) - num_inducing = config.getint("common", "num_inducing", fallback=99) - fallback_allocator = config.getobj( - name, "fallback_allocator", fallback=DummyAllocator(bounds=bounds) - ) - points = config.gettensor( - name, - "points", - fallback=fallback_allocator.allocate_inducing_points( - num_inducing=num_inducing - ), - ) - return {"points": points, "bounds": bounds} - - -class GreedyVarianceReduction(BaseGreedyVarianceReduction, ConfigurableMixin): - def __init__(self, bounds: Optional[torch.Tensor] = None) -> None: - """Initialize the GreedyVarianceReduction with bounds. - - Args: - bounds (torch.Tensor, optional): Bounds for allocating points. Should be of shape (2, d). - """ - super().__init__() - - self.bounds = bounds - if bounds is not None: - self.dummy_allocator = DummyAllocator(bounds) - self.dim = self._initialize_dim() - - def _initialize_dim(self) -> Optional[int]: - """Initialize the dimension `dim` based on the bounds, if available. - - Returns: - int: The dimension `d` if bounds are provided, or None otherwise. - """ - if self.bounds is not None: - assert self.bounds.shape[0] == 2, "Bounds must have shape (2, d)!" - lb, ub = self.bounds[0], self.bounds[1] - for i, (l, u) in enumerate(zip(lb, ub)): - assert ( - l <= u - ), f"Lower bound {l} is not less than or equal to upper bound {u} on dimension {i}!" - return self.bounds.shape[1] - return None - - def allocate_inducing_points( - self, - inputs: Optional[torch.Tensor] = None, - covar_module: Optional[torch.nn.Module] = None, - num_inducing: int = 10, - input_batch_shape: torch.Size = torch.Size([]), - ) -> torch.Tensor: - """Allocate inducing points using the GreedyVarianceReduction strategy. - - Args: - inputs (torch.Tensor): Input tensor, not required for GreedyVarianceReduction. - covar_module (torch.nn.Module, optional): Kernel covariance module; included for API compatibility, but not used here. - num_inducing (int, optional): The number of inducing points to generate. Defaults to 10. - input_batch_shape (torch.Size, optional): Batch shape, defaults to an empty size; included for API compatibility, but not used here. - - Returns: - torch.Tensor: The allocated inducing points. - """ - if inputs is None and self.bounds is not None: - self.allocator_used = self.dummy_allocator.__class__.__name__ - return self.dummy_allocator.allocate_inducing_points( - inputs=inputs, - covar_module=covar_module, - num_inducing=num_inducing, - input_batch_shape=input_batch_shape, - ) - elif inputs is None and self.bounds is None: - raise ValueError("Either inputs or bounds must be provided.") - else: - self.allocator_used = self.__class__.__name__ - return super().allocate_inducing_points( - inputs=inputs, - covar_module=covar_module, - num_inducing=num_inducing, - input_batch_shape=input_batch_shape, - ) - - @classmethod - def get_config_options( - cls, - config: Config, - name: Optional[str] = None, - options: Optional[Dict[str, Any]] = None, - ) -> Dict[str, Any]: - """Get configuration options for the GreedyVarianceReduction allocator. - - Args: - config (Config): Configuration object. - name (str, optional): Name of the allocator, defaults to None. - options (Dict[str, Any], optional): Additional options, defaults to None. - - Returns: - Dict[str, Any]: Configuration options for the GreedyVarianceReduction allocator. - """ - if name is None: - name = cls.__name__ - lb = config.gettensor("common", "lb") - ub = config.gettensor("common", "ub") - bounds = torch.stack((lb, ub)) - return {"bounds": bounds} diff --git a/aepsych/models/inducing_points/__init__.py b/aepsych/models/inducing_points/__init__.py new file mode 100644 index 000000000..1f0c8628a --- /dev/null +++ b/aepsych/models/inducing_points/__init__.py @@ -0,0 +1,23 @@ +#!/usr/bin/env python3 +# Copyright (c) Meta Platforms, Inc. and its affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import sys + +from ...config import Config +from .fixed import FixedAllocator +from .greedy_variance_reduction import GreedyVarianceReduction +from .kmeans import KMeansAllocator +from .sobol import SobolAllocator + +__all__ = [ + "FixedAllocator", + "GreedyVarianceReduction", + "KMeansAllocator", + "SobolAllocator", +] + +Config.register_module(sys.modules[__name__]) diff --git a/aepsych/models/inducing_points/base.py b/aepsych/models/inducing_points/base.py new file mode 100644 index 000000000..66723baad --- /dev/null +++ b/aepsych/models/inducing_points/base.py @@ -0,0 +1,83 @@ +from abc import abstractmethod +from typing import Any, Dict, Optional + +import torch +from aepsych.config import Config, ConfigurableMixin +from aepsych.utils import get_dims +from botorch.models.utils.inducing_point_allocators import InducingPointAllocator + + +class BaseAllocator(InducingPointAllocator, ConfigurableMixin): + """Base class for inducing point allocators.""" + + def __init__(self, dim: int, *args, **kwargs) -> None: + """ + Initialize the allocator with optional bounds. + + Args: + dim (int): Dimensionality of the search space. + *args, **kwargs: Other allocator specific arguments. + """ + self.dim = dim + self.last_allocator_used: Optional[InducingPointAllocator] = None + + @abstractmethod + def allocate_inducing_points( + self, + inputs: Optional[torch.Tensor] = None, + covar_module: Optional[torch.nn.Module] = None, + num_inducing: int = 100, + input_batch_shape: torch.Size = torch.Size([]), + ) -> torch.Tensor: + """ + Abstract method for allocating inducing points. Must replace the + last_allocator_used attribute for what was actually used to produce the + inducing points. Dummy points should be made when it is not possible to create + inducing points (e.g., inputs is None). + + Args: + inputs (torch.Tensor, optional): Input tensor, implementation-specific. + covar_module (torch.nn.Module, optional): Kernel covariance module. + num_inducing (int): Number of inducing points to allocate. + input_batch_shape (torch.Size): Shape of the input batch. + + Returns: + torch.Tensor: Allocated inducing points. + """ + + def _allocate_dummy_points(self, num_inducing: int = 100) -> torch.Tensor: + """Return dummy inducing points with the correct dimensionality. + + Args: + num_inducing (int): Number of inducing points to make, defaults to 100. + """ + self.last_allocator_used = None + return torch.zeros(num_inducing, self.dim) + + def _get_quality_function(self): + return super()._get_quality_function() + + @classmethod + def get_config_options( + cls, + config: Config, + name: Optional[str] = None, + options: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Get configuration options for the allocator. + + Args: + config (Config): Configuration object. + name (str, optional): Name of the allocator, defaults to None. Ignored. + options (Dict[str, Any], optional): Additional options, defaults to None. + + Returns: + Dict[str, Any]: Configuration options for the DummyAllocator. + """ + if options is None: + options = {} + + if "dim" not in options: + options["dim"] = get_dims(config) + + return options diff --git a/aepsych/models/inducing_points/fixed.py b/aepsych/models/inducing_points/fixed.py new file mode 100644 index 000000000..3932b0169 --- /dev/null +++ b/aepsych/models/inducing_points/fixed.py @@ -0,0 +1,70 @@ +from typing import Any, Dict, Optional + +import torch +from aepsych.config import Config +from aepsych.models.inducing_points.base import BaseAllocator + + +class FixedAllocator(BaseAllocator): + def __init__( + self, + dim: int, + points: torch.Tensor, + ) -> None: + """Initialize the FixedAllocator with inducing points to use and bounds. + + Args: + dim (int): Dimensionality of the search space. + points (torch.Tensor): Inducing points to use (should be n, d). + """ + super().__init__(dim=dim) + self.points = points + + def allocate_inducing_points( + self, + inputs: Optional[torch.Tensor] = None, + covar_module: Optional[torch.nn.Module] = None, + num_inducing: int = 100, + input_batch_shape: torch.Size = torch.Size([]), + ) -> torch.Tensor: + """Allocate inducing points by returning the fixed inducing points. + + Args: + inputs (torch.Tensor): Input tensor, not required for FixedAllocator. + covar_module (torch.nn.Module, optional): Kernel covariance module; included for API compatibility, but not used here. + num_inducing (int, optional): The number of inducing points to generate. Defaults to 10. + input_batch_shape (torch.Size, optional): Batch shape, defaults to an empty size; included for API compatibility, but not used here. + + Returns: + torch.Tensor: The fixed inducing points. + """ + # TODO: Usually, these are initialized such that the transforms are applied to + # points already, this means that if the transforms change over training, the inducing + # points aren't in the space. However, we don't have any changing transforms + # right now. + + self.last_allocator_used = self.__class__ + return self.points + + @classmethod + def get_config_options( + cls, + config: Config, + name: Optional[str] = None, + options: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Get configuration options for the FixedAllocator. + + Args: + config (Config): Configuration object. + name (str, optional): Name of the allocator, defaults to None. + options (Dict[str, Any], optional): Additional options, defaults to None. + + Returns: + Dict[str, Any]: Configuration options for the FixedAllocator. + """ + options = super().get_config_options(config=config, name=name, options=options) + + options["points"] = config.gettensor("FixedAllocator", "points") + + return options diff --git a/aepsych/models/inducing_points/greedy_variance_reduction.py b/aepsych/models/inducing_points/greedy_variance_reduction.py new file mode 100644 index 000000000..bb5558219 --- /dev/null +++ b/aepsych/models/inducing_points/greedy_variance_reduction.py @@ -0,0 +1,53 @@ +from typing import Optional + +import torch +from aepsych.models.inducing_points.base import BaseAllocator +from botorch.models.utils.inducing_point_allocators import ( + GreedyVarianceReduction as BaseGreedyVarianceReduction, +) + + +class GreedyVarianceReduction(BaseGreedyVarianceReduction, BaseAllocator): + def allocate_inducing_points( + self, + inputs: Optional[torch.Tensor] = None, + covar_module: Optional[torch.nn.Module] = None, + num_inducing: int = 100, + input_batch_shape: torch.Size = torch.Size([]), + ) -> torch.Tensor: + """Allocate inducing points using the GreedyVarianceReduction strategy. This is + a thin wrapper around BoTorch's GreedyVarianceRedution inducing point allocator. + + Args: + inputs (torch.Tensor): Input tensor, not required for GreedyVarianceReduction. + covar_module (torch.nn.Module, optional): Kernel covariance module; included for API compatibility, but not used here. + num_inducing (int, optional): The number of inducing points to generate. Defaults to 10. + input_batch_shape (torch.Size, optional): Batch shape, defaults to an empty size; included for API compatibility, but not used here. + + Returns: + torch.Tensor: The allocated inducing points. + """ + if inputs is None: # Dummy points + return self._allocate_dummy_points(num_inducing=num_inducing) + else: + if covar_module is None: + raise ValueError( + "covar_module must be set for the GreedyVarianceReduction" + ) + + self.last_allocator_used = self.__class__ + + points = BaseGreedyVarianceReduction.allocate_inducing_points( + self, + inputs=inputs, + covar_module=covar_module, + num_inducing=num_inducing, + input_batch_shape=input_batch_shape, + ) + + if points.shape[1] != self.dim: + # We assume if the shape doesn't match the dim, it's because the points + # were augmented by adding it to be end of the shape + points = points[:, : self.dim, ...] + + return points diff --git a/aepsych/models/inducing_points/kmeans.py b/aepsych/models/inducing_points/kmeans.py new file mode 100644 index 000000000..a8d392631 --- /dev/null +++ b/aepsych/models/inducing_points/kmeans.py @@ -0,0 +1,51 @@ +from typing import Optional + +import torch +from aepsych.models.inducing_points.base import BaseAllocator +from scipy.cluster.vq import kmeans2 + + +class KMeansAllocator(BaseAllocator): + """An inducing point allocator that uses k-means++ to allocate inducing points.""" + + def allocate_inducing_points( + self, + inputs: Optional[torch.Tensor] = None, + covar_module: Optional[torch.nn.Module] = None, + num_inducing: int = 100, + input_batch_shape: torch.Size = torch.Size([]), + ) -> torch.Tensor: + """ + Generates `num_inducing` inducing points using k-means++ initialization on the input data. + + Args: + inputs (torch.Tensor): A tensor of shape (n, d) containing the input data. + covar_module (torch.nn.Module, optional): Kernel covariance module; included for API compatibility, but not used here. + num_inducing (int, optional): The number of inducing points to generate. Defaults to 100. + input_batch_shape (torch.Size, optional): Batch shape, defaults to an empty size; included for API compatibility, but not used here. + + Returns: + torch.Tensor: A (num_inducing, d)-dimensional tensor of inducing points selected via k-means++. + """ + if inputs is None: # Dummy points + return self._allocate_dummy_points(num_inducing=num_inducing) + + if inputs.shape[1] != self.dim: + # The inputs were augmented somehow, assuming it was added to the end of dims + inputs = inputs[:, : self.dim, ...] + + self.last_allocator_used = self.__class__ + + # Ensure inputs are unique to avoid duplication issues with k-means++ + unique_inputs = torch.unique(inputs, dim=0) + + # If unique inputs are less than or equal to the required inducing points, return them directly + if unique_inputs.shape[0] <= num_inducing: + return unique_inputs + + # Run k-means++ on the unique inputs to select inducing points + inducing_points = torch.tensor( + kmeans2(unique_inputs.cpu().numpy(), num_inducing, minit="++")[0] + ) + + return inducing_points diff --git a/aepsych/models/inducing_points/sobol.py b/aepsych/models/inducing_points/sobol.py new file mode 100644 index 000000000..7734a8628 --- /dev/null +++ b/aepsych/models/inducing_points/sobol.py @@ -0,0 +1,89 @@ +from typing import Any, Dict, Optional + +import torch +from aepsych.config import Config +from aepsych.models.inducing_points.base import BaseAllocator +from botorch.utils.sampling import draw_sobol_samples + + +class SobolAllocator(BaseAllocator): + """An inducing point allocator that uses Sobol sequences to allocate inducing points.""" + + def __init__(self, dim: int, bounds: torch.Tensor) -> None: + """ + Initializes a Sobol Allocator. This allocator must have bounds. + + Args: + dim (int): Dimensionality of the search space. + bounds (torch.Tensor): Bounds for allocating points. Should be of shape + (2, d). + """ + # Make sure bounds are the right type so the outputs are the right type + self.bounds = bounds.to(torch.float64) + super().__init__( + dim=dim, + ) + + def allocate_inducing_points( + self, + inputs: Optional[torch.Tensor] = None, + covar_module: Optional[torch.nn.Module] = None, + num_inducing: int = 100, + input_batch_shape: torch.Size = torch.Size([]), + ) -> torch.Tensor: + """ + Generates `num_inducing` inducing points within the specified bounds using Sobol sampling. + + Args: + inputs (torch.Tensor): Input tensor, ignored for Sobol points. + covar_module (torch.nn.Module, optional): Kernel covariance module; included for API compatibility, but not used here. + num_inducing (int, optional): The number of inducing points to generate. Defaults to 100. + input_batch_shape (torch.Size, optional): Batch shape, defaults to an empty size; included for API compatibility, but not used here. + + + Returns: + torch.Tensor: A (num_inducing, d)-dimensional tensor of inducing points within the specified bounds. + """ + # TODO: Usually, these are initialized such that the transforms are applied to + # bounds, this means that if the transforms change over training, the inducing + # points aren't in the space. However, we don't have any changing transforms + # right now. + + # Generate Sobol samples within the unit cube [0,1]^d and rescale to [bounds[0], bounds[1]] + inducing_points = draw_sobol_samples( + bounds=self.bounds, n=num_inducing, q=1 + ).squeeze() + + # Ensure correct shape in case Sobol sampling returns a 1D tensor + if inducing_points.ndim == 1: + inducing_points = inducing_points.view(-1, 1) + + self.last_allocator_used = self.__class__ + + return inducing_points + + @classmethod + def get_config_options( + cls, + config: Config, + name: Optional[str] = None, + options: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Get configuration options for the FixedAllocator. + + Args: + config (Config): Configuration object. + name (str, optional): Name of the allocator, defaults to None. + options (Dict[str, Any], optional): Additional options, defaults to None. + + Returns: + Dict[str, Any]: Configuration options for the FixedAllocator. + """ + options = super().get_config_options(config=config, name=name, options=options) + + if "bounds" not in options: + lb = config.gettensor("common", "lb") + ub = config.gettensor("common", "ub") + options["bounds"] = torch.stack((lb, ub)) + + return options diff --git a/aepsych/models/monotonic_projection_gp.py b/aepsych/models/monotonic_projection_gp.py index f6e5d0cda..011fa0f3c 100644 --- a/aepsych/models/monotonic_projection_gp.py +++ b/aepsych/models/monotonic_projection_gp.py @@ -15,9 +15,9 @@ from aepsych.config import Config from aepsych.factory.default import default_mean_covar_factory from aepsych.models.gp_classification import GPClassificationModel -from aepsych.models.inducing_point_allocators import AutoAllocator -from aepsych.utils import get_optimizer_options -from botorch.models.utils.inducing_point_allocators import InducingPointAllocator +from aepsych.models.inducing_points import GreedyVarianceReduction +from aepsych.models.inducing_points.base import InducingPointAllocator +from aepsych.utils import get_dims, get_optimizer_options from botorch.posteriors.gpytorch import GPyTorchPosterior from gpytorch.likelihoods import Likelihood from statsmodels.stats.moment_helpers import corr2cov, cov2corr @@ -97,36 +97,36 @@ def __init__( self, lb: torch.Tensor, ub: torch.Tensor, - inducing_point_method: InducingPointAllocator, + dim: int, monotonic_dims: List[int], monotonic_grid_size: int = 20, min_f_val: Optional[float] = None, - dim: Optional[int] = None, mean_module: Optional[gpytorch.means.Mean] = None, covar_module: Optional[gpytorch.kernels.Kernel] = None, likelihood: Optional[Likelihood] = None, - inducing_size: Optional[int] = None, + inducing_point_method: Optional[InducingPointAllocator] = None, + inducing_size: int = 100, max_fit_time: Optional[float] = None, optimizer_options: Optional[Dict[str, Any]] = None, ) -> None: - """Initialize the MonotonicProjectionGP model. + """Initialize the MonotonicProjectionGP model. Unlike other models, this model needs bounds. Args: lb (torch.Tensor): Lower bounds of the parameters. ub (torch.Tensor): Upper bounds of the parameters. - inducing_point_method (InducingPointAllocator): The method for allocating inducing points. + dim (int, optional): The number of dimensions in the parameter space. monotonic_dims (List[int]): A list of the dimensions on which monotonicity should be enforced. monotonic_grid_size (int): The size of the grid, s, in 1. above. Defaults to 20. min_f_val (float, optional): If provided, maintains this minimum in the projection in 5. Defaults to None. - dim (int, optional): The number of dimensions in the parameter space. If None, it is inferred from the size - of lb and ub. Defaults to None. mean_module (gpytorch.means.Mean, optional): GP mean class. Defaults to a constant with a normal prior. Defaults to None. covar_module (gpytorch.kernels.Kernel, optional): GP covariance kernel class. Defaults to scaled RBF with a gamma prior. Defaults to None. likelihood (Likelihood, optional): The likelihood function to use. If None defaults to Gaussian likelihood. Defaults to None. - inducing_size (int, optional): The number of inducing points to use. Defaults to None. + inducing_point_method (InducingPointAllocator, optional): The method to use for selecting inducing points. + If not set, a GreedyVarianceReduction is made. + inducing_size (int): The number of inducing points to use. Defaults to 100. max_fit_time (float, optional): The maximum amount of time, in seconds, to spend fitting the model. If None, there is no limit to the fitting time. Defaults to None. """ @@ -134,10 +134,10 @@ def __init__( self.monotonic_dims = [int(d) for d in monotonic_dims] self.mon_grid_size = monotonic_grid_size self.min_f_val = min_f_val - self.inducing_point_method = inducing_point_method + self.lb = lb + self.ub = ub + super().__init__( - lb=lb, - ub=ub, dim=dim, mean_module=mean_module, covar_module=covar_module, @@ -233,12 +233,15 @@ def from_config(cls, config: Config) -> MonotonicProjectionGP: """ classname = cls.__name__ - inducing_size = config.getint(classname, "inducing_size", fallback=None) + inducing_size = config.getint(classname, "inducing_size", fallback=100) lb = config.gettensor(classname, "lb") ub = config.gettensor(classname, "ub") dim = config.getint(classname, "dim", fallback=None) + if dim is None: + dim = get_dims(config) + mean_covar_factory = config.getobj( classname, "mean_covar_factory", fallback=default_mean_covar_factory ) @@ -247,7 +250,7 @@ def from_config(cls, config: Config) -> MonotonicProjectionGP: max_fit_time = config.getfloat(classname, "max_fit_time", fallback=None) inducing_point_method_class = config.getobj( - classname, "inducing_point_method", fallback=AutoAllocator + classname, "inducing_point_method", fallback=GreedyVarianceReduction ) # Check if allocator class has a `from_config` method if hasattr(inducing_point_method_class, "from_config"): diff --git a/aepsych/models/monotonic_rejection_gp.py b/aepsych/models/monotonic_rejection_gp.py index 5025ad0a9..de61f9468 100644 --- a/aepsych/models/monotonic_rejection_gp.py +++ b/aepsych/models/monotonic_rejection_gp.py @@ -11,7 +11,6 @@ from typing import Any, Dict, List, Optional, Sequence, Tuple import gpytorch -import numpy as np import torch from aepsych.acquisition.rejection_sampler import RejectionSampler from aepsych.config import Config @@ -19,14 +18,10 @@ from aepsych.kernels.rbf_partial_grad import RBFKernelPartialObsGrad from aepsych.means.constant_partial_grad import ConstantMeanPartialObsGrad from aepsych.models.base import AEPsychMixin -from aepsych.models.inducing_point_allocators import AutoAllocator, SobolAllocator -from aepsych.models.utils import select_inducing_points -from aepsych.utils import _process_bounds, get_optimizer_options, promote_0d +from aepsych.models.inducing_points import GreedyVarianceReduction, SobolAllocator +from aepsych.models.inducing_points.base import InducingPointAllocator +from aepsych.utils import _process_bounds, get_dims, get_optimizer_options, promote_0d from botorch.fit import fit_gpytorch_mll -from botorch.models.utils.inducing_point_allocators import ( - GreedyVarianceReduction, - InducingPointAllocator, -) from gpytorch.kernels import Kernel from gpytorch.likelihoods import BernoulliLikelihood, Likelihood from gpytorch.means import Mean @@ -34,7 +29,6 @@ from gpytorch.models import ApproximateGP from gpytorch.variational import CholeskyVariationalDistribution, VariationalStrategy from scipy.stats import norm -from torch import Tensor class MonotonicRejectionGP(AEPsychMixin, ApproximateGP): @@ -67,7 +61,7 @@ def __init__( num_induc: int = 25, num_samples: int = 250, num_rejection_samples: int = 5000, - inducing_point_method: InducingPointAllocator = AutoAllocator(), + inducing_point_method: Optional[InducingPointAllocator] = None, optimizer_options: Optional[Dict[str, Any]] = None, ) -> None: """Initialize MonotonicRejectionGP. @@ -88,7 +82,8 @@ def __init__( num_samples (int): Number of samples for estimating posterior on preDict or acquisition function evaluation. Defaults to 250. num_rejection_samples (int): Number of samples used for rejection sampling. Defaults to 4096. - inducing_point_method (InducingPointAllocator): Method for selecting inducing points. Defaults to AutoAllocator(). + inducing_point_method (InducingPointAllocator, optional): Method for selecting inducing points. If not set, + a GreedyVarianceReduction is created. optimizer_options (Dict[str, Any], optional): Optimizer options to pass to the SciPy optimizer during fitting. Assumes we are using L-BFGS-B. """ @@ -97,13 +92,17 @@ def __init__( likelihood = BernoulliLikelihood() self.inducing_size = num_induc - self.inducing_point_method = inducing_point_method - - inducing_points = select_inducing_points( - allocator=SobolAllocator(bounds=torch.stack((self.lb, self.ub))), - inducing_size=self.inducing_size, + self.inducing_point_method = inducing_point_method or GreedyVarianceReduction( + dim=self.dim ) + # TODO: This allocator *must* be SobolAllocator and not the set one. This + # suggests that this model doesn't actually properly use data for inducing + # points properly. + inducing_points = SobolAllocator( + bounds=torch.stack([lb, ub]), dim=self.dim + ).allocate_inducing_points(num_inducing=self.inducing_size) + inducing_points_aug = self._augment_with_deriv_index(inducing_points, 0) variational_distribution = CholeskyVariationalDistribution( inducing_points_aug.size(0) @@ -167,12 +166,10 @@ def fit(self, train_x: torch.Tensor, train_y: torch.Tensor, **kwargs) -> None: """ self.set_train_data(train_x, train_y) - self.inducing_points = select_inducing_points( - allocator=self.inducing_point_method, - inducing_size=self.inducing_size, + self.inducing_points = self.inducing_point_method.allocate_inducing_points( + num_inducing=self.inducing_size, covar_module=self.covar_module, - X=self.train_inputs[0], - bounds=self.bounds, + inputs=self._augment_with_deriv_index(self.train_inputs[0], 0), ) self._set_model(train_x, train_y) @@ -356,7 +353,7 @@ def from_config(cls, config: Config) -> MonotonicRejectionGP: lb = config.gettensor(classname, "lb") ub = config.gettensor(classname, "ub") - dim = config.getint(classname, "dim", fallback=None) + dim = get_dims(config) mean_covar_factory = config.getobj( classname, "mean_covar_factory", fallback=monotonic_mean_covar_factory diff --git a/aepsych/models/pairwise_probit.py b/aepsych/models/pairwise_probit.py index e8c22144c..f967e3594 100644 --- a/aepsych/models/pairwise_probit.py +++ b/aepsych/models/pairwise_probit.py @@ -12,7 +12,7 @@ from aepsych.config import Config from aepsych.factory import default_mean_covar_factory from aepsych.models.base import AEPsychMixin -from aepsych.utils import _process_bounds, get_optimizer_options, promote_0d +from aepsych.utils import _process_bounds, get_dims, get_optimizer_options, promote_0d from aepsych.utils_logging import getLogger from botorch.fit import fit_gpytorch_mll from botorch.models import PairwiseGP, PairwiseLaplaceMarginalLogLikelihood @@ -292,7 +292,7 @@ def from_config(cls, config: Config) -> "PairwiseProbitModel": lb = config.gettensor(classname, "lb") ub = config.gettensor(classname, "ub") - dim = lb.shape[0] + dim = get_dims(config) max_fit_time = config.getfloat(classname, "max_fit_time", fallback=None) diff --git a/aepsych/models/semi_p.py b/aepsych/models/semi_p.py index 9b2c8447d..9fd1adadc 100644 --- a/aepsych/models/semi_p.py +++ b/aepsych/models/semi_p.py @@ -7,6 +7,7 @@ from __future__ import annotations +import warnings from copy import deepcopy from typing import Any, Dict, Optional, Tuple @@ -18,11 +19,11 @@ from aepsych.config import Config from aepsych.likelihoods import BernoulliObjectiveLikelihood, LinearBernoulliLikelihood from aepsych.models import GPClassificationModel -from aepsych.models.inducing_point_allocators import AutoAllocator -from aepsych.utils import _process_bounds, get_optimizer_options, promote_0d +from aepsych.models.inducing_points import GreedyVarianceReduction +from aepsych.models.inducing_points.base import InducingPointAllocator +from aepsych.utils import get_dims, get_optimizer_options, promote_0d from aepsych.utils_logging import getLogger from botorch.acquisition.objective import PosteriorTransform -from botorch.models.utils.inducing_point_allocators import InducingPointAllocator from botorch.optim.fit import fit_gpytorch_mll_scipy from botorch.posteriors import GPyTorchPosterior from gpytorch.distributions import MultivariateNormal @@ -251,27 +252,21 @@ class SemiParametricGPModel(GPClassificationModel): def __init__( self, - lb: torch.Tensor, - ub: torch.Tensor, - inducing_point_method: InducingPointAllocator, - dim: Optional[int] = None, + dim: int, stim_dim: int = 0, mean_module: Optional[gpytorch.means.Mean] = None, covar_module: Optional[gpytorch.kernels.Kernel] = None, likelihood: Optional[Any] = None, slope_mean: float = 2, - inducing_size: Optional[int] = None, + inducing_point_method: Optional[InducingPointAllocator] = None, + inducing_size: int = 100, max_fit_time: Optional[float] = None, optimizer_options: Optional[Dict[str, Any]] = None, ) -> None: """ Initialize SemiParametricGP. Args: - lb (torch.Tensor): Lower bounds of the parameters. - ub (torch.Tensor): Upper bounds of the parameters. - inducing_point_method (InducingPointAllocator): The method to use to select the inducing points. - dim (int, optional): The number of dimensions in the parameter space. If None, it is inferred from the size - of lb and ub. Defaults to None. + dim (int, optional): The number of dimensions in the parameter space. stim_dim (int): Index of the intensity (monotonic) dimension. Defaults to 0. mean_module (gpytorch.means.Mean, optional): GP mean class. Defaults to a constant with a normal prior. covar_module (gpytorch.kernels.Kernel, optional): GP covariance kernel class. Defaults to scaled RBF with a @@ -279,16 +274,18 @@ def __init__( likelihood (gpytorch.likelihood.Likelihood, optional): The likelihood function to use. If None defaults to linear-Bernouli likelihood with probit link. slope_mean (float): The mean of the slope. Defaults to 2. - inducing_size (int, optional): Number of inducing points. Defaults to 99. + inducing_point_method (InducingPointAllocator, optional): The method to use for selecting inducing points. + If not set, a GreedyVarianceReduction is made. + inducing_size (int): Number of inducing points. Defaults to 100. max_fit_time (float, optional): The maximum amount of time, in seconds, to spend fitting the model. If None, there is no limit to the fitting time. optimizer_options (Dict[str, Any], optional): Optimizer options to pass to the SciPy optimizer during fitting. Assumes we are using L-BFGS-B. """ - lb, ub, dim = _process_bounds(lb, ub, dim) + self.dim = dim self.stim_dim = stim_dim - self.context_dims = list(range(dim)) + self.context_dims = list(range(self.dim)) self.context_dims.pop(stim_dim) if mean_module is None: @@ -301,7 +298,7 @@ def __init__( if covar_module is None: covar_module = ScaleKernel( RBFKernel( - ard_num_dims=dim - 1, + ard_num_dims=self.dim - 1, lengthscale_prior=GammaPrior(3, 6), active_dims=self.context_dims, # Operate only on x_s batch_shape=torch.Size([2]), @@ -313,11 +310,8 @@ def __init__( assert isinstance( likelihood, LinearBernoulliLikelihood ), "SemiP model only supports linear Bernoulli likelihoods!" - self.inducing_point_method = inducing_point_method super().__init__( - lb=lb, - ub=ub, dim=dim, mean_module=mean_module, covar_module=covar_module, @@ -343,16 +337,17 @@ def from_config(cls, config: Config) -> SemiParametricGPModel: """ classname = cls.__name__ - inducing_size = config.getint(classname, "inducing_size", fallback=None) + inducing_size = config.getint(classname, "inducing_size", fallback=100) - lb = config.gettensor(classname, "lb") - ub = config.gettensor(classname, "ub") dim = config.getint(classname, "dim", fallback=None) + if dim is None: + dim = get_dims(config) + max_fit_time = config.getfloat(classname, "max_fit_time", fallback=None) inducing_point_method_class = config.getobj( - classname, "inducing_point_method", fallback=AutoAllocator + classname, "inducing_point_method", fallback=GreedyVarianceReduction ) # Check if allocator class has a `from_config` method if hasattr(inducing_point_method_class, "from_config"): @@ -375,8 +370,6 @@ def from_config(cls, config: Config) -> SemiParametricGPModel: optimizer_options = get_optimizer_options(config, classname) return cls( - lb=lb, - ub=ub, stim_dim=stim_dim, dim=dim, likelihood=likelihood, @@ -519,10 +512,7 @@ class HadamardSemiPModel(GPClassificationModel): def __init__( self, - lb: torch.Tensor, - ub: torch.Tensor, - inducing_point_method: InducingPointAllocator, - dim: Optional[int] = None, + dim: int, stim_dim: int = 0, slope_mean_module: Optional[gpytorch.means.Mean] = None, slope_covar_module: Optional[gpytorch.kernels.Kernel] = None, @@ -530,18 +520,15 @@ def __init__( offset_covar_module: Optional[gpytorch.kernels.Kernel] = None, likelihood: Optional[Likelihood] = None, slope_mean: float = 2, - inducing_size: Optional[int] = None, + inducing_point_method: Optional[InducingPointAllocator] = None, + inducing_size: int = 100, max_fit_time: Optional[float] = None, optimizer_options: Optional[Dict[str, Any]] = None, ) -> None: """ Initialize HadamardSemiPModel. Args: - lb (torch.Tensor): Lower bounds of the parameters. - ub (torch.Tensor): Upper bounds of the parameters. - inducing_point_method (InducingPointAllocator): The method to use to select the inducing points. - dim (int, optional): The number of dimensions in the parameter space. If None, it is inferred from the size - of lb and ub. + dim (int): The number of dimensions in the parameter space. stim_dim (int): Index of the intensity (monotonic) dimension. Defaults to 0. slope_mean_module (gpytorch.means.Mean, optional): Mean module to use (default: constant mean) for slope. slope_covar_module (gpytorch.kernels.Kernel, optional): Covariance kernel to use (default: scaled RBF) for slope. @@ -549,16 +536,15 @@ def __init__( offset_covar_module (gpytorch.kernels.Kernel, optional): Covariance kernel to use (default: scaled RBF) for offset. likelihood (gpytorch.likelihood.Likelihood, optional)): defaults to bernoulli with logistic input and a floor of .5 slope_mean (float): The mean of the slope. Defaults to 2. - inducing_size (int, optional): Number of inducing points. Defaults to 99. + inducing_point_method (InducingPointAllocator, optional): The method to use for selecting inducing points. + If not set, a GreedyVarianceReduction is made. + inducing_size (int): Number of inducing points. Defaults to 100. max_fit_time (float, optional): The maximum amount of time, in seconds, to spend fitting the model. If None, there is no limit to the fitting time. optimizer_options (Dict[str, Any], optional): Optimizer options to pass to the SciPy optimizer during fitting. Assumes we are using L-BFGS-B. """ - self.inducing_point_method = inducing_point_method super().__init__( - lb=lb, - ub=ub, dim=dim, inducing_size=inducing_size, max_fit_time=max_fit_time, @@ -656,11 +642,11 @@ def from_config(cls, config: Config) -> HadamardSemiPModel: """ classname = cls.__name__ - inducing_size = config.getint(classname, "inducing_size", fallback=None) + inducing_size = config.getint(classname, "inducing_size", fallback=100) - lb = config.gettensor(classname, "lb") - ub = config.gettensor(classname, "ub") dim = config.getint(classname, "dim", fallback=None) + if dim is None: + dim = get_dims(config) slope_mean_module = config.getobj(classname, "slope_mean_module", fallback=None) slope_covar_module = config.getobj( @@ -676,7 +662,7 @@ def from_config(cls, config: Config) -> HadamardSemiPModel: max_fit_time = config.getfloat(classname, "max_fit_time", fallback=None) inducing_point_method_class = config.getobj( - classname, "inducing_point_method", fallback=AutoAllocator + classname, "inducing_point_method", fallback=GreedyVarianceReduction ) # Check if allocator class has a `from_config` method if hasattr(inducing_point_method_class, "from_config"): @@ -696,8 +682,6 @@ def from_config(cls, config: Config) -> HadamardSemiPModel: optimizer_options = get_optimizer_options(config, classname) return cls( - lb=lb, - ub=ub, stim_dim=stim_dim, dim=dim, slope_mean_module=slope_mean_module, diff --git a/aepsych/models/utils.py b/aepsych/models/utils.py index 28215f568..8677221c8 100644 --- a/aepsych/models/utils.py +++ b/aepsych/models/utils.py @@ -11,21 +11,18 @@ import numpy as np import torch -from aepsych.models.inducing_point_allocators import GreedyVarianceReduction +from aepsych.models.base import ModelProtocol +from aepsych.utils import dim_grid, get_jnd_multid from botorch.acquisition import PosteriorMean from botorch.acquisition.objective import ( PosteriorTransform, ScalarizedPosteriorTransform, ) from botorch.models.model import Model -from botorch.models.utils.inducing_point_allocators import InducingPointAllocator from botorch.optim import optimize_acqf from botorch.posteriors import GPyTorchPosterior -from botorch.utils.sampling import draw_sobol_samples from gpytorch.distributions import MultivariateNormal -from gpytorch.kernels import Kernel from gpytorch.likelihoods import BernoulliLikelihood, Likelihood -from scipy.cluster.vq import kmeans2 from scipy.special import owens_t from scipy.stats import norm from torch.distributions import Normal @@ -60,84 +57,6 @@ def compute_p_quantile( return norm.cdf(f_mean + norm.icdf(alpha) * f_std) -def select_inducing_points( - inducing_size: int, - allocator: Union[str, InducingPointAllocator], - covar_module: Optional[torch.nn.Module] = None, - X: Optional[torch.Tensor] = None, - bounds: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """ - Select inducing points using a specified allocator instance or legacy method. - - Args: - inducing_size (int): Number of inducing points. - allocator (Union[str, InducingPointAllocator]): An inducing point allocator or a legacy string indicating method. - covar_module (torch.nn.Module, optional): Covariance module, required for some allocators. - X (torch.Tensor, optional): Input data tensor, required for most allocators. - bounds (torch.Tensor, optional): Bounds for Sobol sampling in legacy mode. - - Returns: - torch.Tensor: Selected inducing points. - """ - # Handle legacy string methods with a deprecation warning - if isinstance(allocator, str): - warnings.warn( - f"Using string '{allocator}' for inducing point method is deprecated. " - "Please use an InducingPointAllocator class instead.", - DeprecationWarning, - ) - - if allocator == "sobol": - assert ( - bounds is not None - ), "Bounds must be provided for Sobol inducing points!" - inducing_points = ( - draw_sobol_samples(bounds=bounds, n=inducing_size, q=1) - .squeeze() - .to(bounds.device) - ) - if inducing_points.ndim == 1: - inducing_points = inducing_points.view(-1, 1) - return inducing_points - - assert X is not None, "Must pass X for non-Sobol inducing point selection!" - - unique_X = torch.unique(X, dim=0) - if allocator == "auto": - if unique_X.shape[0] <= inducing_size: - return unique_X - else: - allocator = "kmeans++" - - if allocator == "pivoted_chol": - inducing_point_allocator = GreedyVarianceReduction() - inducing_points = inducing_point_allocator.allocate_inducing_points( - inputs=X, - covar_module=covar_module, - num_inducing=inducing_size, - input_batch_shape=torch.Size([]), - ).to(X.device) - elif allocator == "kmeans++": - inducing_points = torch.tensor( - kmeans2(unique_X.cpu().numpy(), inducing_size, minit="++")[0], - dtype=X.dtype, - ).to(X.device) - - return inducing_points - - # Call allocate_inducing_points with allocator instance - if isinstance(allocator, InducingPointAllocator): - inducing_points = allocator.allocate_inducing_points( - inputs=X, - covar_module=covar_module, - num_inducing=inducing_size, - input_batch_shape=torch.Size([]), - ) - - return inducing_points - - def get_probability_space( likelihood: Likelihood, posterior: GPyTorchPosterior ) -> Tuple[torch.Tensor, torch.Tensor]: @@ -239,8 +158,77 @@ def get_extremum( return best_val, best_point.squeeze(0) +def get_min( + model: ModelProtocol, + bounds: torch.Tensor, + locked_dims: Optional[Mapping[int, float]] = None, + probability_space: bool = False, + n_samples: int = 1000, + max_time: Optional[float] = None, +) -> Tuple[float, torch.Tensor]: + """Return the minimum of the modeled function, subject to constraints + Args: + model (ModelProtocol): AEPsychModel to get the minimum of. + bounds (torch.Tensor): Bounds of the space to find the minimum. + locked_dims (Mapping[int, float], optional): Dimensions to fix, so that the + inverse is along a slice of the full surface. + probability_space (bool): Is y (and therefore the returned nearest_y) in + probability space instead of latent function space? Defaults to False. + n_samples (int): number of coarse grid points to sample for optimization estimate. + max_time (float, optional): Maximum time to spend optimizing. Defaults to None. + + Returns: + Tuple[float, torch.Tensor]: Tuple containing the min and its location (argmin). + """ + _, _arg = get_extremum( + model, "min", bounds, locked_dims, n_samples, max_time=max_time + ) + arg = torch.tensor(_arg.reshape(1, bounds.shape[1])) + if probability_space: + val, _ = model.predict_probability(arg) + else: + val, _ = model.predict(arg) + + return float(val.item()), arg + + +def get_max( + model: ModelProtocol, + bounds: torch.Tensor, + locked_dims: Optional[Mapping[int, float]] = None, + probability_space: bool = False, + n_samples: int = 1000, + max_time: Optional[float] = None, +) -> Tuple[float, torch.Tensor]: + """Return the maximum of the modeled function, subject to constraints + + Args: + model (ModelProtocol): AEPsychModel to get the maximum of. + bounds (torch.Tensor): Bounds of the space to find the maximum. + locked_dims (Mapping[int, float], optional): Dimensions to fix, so that the + inverse is along a slice of the full surface. Defaults to None. + probability_space (bool): Is y (and therefore the returned nearest_y) in + probability space instead of latent function space? Defaults to False. + n_samples (int): number of coarse grid points to sample for optimization estimate. + max_time (float, optional): Maximum time to spend optimizing. Defaults to None. + + Returns: + Tuple[float, torch.Tensor]: Tuple containing the max and its location (argmax). + """ + _, _arg = get_extremum( + model, "max", bounds, locked_dims, n_samples, max_time=max_time + ) + arg = torch.tensor(_arg.reshape(1, bounds.shape[1])) + if probability_space: + val, _ = model.predict_probability(arg) + else: + val, _ = model.predict(arg) + + return float(val.item()), arg + + def inv_query( - model: Model, + model: ModelProtocol, y: Union[float, torch.Tensor], bounds: torch.Tensor, locked_dims: Optional[Mapping[int, float]] = None, @@ -253,9 +241,10 @@ def inv_query( Return nearest x such that f(x) = queried y, and also return the value of f at that point. Args: + model (ModelProtocol): AEPsychModel to get the find the inverse from y. y (Union[float, torch.Tensor]): Points at which to find the inverse. bounds (torch.Tensor): Lower and upper bounds of the search space. - locked_dims (Mapping[int, List[float]], optional): Dimensions to fix, so that the + locked_dims (Mapping[int, float], optional): Dimensions to fix, so that the inverse is along a slice of the full surface. Defaults to None. probability_space (bool): Is y (and therefore the returned nearest_y) in probability space instead of latent @@ -268,9 +257,9 @@ def inv_query( nearest to queried y and the x position of this value. """ locked_dims = locked_dims or {} - if model.num_outputs > 1: + if model._num_outputs > 1: if weights is None: - weights = torch.Tensor([1] * model.num_outputs) + weights = torch.Tensor([1] * model._num_outputs) if probability_space: warnings.warn( "Inverse querying with probability_space=True assumes that the model uses Probit-Bernoulli likelihood!" @@ -278,7 +267,7 @@ def inv_query( posterior_transform = TargetProbabilityDistancePosteriorTransform(y, weights) else: posterior_transform = TargetDistancePosteriorTransform(y, weights) - val, arg = get_extremum( + _, _arg = get_extremum( model, "min", bounds, @@ -288,7 +277,115 @@ def inv_query( max_time, weights, ) - return val, arg + + arg = torch.tensor(_arg.reshape(1, bounds.shape[1])) + if probability_space: + val, _ = model.predict_probability(arg) + else: + val, _ = model.predict(arg) + + return float(val.item()), arg + + +def get_jnd( + model: ModelProtocol, + lb: torch.Tensor, + ub: torch.Tensor, + dim: int, + grid: Optional[Union[np.ndarray, torch.Tensor]] = None, + cred_level: Optional[float] = None, + intensity_dim: int = -1, + confsamps: int = 500, + method: str = "step", +) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + """Calculate the JND. + + Note that JND can have multiple plausible definitions + outside of the linear case, so we provide options for how to compute it. + For method="step", we report how far one needs to go over in stimulus + space to move 1 unit up in latent space (this is a lot of people's + conventional understanding of the JND). + For method="taylor", we report the local derivative, which also maps to a + 1st-order Taylor expansion of the latent function. This is a formal + generalization of JND as defined in Weber's law. + Both definitions are equivalent for linear psychometric functions. + + Args: + model (ModelProtocol): Model to use for prediction. + lb (torch.Tensor): Lower bounds of the input space. + ub (torch.Tensor): Upper bounds of the input space. + dim (int): Dimensionality of the input space. + grid (Optional[np.ndarray], optional): Mesh grid over which to find the JND. + Defaults to a square grid of size as determined by aepsych.utils.dim_grid + cred_level (float, optional): Credible level for computing an interval. + Defaults to None, computing no interval. + intensity_dim (int, optional): Dimension over which to compute the JND. + Defaults to -1. + confsamps (int, optional): Number of posterior samples to use for + computing the credible interval. Defaults to 500. + method (str, optional): "taylor" or "step" method (see docstring). + Defaults to "step". + + Raises: + RuntimeError: for passing an unknown method. + + Returns: + Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: either the + mean JND, or a median, lower, upper tuple of the JND posterior. + """ + if grid is None: + grid = dim_grid(lower=lb, upper=ub, gridsize=30, slice_dims=None) + elif isinstance(grid, np.ndarray): + grid = torch.tensor(grid) + + # this is super awkward, back into intensity dim grid assuming a square grid + gridsize = int(grid.shape[0] ** (1 / grid.shape[1])) + coords = torch.linspace( + lb[intensity_dim].item(), ub[intensity_dim].item(), gridsize + ) + + if cred_level is None: + fmean, _ = model.predict(grid) + fmean = fmean.reshape(*[gridsize for i in range(dim)]) + + if method == "taylor": + return torch.tensor(1 / np.gradient(fmean, coords, axis=intensity_dim)) + elif method == "step": + return torch.clip( + get_jnd_multid( + fmean, + coords, + mono_dim=intensity_dim, + ), + 0, + np.inf, + ) + + alpha = 1 - cred_level # type: ignore + qlower = alpha / 2 + qupper = 1 - alpha / 2 + + fsamps = model.sample(grid, confsamps) + if method == "taylor": + jnds = torch.tensor( + 1 + / np.gradient( + fsamps.reshape(confsamps, *[gridsize for i in range(dim)]), + coords, + axis=intensity_dim, + ) + ) + elif method == "step": + samps = [s.reshape((gridsize,) * dim) for s in fsamps] + jnds = torch.stack( + [get_jnd_multid(s, coords, mono_dim=intensity_dim) for s in samps] + ) + else: + raise RuntimeError(f"Unknown method {method}!") + upper = torch.clip(torch.quantile(jnds, qupper, axis=0), 0, np.inf) # type: ignore + lower = torch.clip(torch.quantile(jnds, qlower, axis=0), 0, np.inf) # type: ignore + median = torch.clip(torch.quantile(jnds, 0.5, axis=0), 0, np.inf) # type: ignore + return median, lower, upper class TargetDistancePosteriorTransform(PosteriorTransform): diff --git a/aepsych/plotting.py b/aepsych/plotting.py index 6d0414811..726d1fdc4 100644 --- a/aepsych/plotting.py +++ b/aepsych/plotting.py @@ -13,7 +13,6 @@ from aepsych.strategy import Strategy from aepsych.utils import get_lse_contour, get_lse_interval, make_scaled_sobol from matplotlib.axes import Axes - from matplotlib.image import AxesImage from scipy.stats import norm diff --git a/aepsych/strategy.py b/aepsych/strategy.py index 70b95cbe3..4a7c2ab39 100644 --- a/aepsych/strategy.py +++ b/aepsych/strategy.py @@ -9,19 +9,7 @@ import time import warnings -from typing import ( - Any, - Callable, - Dict, - List, - Literal, - Mapping, - Optional, - Sequence, - Tuple, - Type, - Union, -) +from typing import Any, Callable, List, Mapping, Optional, Tuple, Union import numpy as np import torch @@ -34,6 +22,7 @@ from aepsych.config import Config from aepsych.generators.base import AEPsychGenerator from aepsych.models.base import AEPsychMixin +from aepsych.models.utils import get_extremum, get_jnd, get_max, get_min, inv_query from aepsych.transforms import ( ParameterTransformedGenerator, ParameterTransformedModel, @@ -267,6 +256,7 @@ def __init__( ) self.name = name + self.bounds = torch.stack([self.lb, self.ub]) def normalize_inputs( self, x: torch.Tensor, y: torch.Tensor @@ -338,7 +328,7 @@ def get_max( probability_space: bool = False, max_time: Optional[float] = None, ) -> Tuple[float, torch.Tensor]: - """Get the maximum value of the acquisition function. + """Return the maximum of the modeled function, subject to constraints Args: constraints (Mapping[int, float], optional): Which parameters to fix at specfic points. Defaults to None. @@ -346,17 +336,23 @@ def get_max( max_time (float, optional): Maximum time to run the optimization. Defaults to None. Returns: - Tuple[float, torch.Tensor]: The maximum value of the acquisition function and the corresponding input. + Tuple[float, torch.Tensor]: Tuple containing the max and its location (argmax). """ - constraints = constraints or {} assert ( self.model is not None ), "model is None! Cannot get the max without a model!" self.model.to(self.model_device) - return self.model.get_max( - constraints, probability_space=probability_space, max_time=max_time + + val, arg = get_max( + self.model, + self.bounds, + locked_dims=constraints, + probability_space=probability_space, + max_time=max_time, ) + return val, arg + @ensure_model_is_fresh def get_min( self, @@ -364,25 +360,28 @@ def get_min( probability_space: bool = False, max_time: Optional[float] = None, ) -> Tuple[float, torch.Tensor]: - """Get the minimum value of the acquisition function. + """Return the minimum of the modeled function, subject to constraints Args: constraints (Mapping[int, float], optional): Which parameters to fix at specific points. Defaults to None. probability_space (bool): Whether to return the min in probability space. Defaults to False. max_time (float, optional): Maximum time to run the optimization. Defaults to None. - - Returns: - Tuple[float, torch.Tensor]: The minimum value of the acquisition function and the corresponding input. """ - constraints = constraints or {} assert ( self.model is not None ), "model is None! Cannot get the min without a model!" self.model.to(self.model_device) - return self.model.get_min( - constraints, probability_space=probability_space, max_time=max_time + + val, arg = get_min( + self.model, + self.bounds, + locked_dims=constraints, + probability_space=probability_space, + max_time=max_time, ) + return val, arg + @ensure_model_is_fresh def inv_query( self, @@ -402,15 +401,22 @@ def inv_query( Returns: Tuple[float, torch.Tensor]: The input that corresponds to the given output value and the corresponding output. """ - constraints = constraints or {} assert ( self.model is not None ), "model is None! Cannot get the inv_query without a model!" self.model.to(self.model_device) - return self.model.inv_query( - y, constraints, probability_space, max_time=max_time + + val, arg = inv_query( + model=self.model, + y=y, + bounds=self.bounds, + locked_dims=constraints, + probability_space=probability_space, + max_time=max_time, ) + return val, arg + @ensure_model_is_fresh def predict(self, x: torch.Tensor, probability_space: bool = False) -> torch.Tensor: """Predict the output value(s) for the given input(s). @@ -439,7 +445,9 @@ def get_jnd( self.model is not None ), "model is None! Cannot get the get jnd without a model!" self.model.to(self.model_device) - return self.model.get_jnd(*args, **kwargs) + return get_jnd( # type: ignore + model=self.model, lb=self.lb, ub=self.ub, dim=self.dim, *args, **kwargs + ) @ensure_model_is_fresh def sample( diff --git a/aepsych/transforms/parameters.py b/aepsych/transforms/parameters.py index b801cc963..f4c13d4d6 100644 --- a/aepsych/transforms/parameters.py +++ b/aepsych/transforms/parameters.py @@ -508,10 +508,6 @@ def __init__( """ # Alternative instantiation method for analysis (and not live) if isinstance(model, type): - if "lb" in kwargs: - kwargs["lb"] = transforms.transform(kwargs["lb"].to(torch.float64)) - if "ub" in kwargs: - kwargs["ub"] = transforms.transform(kwargs["ub"].to(torch.float64)) _base_obj = model(**kwargs) else: _base_obj = model @@ -612,20 +608,6 @@ def posterior(self, X: torch.Tensor, **kwargs) -> Posterior: X = torch.Tensor(X) return self._base_obj.posterior(X=X, **kwargs) - def dim_grid(self, gridsize: int = 30) -> torch.Tensor: - """Returns an untransformed grid based on the model's bounds and dimensionality. - - Args: - gridsize (int): How many points to form the grid with, defaults to - 30. - - Returns: - Tensor: A grid based on the model's bounds and dimensionality with the - number of points requested, which will be untransformed. - """ - grid = self._base_obj.dim_grid(gridsize) - return self.transforms.untransform(grid) - @_promote_1d def fit(self, train_x: torch.Tensor, train_y: torch.Tensor, **kwargs) -> None: """Fit underlying model. @@ -654,198 +636,6 @@ def update(self, train_x: torch.Tensor, train_y: torch.Tensor, **kwargs) -> None train_x = self.transforms.transform(train_x) self._base_obj.update(train_x, train_y, **kwargs) - def get_max( - self, - locked_dims: Optional[Mapping[int, float]] = None, - probability_space: bool = False, - n_samples: int = 1000, - max_time: Optional[float] = None, - ) -> Tuple[float, torch.Tensor]: - """Return the maximum of the modeled function, subject to constraints - - Args: - locked_dims (Mapping[int, float], optional): Dimensions to fix, so that the - max is along a slice of the full surface. Defaults to None. - probability_space (bool): Is y (and therefore the returned nearest_y) in - probability space instead of latent function space? Defaults to False. - n_samples (int): number of coarse grid points to sample for optimization estimate. - Defaults to 1000. - max_time (float, optional): Maximum time to spend optimizing. Defaults to None. - - Returns: - Tuple[float, torch.Tensor]: Tuple containing the max and its untransformed - location (argmax). - """ - locked_dims = locked_dims or {} - - # Transform locked dims - tmp = {} - for key, value in locked_dims.items(): - dims = list(self.transforms.values())[0]._d - tensor = torch.zeros(dims) - tensor[key] = value - tensor = self.transforms.transform(tensor) - tmp[key] = tensor[key].item() - locked_dims = tmp - - max_, loc = self._base_obj.get_max( # type: ignore - locked_dims=locked_dims, - probability_space=probability_space, - n_samples=n_samples, - max_time=max_time, - ) - loc = self.transforms.untransform(loc) - - return max_, loc - - def get_min( - self, - locked_dims: Optional[Mapping[int, float]] = None, - probability_space: bool = False, - n_samples: int = 1000, - max_time: Optional[float] = None, - ) -> Tuple[float, torch.Tensor]: - """Return the minimum of the modeled function, subject to constraints - - Args: - locked_dims (Mapping[int, float], optional): Dimensions to fix, so that the - min is along a slice of the full surface. - probability_space (bool): Is y (and therefore the returned nearest_y) in - probability space instead of latent function space? Defaults to False. - n_samples (int): number of coarse grid points to sample for optimization estimate. - Defaults to 1000. - max_time (float, optional): Maximum time to spend optimizing. Defaults to None. - - Returns: - Tuple[float, torch.Tensor]: Tuple containing the min and its untransformed location (argmin). - """ - locked_dims = locked_dims or {} - - # Transform locked dims - tmp = {} - for key, value in locked_dims.items(): - dims = list(self.transforms.values())[0]._d - tensor = torch.zeros(dims) - tensor[key] = value - tensor = self.transforms.transform(tensor) - tmp[key] = tensor[key].item() - locked_dims = tmp - - min_, loc = self._base_obj.get_min( # type: ignore - locked_dims=locked_dims, - probability_space=probability_space, - n_samples=n_samples, - max_time=max_time, - ) - loc = self.transforms.untransform(loc) - - return min_, loc - - def inv_query( - self, - y: float, - locked_dims: Optional[Mapping[int, float]] = None, - probability_space: bool = False, - n_samples: int = 1000, - max_time: Optional[float] = None, - weights: Optional[torch.Tensor] = None, - ) -> Tuple[float, torch.Tensor]: - """Query the model inverse. - - Return nearest untransformed x such that f(x) = queried y, and also return the - value of f at that point. - - Args: - y (float): Points at which to find the inverse. - locked_dims (Mapping[int, float], optional): Dimensions to fix, so that the - inverse is along a slice of the full surface. - probability_space (bool): Is y (and therefore the returned nearest_y) in - probability space instead of latent function space? Defaults to False. - n_samples (int): number of coarse grid points to sample for optimization estimate. Defaults to 1000. - max_time (float, optional): Maximum time to spend optimizing. Defaults to None. - weights (torch.Tensor, optional): Weights for the optimization. Defaults to None. - - Returns: - Tuple[float, torch.Tensor]: Tuple containing the value of f - nearest to queried y and the untransformed x position of this value. - """ - locked_dims = locked_dims or {} - - # Transform locked dims - tmp = {} - for key, value in locked_dims.items(): - dims = list(self.transforms.values())[0]._d - tensor = torch.zeros(dims) - tensor[key] = value - tensor = self.transforms.transform(tensor) - tmp[key] = tensor[key].item() - locked_dims = tmp - - val, loc = self._base_obj.inv_query( # type: ignore - y=y, - locked_dims=locked_dims, - probability_space=probability_space, - n_samples=n_samples, - max_time=max_time, - weights=weights, - ) - - loc = self.transforms.untransform(loc) - return val, loc - - def get_jnd( - self, - grid: Optional[Union[np.ndarray, torch.Tensor]] = None, - cred_level: Optional[float] = None, - intensity_dim: int = -1, - confsamps: int = 500, - method: str = "step", - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: - """Calculate the JND. - - Note that JND can have multiple plausible definitions - outside of the linear case, so we provide options for how to compute it. - For method="step", we report how far one needs to go over in stimulus - space to move 1 unit up in latent space (this is a lot of people's - conventional understanding of the JND). - For method="taylor", we report the local derivative, which also maps to a - 1st-order Taylor expansion of the latent function. This is a formal - generalization of JND as defined in Weber's law. - Both definitions are equivalent for linear psychometric functions. - - Args: - grid (torch.Tensor, optional): Untransformed mesh grid over which to find the JND. - Defaults to a square grid of size as determined by aepsych.utils.dim_grid. - cred_level (float, optional): Credible level for computing an interval. - Defaults to None, computing no interval. - intensity_dim (int): Dimension over which to compute the JND. - Defaults to -1. - confsamps (int): Number of posterior samples to use for - computing the credible interval. Defaults to 500. - method (str): "taylor" or "step" method (see docstring). - Defaults to "step". - - Returns: - Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: either the - mean JND, or a median, lower, upper tuple of the JND posterior. All values - are in the untransformed space. - """ - jnds = self._base_obj.get_jnd( # type: ignore - grid=grid, - cred_level=cred_level, - intensity_dim=intensity_dim, - confsamps=confsamps, - method=method, - ) - - if isinstance(jnds, torch.Tensor): - jnds = self.transforms.untransform(jnds) - else: - jnds = [self.transforms.untransform(jnd) for jnd in jnds] - jnds = tuple(jnds) - - return jnds - def p_below_threshold( self, x: torch.Tensor, f_thresh: torch.Tensor ) -> torch.Tensor: diff --git a/aepsych/utils.py b/aepsych/utils.py index c0f723e83..c1474e711 100644 --- a/aepsych/utils.py +++ b/aepsych/utils.py @@ -405,3 +405,29 @@ def get_optimizer_options(config: Config, name: str) -> Dict[str, Any]: # Filter all the nones out, which could just come back as an empty dict options = {key: value for key, value in options.items() if value is not None} return options + + +def get_dims(config: Config) -> int: + """Return the number of dimensions in the parameter space. This accounts for any + transforms that may modify the the parameter space for the model (e.g., Fixed + parameters will not be included). + + Args: + config (Config): The config to look for the number of dimensions. + + Return: + int: The number of dimensions in the search space. + """ + parnames = config.getlist("common", "parnames", element_type=str) + + # This is pretty weak but fixed is currently the only thing that will changed the + # search space dims, when categorial transforms go in, this function needs to be + # smarter. + try: + valid_pars = [ + parname for parname in parnames if config[parname]["par_type"] != "fixed" + ] + return len(valid_pars) + except KeyError: + # Likely old style of parameter definition, fallback to looking at a bound + return len(config.getlist("common", "lb", element_type=float)) diff --git a/pubs/owenetal/code/prior_plots.py b/pubs/owenetal/code/prior_plots.py index 2441a5a56..efa125f76 100644 --- a/pubs/owenetal/code/prior_plots.py +++ b/pubs/owenetal/code/prior_plots.py @@ -80,6 +80,8 @@ def plot_prior_samps_1d(): mono_mean, mono_covar = monotonic_mean_covar_factory(config) mono_model = MonotonicRejectionGP( + lb=lb, + ub=ub, likelihood="probit-bernoulli", monotonic_idxs=[0], mean_module=mono_mean, @@ -176,6 +178,8 @@ def plot_prior_samps_2d(): mono_mean, mono_covar = monotonic_mean_covar_factory(config) mono_model = MonotonicRejectionGP( + lb=lb, + ub=ub, likelihood="probit-bernoulli", monotonic_idxs=[1], mean_module=mono_mean, diff --git a/tests/acquisition/test_mi.py b/tests/acquisition/test_mi.py index 316571519..bd1ecc9ad 100644 --- a/tests/acquisition/test_mi.py +++ b/tests/acquisition/test_mi.py @@ -21,8 +21,6 @@ SobolGenerator, ) from aepsych.models import GPClassificationModel, MonotonicRejectionGP - -from aepsych.models.inducing_point_allocators import AutoAllocator from aepsych.strategy import SequentialStrategy, Strategy from gpytorch.kernels import LinearKernel from gpytorch.means import ConstantMean @@ -36,8 +34,10 @@ def test_1d_monotonic_single_probit(self): np.random.seed(seed) n_init = 15 n_opt = 1 - lb = -4.0 - ub = 4.0 + lb = torch.tensor([-4.0]) + ub = torch.tensor([4.0]) + inducing_size = 10 + acqf = MonotonicBernoulliMCMutualInformation acqf_kwargs = {"objective": ProbitObjective()} model_list = [ @@ -56,13 +56,12 @@ def test_1d_monotonic_single_probit(self): model=MonotonicRejectionGP( lb=lb, ub=ub, - dim=1, monotonic_idxs=[0], - inducing_point_method=AutoAllocator( - bounds=torch.stack((torch.Tensor([lb]), torch.Tensor([ub]))) - ), + num_induc=inducing_size, + ), + generator=MonotonicRejectionGenerator( + lb=lb, ub=ub, acqf=acqf, acqf_kwargs=acqf_kwargs ), - generator=MonotonicRejectionGenerator(acqf, acqf_kwargs), stimuli_per_trial=1, outcome_types=["binary"], ), @@ -91,8 +90,10 @@ def test_1d_single_probit(self): np.random.seed(seed) n_init = 15 n_opt = 20 - lb = -4.0 - ub = 4.0 + lb = torch.tensor([-4.0]) + ub = torch.tensor([4.0]) + inducing_size = 10 + acqf = BernoulliMCMutualInformation extra_acqf_args = {"objective": ProbitObjective()} @@ -109,15 +110,12 @@ def test_1d_single_probit(self): lb=lb, ub=ub, model=GPClassificationModel( - lb=lb, - ub=ub, + inducing_size=inducing_size, dim=1, - inducing_size=10, - inducing_point_method=AutoAllocator( - bounds=torch.stack((torch.Tensor([lb]), torch.Tensor([ub]))) - ), ), - generator=OptimizeAcqfGenerator(acqf, extra_acqf_args), + generator=OptimizeAcqfGenerator( + lb=lb, ub=ub, acqf=acqf, acqf_kwargs=extra_acqf_args + ), min_asks=n_opt, stimuli_per_trial=1, outcome_types=["binary"], @@ -144,15 +142,15 @@ def test_1d_single_probit(self): def test_mi_acqf(self): mean = ConstantMean().initialize(constant=1.2) covar = LinearKernel().initialize(variance=1.0) + lb = torch.tensor([0.0]) + ub = torch.tensor([1.0]) + inducing_size = 10 + model = GPClassificationModel( - lb=torch.Tensor([0]), - ub=torch.Tensor([1]), - inducing_size=10, + dim=1, + inducing_size=inducing_size, mean_module=mean, covar_module=covar, - inducing_point_method=AutoAllocator( - bounds=torch.stack((torch.Tensor([0]), torch.Tensor([1]))) - ), ) x = torch.rand(size=(10, 1)) acqf = BernoulliMCMutualInformation(model=model, objective=ProbitObjective()) diff --git a/tests/generators/test_epsilon_greedy_generator.py b/tests/generators/test_epsilon_greedy_generator.py index d418c378f..4dfb31404 100644 --- a/tests/generators/test_epsilon_greedy_generator.py +++ b/tests/generators/test_epsilon_greedy_generator.py @@ -22,13 +22,17 @@ def test_epsilon_greedy(self): np.random.seed(seed) total_trials = 2000 extra_acqf_args = {"target": 0.75, "beta": 1.96} + lb = torch.tensor([0.0]) + ub = torch.tensor([1.0]) for epsilon in (0.1, 0.5): gen = EpsilonGreedyGenerator( subgenerator=MonotonicRejectionGenerator( - acqf=MonotonicMCLSE, acqf_kwargs=extra_acqf_args + acqf=MonotonicMCLSE, acqf_kwargs=extra_acqf_args, lb=lb, ub=ub ), epsilon=epsilon, + lb=lb, + ub=ub, ) model = MagicMock() gen.subgenerator.gen = MagicMock() @@ -44,6 +48,8 @@ def test_greedyepsilon_config(self): config_str = """ [common] acqf = MonotonicMCLSE + lb = [0] + ub = [1] [EpsilonGreedyGenerator] subgenerator = MonotonicRejectionGenerator epsilon = .5 diff --git a/tests/generators/test_optimize_acqf_generator.py b/tests/generators/test_optimize_acqf_generator.py index fc4b97f63..fbb53f7ac 100644 --- a/tests/generators/test_optimize_acqf_generator.py +++ b/tests/generators/test_optimize_acqf_generator.py @@ -14,8 +14,6 @@ from aepsych.config import Config from aepsych.generators import OptimizeAcqfGenerator from aepsych.models import GPClassificationModel, PairwiseProbitModel - -from aepsych.models.inducing_point_allocators import AutoAllocator from botorch.acquisition.preference import AnalyticExpectedUtilityOfBestOption from sklearn.datasets import make_classification @@ -35,20 +33,23 @@ def test_time_limits(self): n_clusters_per_class=4, ) X, y = torch.Tensor(X), torch.Tensor(y) + lb = -3 * torch.ones(8) + ub = 3 * torch.ones(8) + inducing_size = 10 + bounds = torch.stack([lb, ub]) model = GPClassificationModel( - lb=-3 * torch.ones(8), - ub=3 * torch.ones(8), + dim=8, max_fit_time=0.5, - inducing_size=10, - inducing_point_method=AutoAllocator( - bounds=torch.stack((-3 * torch.ones(8), 3 * torch.ones(8))) - ), + inducing_size=inducing_size, ) model.fit(X, y) generator = OptimizeAcqfGenerator( - acqf=MCLevelSetEstimation, acqf_kwargs={"beta": 1.96, "target": 0.5} + acqf=MCLevelSetEstimation, + acqf_kwargs={"beta": 1.96, "target": 0.5}, + lb=lb, + ub=ub, ) start = time.time() @@ -59,6 +60,8 @@ def test_time_limits(self): acqf=MCLevelSetEstimation, acqf_kwargs={"beta": 1.96, "target": 0.5}, max_gen_time=0.1, + lb=lb, + ub=ub, ) start = time.time() @@ -73,6 +76,10 @@ def test_time_limits(self): def test_instantiate_eubo(self): config = """ + [common] + lb = [-1] + ub = [1] + [OptimizeAcqfGenerator] acqf = AnalyticExpectedUtilityOfBestOption stimuli_per_trial = 2 diff --git a/tests/models/test_gp_classification.py b/tests/models/test_gp_classification.py index 58b288426..fa3a25d60 100644 --- a/tests/models/test_gp_classification.py +++ b/tests/models/test_gp_classification.py @@ -10,8 +10,6 @@ import torch -from aepsych.models.inducing_point_allocators import AutoAllocator, SobolAllocator - # run on single threads to keep us from deadlocking weirdly in CI if "CI" in os.environ or "SANDCASTLE" in os.environ: torch.set_num_threads(1) @@ -120,14 +118,11 @@ def test_1d_classification(self): Just see if we memorize the training set """ X, y = self.X, self.y + inducing_size = 10 model = GPClassificationModel( - torch.Tensor([-3]), - torch.Tensor([3]), - inducing_size=10, - inducing_point_method=AutoAllocator( - bounds=torch.stack((torch.tensor([-3]), torch.tensor([3]))) - ), + dim=1, + inducing_size=inducing_size, ) model.fit(X[:50], y[:50]) @@ -160,13 +155,11 @@ def test_1d_classification_pytorchopt(self): Just see if we memorize the training set """ X, y = self.X, self.y + inducing_size = 10 + model = GPClassificationModel( - torch.Tensor([-3]), - torch.Tensor([3]), - inducing_size=10, - inducing_point_method=SobolAllocator( - bounds=torch.stack((torch.tensor([-3]), torch.tensor([3]))) - ), + dim=1, + inducing_size=inducing_size, ) model.fit( @@ -223,19 +216,18 @@ def test_1d_classification_different_scales(self): X, y = torch.Tensor(X), torch.Tensor(y) X[:, 0] = X[:, 0] * 1000 X[:, 1] = X[:, 1] / 1000 - lb = torch.tensor([-3000, -0.003]) - ub = torch.tensor([3000, 0.003]) + lb = torch.tensor([-3000.0, -0.003]) + ub = torch.tensor([3000.0, 0.003]) + inducing_size = 20 transforms = ParameterTransforms( normalize=NormalizeScale(2, bounds=torch.stack((lb, ub))) ) model = ParameterTransformedModel( model=GPClassificationModel, - lb=lb, - ub=ub, - inducing_size=20, + inducing_size=inducing_size, transforms=transforms, - inducing_point_method=AutoAllocator(bounds=torch.stack((lb, ub))), + dim=2, ) model.fit(X[:50], y[:50]) @@ -264,12 +256,8 @@ def test_1d_classification_different_scales(self): def test_reset_hyperparams(self): model = GPClassificationModel( - lb=[-3], - ub=[3], + dim=1, inducing_size=20, - inducing_point_method=SobolAllocator( - bounds=torch.stack((torch.tensor([-3]), torch.tensor([3]))) - ), ) ls_before = model.covar_module.lengthscale.clone().detach().numpy() @@ -289,12 +277,8 @@ def test_reset_hyperparams(self): def test_reset_variational_strategy(self): model = GPClassificationModel( - lb=[-3], - ub=[3], + dim=1, inducing_size=20, - inducing_point_method=SobolAllocator( - bounds=torch.stack((torch.tensor([-3]), torch.tensor([3]))) - ), ) variational_params_before = [ @@ -319,25 +303,28 @@ def test_reset_variational_strategy(self): # before should be different from after and after should be different # from reset self.assertFalse(np.allclose(induc_before, induc_after)) - self.assertFalse(np.allclose(induc_after, induc_reset)) + if ( + induc_after.shape == induc_reset.shape + ): # If they're not the same shape, we definitely can't fail the assertFalse + self.assertFalse(np.allclose(induc_after, induc_reset)) + for before, after in zip(variational_params_before, variational_params_after): self.assertFalse(np.allclose(before, after)) for after, reset in zip(variational_params_after, variational_params_reset): - self.assertFalse(np.allclose(after, reset)) + if after.shape == reset.shape: # Same as above + self.assertFalse(np.allclose(after, reset)) def test_predict_p(self): """ Verify analytic p-space mean and var is correct. """ X, y = self.X, self.y + inducing_size = 10 + model = GPClassificationModel( - torch.Tensor([-3]), - torch.Tensor([3]), - inducing_size=10, - inducing_point_method=SobolAllocator( - bounds=torch.stack((torch.tensor([-3]), torch.tensor([3]))) - ), + dim=1, + inducing_size=inducing_size, ) model.fit(X, y) @@ -359,8 +346,9 @@ def test_1d_single_probit_new_interface(self): np.random.seed(seed) n_init = 50 n_opt = 1 - lb = -4.0 - ub = 4.0 + lb = torch.tensor([-4.0]) + ub = torch.tensor([4.0]) + inducing_size = 10 model_list = [ Strategy( @@ -375,15 +363,11 @@ def test_1d_single_probit_new_interface(self): lb=lb, ub=ub, model=GPClassificationModel( - lb=lb, - ub=ub, - inducing_size=10, - inducing_point_method=SobolAllocator( - bounds=torch.stack((torch.tensor([-4]), torch.tensor([4]))) - ), + dim=1, + inducing_size=inducing_size, ), generator=OptimizeAcqfGenerator( - qUpperConfidenceBound, acqf_kwargs={"beta": 1.96} + acqf=qUpperConfidenceBound, acqf_kwargs={"beta": 1.96}, lb=lb, ub=ub ), min_asks=n_opt, stimuli_per_trial=1, @@ -412,8 +396,9 @@ def test_1d_single_probit_batched(self): np.random.seed(seed) n_init = 50 n_opt = 2 - lb = -4.0 - ub = 4.0 + lb = torch.tensor([-4.0]) + ub = torch.tensor([4.0]) + inducing_size = 10 model_list = [ Strategy( @@ -428,15 +413,11 @@ def test_1d_single_probit_batched(self): lb=lb, ub=ub, model=GPClassificationModel( - lb=lb, - ub=ub, - inducing_size=10, - inducing_point_method=SobolAllocator( - bounds=torch.stack((torch.tensor([-4]), torch.tensor([4]))) - ), + dim=1, + inducing_size=inducing_size, ), generator=OptimizeAcqfGenerator( - qUpperConfidenceBound, acqf_kwargs={"beta": 1.96} + acqf=qUpperConfidenceBound, acqf_kwargs={"beta": 1.96}, lb=lb, ub=ub ), min_asks=n_opt, stimuli_per_trial=1, @@ -464,8 +445,9 @@ def test_1d_single_probit(self): np.random.seed(seed) n_init = 50 n_opt = 1 - lb = -4.0 - ub = 4.0 + lb = torch.tensor([-4.0]) + ub = torch.tensor([4.0]) + inducing_size = 10 model_list = [ Strategy( @@ -480,15 +462,11 @@ def test_1d_single_probit(self): lb=lb, ub=ub, model=GPClassificationModel( - lb=lb, - ub=ub, - inducing_size=10, - inducing_point_method=SobolAllocator( - bounds=torch.stack((torch.tensor([-4]), torch.tensor([4]))) - ), + dim=1, + inducing_size=inducing_size, ), generator=OptimizeAcqfGenerator( - qUpperConfidenceBound, acqf_kwargs={"beta": 1.96} + acqf=qUpperConfidenceBound, acqf_kwargs={"beta": 1.96}, lb=lb, ub=ub ), min_asks=n_opt, stimuli_per_trial=1, @@ -515,8 +493,9 @@ def test_1d_single_probit_pure_exploration(self): np.random.seed(seed) n_init = 50 n_opt = 1 - lb = -4.0 - ub = 4.0 + lb = torch.tensor([-4.0]) + ub = torch.tensor([4.0]) + inducing_size = 10 strat_list = [ Strategy( @@ -531,15 +510,11 @@ def test_1d_single_probit_pure_exploration(self): lb=lb, ub=ub, model=GPClassificationModel( - lb=lb, - ub=ub, - inducing_size=10, - inducing_point_method=SobolAllocator( - bounds=torch.stack((torch.tensor([-4]), torch.tensor([4]))) - ), + dim=1, + inducing_size=inducing_size, ), generator=OptimizeAcqfGenerator( - qUpperConfidenceBound, acqf_kwargs={"beta": 1.96} + acqf=qUpperConfidenceBound, acqf_kwargs={"beta": 1.96}, lb=lb, ub=ub ), min_asks=n_opt, stimuli_per_trial=1, @@ -571,8 +546,9 @@ def test_2d_single_probit_pure_exploration(self): np.random.seed(seed) n_init = 50 n_opt = 1 - lb = [-1, -1] - ub = [1, 1] + lb = torch.tensor([-1.0, -1.0]) + ub = torch.tensor([1.0, 1.0]) + inducing_size = 10 strat_list = [ Strategy( @@ -587,17 +563,11 @@ def test_2d_single_probit_pure_exploration(self): lb=lb, ub=ub, model=GPClassificationModel( - lb=lb, - ub=ub, - inducing_size=10, - inducing_point_method=SobolAllocator( - bounds=torch.stack( - (torch.tensor([-1, -1]), torch.tensor([1, 1])) - ) - ), + dim=2, + inducing_size=inducing_size, ), generator=OptimizeAcqfGenerator( - qUpperConfidenceBound, acqf_kwargs={"beta": 1.96} + acqf=qUpperConfidenceBound, acqf_kwargs={"beta": 1.96}, lb=lb, ub=ub ), min_asks=n_opt, stimuli_per_trial=1, @@ -627,8 +597,9 @@ def test_1d_single_targeting(self): np.random.seed(seed) n_init = 50 n_opt = 1 - lb = -4.0 - ub = 4.0 + lb = torch.tensor([-4.0]) + ub = torch.tensor([4.0]) + inducing_size = 10 target = 0.75 @@ -648,15 +619,11 @@ def obj(x): lb=lb, ub=ub, model=GPClassificationModel( - lb=lb, - ub=ub, - inducing_size=10, - inducing_point_method=SobolAllocator( - bounds=torch.stack((torch.tensor([-4]), torch.tensor([4]))) - ), + dim=1, + inducing_size=inducing_size, ), generator=OptimizeAcqfGenerator( - qUpperConfidenceBound, acqf_kwargs={"beta": 1.96} + acqf=qUpperConfidenceBound, acqf_kwargs={"beta": 1.96}, lb=lb, ub=ub ), min_asks=n_opt, stimuli_per_trial=1, @@ -685,8 +652,9 @@ def test_1d_jnd(self): np.random.seed(seed) n_init = 150 n_opt = 1 - lb = -4.0 - ub = 4.0 + lb = torch.tensor([-4.0]) + ub = torch.tensor([4.0]) + inducing_size = 10 target = 0.5 @@ -706,15 +674,11 @@ def obj(x): lb=lb, ub=ub, model=GPClassificationModel( - lb=lb, - ub=ub, - inducing_size=10, - inducing_point_method=SobolAllocator( - bounds=torch.stack((torch.tensor([-4]), torch.tensor([4]))) - ), + dim=1, + inducing_size=inducing_size, ), generator=OptimizeAcqfGenerator( - qUpperConfidenceBound, acqf_kwargs={"beta": 1.96} + acqf=qUpperConfidenceBound, acqf_kwargs={"beta": 1.96}, lb=lb, ub=ub ), min_asks=n_opt, stimuli_per_trial=1, @@ -750,8 +714,9 @@ def test_1d_single_lse(self): np.random.seed(seed) n_init = 50 n_opt = 1 - lb = -4.0 - ub = 4.0 + lb = torch.tensor([-4.0]) + ub = torch.tensor([4.0]) + inducing_size = 10 # target is in z space not phi(z) space, maybe that's # weird @@ -770,16 +735,12 @@ def test_1d_single_lse(self): lb=lb, ub=ub, model=GPClassificationModel( - lb=lb, - ub=ub, - inducing_size=10, - inducing_point_method=SobolAllocator( - bounds=torch.stack((torch.tensor([-4]), torch.tensor([4]))) - ), + dim=1, + inducing_size=inducing_size, ), min_asks=n_opt, generator=OptimizeAcqfGenerator( - MCLevelSetEstimation, acqf_kwargs=extra_acqf_args + acqf=MCLevelSetEstimation, acqf_kwargs=extra_acqf_args, lb=lb, ub=ub ), stimuli_per_trial=1, outcome_types=["binary"], @@ -806,8 +767,9 @@ def test_2d_single_probit(self): np.random.seed(seed) n_init = 150 n_opt = 1 - lb = [-1, -1] - ub = [1, 1] + lb = torch.tensor([-1.0, -1.0]) + ub = torch.tensor([1.0, 1.0]) + inducing_size = 20 strat_list = [ Strategy( @@ -822,15 +784,11 @@ def test_2d_single_probit(self): lb=lb, ub=ub, model=GPClassificationModel( - lb=lb, - ub=ub, - inducing_size=20, - inducing_point_method=AutoAllocator( - bounds=torch.stack((torch.tensor([-1]), torch.tensor([1]))) - ), + dim=2, + inducing_size=inducing_size, ), generator=OptimizeAcqfGenerator( - qUpperConfidenceBound, acqf_kwargs={"beta": 1.96} + acqf=qUpperConfidenceBound, acqf_kwargs={"beta": 1.96}, lb=lb, ub=ub ), min_asks=n_opt, stimuli_per_trial=1, @@ -856,8 +814,9 @@ def test_extra_ask_warns(self): np.random.seed(seed) n_init = 3 n_opt = 1 - lb = -4.0 - ub = 4.0 + lb = torch.tensor([-4.0]) + ub = torch.tensor([4.0]) + inducing_size = 10 model_list = [ Strategy( @@ -872,15 +831,11 @@ def test_extra_ask_warns(self): lb=lb, ub=ub, model=GPClassificationModel( - lb=lb, - ub=ub, - inducing_size=10, - inducing_point_method=SobolAllocator( - bounds=torch.stack((torch.tensor([-4]), torch.tensor([4]))) - ), + dim=1, + inducing_size=inducing_size, ), generator=OptimizeAcqfGenerator( - qUpperConfidenceBound, acqf_kwargs={"beta": 1.96} + acqf=qUpperConfidenceBound, acqf_kwargs={"beta": 1.96}, lb=lb, ub=ub ), min_asks=n_opt, stimuli_per_trial=1, @@ -899,13 +854,9 @@ def test_extra_ask_warns(self): def test_hyperparam_consistency(self): # verify that creating the model `from_config` or with `__init__` has the same hyperparams - m1 = GPClassificationModel( - lb=[1, 2], - ub=[3, 4], - inducing_point_method=SobolAllocator( - bounds=torch.stack((torch.tensor([1, 2]), torch.tensor([3, 4]))) - ), + dim=2, + inducing_size=2, ) config = Config( @@ -914,6 +865,7 @@ def test_hyperparam_consistency(self): "parnames": ["par1", "par2"], "lb": "[1, 2]", "ub": "[3, 4]", + "inducing_size": 2, }, "par1": {"value_type": "float"}, "par2": {"value_type": "float"}, diff --git a/tests/models/test_gp_regression.py b/tests/models/test_gp_regression.py index 263a1e58c..0198f34d0 100644 --- a/tests/models/test_gp_regression.py +++ b/tests/models/test_gp_regression.py @@ -59,6 +59,7 @@ def setUp(self): [GPRegressionModel] likelihood = GaussianLikelihood max_fit_time = 1 + dim = 1 """ self.server = AEPsychServer(database_path=dbname) configure(self.server, config_str=config) @@ -88,9 +89,9 @@ def test_extremum(self): def test_from_config(self): model = self.server.strat.model - npt.assert_allclose(model.transforms.untransform(model.lb), [-1.0]) - npt.assert_allclose(model.transforms.untransform(model.ub), [3.0]) - self.assertEqual(model.dim, 1) + generator = self.server.strat.generator + npt.assert_allclose(model.transforms.untransform(generator.lb), [-1.0]) + npt.assert_allclose(model.transforms.untransform(generator.ub), [3.0]) self.assertIsInstance(model.likelihood, GaussianLikelihood) self.assertEqual(model.max_fit_time, 1) diff --git a/tests/models/test_monotonic_rejection_gp.py b/tests/models/test_monotonic_rejection_gp.py index 083b8023b..b48f445af 100644 --- a/tests/models/test_monotonic_rejection_gp.py +++ b/tests/models/test_monotonic_rejection_gp.py @@ -9,13 +9,10 @@ import torch -from aepsych.models.inducing_point_allocators import AutoAllocator - # run on single threads to keep us from deadlocking weirdly in CI if "CI" in os.environ or "SANDCASTLE" in os.environ: torch.set_num_threads(1) - from aepsych.acquisition.monotonic_rejection import MonotonicMCLSE from aepsych.acquisition.objective import ProbitObjective from aepsych.generators import MonotonicRejectionGenerator @@ -32,18 +29,19 @@ def testRegression(self): # Init target = 1.5 model_gen_options = {"num_restarts": 1, "raw_samples": 3, "epochs": 5} - lb = torch.tensor([0, 0]) - ub = torch.tensor([4, 4]) + lb = torch.tensor([0.0, 0.0]) + ub = torch.tensor([4.0, 4.0]) + inducing_size = 2 + m = MonotonicRejectionGP( lb=lb, ub=ub, likelihood=GaussianLikelihood(), fixed_prior_mean=target, monotonic_idxs=[1], - num_induc=2, + num_induc=inducing_size, num_samples=3, num_rejection_samples=4, - inducing_point_method=AutoAllocator(bounds=torch.stack((lb, ub))), ) strat = Strategy( lb=lb, @@ -53,6 +51,8 @@ def testRegression(self): MonotonicMCLSE, acqf_kwargs={"target": target}, model_gen_options=model_gen_options, + lb=lb, + ub=ub, ), min_asks=1, stimuli_per_trial=1, @@ -85,18 +85,19 @@ def testClassification(self): # Init target = 0.75 model_gen_options = {"num_restarts": 1, "raw_samples": 3, "epochs": 5} - lb = torch.tensor([0, 0]) - ub = torch.tensor([4, 4]) + lb = torch.tensor([0.0, 0.0]) + ub = torch.tensor([4.0, 4.0]) + inducing_size = 2 + m = MonotonicRejectionGP( lb=lb, ub=ub, likelihood=BernoulliLikelihood(), fixed_prior_mean=target, monotonic_idxs=[1], - num_induc=2, + num_induc=inducing_size, num_samples=3, num_rejection_samples=4, - inducing_point_method=AutoAllocator(bounds=torch.stack((lb, ub))), ) strat = Strategy( lb=lb, @@ -106,6 +107,8 @@ def testClassification(self): MonotonicMCLSE, acqf_kwargs={"target": target, "objective": ProbitObjective()}, model_gen_options=model_gen_options, + lb=lb, + ub=ub, ), min_asks=1, stimuli_per_trial=1, diff --git a/tests/models/test_multitask_regression.py b/tests/models/test_multitask_regression.py index e4cfbaa55..9bd54d164 100644 --- a/tests/models/test_multitask_regression.py +++ b/tests/models/test_multitask_regression.py @@ -19,8 +19,19 @@ torch.set_num_threads(1) models = [ - (MultitaskGPRModel(num_outputs=2, rank=2, lb=[-1], ub=[3]),), - (IndependentMultitaskGPRModel(num_outputs=2, lb=[-1], ub=[3]),), + ( + MultitaskGPRModel( + num_outputs=2, + rank=2, + dim=1, + ), + ), + ( + IndependentMultitaskGPRModel( + num_outputs=2, + dim=1, + ), + ), ] diff --git a/tests/models/test_pairwise_probit.py b/tests/models/test_pairwise_probit.py index ce7fb15a0..de86a47af 100644 --- a/tests/models/test_pairwise_probit.py +++ b/tests/models/test_pairwise_probit.py @@ -28,6 +28,7 @@ ParameterTransforms, ) from aepsych.transforms.ops import NormalizeScale +from aepsych.utils import dim_grid from botorch.acquisition import qUpperConfidenceBound from botorch.acquisition.active_learning import PairwiseMCPosteriorVariance from scipy.stats import bernoulli, norm, pearsonr @@ -93,8 +94,8 @@ def test_pairwise_probit_batched(self): np.random.seed(seed) n_init = 20 n_opt = 1 - lb = [-4.0, 1e-5] - ub = [-1e-5, 4.0] + lb = torch.tensor([-4.0, 1e-5]) + ub = torch.tensor([-1e-5, 4.0]) extra_acqf_args = {"beta": 3.84} model_list = [ Strategy( @@ -113,6 +114,8 @@ def test_pairwise_probit_batched(self): acqf=qUpperConfidenceBound, acqf_kwargs=extra_acqf_args, stimuli_per_trial=2, + lb=lb, + ub=ub, ), min_asks=n_opt, stimuli_per_trial=2, @@ -136,7 +139,7 @@ def test_pairwise_probit_batched(self): ), ) - xgrid = strat.model.dim_grid(gridsize=10) + xgrid = dim_grid(lb, ub, gridsize=10) zhat, _ = strat.predict(xgrid) # true max is 0, very loose test @@ -218,6 +221,8 @@ def test_1d_pairwise_probit(self): acqf_kwargs=extra_acqf_args, stimuli_per_trial=2, transforms=transforms, + lb=lb, + ub=ub, ) probit_model = ParameterTransformedModel( model=PairwiseProbitModel, lb=lb, ub=ub, transforms=transforms @@ -264,8 +269,8 @@ def test_1d_pairwise_probit_pure_exploration(self): np.random.seed(seed) n_init = 50 n_opt = 1 - lb = -2.0 - ub = 2.0 + lb = torch.tensor([-2.0]) + ub = torch.tensor([2.0]) acqf = PairwiseMCPosteriorVariance extra_acqf_args = {"objective": ProbitObjective()} @@ -284,7 +289,11 @@ def test_1d_pairwise_probit_pure_exploration(self): ub=ub, model=PairwiseProbitModel(lb=lb, ub=ub), generator=OptimizeAcqfGenerator( - acqf=acqf, acqf_kwargs=extra_acqf_args, stimuli_per_trial=2 + acqf=acqf, + acqf_kwargs=extra_acqf_args, + stimuli_per_trial=2, + lb=lb, + ub=ub, ), min_asks=n_opt, stimuli_per_trial=2, @@ -328,8 +337,8 @@ def test_2d_pairwise_probit(self): np.random.seed(seed) n_init = 20 n_opt = 1 - lb = np.r_[-1, -1] - ub = np.r_[1, 1] + lb = torch.tensor([-1.0, -1.0]) + ub = torch.tensor([1.0, 1.0]) extra_acqf_args = {"beta": 3.84} model_list = [ @@ -349,6 +358,8 @@ def test_2d_pairwise_probit(self): acqf=qUpperConfidenceBound, acqf_kwargs=extra_acqf_args, stimuli_per_trial=2, + lb=lb, + ub=ub, ), min_asks=n_opt, stimuli_per_trial=2, @@ -377,8 +388,8 @@ def test_2d_pairwise_probit_pure_exploration(self): np.random.seed(seed) n_init = 20 n_opt = 1 - lb = np.r_[-1, -1] - ub = np.r_[1, 1] + lb = torch.tensor([-1.0, -1.0]) + ub = torch.tensor([1.0, 1.0]) acqf = PairwiseMCPosteriorVariance extra_acqf_args = {"objective": ProbitObjective()} @@ -396,7 +407,11 @@ def test_2d_pairwise_probit_pure_exploration(self): ub=ub, model=PairwiseProbitModel(lb=lb, ub=ub), generator=OptimizeAcqfGenerator( - acqf=acqf, acqf_kwargs=extra_acqf_args, stimuli_per_trial=2 + acqf=acqf, + acqf_kwargs=extra_acqf_args, + stimuli_per_trial=2, + lb=lb, + ub=ub, ), min_asks=n_opt, stimuli_per_trial=2, @@ -454,7 +469,11 @@ def test_hyperparam_consistency(self): m1 = PairwiseProbitModel(lb=[1, 2], ub=[3, 4]) m2 = PairwiseProbitModel.from_config( - config=Config(config_dict={"common": {"lb": "[1,2]", "ub": "[3,4]"}}) + config=Config( + config_dict={ + "common": {"lb": "[1,2]", "ub": "[3,4]", "parnames": "[par1, par2]"} + } + ) ) self.assertTrue(isinstance(m1.covar_module, type(m2.covar_module))) diff --git a/tests/models/test_semi_p.py b/tests/models/test_semi_p.py index 4c4e6f610..1f2caf12f 100644 --- a/tests/models/test_semi_p.py +++ b/tests/models/test_semi_p.py @@ -26,51 +26,37 @@ from aepsych.likelihoods import BernoulliObjectiveLikelihood from aepsych.likelihoods.semi_p import LinearBernoulliLikelihood from aepsych.models import HadamardSemiPModel, SemiParametricGPModel -from aepsych.models.inducing_point_allocators import AutoAllocator from aepsych.models.semi_p import _hadamard_mvn_approx, semi_p_posterior_transform from aepsych.strategy import SequentialStrategy, Strategy +from aepsych.utils import make_scaled_sobol from gpytorch.distributions import MultivariateNormal from parameterized import parameterized def _hadamard_model_constructor( - lb, - ub, stim_dim, floor, objective=FloorLogitObjective, - inducing_point_method=AutoAllocator, ): return HadamardSemiPModel( - lb=lb, - ub=ub, + dim=2, stim_dim=stim_dim, likelihood=BernoulliObjectiveLikelihood(objective=objective(floor=floor)), inducing_size=10, - inducing_point_method=AutoAllocator( - bounds=torch.stack((torch.tensor(lb), torch.tensor(ub))) - ), max_fit_time=0.5, ) def _semip_model_constructor( - lb, - ub, stim_dim, floor, objective=FloorLogitObjective, - inducing_point_method=AutoAllocator, ): return SemiParametricGPModel( - lb=lb, - ub=ub, + dim=2, stim_dim=stim_dim, likelihood=LinearBernoulliLikelihood(objective=objective(floor=floor)), inducing_size=10, - inducing_point_method=AutoAllocator( - bounds=torch.stack((torch.tensor(lb), torch.tensor(ub))) - ), ) @@ -100,9 +86,10 @@ def setUp(self): self.f = torch.Tensor(slope * (intercept + xintensity)).unsqueeze(-1) X[:, 0] = X[:, 0] * 100 X[:, 1] = X[:, 1] / 100 - self.lb = [-100, -0.01] - self.ub = [100, 0.01] + self.lb = torch.tensor([-100.0, -0.01]) + self.ub = torch.tensor([100.0, 0.01]) self.X = torch.Tensor(X) + self.inducing_size = 10 @parameterized.expand( [(SemiPThresholdObjective(target=0.75),), (SemiPProbabilityObjective(),)] @@ -111,20 +98,18 @@ def test_mc_generation(self, objective): # no objective here, the objective for `gen` is not the same as the objective # for the likelihood in this case model = SemiParametricGPModel( - lb=self.lb, - ub=self.ub, + dim=2, stim_dim=self.stim_dim, likelihood=LinearBernoulliLikelihood(), inducing_size=10, - inducing_point_method=AutoAllocator( - bounds=torch.stack((torch.tensor(self.lb), torch.tensor(self.ub))) - ), ) generator = OptimizeAcqfGenerator( acqf=MCPosteriorVariance, acqf_kwargs={"objective": objective}, max_gen_time=0.1, + lb=self.lb, + ub=self.ub, ) y = torch.bernoulli(model.likelihood.objective(self.f)) @@ -144,8 +129,6 @@ def test_analytic_lookahead_generation(self): floor = 0 objective = FloorProbitObjective model = _semip_model_constructor( - lb=self.lb, - ub=self.ub, stim_dim=self.stim_dim, floor=floor, objective=objective, @@ -157,8 +140,11 @@ def test_analytic_lookahead_generation(self): "posterior_transform": semi_p_posterior_transform, "target": 0.75, "query_set_size": 100, + "Xq": make_scaled_sobol(self.lb, self.ub, 100), }, max_gen_time=0.2, + lb=self.lb, + ub=self.ub, ) link = objective(floor=floor) y = torch.bernoulli(link(self.f)) @@ -193,8 +179,7 @@ def test_memorize_data(self, objective, floor, model_constructor): y = torch.bernoulli(link(self.f)) model = model_constructor( - lb=self.lb, - ub=self.ub, + inducing_point_method=self.inducing_point_method, stim_dim=self.stim_dim, floor=floor, objective=objective, @@ -215,11 +200,14 @@ def test_memorize_data(self, objective, floor, model_constructor): @parameterized.expand([(_semip_model_constructor,), (_hadamard_model_constructor,)]) def test_prediction_shapes(self, model_constructor): n_opt = 1 - lb = [-1, -1] - ub = [1, 1] + lb = torch.tensor([-1.0, -1.0]) + ub = torch.tensor([1.0, 1.0]) with self.subTest(model_constructor=model_constructor): - model = model_constructor(lb=lb, ub=ub, stim_dim=self.stim_dim, floor=0) + model = model_constructor( + stim_dim=self.stim_dim, + floor=0, + ) strat_list = [ Strategy( @@ -276,18 +264,11 @@ def test_prediction_shapes(self, model_constructor): @parameterized.expand([(_semip_model_constructor,), (_hadamard_model_constructor,)]) def test_reset_variational_strategy(self, model_constructor): - lb = [-3, -3] - ub = [3, 3] stim_dim = 0 with self.subTest(model_constructor=model_constructor): model = model_constructor( - lb=lb, - ub=ub, stim_dim=stim_dim, floor=0, - inducing_point_method=AutoAllocator( - bounds=torch.stack((torch.tensor(lb), torch.tensor(ub))) - ), ) link = FloorLogitObjective(floor=0) y = torch.bernoulli(link(self.f)) @@ -309,16 +290,16 @@ def test_reset_variational_strategy(self, model_constructor): variational_params_reset = [ v.clone().detach().numpy() for v in model.variational_parameters() ] - induc_reset = model.variational_strategy.inducing_points - # before should be different from after and after should be different - # from reset - self.assertFalse(np.allclose(induc_before, induc_after)) - self.assertFalse(np.allclose(induc_after, induc_reset)) + # before should be different from after + if induc_before.shape == induc_after.shape: # Not same can't fail + self.assertFalse(np.allclose(induc_before, induc_after)) + for before, after in zip( variational_params_before, variational_params_after ): - self.assertFalse(np.allclose(before, after)) + if before.shape == after.shape: # Not same can't fail + self.assertFalse(np.allclose(before, after)) for after, reset in zip(variational_params_after, variational_params_reset): self.assertFalse(np.allclose(after, reset)) @@ -326,28 +307,20 @@ def test_reset_variational_strategy(self, model_constructor): def test_slope_mean_setting(self): for slope_mean in (2, 4): model = SemiParametricGPModel( - lb=self.lb, - ub=self.ub, + dim=2, stim_dim=self.stim_dim, likelihood=LinearBernoulliLikelihood(), - inducing_size=10, + inducing_size=self.inducing_size, slope_mean=slope_mean, - inducing_point_method=AutoAllocator( - bounds=torch.stack((torch.tensor(self.lb), torch.tensor(self.ub))) - ), ) with self.subTest(model=model, slope_mean=slope_mean): self.assertEqual(model.mean_module.constant[-1], slope_mean) model = HadamardSemiPModel( - lb=self.lb, - ub=self.ub, + dim=2, stim_dim=self.stim_dim, likelihood=BernoulliObjectiveLikelihood(objective=ProbitObjective()), - inducing_size=10, + inducing_size=self.inducing_size, slope_mean=slope_mean, - inducing_point_method=AutoAllocator( - bounds=torch.stack((torch.tensor(self.lb), torch.tensor(self.ub))) - ), ) with self.subTest(model=model, slope_mean=slope_mean): self.assertEqual(model.slope_mean_module.constant.item(), slope_mean) @@ -365,12 +338,8 @@ def setUp(self): def test_reset_hyperparams(self): model = HadamardSemiPModel( - lb=[-3, -3], - ub=[3, 3], + dim=2, inducing_size=20, - inducing_point_method=AutoAllocator( - bounds=torch.stack((torch.tensor([-3, -3]), torch.tensor([3, 3]))) - ), ) slope_os_before = model.slope_covar_module.outputscale.clone().detach().numpy() diff --git a/tests/models/test_utils.py b/tests/models/test_utils.py deleted file mode 100644 index 3226b1415..000000000 --- a/tests/models/test_utils.py +++ /dev/null @@ -1,159 +0,0 @@ -#!/usr/bin/env python3 -# Copyright (c) Facebook, Inc. and its affiliates. -# All rights reserved. - -# This source code is licensed under the license found in the -# LICENSE file in the root directory of this source tree. - -import unittest - -import numpy as np -import torch -from aepsych.models import GPClassificationModel -from aepsych.models.inducing_point_allocators import ( - AutoAllocator, - DummyAllocator, - FixedAllocator, - GreedyVarianceReduction, - KMeansAllocator, - SobolAllocator, -) -from aepsych.models.utils import select_inducing_points - -from sklearn.datasets import make_classification - - -class UtilsTestCase(unittest.TestCase): - def test_select_inducing_points(self): - """Verify that when we have n_induc > data size, we use data as inducing, - and otherwise we correctly select inducing points.""" - X, y = make_classification( - n_samples=100, - n_features=1, - n_redundant=0, - n_informative=1, - random_state=1, - n_clusters_per_class=1, - ) - X, y = torch.Tensor(X), torch.Tensor(y) - inducing_size = 20 - - model = GPClassificationModel( - torch.Tensor([-3]), - torch.Tensor([3]), - inducing_size=inducing_size, - inducing_point_method=AutoAllocator( - bounds=torch.stack((torch.Tensor([-3]), torch.Tensor([3]))) - ), - ) - model.set_train_data(X[:10, ...], y[:10]) - - # (inducing point selection sorts the inputs so we sort X to verify) - self.assertTrue( - np.allclose( - select_inducing_points( - allocator=AutoAllocator(), - inducing_size=inducing_size, - covar_module=model.covar_module, - X=model.train_inputs[0], - bounds=model.bounds, - ), - X[:10].sort(0).values, - ) - ) - - model.set_train_data(X, y) - - self.assertTrue( - len( - select_inducing_points( - allocator=AutoAllocator(), - inducing_size=inducing_size, - covar_module=model.covar_module, - X=model.train_inputs[0], - bounds=model.bounds, - ) - ) - <= 20 - ) - - self.assertTrue( - len( - select_inducing_points( - allocator=GreedyVarianceReduction(), - inducing_size=inducing_size, - covar_module=model.covar_module, - X=model.train_inputs[0], - ) - ) - <= 20 - ) - - self.assertEqual( - len( - select_inducing_points( - allocator=KMeansAllocator(), - inducing_size=inducing_size, - covar_module=model.covar_module, - X=model.train_inputs[0], - bounds=model.bounds, - ) - ), - 20, - ) - self.assertTrue( - len( - select_inducing_points( - allocator="auto", - inducing_size=inducing_size, - covar_module=model.covar_module, - X=model.train_inputs[0], - bounds=model.bounds, - ) - ) - <= 20 - ) - self.assertTrue( - len( - select_inducing_points( - allocator=SobolAllocator( - bounds=torch.stack([torch.tensor([0]), torch.tensor([1])]) - ), - inducing_size=inducing_size, - covar_module=model.covar_module, - X=model.train_inputs[0], - bounds=model.bounds, - ) - ) - <= 20 - ) - self.assertTrue( - len( - select_inducing_points( - allocator=DummyAllocator( - bounds=torch.stack([torch.tensor([0]), torch.tensor([1])]) - ), - inducing_size=inducing_size, - covar_module=model.covar_module, - X=model.train_inputs[0], - bounds=model.bounds, - ) - ) - <= 20 - ) - self.assertTrue( - len( - select_inducing_points( - allocator=FixedAllocator(points=torch.tensor([0, 1, 2, 3])), - inducing_size=inducing_size, - covar_module=model.covar_module, - X=model.train_inputs[0], - bounds=model.bounds, - ) - ) - <= 20 - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/test_benchmark.py b/tests/test_benchmark.py index 1a8d1d23f..c377cf2ca 100644 --- a/tests/test_benchmark.py +++ b/tests/test_benchmark.py @@ -21,9 +21,7 @@ Problem, ) from aepsych.models import GPClassificationModel - -from aepsych.models.inducing_point_allocators import AutoAllocator -from scipy.stats import norm +from aepsych.models.inducing_points import GreedyVarianceReduction torch.set_num_threads(1) torch.set_num_interop_threads(1) @@ -74,10 +72,10 @@ def setUp(self): self.n_thresholds = 5 self.thresholds = torch.linspace(0.55, 0.95, self.n_thresholds) self.test_problem = example_problems.DiscrimLowDim(thresholds=self.thresholds) + self.model = GPClassificationModel( - lb=self.test_problem.lb, - ub=self.test_problem.ub, - inducing_point_method=AutoAllocator(bounds=self.test_problem.bounds), + dim=2, + inducing_point_method=GreedyVarianceReduction(dim=2), ) def unvectorized_p_below_threshold(self, x, f_thresh) -> torch.Tensor: @@ -420,7 +418,6 @@ def test_monotonic_single_lse_eval(self): "inducing_size": 10, "mean_covar_factory": "monotonic_mean_covar_factory", "monotonic_idxs": "[1]", - "inducing_point_method": "SobolAllocator", }, "MonotonicRejectionGenerator": { "model_gen_options": { diff --git a/tests/test_config.py b/tests/test_config.py index 24df57d74..7c0039092 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -27,7 +27,7 @@ MonotonicRejectionGP, PairwiseProbitModel, ) -from aepsych.models.inducing_point_allocators import SobolAllocator +from aepsych.models.inducing_points import SobolAllocator from aepsych.server import AEPsychServer from aepsych.server.message_handlers.handle_setup import configure from aepsych.strategy import SequentialStrategy, Strategy @@ -126,18 +126,17 @@ def test_single_probit_config(self): self.assertTrue(torch.all(strat.strat_list[0].lb == strat.strat_list[1].lb)) self.assertTrue( torch.all( - strat.transforms.untransform(strat.strat_list[1].model.lb) + strat.transforms.untransform(strat.strat_list[1].generator.lb) == torch.Tensor([1, -1]) ) ) self.assertTrue(torch.all(strat.strat_list[0].ub == strat.strat_list[1].ub)) self.assertTrue( torch.all( - strat.transforms.untransform(strat.strat_list[1].model.ub) + strat.transforms.untransform(strat.strat_list[1].generator.ub) == torch.Tensor([10, 1]) ) ) - self.assertEqual(strat.strat_list[0].min_total_outcome_occurrences, 5) self.assertEqual(strat.strat_list[0].min_post_range, None) self.assertEqual(strat.strat_list[0].keep_most_recent, None) @@ -171,7 +170,8 @@ def test_single_probit_config_file(self): ) self.assertTrue(strat.strat_list[1].generator.acqf is EAVC) self.assertTrue( - set(strat.strat_list[1].generator.acqf_kwargs.keys()) == {"target"} + set(strat.strat_list[1].generator.acqf_kwargs.keys()) + == {"lb", "ub", "target"} ) self.assertTrue(strat.strat_list[1].generator.acqf_kwargs["target"] == 0.75) self.assertTrue(strat.strat_list[1].generator.samps == 1000) @@ -180,9 +180,13 @@ def test_single_probit_config_file(self): self.assertTrue(strat.strat_list[0].outcome_types == ["binary"]) self.assertTrue(strat.strat_list[1].min_asks == 20) self.assertTrue(torch.all(strat.strat_list[0].lb == strat.strat_list[1].lb)) - self.assertTrue(torch.all(strat.strat_list[1].model.lb == torch.Tensor([0, 0]))) + self.assertTrue( + torch.all(strat.strat_list[1].generator.lb == torch.Tensor([0, 0])) + ) self.assertTrue(torch.all(strat.strat_list[0].ub == strat.strat_list[1].ub)) - self.assertTrue(torch.all(strat.strat_list[1].model.ub == torch.Tensor([1, 1]))) + self.assertTrue( + torch.all(strat.strat_list[1].generator.ub == torch.Tensor([1, 1])) + ) def test_nonmonotonic_optimization_config_file(self): config_file = "../configs/nonmonotonic_optimization_example.ini" @@ -216,9 +220,13 @@ def test_nonmonotonic_optimization_config_file(self): self.assertTrue(strat.strat_list[0].outcome_types == ["binary"]) self.assertTrue(strat.strat_list[1].min_asks == 20) self.assertTrue(torch.all(strat.strat_list[0].lb == strat.strat_list[1].lb)) - self.assertTrue(torch.all(strat.strat_list[1].model.lb == torch.Tensor([0, 0]))) + self.assertTrue( + torch.all(strat.strat_list[1].generator.lb == torch.Tensor([0, 0])) + ) self.assertTrue(torch.all(strat.strat_list[0].ub == strat.strat_list[1].ub)) - self.assertTrue(torch.all(strat.strat_list[1].model.ub == torch.Tensor([1, 1]))) + self.assertTrue( + torch.all(strat.strat_list[1].generator.ub == torch.Tensor([1, 1])) + ) def test_name_conflict_warns(self): class DummyMod: @@ -262,7 +270,6 @@ def test_multiple_models_and_strats(self): min_asks = 1 model = MonotonicRejectionGP acqf = MonotonicMCLSE - inducing_point_method=AutoAllocator """ @@ -519,9 +526,13 @@ def test_pairwise_probit_config(self): self.assertTrue(strat.strat_list[0].outcome_types == ["binary"]) self.assertTrue(strat.strat_list[1].min_asks == 20) self.assertTrue(torch.all(strat.strat_list[0].lb == strat.strat_list[1].lb)) - self.assertTrue(torch.all(strat.strat_list[1].model.lb == torch.Tensor([0, 0]))) + self.assertTrue( + torch.all(strat.strat_list[1].generator.lb == torch.Tensor([0, 0])) + ) self.assertTrue(torch.all(strat.strat_list[0].ub == strat.strat_list[1].ub)) - self.assertTrue(torch.all(strat.strat_list[1].model.ub == torch.Tensor([1, 1]))) + self.assertTrue( + torch.all(strat.strat_list[1].generator.ub == torch.Tensor([1, 1])) + ) def test_pairwise_probit_config_file(self): config_file = "../configs/pairwise_al_example.ini" @@ -557,9 +568,13 @@ def test_pairwise_probit_config_file(self): self.assertTrue(strat.strat_list[0].outcome_types == ["binary"]) self.assertTrue(strat.strat_list[1].min_asks == 20) self.assertTrue(torch.all(strat.strat_list[0].lb == strat.strat_list[1].lb)) - self.assertTrue(torch.all(strat.strat_list[1].model.lb == torch.Tensor([0, 0]))) + self.assertTrue( + torch.all(strat.strat_list[1].generator.lb == torch.Tensor([0, 0])) + ) self.assertTrue(torch.all(strat.strat_list[0].ub == strat.strat_list[1].ub)) - self.assertTrue(torch.all(strat.strat_list[1].model.ub == torch.Tensor([1, 1]))) + self.assertTrue( + torch.all(strat.strat_list[1].generator.ub == torch.Tensor([1, 1])) + ) def test_pairwise_al_config_file(self): # random datebase path name without dashes @@ -598,9 +613,13 @@ def test_pairwise_al_config_file(self): self.assertTrue(strat.strat_list[0].outcome_types == ["binary"]) self.assertTrue(strat.strat_list[1].min_asks == 20) self.assertTrue(torch.all(strat.strat_list[0].lb == strat.strat_list[1].lb)) - self.assertTrue(torch.all(strat.strat_list[1].model.lb == torch.Tensor([0, 0]))) + self.assertTrue( + torch.all(strat.strat_list[1].generator.lb == torch.Tensor([0, 0])) + ) self.assertTrue(torch.all(strat.strat_list[0].ub == strat.strat_list[1].ub)) - self.assertTrue(torch.all(strat.strat_list[1].model.ub == torch.Tensor([1, 1]))) + self.assertTrue( + torch.all(strat.strat_list[1].generator.ub == torch.Tensor([1, 1])) + ) # cleanup the db if server.db is not None: server.db.delete_db() @@ -640,9 +659,13 @@ def test_pairwise_opt_config(self): self.assertTrue(strat.strat_list[0].outcome_types == ["binary"]) self.assertTrue(strat.strat_list[1].min_asks == 20) self.assertTrue(torch.all(strat.strat_list[0].lb == strat.strat_list[1].lb)) - self.assertTrue(torch.all(strat.strat_list[1].model.lb == torch.Tensor([0, 0]))) + self.assertTrue( + torch.all(strat.strat_list[1].generator.lb == torch.Tensor([0, 0])) + ) self.assertTrue(torch.all(strat.strat_list[0].ub == strat.strat_list[1].ub)) - self.assertTrue(torch.all(strat.strat_list[1].model.ub == torch.Tensor([1, 1]))) + self.assertTrue( + torch.all(strat.strat_list[1].generator.ub == torch.Tensor([1, 1])) + ) # cleanup the db if server.db is not None: server.db.delete_db() @@ -1024,11 +1047,11 @@ def test_semip_config(self): strat = SequentialStrategy.from_config(config) opt_strat = strat.strat_list[1] model = opt_strat.model + generator = opt_strat.generator self.assertTrue(isinstance(model, HadamardSemiPModel)) - self.assertTrue(torch.all(model.lb == torch.Tensor([0, 0]))) - self.assertTrue(torch.all(model.ub == torch.Tensor([1, 1]))) - self.assertTrue(model.dim == 2) + self.assertTrue(torch.all(generator.lb == torch.Tensor([0, 0]))) + self.assertTrue(torch.all(generator.ub == torch.Tensor([1, 1]))) self.assertTrue(model.inducing_size == 10) self.assertTrue(model.stim_dim == 1) @@ -1079,10 +1102,9 @@ def test_derived_bounds(self): strat = SequentialStrategy.from_config(config) opt_strat = strat.strat_list[1] - model = opt_strat.model - self.assertTrue(torch.all(model.lb == torch.Tensor([0, -10]))) - self.assertTrue(torch.all(model.ub == torch.Tensor([1, 10]))) + self.assertTrue(torch.all(opt_strat.lb == torch.Tensor([0, -10]))) + self.assertTrue(torch.all(opt_strat.ub == torch.Tensor([1, 10]))) def test_ignore_common_bounds(self): config_str = """ @@ -1123,10 +1145,9 @@ def test_ignore_common_bounds(self): strat = SequentialStrategy.from_config(config) opt_strat = strat.strat_list[1] - model = opt_strat.model - self.assertTrue(torch.all(model.lb == torch.Tensor([0, -5]))) - self.assertTrue(torch.all(model.ub == torch.Tensor([1, 1]))) + self.assertTrue(torch.all(opt_strat.lb == torch.Tensor([0, -5]))) + self.assertTrue(torch.all(opt_strat.ub == torch.Tensor([1, 1]))) def test_common_fallback_bounds(self): config_str = """ @@ -1167,10 +1188,9 @@ def test_common_fallback_bounds(self): strat = SequentialStrategy.from_config(config) opt_strat = strat.strat_list[1] - model = opt_strat.model - self.assertTrue(torch.all(model.lb == torch.Tensor([0, 0]))) - self.assertTrue(torch.all(model.ub == torch.Tensor([1, 100]))) + self.assertTrue(torch.all(opt_strat.lb == torch.Tensor([0, 0]))) + self.assertTrue(torch.all(opt_strat.ub == torch.Tensor([1, 100]))) def test_parameter_setting_block_validation(self): config_str = """ diff --git a/tests/test_points_allocators.py b/tests/test_points_allocators.py index c0d0d2eb6..f97b7f3e5 100644 --- a/tests/test_points_allocators.py +++ b/tests/test_points_allocators.py @@ -1,87 +1,26 @@ import unittest +import gpytorch +import numpy as np import torch from aepsych.config import Config +from aepsych.kernels import RBFKernelPartialObsGrad from aepsych.models.gp_classification import GPClassificationModel -from aepsych.models.inducing_point_allocators import ( - AutoAllocator, - DummyAllocator, +from aepsych.models.inducing_points import ( FixedAllocator, GreedyVarianceReduction, KMeansAllocator, SobolAllocator, ) -from aepsych.models.utils import select_inducing_points - from aepsych.strategy import Strategy from aepsych.transforms.parameters import ParameterTransforms, transform_options -from botorch.models.utils.inducing_point_allocators import GreedyImprovementReduction -from botorch.utils.sampling import draw_sobol_samples from sklearn.datasets import make_classification class TestInducingPointAllocators(unittest.TestCase): - def test_sobol_allocator_from_config(self): - config_str = """ - [common] - parnames = [par1] - - [par1] - par_type = continuous - lower_bound = 0.0 - upper_bound = 1.0 - log_scale = true - - """ - config = Config() - config.update(config_str=config_str) - allocator = SobolAllocator.from_config(config) - - # Check if bounds are correctly loaded - expected_bounds = torch.tensor([[0.0], [1.0]]) - self.assertTrue(torch.equal(allocator.bounds, expected_bounds)) - - def test_kmeans_allocator_from_config(self): - config_str = """ - [common] - parnames = [par1] - - [par1] - par_type = continuous - lower_bound = 0.0 - upper_bound = 1.0 - log_scale = true - - [KMeansAllocator] - """ - config = Config() - config.update(config_str=config_str) - allocator = KMeansAllocator.from_config(config) - # No specific configuration to check, just test instantiation - self.assertTrue(isinstance(allocator, KMeansAllocator)) - - def test_auto_allocator_from_config_with_fallback(self): - config_str = """ - [common] - parnames = [par1] - - [par1] - par_type = continuous - lower_bound = 0.0 - upper_bound = 1.0 - log_scale = true - - """ - config = Config() - config.update(config_str=config_str) - allocator = AutoAllocator.from_config(config) - - # Check if fallback allocator is an instance of SobolAllocator with correct bounds - self.assertTrue(isinstance(allocator.fallback_allocator, KMeansAllocator)) - def test_sobol_allocator_allocate_inducing_points(self): bounds = torch.tensor([[0.0], [1.0]]) - allocator = SobolAllocator(bounds=bounds) + allocator = SobolAllocator(bounds=bounds, dim=1) inducing_points = allocator.allocate_inducing_points(num_inducing=5) # Check shape and bounds of inducing points @@ -90,112 +29,32 @@ def test_sobol_allocator_allocate_inducing_points(self): torch.all(inducing_points >= bounds[0]) and torch.all(inducing_points <= bounds[1]) ) - - def test_kmeans_allocator_allocate_inducing_points(self): - inputs = torch.rand(100, 2) # 100 points in 2D - allocator = KMeansAllocator() - inducing_points = allocator.allocate_inducing_points( - inputs=inputs, num_inducing=10 - ) - - # Check shape of inducing points - self.assertEqual(inducing_points.shape, (10, 2)) - - def test_auto_allocator_with_kmeans_fallback(self): - inputs = torch.rand(50, 2) - fallback_allocator = KMeansAllocator() - allocator = AutoAllocator(fallback_allocator=fallback_allocator) - inducing_points = allocator.allocate_inducing_points( - inputs=inputs, num_inducing=10 - ) - - # Check shape of inducing points and that fallback allocator is used - self.assertEqual(inducing_points.shape, (10, 2)) - - def test_select_inducing_points_legacy(self): - with self.assertWarns(DeprecationWarning): - # Call select_inducing_points directly with a string for allocator to trigger the warning - bounds = torch.tensor([[0.0], [1.0]]) - points = select_inducing_points( - inducing_size=5, - allocator="sobol", # Legacy string argument to trigger DeprecationWarning - bounds=bounds, - ) - self.assertEqual(points.shape, (5, 1)) - - def test_auto_allocator_allocate_inducing_points(self): - train_X = torch.tensor([[0.0, 0.0], [1.0, 1.0], [2.0, 2.0]]) - train_Y = torch.tensor([[1.0], [2.0], [3.0]]) - lb = torch.tensor([0, 0]) - ub = torch.tensor([4, 4]) - bounds = torch.stack([lb, ub]) - model = GPClassificationModel( - lb=lb, - ub=ub, - inducing_point_method=AutoAllocator(bounds=bounds), - inducing_size=3, - ) - self.assertTrue(model.last_inducing_points_method == "DummyAllocator") - auto_inducing_points = AutoAllocator(bounds=bounds).allocate_inducing_points( - inputs=train_X, - covar_module=model.covar_module, - num_inducing=model.inducing_size, - ) - inital_inducing_points = DummyAllocator(bounds=bounds).allocate_inducing_points( - inputs=train_X, - covar_module=model.covar_module, - num_inducing=model.inducing_size, - ) - - # Should be different from the initial inducing points - self.assertFalse( - torch.allclose( - auto_inducing_points, model.variational_strategy.inducing_points - ) - ) - self.assertTrue( - torch.allclose( - inital_inducing_points, model.variational_strategy.inducing_points - ) - ) - - model.fit(train_X, train_Y) - self.assertTrue(model.last_inducing_points_method == "AutoAllocator") - self.assertEqual( - model.variational_strategy.inducing_points.shape, auto_inducing_points.shape - ) - - # Check that inducing points are updated after fitting - self.assertTrue( - torch.allclose( - auto_inducing_points, model.variational_strategy.inducing_points - ) - ) + self.assertIs(allocator.last_allocator_used, SobolAllocator) def test_sobol_allocator_from_model_config(self): config_str = """ - [common] - parnames = [par1] - stimuli_per_trial = 1 - outcome_types = [binary] - strategy_names = [init_strat] - - [par1] - par_type = continuous - lower_bound = 10 - upper_bound = 1000 - log_scale = True - - [init_strat] - generator = OptimizeAcqfGenerator - acqf = MCLevelSetEstimation - min_asks = 2 - model = GPClassificationModel - - [GPClassificationModel] - inducing_point_method = SobolAllocator - inducing_size = 2 - """ + [common] + parnames = [par1] + stimuli_per_trial = 1 + outcome_types = [binary] + strategy_names = [init_strat] + + [par1] + par_type = continuous + lower_bound = 10 + upper_bound = 1000 + log_scale = True + + [init_strat] + generator = OptimizeAcqfGenerator + acqf = MCLevelSetEstimation + min_asks = 2 + model = GPClassificationModel + + [GPClassificationModel] + inducing_point_method = SobolAllocator + inducing_size = 2 + """ config = Config() config.update(config_str=config_str) @@ -206,7 +65,7 @@ def test_sobol_allocator_from_model_config(self): # check that the bounds are scaled correctly self.assertFalse( torch.equal( - strat.model.inducing_point_method.bounds, torch.tensor([[10], [100]]) + strat.model.inducing_point_method.bounds, torch.tensor([[10], [1000]]) ) ) self.assertTrue( @@ -215,262 +74,218 @@ def test_sobol_allocator_from_model_config(self): ) ) + def test_kmeans_allocator_allocate_inducing_points(self): + # Mock data for testing + train_X = torch.randint(low=0, high=100, size=(100, 2), dtype=torch.float64) + train_Y = torch.rand(100, 1) + model = GPClassificationModel( + inducing_point_method=KMeansAllocator(dim=2), + inducing_size=10, + dim=2, + ) + + # Check if model has dummy points + self.assertIsNone(model.inducing_point_method.last_allocator_used) + self.assertTrue(torch.all(model.variational_strategy.inducing_points == 0)) + self.assertTrue(model.variational_strategy.inducing_points.shape == (10, 2)) + + # Fit with small data leess than inducing_size + model.fit(train_X[:9], train_Y[:9]) + + self.assertIs(model.inducing_point_method.last_allocator_used, KMeansAllocator) + inducing_points = model.variational_strategy.inducing_points + self.assertTrue(inducing_points.shape == (9, 2)) + # We made ints, so mod 1 should be 0s, so we know these were the original inputs + self.assertTrue(torch.all(inducing_points % 1 == 0)) + + # Then fit the model and check that the inducing points are updated + model.fit(train_X, train_Y) + + self.assertIs(model.inducing_point_method.last_allocator_used, KMeansAllocator) + inducing_points = model.variational_strategy.inducing_points + self.assertTrue(inducing_points.shape == (10, 2)) + # It's highly unlikely clustered will all be integers, so check against extents too + self.assertFalse(torch.all(inducing_points % 1 == 0)) + self.assertTrue(torch.all((inducing_points >= 0) & (inducing_points <= 100))) + def test_kmeans_allocator_from_model_config(self): config_str = """ - [common] - parnames = [par1] - stimuli_per_trial = 1 - outcome_types = [binary] - strategy_names = [init_strat] - - [par1] - par_type = continuous - lower_bound = 10 - upper_bound = 1000 - log_scale = True - - [init_strat] - generator = OptimizeAcqfGenerator - acqf = MCLevelSetEstimation - min_asks = 2 - model = GPClassificationModel - - [GPClassificationModel] - inducing_point_method = KMeansAllocator - inducing_size = 2 - """ + [common] + parnames = [par1, par2] + stimuli_per_trial = 1 + outcome_types = [binary] + strategy_names = [init_strat] + + [par1] + par_type = continuous + lower_bound = 10 + upper_bound = 1000 + log_scale = True + + [par2] + par_type = integer + lower_bound = 0 + upper_bound = 1 + + [init_strat] + generator = OptimizeAcqfGenerator + acqf = MCLevelSetEstimation + min_asks = 2 + model = GPClassificationModel + + [GPClassificationModel] + inducing_point_method = KMeansAllocator + inducing_size = 2 + """ config = Config() config.update(config_str=config_str) strat = Strategy.from_config(config, "init_strat") self.assertTrue(isinstance(strat.model.inducing_point_method, KMeansAllocator)) - # check that the bounds are scaled correctly - self.assertFalse( - torch.equal( - strat.model.inducing_point_method.bounds, torch.tensor([[10], [100]]) - ) - ) - self.assertTrue( - torch.equal( - strat.model.inducing_point_method.bounds, torch.tensor([[0], [1]]) - ) - ) + self.assertTrue(strat.model.inducing_point_method.dim == 2) - def test_auto_allocator_from_model_config(self): - config_str = """ - [common] - parnames = [par1] - stimuli_per_trial = 1 - outcome_types = [binary] - strategy_names = [init_strat] - - [par1] - par_type = continuous - lower_bound = 10 - upper_bound = 1000 - log_scale = True - - [init_strat] - generator = OptimizeAcqfGenerator - acqf = MCLevelSetEstimation - min_asks = 2 - model = GPClassificationModel - - [GPClassificationModel] - inducing_point_method = AutoAllocator - inducing_size = 2 - """ + def test_kmeans_shape_handling(self): + allocator = KMeansAllocator(dim=1) - config = Config() - config.update(config_str=config_str) - strat = Strategy.from_config(config, "init_strat") - self.assertTrue(isinstance(strat.model.inducing_point_method, AutoAllocator)) + inputs = torch.tensor([[1], [2], [3]]) - # check that the bounds are scaled correctly - self.assertFalse( - torch.equal( - strat.model.inducing_point_method.bounds, torch.tensor([[10], [100]]) - ) - ) - self.assertTrue( - torch.equal( - strat.model.inducing_point_method.bounds, torch.tensor([[0], [1]]) - ) - ) + inputs_aug = torch.hstack([inputs, torch.zeros(size=[3, 1])]) - def test_dummy_allocator_from_model_config(self): - config_str = """ - [common] - parnames = [par1] - stimuli_per_trial = 1 - outcome_types = [binary] - strategy_names = [init_strat] - - [par1] - par_type = continuous - lower_bound = 10 - upper_bound = 1000 - log_scale = True - - [init_strat] - generator = OptimizeAcqfGenerator - acqf = MCLevelSetEstimation - min_asks = 2 - model = GPClassificationModel - - [GPClassificationModel] - inducing_point_method = DummyAllocator - inducing_size = 2 - """ + points = allocator.allocate_inducing_points(inputs=inputs_aug, num_inducing=2) + self.assertTrue(points.shape == (2, 1)) - config = Config() - config.update(config_str=config_str) - strat = Strategy.from_config(config, "init_strat") - self.assertTrue(isinstance(strat.model.inducing_point_method, DummyAllocator)) + points = allocator.allocate_inducing_points(inputs=inputs_aug, num_inducing=100) + self.assertTrue(torch.equal(points, inputs)) - # check that the bounds are scaled correctly - self.assertFalse( - torch.equal( - strat.model.inducing_point_method.bounds, torch.tensor([[10], [100]]) + def test_greedy_variance_allocator_no_covar_raise(self): + allocator = GreedyVarianceReduction(dim=2) + + with self.assertRaises(ValueError): + _ = allocator.allocate_inducing_points( + inputs=torch.zeros((30, 1)), num_inducing=10 ) + + def test_greedy_variance_reduction_allocate_inducing_points(self): + # Mock data for testing + train_X = torch.randint(low=0, high=100, size=(100, 2), dtype=torch.float64) + train_Y = torch.rand(100, 1) + model = GPClassificationModel( + inducing_point_method=GreedyVarianceReduction(dim=2), + inducing_size=10, + dim=2, ) - self.assertTrue( - torch.equal( - strat.model.inducing_point_method.bounds, torch.tensor([[0], [1]]) - ) + + # Check if model has dummy points + self.assertIsNone(model.inducing_point_method.last_allocator_used) + self.assertTrue(torch.all(model.variational_strategy.inducing_points == 0)) + self.assertTrue(model.variational_strategy.inducing_points.shape == (10, 2)) + + # Then fit the model and check that the inducing points are updated + model.fit(train_X, train_Y) + + self.assertIs( + model.inducing_point_method.last_allocator_used, GreedyVarianceReduction ) + inducing_points = model.variational_strategy.inducing_points + self.assertTrue(inducing_points.shape == (10, 2)) + self.assertTrue(torch.all((inducing_points >= 0) & (inducing_points <= 100))) - def test_inducing_point_before_and_after_auto(self): + def test_greedy_variance_from_config(self): config_str = """ - [common] - parnames = [par1] - stimuli_per_trial = 1 - outcome_types = [binary] - strategy_names = [init_strat] - - [par1] - par_type = continuous - lower_bound = 10 - upper_bound = 1000 - log_scale = True - - [init_strat] - generator = OptimizeAcqfGenerator - acqf = MCLevelSetEstimation - min_asks = 2 - model = GPClassificationModel - - [GPClassificationModel] - inducing_point_method = AutoAllocator - inducing_size = 2 - """ + [common] + parnames = [par1] + stimuli_per_trial = 1 + outcome_types = [binary] + strategy_names = [init_strat] + + [par1] + par_type = continuous + lower_bound = 10 + upper_bound = 1000 + log_scale = True + + [init_strat] + generator = OptimizeAcqfGenerator + acqf = MCLevelSetEstimation + min_asks = 2 + model = GPClassificationModel + + [GPClassificationModel] + inducing_point_method = GreedyVarianceReduction + inducing_size = 2 + """ config = Config() config.update(config_str=config_str) strat = Strategy.from_config(config, "init_strat") - self.assertTrue(isinstance(strat.model.inducing_point_method, AutoAllocator)) - - # check that the bounds are scaled correctly - self.assertFalse( - torch.equal( - strat.model.inducing_point_method.bounds, torch.tensor([[10], [100]]) - ) - ) self.assertTrue( - torch.equal( - strat.model.inducing_point_method.bounds, torch.tensor([[0], [1]]) - ) + isinstance(strat.model.inducing_point_method, GreedyVarianceReduction) ) - train_X = torch.tensor([[0.0], [1.0]]) - train_Y = torch.tensor([[1.0], [0.0]]) + def test_greedy_variance_shape_handling(self): + allocator = GreedyVarianceReduction(dim=1) - auto_inducing_points = AutoAllocator( - bounds=torch.stack([torch.tensor([0]), torch.tensor([1])]) - ).allocate_inducing_points( - inputs=train_X, - covar_module=strat.model.covar_module, - num_inducing=strat.model.inducing_size, - ) - inital_inducing_points = DummyAllocator( - bounds=torch.stack([torch.tensor([0]), torch.tensor([1])]) - ).allocate_inducing_points( - inputs=train_X, - covar_module=strat.model.covar_module, - num_inducing=strat.model.inducing_size, - ) + inputs = torch.tensor([[1], [2], [3]]) + inputs_aug = torch.hstack([inputs, torch.zeros(size=[3, 1])]) - # Should be different from the initial inducing points - self.assertFalse( - torch.allclose( - auto_inducing_points, strat.model.variational_strategy.inducing_points - ) + ls_prior = gpytorch.priors.GammaPrior( + concentration=4.6, rate=1.0, transform=lambda x: 1 / x ) - # Should be the same as the initial inducing points - self.assertTrue( - torch.allclose( - inital_inducing_points, strat.model.variational_strategy.inducing_points - ) + ls_prior_mode = ls_prior.rate / (ls_prior.concentration + 1) + ls_constraint = gpytorch.constraints.GreaterThan( + lower_bound=1e-4, transform=None, initial_value=ls_prior_mode + ) + covar_module = gpytorch.kernels.ScaleKernel( + RBFKernelPartialObsGrad( + lengthscale_prior=ls_prior, + lengthscale_constraint=ls_constraint, + ard_num_dims=1, + ), + outputscale_prior=gpytorch.priors.SmoothedBoxPrior(a=1, b=4), ) - # Fit the model and check that the inducing points are updated - strat.add_data(train_X, train_Y) - strat.fit() - self.assertEqual( - strat.model.variational_strategy.inducing_points.shape, - auto_inducing_points.shape, + points = allocator.allocate_inducing_points( + inputs=inputs_aug, covar_module=covar_module, num_inducing=2 ) + self.assertTrue(points.shape[1] == 1) def test_fixed_allocator_allocate_inducing_points(self): config_str = """ - [common] - parnames = [par1] - stimuli_per_trial = 1 - outcome_types = [binary] - strategy_names = [init_strat] - - [par1] - par_type = continuous - lower_bound = 10 - upper_bound = 1000 - log_scale = True - - [init_strat] - generator = OptimizeAcqfGenerator - acqf = MCLevelSetEstimation - min_asks = 2 - model = GPClassificationModel - - [GPClassificationModel] - inducing_point_method = FixedAllocator - inducing_size = 2 - - [FixedAllocator] - points = [[0.1], [0.2]] - - """ + [common] + parnames = [par1] + stimuli_per_trial = 1 + outcome_types = [binary] + strategy_names = [init_strat] + + [par1] + par_type = continuous + lower_bound = 10 + upper_bound = 1000 + log_scale = True + + [init_strat] + generator = OptimizeAcqfGenerator + acqf = MCLevelSetEstimation + min_asks = 2 + model = GPClassificationModel + + [GPClassificationModel] + inducing_point_method = FixedAllocator + inducing_size = 2 + + [FixedAllocator] + points = [[0.1], [0.2]] + """ config = Config() config.update(config_str=config_str) strat = Strategy.from_config(config, "init_strat") self.assertTrue(isinstance(strat.model.inducing_point_method, FixedAllocator)) - # check that the bounds are scaled correctly - self.assertFalse( - torch.equal( - strat.model.inducing_point_method.bounds, torch.tensor([[10], [100]]) - ) - ) - self.assertTrue( - torch.equal( - strat.model.inducing_point_method.bounds, torch.tensor([[0], [1]]) - ) - ) - # Check that the inducing points are the same as the fixed points (pre-transformation) inducing_points_pre_transform = FixedAllocator( - points=torch.tensor([[0.1], [0.2]]) + points=torch.tensor([[0.1], [0.2]]), dim=1 ).allocate_inducing_points(num_inducing=2) self.assertTrue( torch.equal(inducing_points_pre_transform, torch.tensor([[0.1], [0.2]])) @@ -506,101 +321,133 @@ def test_fixed_allocator_allocate_inducing_points(self): ) ) + def test_allocator_model_fit(self): + config_str = """ + [common] + parnames = [par1, par2] + stimuli_per_trial = 1 + outcome_types = [binary] + strategy_names = [init_strat] + + [par1] + par_type = continuous + lower_bound = 10 + upper_bound = 1000 + log_scale = True -class TestGreedyAllocators(unittest.TestCase): - def test_greedy_variance_reduction_allocate_inducing_points(self): - # Mock data for testing - train_X = torch.rand(100, 1) - train_Y = torch.rand(100, 1) - lb = torch.tensor([0]) - ub = torch.tensor([1]) - bounds = torch.stack([lb, ub]) - model = GPClassificationModel( - lb=lb, - ub=ub, - inducing_point_method=GreedyVarianceReduction(bounds=bounds), - inducing_size=10, + [par2] + par_type = continuous + lower_bound = 0 + upper_bound = 1 + + [init_strat] + generator = OptimizeAcqfGenerator + acqf = MCLevelSetEstimation + min_asks = 2 + model = GPClassificationModel + + [GPClassificationModel] + inducing_size = 2 + """ + + config = Config() + config.update(config_str=config_str) + strat = Strategy.from_config(config, "init_strat") + self.assertIsInstance( + strat.model.inducing_point_method, GreedyVarianceReduction ) + self.assertIsNone(strat.model.inducing_point_method.last_allocator_used) - # Instantiate GreedyVarianceReduction allocator - allocator = GreedyVarianceReduction(bounds=bounds) + train_X = torch.tensor([[12.0, 0.0], [600.0, 1.0]]) + train_Y = torch.tensor([[1.0], [0.0]]) - # Allocate inducing points and verify output shape - inducing_points = allocator.allocate_inducing_points( - inputs=train_X, - covar_module=model.covar_module, - num_inducing=10, - input_batch_shape=torch.Size([]), + # Fit the model and check that the inducing points are updated + strat.add_data(train_X, train_Y) + strat.fit() + self.assertIs( + strat.model.inducing_point_method.last_allocator_used, + GreedyVarianceReduction, ) - inital_inducing_points = DummyAllocator( - bounds=torch.stack([torch.tensor([0]), torch.tensor([1])]) - ).allocate_inducing_points( - inputs=train_X, - covar_module=model.covar_module, - num_inducing=10, - input_batch_shape=torch.Size([]), + self.assertEqual( + strat.model.variational_strategy.inducing_points.shape, train_X.shape ) - self.assertEqual(inducing_points.shape, (10, 1)) - # Should be different from the initial inducing points - self.assertFalse( - torch.allclose(inducing_points, model.variational_strategy.inducing_points) + + def test_select_inducing_points(self): + """Verify that when we have n_induc > data size, we use data as inducing, + and otherwise we correctly select inducing points.""" + X, y = make_classification( + n_samples=100, + n_features=1, + n_redundant=0, + n_informative=1, + random_state=1, + n_clusters_per_class=1, ) - self.assertTrue( - torch.allclose( - inital_inducing_points, model.variational_strategy.inducing_points - ) + X, y = torch.Tensor(X), torch.Tensor(y) + inducing_size = 20 + + model = GPClassificationModel( + dim=1, + inducing_size=inducing_size, + inducing_point_method=GreedyVarianceReduction(dim=1), ) - # Then fit the model and check that the inducing points are updated - model.fit(train_X, train_Y) + # Test dummy + points = model.inducing_point_method.allocate_inducing_points( + inputs=None, + covar_module=model.covar_module, + num_inducing=inducing_size, + ) + self.assertTrue(torch.all(points == 0)) - self.assertTrue( - torch.allclose(inducing_points, model.variational_strategy.inducing_points) + model.set_train_data(X, y) + + points = model.inducing_point_method.allocate_inducing_points( + inputs=model.train_inputs[0], + covar_module=model.covar_module, + num_inducing=inducing_size, ) + self.assertTrue(len(points) <= 20) - def test_greedy_variance_from_config(self): - config_str = """ - [common] - parnames = [par1] - stimuli_per_trial = 1 - outcome_types = [binary] - strategy_names = [init_strat] - - [par1] - par_type = continuous - lower_bound = 10 - upper_bound = 1000 - log_scale = True - - [init_strat] - generator = OptimizeAcqfGenerator - acqf = MCLevelSetEstimation - min_asks = 2 - model = GPClassificationModel - - [GPClassificationModel] - inducing_point_method = GreedyVarianceReduction - inducing_size = 2 - """ + allocator = GreedyVarianceReduction(dim=1) + points = allocator.allocate_inducing_points( + inputs=model.train_inputs[0], + num_inducing=inducing_size, + covar_module=model.covar_module, + ) + self.assertTrue(len(points) <= 20) - config = Config() - config.update(config_str=config_str) - strat = Strategy.from_config(config, "init_strat") - self.assertTrue( - isinstance(strat.model.inducing_point_method, GreedyVarianceReduction) + allocator = KMeansAllocator(dim=1) + points = allocator.allocate_inducing_points( + inputs=model.train_inputs[0], + num_inducing=inducing_size, + covar_module=model.covar_module, ) + self.assertEqual(len(points), 20) - # check that the bounds are scaled correctly - self.assertFalse( - torch.equal( - strat.model.inducing_point_method.bounds, torch.tensor([[10], [100]]) - ) + allocator = SobolAllocator( + bounds=torch.stack([torch.tensor([0]), torch.tensor([1])]), dim=1 ) - self.assertTrue( - torch.equal( - strat.model.inducing_point_method.bounds, torch.tensor([[0], [1]]) - ) + points = allocator.allocate_inducing_points( + inputs=model.train_inputs[0], + num_inducing=inducing_size, + covar_module=model.covar_module, ) + self.assertTrue(len(points) <= 20) + + allocator = FixedAllocator(points=torch.tensor([[0], [1], [2], [3]]), dim=1) + points = allocator.allocate_inducing_points( + inputs=model.train_inputs[0], + num_inducing=inducing_size, + covar_module=model.covar_module, + ) + self.assertTrue(len(points) <= 20) + + def test_model_default_allocator(self): + model = GPClassificationModel(dim=2) + + self.assertIsInstance(model.inducing_point_method, GreedyVarianceReduction) + self.assertTrue(model.inducing_point_method.dim == 2) if __name__ == "__main__": diff --git a/tests/test_strategy.py b/tests/test_strategy.py index a95d5df6b..814eeb2e0 100644 --- a/tests/test_strategy.py +++ b/tests/test_strategy.py @@ -18,7 +18,6 @@ SobolGenerator, ) from aepsych.models.gp_classification import GPClassificationModel -from aepsych.models.inducing_point_allocators import AutoAllocator, SobolAllocator from aepsych.models.monotonic_rejection_gp import MonotonicRejectionGP from aepsych.strategy import SequentialStrategy, Strategy from aepsych.transforms import ( @@ -26,7 +25,6 @@ ParameterTransformedModel, ParameterTransforms, ) - from aepsych.transforms.ops import NormalizeScale @@ -47,11 +45,9 @@ def setUp(self): self.strat = Strategy( model=ParameterTransformedModel( MonotonicRejectionGP, - transforms=transforms, - dim=2, lb=lb, ub=ub, - inducing_point_method=SobolAllocator(bounds=torch.stack((lb, ub))), + transforms=transforms, monotonic_idxs=[1], ), generator=ParameterTransformedGenerator( @@ -59,6 +55,8 @@ def setUp(self): transforms=transforms, acqf=MonotonicMCLSE, acqf_kwargs=extra_acqf_args, + lb=lb, + ub=ub, ), min_asks=50, lb=lb, @@ -131,7 +129,6 @@ def test_finish_criteria(self): self.strat.gen() self.strat.add_data(np.r_[1.0, 1.0], [1]) self.assertFalse(self.strat.finished) - self.strat.gen() self.strat.add_data(np.r_[1.0, 1.0], [1]) self.assertFalse(self.strat.finished) # not enough "no" trials @@ -142,10 +139,36 @@ def test_finish_criteria(self): self.strat.finished ) # not enough difference between posterior min/max - for _ in range(50): - self.strat.gen() - self.strat.add_data(np.r_[0.0, 0.0], [0]) - self.assertTrue(self.strat.finished) + def test_min_post_range(self): + seed = 1 + torch.manual_seed(seed) + np.random.seed(seed) + lb = [0] + ub = [1] + + self.strat = Strategy( + model=GPClassificationModel( + dim=1, + ), + generator=SobolGenerator(lb=lb, ub=ub), + min_asks=10, + lb=lb, + ub=ub, + stimuli_per_trial=1, + outcome_types=["binary"], + min_post_range=0.3, + ) + + loops = 0 + while not self.strat.finished: + points = self.strat.gen(1) + response = int(np.random.rand() < points) + self.strat.add_data(points, torch.tensor([response])) + loops += 1 + + if loops > 50: + self.fail("min_post_range didn't finish even after 50 loops.") + break def test_max_asks(self): self.strat.max_asks = 50 @@ -167,11 +190,7 @@ def test_keep_most_recent(self): self.strat = Strategy( model=GPClassificationModel( - lb=lb, - ub=ub, - inducing_point_method=SobolAllocator( - bounds=torch.stack([torch.tensor(lb), torch.tensor(ub)]) - ), + dim=2, ), generator=SobolGenerator(lb=lb, ub=ub), min_asks=50, @@ -193,17 +212,13 @@ def test_keep_most_recent(self): ) def test_run_indefinitely(self): - lb = [-1.0, -1.0] - ub = [1.0, 1.0] + lb = torch.tensor([-1.0, -1.0]) + ub = torch.tensor([1.0, 1.0]) with self.assertWarns(UserWarning): self.strat = Strategy( model=GPClassificationModel( - lb=lb, - ub=ub, - inducing_point_method=SobolAllocator( - bounds=torch.stack([torch.tensor(lb), torch.tensor(ub)]) - ), + dim=2, ), generator=SobolGenerator(lb=lb, ub=ub), lb=lb, @@ -375,11 +390,7 @@ def test_no_gpu_model_warn(self): stimuli_per_trial=1, outcome_types=["binary"], model=GPClassificationModel( - lb=[0], - ub=[1], - inducing_point_method=AutoAllocator( - bounds=torch.stack([torch.tensor([0]), torch.tensor([1])]) - ), + dim=1, ), generator=SobolGenerator(lb=[0], ub=[1]), use_gpu_modeling=True, @@ -396,13 +407,11 @@ def test_no_gpu_generator_warn(self): stimuli_per_trial=1, outcome_types=["binary"], model=GPClassificationModel( - lb=[0], - ub=[1], - inducing_point_method=AutoAllocator( - bounds=torch.stack([torch.tensor([0]), torch.tensor([1])]) - ), + dim=1, + ), + generator=OptimizeAcqfGenerator( + lb=[0], ub=[1], acqf=MCLevelSetEstimation ), - generator=OptimizeAcqfGenerator(acqf=MCLevelSetEstimation), use_gpu_generating=True, ) diff --git a/tests/test_transforms.py b/tests/test_transforms.py index 55fcf6e8c..e48cf8797 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -11,7 +11,6 @@ from aepsych.config import Config, ParameterConfigError from aepsych.generators import SobolGenerator from aepsych.models import GPClassificationModel -from aepsych.models.inducing_point_allocators import AutoAllocator from aepsych.strategy import SequentialStrategy from aepsych.transforms import ( ParameterTransformedGenerator, @@ -69,8 +68,8 @@ def test_generator_init_equivalent(self): class_gen = ParameterTransformedGenerator( generator=SobolGenerator, - lb=torch.tensor([1, 1]), - ub=torch.tensor([100, 100]), + lb=torch.tensor([1.0, 1.0]), + ub=torch.tensor([100.0, 100.0]), seed=12345, transforms=self.strat.strat_list[0].transforms, ) @@ -89,21 +88,23 @@ def test_generator_init_equivalent(self): def test_model_init_equivalent(self): config_model = self.strat.strat_list[1].model + config_generator = self.strat.strat_list[1].generator obj_model = ParameterTransformedModel( model=GPClassificationModel, - lb=torch.tensor([1, 1]), - ub=torch.tensor([100, 100]), - inducing_point_method=AutoAllocator( - bounds=torch.stack((torch.tensor([1, 1]), torch.tensor([100, 100]))) - ), + transforms=self.strat.strat_list[1].transforms, + dim=2, + ) + obj_generator = ParameterTransformedGenerator( + generator=SobolGenerator, + lb=torch.tensor([1.0, 1.0]), + ub=torch.tensor([100.0, 100.0]), transforms=self.strat.strat_list[1].transforms, ) self.assertTrue(type(config_model._base_obj) is type(obj_model._base_obj)) - self.assertTrue(torch.equal(config_model.bounds, obj_model.bounds)) - self.assertTrue(torch.equal(config_model.bounds, obj_model.bounds)) - + self.assertTrue(torch.equal(config_generator.lb, obj_generator.lb)) + self.assertTrue(torch.equal(config_generator.ub, obj_generator.ub)) self.assertEqual( len(config_model.transforms.values()), len(obj_model.transforms.values()) ) @@ -209,8 +210,8 @@ def test_transform_manual_generator(self): class TransformsLog10Test(unittest.TestCase): def test_transform_reshape3D(self): - lb = torch.tensor([-1, 0, 10]) - ub = torch.tensor([-1e-6, 9, 99]) + lb = torch.tensor([-1.0, 0.0, 10.0]) + ub = torch.tensor([-1e-6, 9.0, 99.0]) x = SobolGenerator(lb=lb, ub=ub, stimuli_per_trial=2).gen(4) transforms = ParameterTransforms( @@ -583,6 +584,7 @@ def test_fixed_from_config(self): self.assertTrue(len(strat.strat_list[0].generator.lb) == 2) self.assertTrue(len(strat.strat_list[0].generator.ub) == 2) + self.assertTrue(strat.strat_list[-1].model.dim == 2) bad_config_str = """ [common] diff --git a/tests/test_utils.py b/tests/test_utils.py index 8b9e661a7..e8de9fe3f 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -10,7 +10,6 @@ import numpy as np import torch from aepsych.models import GPClassificationModel -from aepsych.models.inducing_point_allocators import AutoAllocator from aepsych.utils import _process_bounds, dim_grid, make_scaled_sobol @@ -33,14 +32,11 @@ def test_dim_grid_model_size(self): dim = 1 gridsize = 10 mb = GPClassificationModel( - lb=lb, - ub=ub, - dim=dim, - inducing_point_method=AutoAllocator( - bounds=torch.stack([torch.tensor([lb]), torch.tensor([ub])]) - ), + dim=1, + ) + grid = dim_grid( + lower=torch.tensor([lb]), upper=torch.tensor([ub]), gridsize=gridsize ) - grid = GPClassificationModel.dim_grid(mb, gridsize=gridsize) self.assertEqual(grid.shape, torch.Size([10, 1])) def test_dim_grid_slice(self): diff --git a/tests_gpu/generators/test_optimize_acqf_generator.py b/tests_gpu/generators/test_optimize_acqf_generator.py index a6d780f19..9caef06e4 100644 --- a/tests_gpu/generators/test_optimize_acqf_generator.py +++ b/tests_gpu/generators/test_optimize_acqf_generator.py @@ -22,7 +22,7 @@ from aepsych.acquisition.mutual_information import BernoulliMCMutualInformation from aepsych.generators import OptimizeAcqfGenerator from aepsych.models import GPClassificationModel -from aepsych.models.inducing_point_allocators import GreedyVarianceReduction +from aepsych.models.inducing_points import GreedyVarianceReduction from aepsych.strategy import Strategy from parameterized import parameterized @@ -49,16 +49,23 @@ class TestOptimizeAcqfGenerator(unittest.TestCase): @parameterized.expand(acqfs) def test_gpu_smoketest(self, acqf, acqf_kwargs): - lb = torch.tensor([0]) - ub = torch.tensor([1]) + lb = torch.tensor([0.0]) + ub = torch.tensor([1.0]) + bounds = torch.stack([lb, ub]) + inducing_size = 10 + model = GPClassificationModel( - lb=lb, - ub=ub, - inducing_size=10, - inducing_point_method=GreedyVarianceReduction(bounds=torch.stack([lb, ub])), + dim=1, + inducing_size=inducing_size, + inducing_point_method=GreedyVarianceReduction(dim=1), ) - generator = OptimizeAcqfGenerator(acqf=acqf, acqf_kwargs=acqf_kwargs) + generator = OptimizeAcqfGenerator( + acqf=acqf, + acqf_kwargs=acqf_kwargs, + lb=torch.tensor([0.0]), + ub=torch.tensor([1.0]), + ) strat = Strategy( lb=torch.tensor([0]), diff --git a/tests_gpu/models/test_gp_classification.py b/tests_gpu/models/test_gp_classification.py index fc5995db6..7204fa30b 100644 --- a/tests_gpu/models/test_gp_classification.py +++ b/tests_gpu/models/test_gp_classification.py @@ -10,8 +10,6 @@ import torch -from aepsych.models.inducing_point_allocators import SobolAllocator - # run on single threads to keep us from deadlocking weirdly in CI if "CI" in os.environ or "SANDCASTLE" in os.environ: torch.set_num_threads(1) @@ -107,14 +105,11 @@ def test_1d_classification_gpu(self): Just see if we memorize the training set but on gpu """ X, y = self.X, self.y + inducing_size = 10 model = GPClassificationModel( - torch.Tensor([-3]), - torch.Tensor([3]), - inducing_size=10, - inducing_point_method=SobolAllocator( - bounds=torch.stack([torch.tensor([-3]), torch.tensor([3])]) - ), + dim=1, + inducing_size=inducing_size, ).to("cuda") model.fit(X[:50], y[:50]) diff --git a/tests_gpu/test_strategy.py b/tests_gpu/test_strategy.py index 99dd51e0a..1fbb98790 100644 --- a/tests_gpu/test_strategy.py +++ b/tests_gpu/test_strategy.py @@ -8,11 +8,9 @@ import unittest import torch - from aepsych.acquisition.monotonic_rejection import MonotonicMCLSE from aepsych.generators import OptimizeAcqfGenerator, SobolGenerator from aepsych.models.gp_classification import GPClassificationModel -from aepsych.models.inducing_point_allocators import AutoAllocator from aepsych.strategy import Strategy @@ -20,8 +18,8 @@ class TestStrategyGPU(unittest.TestCase): def test_gpu_no_model_generator_warn(self): with self.assertWarns(UserWarning): Strategy( - lb=[0], - ub=[1], + lb=[0.0], + ub=[1.0], stimuli_per_trial=1, outcome_types=["binary"], min_asks=1, @@ -32,19 +30,15 @@ def test_gpu_no_model_generator_warn(self): def test_no_gpu_acqf(self): with self.assertWarns(UserWarning): Strategy( - lb=[0], - ub=[1], + lb=[0.0], + ub=[1.0], stimuli_per_trial=1, outcome_types=["binary"], min_asks=1, model=GPClassificationModel( - lb=[0], - ub=[1], - inducing_point_method=AutoAllocator( - bounds=torch.stack([torch.tensor([0]), torch.tensor([1])]) - ), + dim=1, ), - generator=OptimizeAcqfGenerator(acqf=MonotonicMCLSE), + generator=OptimizeAcqfGenerator(acqf=MonotonicMCLSE, lb=[0], ub=[1]), use_gpu_generating=True, )