Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Revert D60472381 #2627

Closed
wants to merge 1 commit into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
92 changes: 1 addition & 91 deletions ax/benchmark/tests/test_benchmark.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,11 +6,9 @@
# pyre-strict

import tempfile
from typing import List
from unittest.mock import patch

import numpy as np
import torch
from ax.benchmark.benchmark import (
_create_benchmark_experiment,
benchmark_multiple_problems_methods,
Expand All @@ -23,10 +21,7 @@
BenchmarkMethod,
get_benchmark_scheduler_options,
)
from ax.benchmark.benchmark_problem import (
BenchmarkProblem,
create_single_objective_problem_from_botorch,
)
from ax.benchmark.benchmark_problem import create_single_objective_problem_from_botorch
from ax.benchmark.benchmark_result import BenchmarkResult
from ax.benchmark.methods.modular_botorch import get_sobol_botorch_modular_acquisition
from ax.benchmark.metrics.base import GroundTruthMetricMixin
Expand Down Expand Up @@ -58,13 +53,6 @@
from botorch.models.gp_regression import SingleTaskGP
from botorch.optim.optimize import optimize_acqf
from botorch.test_functions.synthetic import Branin
from pyre_extensions import none_throws


# Branin is subclassed here for convenience.
class ThreeProblem(Branin):
def evaluate_true(self, X: torch.Tensor) -> torch.Tensor:
return torch.tensor(3.0, dtype=torch.float64)


class TestBenchmark(TestCase):
Expand Down Expand Up @@ -102,84 +90,6 @@ def test_batch(self) -> None:
)
self.assertEqual(mock_optimize_acqf.call_args.kwargs["q"], batch_size)

def test_noisy_property_test(self) -> None:
problem = create_single_objective_problem_from_botorch(
test_problem_class=ThreeProblem,
test_problem_kwargs={"noise_std": 1.0},
lower_is_better=True,
num_trials=5,
)

res = benchmark_replication(
problem=problem,
method=get_sobol_benchmark_method(),
seed=0,
)
self.assertTrue((res.optimization_trace == 3).all())
df = none_throws(res.experiment).fetch_data().df
self.assertTrue(
(df.loc[df["metric_name"] == "ThreeProblem", "mean"] != 3).all()
)

def test_characterization_noisy_soo(self) -> None:
soo_problem = get_single_objective_benchmark_problem(
test_problem_kwargs={"noise_std": 1.0}
)
expected = [
104.36542659448965,
90.20076516036951,
56.18382601460175,
7.5867050504430775,
]
self._noisy_characterization_test(problem=soo_problem, expected_result=expected)

def test_characterization_noisy_moo(self) -> None:
moo_problem = get_multi_objective_benchmark_problem(
test_problem_kwargs={"noise_std": 1.0}
)
expected = [0.0, 0.0, 6.952682254445629, 6.952682254445629]
self._noisy_characterization_test(problem=moo_problem, expected_result=expected)

def test_characterization_noisy_soo_surrogate(self) -> None:
# Surrogate construction shouldn't be random
torch.manual_seed(0)
soo_surrogate_problem = get_soo_surrogate(noise_stds=1.0)
expected = [
104.36542659448779,
104.36542659448779,
104.36542659448779,
104.36542659448779,
]
self._noisy_characterization_test(
problem=soo_surrogate_problem,
expected_result=expected,
)

def test_characterization_noisy_moo_surrogate(self) -> None:
# Surrogate construction shouldn't be random
torch.manual_seed(0)
moo_surrogate_problem = get_moo_surrogate(noise_stds=1.0)
expected = [
24.999999999997655,
24.999999999997655,
24.999999999997655,
24.999999999997655,
]
self._noisy_characterization_test(
problem=moo_surrogate_problem,
expected_result=expected,
)

def _noisy_characterization_test(
self, problem: BenchmarkProblem, expected_result: List[float]
) -> None:
res = benchmark_replication(
problem=problem,
method=get_sobol_benchmark_method(),
seed=0,
)
self.assertTrue((res.optimization_trace == expected_result).all())

def test_storage(self) -> None:
problem = get_single_objective_benchmark_problem()
res = benchmark_replication(
Expand Down
84 changes: 22 additions & 62 deletions ax/utils/testing/benchmark_stubs.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,26 +17,24 @@
MultiObjectiveBenchmarkProblem,
)
from ax.benchmark.benchmark_result import AggregatedBenchmarkResult, BenchmarkResult
from ax.benchmark.metrics.benchmark import BenchmarkMetric
from ax.benchmark.problems.surrogate import (
MOOSurrogateBenchmarkProblem,
SOOSurrogateBenchmarkProblem,
)
from ax.benchmark.runners.surrogate import SurrogateRunner
from ax.core.experiment import Experiment
from ax.core.objective import MultiObjective, Objective
from ax.core.optimization_config import (
MultiObjectiveOptimizationConfig,
OptimizationConfig,
)
from ax.core.outcome_constraint import ObjectiveThreshold
from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy
from ax.modelbridge.registry import Models
from ax.modelbridge.torch import TorchModelBridge
from ax.models.torch.botorch_modular.model import BoTorchModel
from ax.models.torch.botorch_modular.surrogate import Surrogate
from ax.service.scheduler import SchedulerOptions
from ax.utils.common.constants import Keys
from ax.utils.common.typeutils import checked_cast
from ax.utils.testing.core_stubs import (
get_branin_experiment,
get_branin_experiment_with_multi_objective,
Expand All @@ -62,13 +60,11 @@ def get_single_objective_benchmark_problem(


def get_multi_objective_benchmark_problem(
observe_noise_sd: bool = False,
num_trials: int = 4,
test_problem_kwargs: Optional[Dict[str, Any]] = None,
observe_noise_sd: bool = False, num_trials: int = 4
) -> MultiObjectiveBenchmarkProblem:
return create_multi_objective_problem_from_botorch(
test_problem_class=BraninCurrin,
test_problem_kwargs={} if test_problem_kwargs is None else test_problem_kwargs,
test_problem_kwargs={},
num_trials=num_trials,
observe_noise_sd=observe_noise_sd,
)
Expand Down Expand Up @@ -98,97 +94,61 @@ def get_sobol_benchmark_method() -> BenchmarkMethod:
)


def get_soo_surrogate(noise_stds: float = 0.0) -> SOOSurrogateBenchmarkProblem:
outcome_name = "branin"
observe_noise_stds = True
def get_soo_surrogate() -> SOOSurrogateBenchmarkProblem:
experiment = get_branin_experiment(with_completed_trial=True)

optimization_config = OptimizationConfig(
objective=Objective(
metric=BenchmarkMetric(
name=outcome_name,
lower_is_better=False,
observe_noise_sd=observe_noise_stds,
),
minimize=False,
)
)
surrogate = Surrogate(botorch_model_class=SingleTaskGP)
model_bridge = TorchModelBridge(
surrogate = TorchModelBridge(
experiment=experiment,
search_space=experiment.search_space,
model=BoTorchModel(surrogate=surrogate),
model=BoTorchModel(surrogate=Surrogate(botorch_model_class=SingleTaskGP)),
data=experiment.lookup_data(),
transforms=[],
)
datasets = surrogate.training_data
runner = SurrogateRunner(
name=outcome_name,
name="test",
search_space=experiment.search_space,
outcome_names=[outcome_name],
get_surrogate_and_datasets=lambda: (model_bridge, datasets),
noise_stds=noise_stds,
outcome_names=["branin"],
get_surrogate_and_datasets=lambda: (surrogate, []),
)
return SOOSurrogateBenchmarkProblem(
name="test",
search_space=experiment.search_space,
optimization_config=optimization_config,
optimization_config=checked_cast(
OptimizationConfig, experiment.optimization_config
),
num_trials=6,
observe_noise_stds=observe_noise_stds,
observe_noise_stds=True,
optimal_value=0.0,
runner=runner,
is_noiseless=runner.is_noiseless,
)


def get_moo_surrogate(noise_stds: float = 0.0) -> MOOSurrogateBenchmarkProblem:
observe_noise_stds = True
outcome_names = ["branin_a", "branin_b"]
# set this to be easy to beat, so hypervolume computations aren't all zero
ref_point = [10.0, 10.0]
surrogate = Surrogate(botorch_model_class=SingleTaskGP)
def get_moo_surrogate() -> MOOSurrogateBenchmarkProblem:
experiment = get_branin_experiment_with_multi_objective(with_completed_trial=True)
model_bridge = TorchModelBridge(
surrogate = TorchModelBridge(
experiment=experiment,
search_space=experiment.search_space,
model=BoTorchModel(surrogate=surrogate),
model=BoTorchModel(surrogate=Surrogate(botorch_model_class=SingleTaskGP)),
data=experiment.lookup_data(),
transforms=[],
)

runner = SurrogateRunner(
name="test",
search_space=experiment.search_space,
outcome_names=outcome_names,
get_surrogate_and_datasets=lambda: (model_bridge, surrogate.training_data),
noise_stds=noise_stds,
outcome_names=["branin_a", "branin_b"],
get_surrogate_and_datasets=lambda: (surrogate, []),
)
metrics = [
BenchmarkMetric(
name=name,
lower_is_better=True,
observe_noise_sd=observe_noise_stds,
)
for name in outcome_names
]
objectives = [Objective(metric=metric) for metric in metrics]
objective_thresholds = [
ObjectiveThreshold(metric=metric, bound=ref_p, relative=False)
for metric, ref_p in zip(metrics, ref_point)
]
optimization_config = MultiObjectiveOptimizationConfig(
objective=MultiObjective(objectives=objectives),
objective_thresholds=objective_thresholds,
)

return MOOSurrogateBenchmarkProblem(
name="test",
search_space=experiment.search_space,
optimization_config=optimization_config,
optimization_config=checked_cast(
MultiObjectiveOptimizationConfig, experiment.optimization_config
),
num_trials=10,
observe_noise_stds=True,
optimal_value=1.0,
reference_point=ref_point,
reference_point=[],
runner=runner,
is_noiseless=runner.is_noiseless,
)
Expand Down