From f4a74333af3a03a21bfb8c3ce9cd46b2b543eb78 Mon Sep 17 00:00:00 2001 From: Sait Cakmak Date: Mon, 29 Jul 2024 10:26:12 -0700 Subject: [PATCH] Delete legacy get posterior mean Summary: Legacy models have been deprecated for a while and they're being removed one by one. This particular one was being used to construct MC-PosteriorMean (qSimpleRegret) and was used only in `compute_posterior_pareto_frontier`. Updated the usage with MBM equivalent and removed the code. Differential Revision: D60396661 --- ax/models/tests/test_posterior_mean.py | 119 ------------------------- ax/models/torch/posterior_mean.py | 84 ----------------- ax/plot/pareto_utils.py | 30 +++---- sphinx/source/models.rst | 8 -- 4 files changed, 10 insertions(+), 231 deletions(-) delete mode 100644 ax/models/tests/test_posterior_mean.py delete mode 100644 ax/models/torch/posterior_mean.py diff --git a/ax/models/tests/test_posterior_mean.py b/ax/models/tests/test_posterior_mean.py deleted file mode 100644 index 66bbc09e484..00000000000 --- a/ax/models/tests/test_posterior_mean.py +++ /dev/null @@ -1,119 +0,0 @@ -#!/usr/bin/env python3 -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# This source code is licensed under the MIT license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -from typing import Any, Dict - -import torch -from ax.core.search_space import SearchSpaceDigest -from ax.models.torch.botorch import BotorchModel -from ax.models.torch.botorch_moo import MultiObjectiveBotorchModel -from ax.models.torch.posterior_mean import get_PosteriorMean -from ax.models.torch_base import TorchOptConfig -from ax.utils.common.testutils import TestCase -from ax.utils.testing.mock import fast_botorch_optimize -from botorch.utils.datasets import SupervisedDataset - - -# TODO (jej): Streamline testing for a simple acquisition function. -class PosteriorMeanTest(TestCase): - def setUp(self) -> None: - self.tkwargs: Dict[str, Any] = { - "device": torch.device("cpu"), - "dtype": torch.double, - } - self.X = torch.tensor([[1.0, 2.0, 3.0], [2.0, 3.0, 4.0]], **self.tkwargs) - self.Y = torch.tensor([[3.0], [4.0]], **self.tkwargs) - self.Yvar = torch.tensor([[0.0], [2.0]], **self.tkwargs) - self.bounds = [(0.0, 1.0), (1.0, 4.0), (2.0, 5.0)] - self.feature_names = ["x1", "x2", "x3"] - self.objective_weights = torch.tensor([1.0], **self.tkwargs) - self.outcome_constraints = ( - torch.tensor([[1.0]], **self.tkwargs), - torch.tensor([[5.0]], **self.tkwargs), - ) - self.search_space_digest = SearchSpaceDigest( - feature_names=self.feature_names, - bounds=self.bounds, - ) - - @fast_botorch_optimize - def test_GetPosteriorMean(self) -> None: - - model = BotorchModel(acqf_constructor=get_PosteriorMean) - dataset = SupervisedDataset( - X=self.X, - Y=self.Y, - Yvar=self.Yvar, - feature_names=self.feature_names, - outcome_names=["y"], - ) - model.fit( - datasets=[dataset], - search_space_digest=self.search_space_digest, - ) - - # test model.gen() with no outcome_constraints. Analytic. - new_X_dummy = torch.rand(1, 1, 3, **self.tkwargs) - gen_results = model.gen( - n=1, - search_space_digest=self.search_space_digest, - torch_opt_config=TorchOptConfig( - objective_weights=self.objective_weights, - ), - ) - self.assertTrue( - torch.equal(gen_results.weights, torch.ones(1, dtype=self.tkwargs["dtype"])) - ) - - # test model.gen() works with outcome_constraints. qSimpleRegret. - new_X_dummy = torch.rand(1, 1, 3, **self.tkwargs) - model.gen( - n=1, - search_space_digest=self.search_space_digest, - torch_opt_config=TorchOptConfig( - objective_weights=self.objective_weights, - outcome_constraints=self.outcome_constraints, - ), - ) - - # test model.gen() works with chebyshev scalarization. - model = MultiObjectiveBotorchModel(acqf_constructor=get_PosteriorMean) - model.fit( - datasets=[dataset, dataset], - search_space_digest=self.search_space_digest, - ) - new_X_dummy = torch.rand(1, 1, 3, **self.tkwargs) - model.gen( - n=1, - search_space_digest=self.search_space_digest, - torch_opt_config=TorchOptConfig( - objective_weights=torch.ones(2, **self.tkwargs), - outcome_constraints=( - torch.tensor([[1.0, 0.0]], **self.tkwargs), - torch.tensor([[5.0]], **self.tkwargs), - ), - objective_thresholds=torch.zeros(2, **self.tkwargs), - model_gen_options={ - "acquisition_function_kwargs": {"chebyshev_scalarization": True} - }, - ), - ) - - # ValueError with empty X_Observed - with self.assertRaises(ValueError): - get_PosteriorMean( - # pyre-fixme[6]: For 1st param expected `Model` but got - # `MultiObjectiveBotorchModel`. - model=model, - objective_weights=self.objective_weights, - X_observed=None, - ) - - # test model.predict() - new_X_dummy = torch.rand(1, 1, 3, **self.tkwargs) - model.predict(new_X_dummy) diff --git a/ax/models/torch/posterior_mean.py b/ax/models/torch/posterior_mean.py deleted file mode 100644 index 250bcb75fa3..00000000000 --- a/ax/models/torch/posterior_mean.py +++ /dev/null @@ -1,84 +0,0 @@ -#!/usr/bin/env python3 -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# This source code is licensed under the MIT license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - - -from typing import Any, Optional, Tuple - -import torch -from ax.models.torch.botorch_defaults import NO_FEASIBLE_POINTS_MESSAGE -from ax.utils.common.typeutils import checked_cast -from botorch.acquisition.acquisition import AcquisitionFunction -from botorch.acquisition.monte_carlo import qSimpleRegret -from botorch.acquisition.objective import ConstrainedMCObjective, GenericMCObjective -from botorch.acquisition.utils import get_infeasible_cost -from botorch.models.model import Model -from botorch.posteriors.gpytorch import GPyTorchPosterior -from botorch.utils import ( - get_objective_weights_transform, - get_outcome_constraint_transforms, -) -from botorch.utils.multi_objective.scalarization import get_chebyshev_scalarization -from torch import Tensor - - -def get_PosteriorMean( - model: Model, - objective_weights: Tensor, - outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None, - X_observed: Optional[Tensor] = None, - X_pending: Optional[Tensor] = None, - **kwargs: Any, -) -> AcquisitionFunction: - r"""Instantiates a PosteriorMean acquisition function. - - Note: If no OutcomeConstraints given, return an analytic acquisition - function. This requires {optimizer_kwargs: {joint_optimization: True}} or an - optimizer that does not assume pending point support. - - Args: - objective_weights: The objective is to maximize a weighted sum of - the columns of f(x). These are the weights. - outcome_constraints: A tuple of (A, b). For k outcome constraints - and m outputs at f(x), A is (k x m) and b is (k x 1) such that - A f(x) <= b. (Not used by single task models) - X_observed: A tensor containing points observed for all objective - outcomes and outcomes that appear in the outcome constraints (if - there are any). - X_pending: A tensor containing points whose evaluation is pending (i.e. - that have been submitted for evaluation) present for all objective - outcomes and outcomes that appear in the outcome constraints (if - there are any). - - Returns: - PosteriorMean: The instantiated acquisition function. - """ - if X_observed is None: - raise ValueError(NO_FEASIBLE_POINTS_MESSAGE) - # construct Objective module - if kwargs.get("chebyshev_scalarization", False): - with torch.no_grad(): - Y = checked_cast(GPyTorchPosterior, model.posterior(X_observed)).mean - obj_tf = get_chebyshev_scalarization(weights=objective_weights, Y=Y) - else: - obj_tf = get_objective_weights_transform(objective_weights) - - # pyre-fixme[53]: Captured variable `obj_tf` is not annotated. - def obj_fn(samples: Tensor, X: Optional[Tensor] = None) -> Tensor: - return obj_tf(samples) - - if outcome_constraints is None: - objective = GenericMCObjective(objective=obj_fn) - else: - con_tfs = get_outcome_constraint_transforms(outcome_constraints) - inf_cost = get_infeasible_cost(X=X_observed, model=model, objective=obj_fn) - objective = ConstrainedMCObjective( - objective=obj_fn, constraints=con_tfs or [], infeasible_cost=inf_cost - ) - # Use qSimpleRegret, not analytic posterior, to handle arbitrary objective fns. - acq_func = qSimpleRegret(model, objective=objective) - return acq_func diff --git a/ax/plot/pareto_utils.py b/ax/plot/pareto_utils.py index 24a63bdcab8..babd7ea8487 100644 --- a/ax/plot/pareto_utils.py +++ b/ax/plot/pareto_utils.py @@ -36,11 +36,11 @@ from ax.modelbridge.registry import Models from ax.modelbridge.torch import TorchModelBridge from ax.modelbridge.transforms.search_space_to_float import SearchSpaceToFloat -from ax.models.torch.posterior_mean import get_PosteriorMean from ax.models.torch_base import TorchModel from ax.utils.common.logger import get_logger from ax.utils.common.typeutils import checked_cast from ax.utils.stats.statstools import relativize +from botorch.acquisition.monte_carlo import qSimpleRegret from botorch.utils.multi_objective import is_non_dominated from botorch.utils.multi_objective.hypervolume import infer_reference_point @@ -347,7 +347,6 @@ def compute_posterior_pareto_frontier( absolute_metrics: Optional[List[str]] = None, num_points: int = 10, trial_index: Optional[int] = None, - chebyshev: bool = True, ) -> ParetoFrontierResults: """Compute the Pareto frontier between two objectives. For experiments with batch trials, a trial index or data object must be provided. @@ -368,16 +367,10 @@ def compute_posterior_pareto_frontier( will be in % relative to status_quo). num_points: The number of points to compute on the Pareto frontier. - chebyshev: Whether to use augmented_chebyshev_scalarization - when computing Pareto Frontier points. Returns: ParetoFrontierResults: A NamedTuple with fields listed in its definition. """ - model_gen_options = { - "acquisition_function_kwargs": {"chebyshev_scalarization": chebyshev} - } - if ( trial_index is None and data is None @@ -426,11 +419,11 @@ def compute_posterior_pareto_frontier( secondary_objective=secondary_objective, outcome_constraints=outcome_constraints, ) - model = Models.MOO( + model = Models.BOTORCH_MODULAR( experiment=experiment, data=data, - acqf_constructor=get_PosteriorMean, optimization_config=oc, + botorch_acqf_class=qSimpleRegret, ) status_quo = experiment.status_quo @@ -472,7 +465,7 @@ def compute_posterior_pareto_frontier( # TODO: (jej) T64002590 Let this serve as a starting point for optimization. # ex. Add global spacing criterion. Implement on BoTorch side. # pyre-fixme [6]: Expected different type for model_gen_options - run = model.gen(1, model_gen_options=model_gen_options, optimization_config=oc) + run = model.gen(1, optimization_config=oc) param_dicts.append(run.arms[0].parameters) # Call predict on points to get their decomposed metrics. @@ -518,6 +511,7 @@ def _extract_pareto_frontier_results( for metric in metrics: if metric not in absolute_metrics and metric in sq_mean: + # pyre-fixme [6]: Expected List[float] but got ndarray. means_out[metric], sems_out[metric] = relativize( means_t=means_out[metric], sems_t=sems_out[metric], @@ -555,18 +549,14 @@ def _validate_outcome_constraints( def _build_new_optimization_config( - # pyre-fixme[2]: Parameter must be annotated. - weights, - # pyre-fixme[2]: Parameter must be annotated. - primary_objective, - # pyre-fixme[2]: Parameter must be annotated. - secondary_objective, - # pyre-fixme[2]: Parameter must be annotated. - outcome_constraints=None, + weights: np.ndarray, + primary_objective: Metric, + secondary_objective: Metric, + outcome_constraints: Optional[List[OutcomeConstraint]] = None, ) -> MultiObjectiveOptimizationConfig: obj = ScalarizedObjective( metrics=[primary_objective, secondary_objective], - weights=weights, + weights=weights.tolist(), minimize=False, ) optimization_config = MultiObjectiveOptimizationConfig( diff --git a/sphinx/source/models.rst b/sphinx/source/models.rst index c7f104a1da5..f3efa5dd433 100644 --- a/sphinx/source/models.rst +++ b/sphinx/source/models.rst @@ -303,14 +303,6 @@ ax.models.torch.fully_bayesian_model_utils module :undoc-members: :show-inheritance: -ax.models.torch.posterior_mean module -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -.. automodule:: ax.models.torch.posterior_mean - :members: - :undoc-members: - :show-inheritance: - ax.models.torch.utils module ~~~~~~~~~~~~~~~~~~~~~~~~~~~~