diff --git a/ax/benchmark/metrics/jenatton.py b/ax/benchmark/metrics/jenatton.py index dd1da3205fd..4be7f5eae1b 100644 --- a/ax/benchmark/metrics/jenatton.py +++ b/ax/benchmark/metrics/jenatton.py @@ -5,78 +5,11 @@ # pyre-strict -from __future__ import annotations +from typing import Optional -from typing import Any, Optional - -import numpy as np -import pandas as pd -from ax.benchmark.metrics.base import BenchmarkMetricBase, GroundTruthMetricMixin -from ax.core.base_trial import BaseTrial -from ax.core.data import Data -from ax.core.metric import MetricFetchE, MetricFetchResult -from ax.utils.common.result import Err, Ok from ax.utils.common.typeutils import not_none -class JenattonMetric(BenchmarkMetricBase): - """Jenatton metric for hierarchical search spaces.""" - - has_ground_truth: bool = True - - def __init__( - self, - name: str = "jenatton", - noise_std: float = 0.0, - observe_noise_sd: bool = False, - ) -> None: - super().__init__(name=name) - self.noise_std = noise_std - self.observe_noise_sd = observe_noise_sd - self.lower_is_better = True - - def fetch_trial_data(self, trial: BaseTrial, **kwargs: Any) -> MetricFetchResult: - try: - mean = [ - jenatton_test_function(**arm.parameters) # pyre-ignore [6] - for _, arm in trial.arms_by_name.items() - ] - if self.noise_std != 0: - mean = [m + self.noise_std * np.random.randn() for m in mean] - df = pd.DataFrame( - { - "arm_name": [name for name, _ in trial.arms_by_name.items()], - "metric_name": self.name, - "mean": mean, - "sem": self.noise_std if self.observe_noise_sd else None, - "trial_index": trial.index, - } - ) - return Ok(value=Data(df=df)) - - except Exception as e: - return Err( - MetricFetchE(message=f"Failed to fetch {self.name}", exception=e) - ) - - def make_ground_truth_metric(self) -> GroundTruthJenattonMetric: - return GroundTruthJenattonMetric(original_metric=self) - - -class GroundTruthJenattonMetric(JenattonMetric, GroundTruthMetricMixin): - def __init__(self, original_metric: JenattonMetric) -> None: - """ - Args: - original_metric: The original JenattonMetric to which this metric - corresponds. - """ - super().__init__( - name=self.get_ground_truth_name(original_metric), - noise_std=0.0, - observe_noise_sd=False, - ) - - def jenatton_test_function( x1: Optional[int] = None, x2: Optional[int] = None, diff --git a/ax/benchmark/problems/synthetic/hss/jenatton.py b/ax/benchmark/problems/synthetic/hss/jenatton.py index f545ac39400..67fa1755372 100644 --- a/ax/benchmark/problems/synthetic/hss/jenatton.py +++ b/ax/benchmark/problems/synthetic/hss/jenatton.py @@ -5,18 +5,52 @@ # pyre-strict +from dataclasses import dataclass +from typing import Optional + +import torch from ax.benchmark.benchmark_problem import BenchmarkProblem -from ax.benchmark.metrics.jenatton import JenattonMetric +from ax.benchmark.metrics.benchmark import BenchmarkMetric +from ax.benchmark.metrics.jenatton import jenatton_test_function +from ax.benchmark.runners.botorch_test import ( + ParamBasedTestProblem, + ParamBasedTestProblemRunner, +) from ax.core.objective import Objective from ax.core.optimization_config import OptimizationConfig from ax.core.parameter import ChoiceParameter, ParameterType, RangeParameter from ax.core.search_space import HierarchicalSearchSpace -from ax.runners.synthetic import SyntheticRunner +from ax.core.types import TParameterization + + +@dataclass(kw_only=True) +class Jenatton(ParamBasedTestProblem): + r"""Jenatton test function for hierarchical search spaces. + + This function is taken from: + + R. Jenatton, C. Archambeau, J. González, and M. Seeger. Bayesian + optimization with tree-structured dependencies. ICML 2017. + """ + + noise_std: Optional[float] = None + negate: bool = False + num_objectives: int = 1 + optimal_value: float = 0.1 + _is_constrained: bool = False + + def evaluate_true(self, params: TParameterization) -> torch.Tensor: + # pyre-fixme: Incompatible parameter type [6]: In call + # `jenatton_test_function`, for 1st positional argument, expected + # `Optional[float]` but got `Union[None, bool, float, int, str]`. + value = jenatton_test_function(**params) + return torch.tensor(value) def get_jenatton_benchmark_problem( num_trials: int = 50, observe_noise_sd: bool = False, + noise_std: float = 0.0, ) -> BenchmarkProblem: search_space = HierarchicalSearchSpace( parameters=[ @@ -55,24 +89,28 @@ def get_jenatton_benchmark_problem( ), ] ) + name = "Jenatton" + ("_observed_noise" if observe_noise_sd else "") optimization_config = OptimizationConfig( objective=Objective( - metric=JenattonMetric(observe_noise_sd=observe_noise_sd), + metric=BenchmarkMetric( + name=name, observe_noise_sd=observe_noise_sd, lower_is_better=True + ), minimize=True, ) ) - - name = "Jenatton" + ("_observed_noise" if observe_noise_sd else "") - return BenchmarkProblem( name=name, search_space=search_space, optimization_config=optimization_config, - runner=SyntheticRunner(), + runner=ParamBasedTestProblemRunner( + test_problem_class=Jenatton, + test_problem_kwargs={"noise_std": noise_std}, + outcome_names=[name], + ), num_trials=num_trials, - is_noiseless=True, + is_noiseless=noise_std == 0.0, observe_noise_stds=observe_noise_sd, has_ground_truth=True, - optimal_value=0.1, + optimal_value=Jenatton.optimal_value, ) diff --git a/ax/benchmark/runners/base.py b/ax/benchmark/runners/base.py index f300742e33f..cfb0eebd6f9 100644 --- a/ax/benchmark/runners/base.py +++ b/ax/benchmark/runners/base.py @@ -6,12 +6,14 @@ # pyre-strict from abc import ABC, abstractmethod +from collections.abc import Iterable from math import sqrt -from typing import Any, Optional, Union +from typing import Any, Union import torch from ax.core.arm import Arm -from ax.core.base_trial import BaseTrial + +from ax.core.base_trial import BaseTrial, TrialStatus from ax.core.batch_trial import BatchTrial from ax.core.runner import Runner from ax.core.trial import Trial @@ -21,45 +23,41 @@ class BenchmarkRunner(Runner, ABC): - - @property - @abstractmethod - def outcome_names(self) -> list[str]: - """The names of the outcomes of the problem (in the order of the outcomes).""" - pass # pragma: no cover + """ + A Runner that produces both observed and ground-truth values. + + Observed values equal ground-truth values plus noise, with the noise added + according to the standard deviations returned by `get_noise_stds()`. + + This runner does require that every benchmark has a ground truth, which + won't necessarily be true for real-world problems. Such problems fall into + two categories: + - If they are deterministic, they can be used with this runner by + viewing them as noiseless problems where the observed values are the + ground truth. The observed values will be used for tracking the + progress of optimization. + - If they are not deterministc, they are not supported. It is not + conceptually clear how to benchmark such problems, so we decided to + not over-engineer for that before such a use case arrives. + """ + + outcome_names: list[str] def get_Y_true(self, arm: Arm) -> Tensor: - """Function returning the ground truth values for a given arm. The - synthetic noise is added as part of the Runner's `run()` method. - For problems that do not have a ground truth, the Runner must - implement the `get_Y_Ystd()` method instead.""" - raise NotImplementedError( - "Must implement method `get_Y_true()` for Runner " - f"{self.__class__.__name__} as it does not implement a " - "`get_Y_Ystd()` method." - ) + """ + Return the ground truth values for a given arm. + + Synthetic noise is added as part of the Runner's `run()` method. + """ + ... + @abstractmethod def get_noise_stds(self) -> Union[None, float, dict[str, float]]: - """Function returning the standard errors for the synthetic noise - to be applied to the observed values. For problems that do not have - a ground truth, the Runner must implement the `get_Y_Ystd()` method - instead.""" - raise NotImplementedError( - "Must implement method `get_Y_Ystd()` for Runner " - f"{self.__class__.__name__} as it does not implement a " - "`get_noise_stds()` method." - ) - - def get_Y_Ystd(self, arm: Arm) -> tuple[Tensor, Optional[Tensor]]: - """Function returning the observed values and their standard errors - for a given arm. This function is unused for problems that have a - ground truth (in this case `get_Y_true()` is used), and is required - for problems that do not have a ground truth.""" - raise NotImplementedError( - "Must implement method `get_Y_Ystd()` for Runner " - f"{self.__class__.__name__} as it does not implement a " - "`get_Y_true()` method." - ) + """ + Return the standard errors for the synthetic noise to be applied to the + observed values. + """ + ... def run(self, trial: BaseTrial) -> dict[str, Any]: """Run the trial by evaluating its parameterization(s). @@ -110,33 +108,32 @@ def run(self, trial: BaseTrial) -> dict[str, Any]: ) for arm in trial.arms: - try: - # Case where we do have a ground truth - Y_true = self.get_Y_true(arm) - Ys_true[arm.name] = Y_true.tolist() - if noise_stds is None: - # No noise, so just return the true outcome. - Ystds[arm.name] = [0.0] * len(Y_true) - Ys[arm.name] = Y_true.tolist() - else: - # We can scale the noise std by the inverse of the relative sample - # budget allocation to each arm. This works b/c (i) we assume that - # observations per unit sample budget are i.i.d. and (ii) the - # normalized weights sum to one. - std = noise_stds_tsr.to(Y_true) / sqrt(nlzd_arm_weights[arm]) - Ystds[arm.name] = std.tolist() - Ys[arm.name] = (Y_true + std * torch.randn_like(Y_true)).tolist() - except NotImplementedError: - # Case where we don't have a ground truth. - Y, Ystd = self.get_Y_Ystd(arm) - Ys[arm.name] = Y.tolist() - Ystds[arm.name] = Ystd.tolist() if Ystd is not None else None + # Case where we do have a ground truth + Y_true = self.get_Y_true(arm) + Ys_true[arm.name] = Y_true.tolist() + if noise_stds is None: + # No noise, so just return the true outcome. + Ystds[arm.name] = [0.0] * len(Y_true) + Ys[arm.name] = Y_true.tolist() + else: + # We can scale the noise std by the inverse of the relative sample + # budget allocation to each arm. This works b/c (i) we assume that + # observations per unit sample budget are i.i.d. and (ii) the + # normalized weights sum to one. + std = noise_stds_tsr.to(Y_true) / sqrt(nlzd_arm_weights[arm]) + Ystds[arm.name] = std.tolist() + Ys[arm.name] = (Y_true + std * torch.randn_like(Y_true)).tolist() run_metadata = { "Ys": Ys, "Ystds": Ystds, "outcome_names": self.outcome_names, + "Ys_true": Ys_true, } - if Ys_true: # only add key if we actually have a ground truth - run_metadata["Ys_true"] = Ys_true return run_metadata + + # This will need to be udpated once asynchronous benchmarks are supported. + def poll_trial_status( + self, trials: Iterable[BaseTrial] + ) -> dict[TrialStatus, set[int]]: + return {TrialStatus.COMPLETED: {t.index for t in trials}} diff --git a/ax/benchmark/runners/botorch_test.py b/ax/benchmark/runners/botorch_test.py index 6796c4b533b..90fddbdefda 100644 --- a/ax/benchmark/runners/botorch_test.py +++ b/ax/benchmark/runners/botorch_test.py @@ -6,42 +6,76 @@ # pyre-strict import importlib -from collections.abc import Iterable +from abc import ABC, abstractmethod +from dataclasses import dataclass from typing import Any, Optional, Union import torch from ax.benchmark.runners.base import BenchmarkRunner from ax.core.arm import Arm -from ax.core.base_trial import BaseTrial, TrialStatus +from ax.core.types import TParameterization from ax.utils.common.base import Base from ax.utils.common.equality import equality_typechecker from ax.utils.common.serialization import TClassDecoderRegistry, TDecoderRegistry -from ax.utils.common.typeutils import checked_cast -from botorch.test_functions.base import BaseTestProblem, ConstrainedBaseTestProblem -from botorch.test_functions.multi_objective import MultiObjectiveTestProblem +from botorch.test_functions.synthetic import ( + ConstrainedSyntheticTestFunction, + SyntheticTestFunction, +) from botorch.utils.transforms import normalize, unnormalize +from pyre_extensions import assert_is_instance from torch import Tensor -class BotorchTestProblemRunner(BenchmarkRunner): - """A Runner for evaluating Botorch BaseTestProblems. +@dataclass(kw_only=True) +class ParamBasedTestProblem(ABC): + """ + Similar to a BoTorch test problem, but evaluated using an Ax + TParameterization rather than a tensor. + """ + + num_objectives: int + optimal_value: float + # Constraints could easily be supported similar to BoTorch test problems, + # but haven't been hooked up. + _is_constrained: bool = False + constraint_noise_std: Optional[Union[float, list[float]]] = None + noise_std: Optional[Union[float, list[float]]] = None + negate: bool = False + + @abstractmethod + def evaluate_true(self, params: TParameterization) -> Tensor: ... + + def evaluate_slack_true(self, params: TParameterization) -> Tensor: + raise NotImplementedError( + f"{self.__class__.__name__} does not support constraints." + ) + + # pyre-fixme: Missing parameter annotation [2]: Parameter `other` must have + # a type other than `Any`. + def __eq__(self, other: Any) -> bool: + if not isinstance(other, type(self)): + return False + return self.__class__.__name__ == other.__class__.__name__ + - Given a trial the Runner will evaluate the BaseTestProblem.forward method for each - arm in the trial, as well as return some metadata about the underlying Botorch - problem such as the noise_std. We compute the full result on the Runner (as opposed - to the Metric as is typical in synthetic test problems) because the BoTorch problem - computes all metrics in one stacked tensor in the MOO case, and we wish to avoid - recomputation per metric. +class SyntheticProblemRunner(BenchmarkRunner, ABC): + """A Runner for evaluating synthetic problems, either BoTorch + `SyntheticTestFunction`s or Ax benchmarking `ParamBasedTestProblem`s. + + Given a trial, the Runner will evaluate the problem noiselessly for each + arm in the trial, as well as return some metadata about the underlying + problem such as the noise_std. """ - test_problem: BaseTestProblem + test_problem: Union[SyntheticTestFunction, ParamBasedTestProblem] _is_constrained: bool - _test_problem_class: type[BaseTestProblem] + _test_problem_class: type[Union[SyntheticTestFunction, ParamBasedTestProblem]] _test_problem_kwargs: Optional[dict[str, Any]] def __init__( self, - test_problem_class: type[BaseTestProblem], + *, + test_problem_class: type[Union[SyntheticTestFunction, ParamBasedTestProblem]], test_problem_kwargs: dict[str, Any], outcome_names: list[str], modified_bounds: Optional[list[tuple[float, float]]] = None, @@ -49,7 +83,8 @@ def __init__( """Initialize the test problem runner. Args: - test_problem_class: The BoTorch test problem class. + test_problem_class: A BoTorch `SyntheticTestFunction` class or Ax + `ParamBasedTestProblem` class. test_problem_kwargs: The keyword arguments used for initializing the test problem. outcome_names: The names of the outcomes returned by the problem. @@ -63,28 +98,27 @@ def __init__( If modified bounds are not provided, the test problem will be evaluated using the raw parameter values. """ - self._test_problem_class = test_problem_class self._test_problem_kwargs = test_problem_kwargs - - # pyre-fixme [45]: Invalid class instantiation - self.test_problem = test_problem_class(**test_problem_kwargs).to( - dtype=torch.double + self.test_problem = ( + # pyre-fixme: Invalid class instantiation [45]: Cannot instantiate + # abstract class with abstract method `evaluate_true`. + test_problem_class(**test_problem_kwargs) ) + if isinstance(self.test_problem, SyntheticTestFunction): + self.test_problem = self.test_problem.to(dtype=torch.double) + # A `ConstrainedSyntheticTestFunction` is a type of `SyntheticTestFunction`; a + # `ParamBasedTestProblem` is never constrained. self._is_constrained: bool = isinstance( - self.test_problem, ConstrainedBaseTestProblem + self.test_problem, ConstrainedSyntheticTestFunction ) - self._is_moo: bool = isinstance(self.test_problem, MultiObjectiveTestProblem) - self._outcome_names = outcome_names + self._is_moo: bool = self.test_problem.num_objectives > 1 + self.outcome_names = outcome_names self._modified_bounds = modified_bounds - @property - def outcome_names(self) -> list[str]: - return self._outcome_names - @equality_typechecker def __eq__(self, other: Base) -> bool: - if not isinstance(other, BotorchTestProblemRunner): + if not isinstance(other, type(self)): return False return ( @@ -129,12 +163,95 @@ def get_noise_stds(self) -> Union[None, float, dict[str, float]]: return noise_std_dict + @classmethod + # pyre-fixme [2]: Parameter `obj` must have a type other than `Any`` + def serialize_init_args(cls, obj: Any) -> dict[str, Any]: + """Serialize the properties needed to initialize the runner. + Used for storage. + """ + runner = assert_is_instance(obj, cls) + + return { + "test_problem_module": runner._test_problem_class.__module__, + "test_problem_class_name": runner._test_problem_class.__name__, + "test_problem_kwargs": runner._test_problem_kwargs, + "outcome_names": runner.outcome_names, + "modified_bounds": runner._modified_bounds, + } + + @classmethod + def deserialize_init_args( + cls, + args: dict[str, Any], + decoder_registry: Optional[TDecoderRegistry] = None, + class_decoder_registry: Optional[TClassDecoderRegistry] = None, + ) -> dict[str, Any]: + """Given a dictionary, deserialize the properties needed to initialize the + runner. Used for storage. + """ + + module = importlib.import_module(args["test_problem_module"]) + + return { + "test_problem_class": getattr(module, args["test_problem_class_name"]), + "test_problem_kwargs": args["test_problem_kwargs"], + "outcome_names": args["outcome_names"], + "modified_bounds": args["modified_bounds"], + } + + +class BotorchTestProblemRunner(SyntheticProblemRunner): + """ + A `SyntheticProblemRunner` for BoTorch `SyntheticTestFunction`s. + + Args: + test_problem_class: A BoTorch `SyntheticTestFunction` class. + test_problem_kwargs: The keyword arguments used for initializing the + test problem. + outcome_names: The names of the outcomes returned by the problem. + modified_bounds: The bounds that are used by the Ax search space + while optimizing the problem. If different from the bounds of the + test problem, we project the parameters into the test problem + bounds before evaluating the test problem. + For example, if the test problem is defined on [0, 1] but the Ax + search space is integers in [0, 10], an Ax parameter value of + 5 will correspond to 0.5 while evaluating the test problem. + If modified bounds are not provided, the test problem will be + evaluated using the raw parameter values. + """ + + def __init__( + self, + *, + test_problem_class: type[SyntheticTestFunction], + test_problem_kwargs: dict[str, Any], + outcome_names: list[str], + modified_bounds: Optional[list[tuple[float, float]]] = None, + ) -> None: + super().__init__( + test_problem_class=test_problem_class, + test_problem_kwargs=test_problem_kwargs, + outcome_names=outcome_names, + modified_bounds=modified_bounds, + ) + self.test_problem: SyntheticTestFunction = self.test_problem.to( + dtype=torch.double + ) + self._is_constrained: bool = isinstance( + self.test_problem, ConstrainedSyntheticTestFunction + ) + def get_Y_true(self, arm: Arm) -> Tensor: - """Converts X to original bounds -- only if modified bounds were provided -- - and evaluates the test problem. See `__init__` docstring for details. + """ + Convert the arm to a tensor and evaluate it on the base test problem. + + Convert the tensor to original bounds -- only if modified bounds were + provided -- and evaluates the test problem. See the docstring for + `modified_bounds` in `BotorchTestProblemRunner.__init__` for details. Args: - X: A `batch_shape x d`-dim tensor of point(s) at which to evaluate the + arm: Arm to evaluate. It will be converted to a + `batch_shape x d`-dim tensor of point(s) at which to evaluate the test problem. Returns: @@ -157,7 +274,7 @@ def get_Y_true(self, arm: Arm) -> Tensor: X = unnormalize(unit_X, self.test_problem.bounds) Y_true = self.test_problem.evaluate_true(X).view(-1) - # `BaseTestProblem.evaluate_true()` does not negate the outcome + # `SyntheticTestFunction.evaluate_true()` does not negate the outcome if self.test_problem.negate: Y_true = -Y_true @@ -171,43 +288,44 @@ def get_Y_true(self, arm: Arm) -> Tensor: return Y_true - def poll_trial_status( - self, trials: Iterable[BaseTrial] - ) -> dict[TrialStatus, set[int]]: - return {TrialStatus.COMPLETED: {t.index for t in trials}} - @classmethod - # pyre-fixme [2]: Parameter `obj` must have a type other than `Any`` - def serialize_init_args(cls, obj: Any) -> dict[str, Any]: - """Serialize the properties needed to initialize the runner. - Used for storage. - """ - runner = checked_cast(BotorchTestProblemRunner, obj) +class ParamBasedTestProblemRunner(SyntheticProblemRunner): + """ + A `SyntheticProblemRunner` for `ParamBasedTestProblem`s. See + `SyntheticProblemRunner` for more information. + """ - return { - "test_problem_module": runner._test_problem_class.__module__, - "test_problem_class_name": runner._test_problem_class.__name__, - "test_problem_kwargs": runner._test_problem_kwargs, - "outcome_names": runner._outcome_names, - "modified_bounds": runner._modified_bounds, - } + # This could easily be supported, but hasn't been hooked up + _is_constrained: bool = False - @classmethod - def deserialize_init_args( - cls, - args: dict[str, Any], - decoder_registry: Optional[TDecoderRegistry] = None, - class_decoder_registry: Optional[TClassDecoderRegistry] = None, - ) -> dict[str, Any]: - """Given a dictionary, deserialize the properties needed to initialize the - runner. Used for storage. - """ + def __init__( + self, + *, + test_problem_class: type[ParamBasedTestProblem], + test_problem_kwargs: dict[str, Any], + outcome_names: list[str], + modified_bounds: Optional[list[tuple[float, float]]] = None, + ) -> None: + if modified_bounds is not None: + raise NotImplementedError( + f"modified_bounds is not supported for {test_problem_class.__name__}" + ) + super().__init__( + test_problem_class=test_problem_class, + test_problem_kwargs=test_problem_kwargs, + outcome_names=outcome_names, + modified_bounds=modified_bounds, + ) + self.test_problem: ParamBasedTestProblem = self.test_problem - module = importlib.import_module(args["test_problem_module"]) + def get_Y_true(self, arm: Arm) -> Tensor: + """Evaluates the test problem. - return { - "test_problem_class": getattr(module, args["test_problem_class_name"]), - "test_problem_kwargs": args["test_problem_kwargs"], - "outcome_names": args["outcome_names"], - "modified_bounds": args["modified_bounds"], - } + Returns: + A `batch_shape x m`-dim tensor of ground truth (noiseless) evaluations. + """ + Y_true = self.test_problem.evaluate_true(arm.parameters).view(-1) + # `ParamBasedTestProblem.evaluate_true()` does not negate the outcome + if self.test_problem.negate: + Y_true = -Y_true + return Y_true diff --git a/ax/benchmark/runners/surrogate.py b/ax/benchmark/runners/surrogate.py index 0054b64dc32..c990d9aa09a 100644 --- a/ax/benchmark/runners/surrogate.py +++ b/ax/benchmark/runners/surrogate.py @@ -6,7 +6,6 @@ # pyre-strict import warnings -from collections.abc import Iterable from typing import Any, Callable, Optional, Union import torch @@ -68,7 +67,7 @@ def __init__( self.get_surrogate_and_datasets = get_surrogate_and_datasets self.name = name self._surrogate = surrogate - self._outcome_names = outcome_names + self.outcome_names = outcome_names self._datasets = datasets self.search_space = search_space self.noise_stds = noise_stds @@ -89,10 +88,6 @@ def datasets(self) -> list[SupervisedDataset]: self.set_surrogate_and_datasets() return none_throws(self._datasets) - @property - def outcome_names(self) -> list[str]: - return self._outcome_names - def get_noise_stds(self) -> Union[None, float, dict[str, float]]: return self.noise_stds @@ -135,11 +130,6 @@ def run(self, trial: BaseTrial) -> dict[str, Any]: run_metadata["outcome_names"] = self.outcome_names return run_metadata - def poll_trial_status( - self, trials: Iterable[BaseTrial] - ) -> dict[TrialStatus, set[int]]: - return {TrialStatus.COMPLETED: {t.index for t in trials}} - @classmethod # pyre-fixme[2]: Parameter annotation cannot be `Any`. def serialize_init_args(cls, obj: Any) -> dict[str, Any]: diff --git a/ax/benchmark/tests/metrics/test_jennaton.py b/ax/benchmark/tests/metrics/test_jennaton.py index f7cb2474a13..3f09fa07e94 100644 --- a/ax/benchmark/tests/metrics/test_jennaton.py +++ b/ax/benchmark/tests/metrics/test_jennaton.py @@ -7,107 +7,216 @@ import math from random import random -from unittest import mock -from ax.benchmark.metrics.jenatton import jenatton_test_function, JenattonMetric +from ax.benchmark.metrics.benchmark import BenchmarkMetric, GroundTruthBenchmarkMetric + +from ax.benchmark.metrics.jenatton import jenatton_test_function +from ax.benchmark.problems.synthetic.hss.jenatton import get_jenatton_benchmark_problem +from ax.benchmark.runners.base import BenchmarkRunner +from ax.benchmark.runners.botorch_test import ParamBasedTestProblemRunner from ax.core.arm import Arm +from ax.core.data import Data +from ax.core.experiment import Experiment from ax.core.trial import Trial +from ax.core.types import TParameterization from ax.utils.common.testutils import TestCase +from pyre_extensions import assert_is_instance -class JenattonMetricTest(TestCase): +class JenattonTest(TestCase): def test_jenatton_test_function(self) -> None: + benchmark_problem = get_jenatton_benchmark_problem() + rand_params = {f"x{i}": random() for i in range(4, 8)} rand_params["r8"] = random() rand_params["r9"] = random() + cases: list[tuple[TParameterization, float]] = [] + for x3 in (0, 1): - self.assertAlmostEqual( - jenatton_test_function( - x1=0, - x2=0, - x3=x3, - **{**rand_params, "x4": 2.0, "r8": 0.05}, + # list of (param dict, expected value) + cases.append( + ( + { + "x1": 0, + "x2": 0, + "x3": x3, + **{**rand_params, "x4": 2.0, "r8": 0.05}, + }, + 4.15, ), - 4.15, ) - self.assertAlmostEqual( - jenatton_test_function( - x1=0, - x2=1, - x3=x3, - **{**rand_params, "x5": 2.0, "r8": 0.05}, - ), - 4.25, + cases.append( + ( + { + "x1": 0, + "x2": 1, + "x3": x3, + **{**rand_params, "x5": 2.0, "r8": 0.05}, + }, + 4.25, + ) ) + for x2 in (0, 1): + cases.append( + ( + { + "x1": 1, + "x2": x2, + "x3": 0, + **{**rand_params, "x6": 2.0, "r9": 0.05}, + }, + 4.35, + ) + ) + cases.append( + ( + { + "x1": 1, + "x2": x2, + "x3": 1, + **{**rand_params, "x7": 2.0, "r9": 0.05}, + }, + 4.45, + ) + ) + + for params, value in cases: + arm = Arm(parameters=params) self.assertAlmostEqual( - jenatton_test_function( - x1=1, - x2=x2, - x3=0, - **{**rand_params, "x6": 2.0, "r9": 0.05}, - ), - 4.35, + # pyre-fixme: Incompatible parameter type [6]: In call + # `jenatton_test_function`, for 1st positional argument, + # expected `Optional[float]` but got `Union[None, bool, float, + # int, str]`. + jenatton_test_function(**params), + value, ) self.assertAlmostEqual( - jenatton_test_function( - x1=1, - x2=x2, - x3=1, - **{**rand_params, "x7": 2.0, "r9": 0.05}, - ), - 4.45, + assert_is_instance(benchmark_problem.runner, BenchmarkRunner) + .get_Y_true(arm) + .item(), + value, + places=6, ) - def test_init(self) -> None: - metric = JenattonMetric() - self.assertEqual(metric.name, "jenatton") + def test_create_problem(self) -> None: + problem = get_jenatton_benchmark_problem() + objective = problem.optimization_config.objective + metric = objective.metric + + self.assertEqual(metric.name, "Jenatton") + self.assertTrue(objective.minimize) self.assertTrue(metric.lower_is_better) - self.assertEqual(metric.noise_std, 0.0) - self.assertFalse(metric.observe_noise_sd) - metric = JenattonMetric(name="nottanej", noise_std=0.1, observe_noise_sd=True) - self.assertEqual(metric.name, "nottanej") + self.assertEqual( + assert_is_instance( + problem.runner, ParamBasedTestProblemRunner + ).test_problem.noise_std, + 0.0, + ) + self.assertTrue(problem.is_noiseless) + self.assertFalse(assert_is_instance(metric, BenchmarkMetric).observe_noise_sd) + + problem = get_jenatton_benchmark_problem( + num_trials=10, noise_std=0.1, observe_noise_sd=True + ) + objective = problem.optimization_config.objective + metric = objective.metric self.assertTrue(metric.lower_is_better) - self.assertEqual(metric.noise_std, 0.1) - self.assertTrue(metric.observe_noise_sd) + self.assertEqual( + assert_is_instance( + problem.runner, ParamBasedTestProblemRunner + ).test_problem.noise_std, + 0.1, + ) + self.assertFalse(problem.is_noiseless) + self.assertTrue(assert_is_instance(metric, BenchmarkMetric).observe_noise_sd) def test_fetch_trial_data(self) -> None: - arm = mock.Mock(spec=Arm) - arm.parameters = {"x1": 0, "x2": 1, "x5": 2.0, "r8": 0.05} - trial = mock.Mock(spec=Trial) - trial.arms_by_name = {"0_0": arm} - trial.index = 0 - - metric = JenattonMetric() - df = metric.fetch_trial_data(trial=trial).value.df # pyre-ignore [16] + problem = get_jenatton_benchmark_problem() + arm = Arm(parameters={"x1": 0, "x2": 1, "x5": 2.0, "r8": 0.05}, name="0_0") + + experiment = Experiment( + search_space=problem.search_space, + name="Jenatton", + optimization_config=problem.optimization_config, + ) + + trial = Trial(experiment=experiment) + trial.add_arm(arm) + metadata = problem.runner.run(trial=trial) + trial.update_run_metadata(metadata) + + expected_metadata = { + "Ys": {"0_0": [4.25]}, + "Ystds": {"0_0": [0.0]}, + "outcome_names": ["Jenatton"], + "Ys_true": {"0_0": [4.25]}, + } + self.assertEqual(metadata, expected_metadata) + + metric = problem.optimization_config.objective.metric + + df = assert_is_instance(metric.fetch_trial_data(trial=trial).value, Data).df self.assertEqual(len(df), 1) res_dict = df.iloc[0].to_dict() self.assertEqual(res_dict["arm_name"], "0_0") - self.assertEqual(res_dict["metric_name"], "jenatton") + self.assertEqual(res_dict["metric_name"], "Jenatton") self.assertEqual(res_dict["mean"], 4.25) self.assertTrue(math.isnan(res_dict["sem"])) self.assertEqual(res_dict["trial_index"], 0) - metric = JenattonMetric(name="nottanej", noise_std=0.1, observe_noise_sd=True) - df = metric.fetch_trial_data(trial=trial).value.df # pyre-ignore [16] + problem = get_jenatton_benchmark_problem(noise_std=0.1, observe_noise_sd=True) + experiment = Experiment( + search_space=problem.search_space, + name="Jenatton", + optimization_config=problem.optimization_config, + ) + + trial = Trial(experiment=experiment) + trial.add_arm(arm) + metadata = problem.runner.run(trial=trial) + trial.update_run_metadata(metadata) + + metric = problem.optimization_config.objective.metric + df = assert_is_instance(metric.fetch_trial_data(trial=trial).value, Data).df self.assertEqual(len(df), 1) res_dict = df.iloc[0].to_dict() self.assertEqual(res_dict["arm_name"], "0_0") - self.assertEqual(res_dict["metric_name"], "nottanej") self.assertNotEqual(res_dict["mean"], 4.25) - self.assertEqual(res_dict["sem"], 0.1) + self.assertAlmostEqual(res_dict["sem"], 0.1) self.assertEqual(res_dict["trial_index"], 0) def test_make_ground_truth_metric(self) -> None: - metric = JenattonMetric() - gt_metric = metric.make_ground_truth_metric() - self.assertIsInstance(gt_metric, JenattonMetric) - self.assertEqual(gt_metric.noise_std, 0.0) - self.assertFalse(gt_metric.observe_noise_sd) - metric = JenattonMetric(noise_std=0.1, observe_noise_sd=True) + problem = get_jenatton_benchmark_problem() + + arm = Arm(parameters={"x1": 0, "x2": 1, "x5": 2.0, "r8": 0.05}, name="0_0") + + experiment = Experiment( + search_space=problem.search_space, + name="Jenatton", + optimization_config=problem.optimization_config, + ) + + trial = Trial(experiment=experiment) + trial.add_arm(arm) + problem.runner.run(trial=trial) + metadata = problem.runner.run(trial=trial) + trial.update_run_metadata(metadata) + + metric = assert_is_instance( + problem.optimization_config.objective.metric, BenchmarkMetric + ) gt_metric = metric.make_ground_truth_metric() - self.assertIsInstance(gt_metric, JenattonMetric) - self.assertEqual(gt_metric.noise_std, 0.0) - self.assertFalse(gt_metric.observe_noise_sd) + self.assertIsInstance(gt_metric, GroundTruthBenchmarkMetric) + runner = assert_is_instance(problem.runner, ParamBasedTestProblemRunner) + self.assertEqual(runner.test_problem.noise_std, 0.0) + self.assertFalse( + assert_is_instance(gt_metric, BenchmarkMetric).observe_noise_sd + ) + + self.assertIsInstance(metric, BenchmarkMetric) + self.assertNotIsInstance(metric, GroundTruthBenchmarkMetric) + self.assertEqual(runner.test_problem.noise_std, 0.0) + self.assertFalse(metric.observe_noise_sd) diff --git a/ax/benchmark/tests/runners/test_botorch_test_problem.py b/ax/benchmark/tests/runners/test_botorch_test_problem.py index 1ca0ab3e24d..ea929723e3d 100644 --- a/ax/benchmark/tests/runners/test_botorch_test_problem.py +++ b/ax/benchmark/tests/runners/test_botorch_test_problem.py @@ -8,35 +8,79 @@ from itertools import product -from typing import Union from unittest.mock import Mock import torch -from ax.benchmark.runners.botorch_test import BotorchTestProblemRunner +from ax.benchmark.runners.botorch_test import ( + BotorchTestProblemRunner, + ParamBasedTestProblemRunner, +) from ax.core.arm import Arm from ax.core.base_trial import TrialStatus from ax.core.trial import Trial from ax.utils.common.testutils import TestCase from ax.utils.common.typeutils import checked_cast -from botorch.test_functions.base import ConstrainedBaseTestProblem +from ax.utils.testing.benchmark_stubs import TestParamBasedTestProblem +from botorch.test_functions.base import BaseTestProblem, ConstrainedBaseTestProblem from botorch.test_functions.synthetic import ConstrainedHartmann, Hartmann from botorch.utils.transforms import normalize - - -class TestBotorchTestProblemRunner(TestCase): - def test_botorch_test_problem_runner(self) -> None: - for test_problem_class, modified_bounds, noise_std in product( - (Hartmann, ConstrainedHartmann), (None, [(0.0, 2.0)] * 6), (None, 0.1) +from pyre_extensions import assert_is_instance + + +class TestSyntheticRunner(TestCase): + def test_synthetic_runner(self) -> None: + botorch_cases = [ + ( + BotorchTestProblemRunner, + test_problem_class, + {"dim": 6}, + modified_bounds, + noise_std, + ) + for test_problem_class, modified_bounds, noise_std in product( + (Hartmann, ConstrainedHartmann), + (None, [(0.0, 2.0)] * 6), + (None, 0.1), + ) + ] + param_based_cases = [ + ( + ParamBasedTestProblemRunner, + TestParamBasedTestProblem, + {"num_objectives": num_objectives, "dim": 6}, + None, + noise_std, + ) + for num_objectives, noise_std in product((1, 2), (None, 0.0, 1.0)) + ] + for ( + runner_cls, + test_problem_class, + test_problem_kwargs, + modified_bounds, + noise_std, + ) in ( + botorch_cases + param_based_cases ): - test_problem = test_problem_class(dim=6).to(dtype=torch.double) - test_problem_kwargs: dict[str, Union[int, float]] = {"dim": 6} if noise_std is not None: + # pyre-fixme[6]: Incompatible parameter type: Expected int, got float test_problem_kwargs["noise_std"] = noise_std - outcome_names = ["objective"] + + num_objectives = ( + test_problem_kwargs["num_objectives"] + if "num_objectives" in test_problem_kwargs + else 1 + ) + outcome_names = [f"objective_{i}" for i in range(num_objectives)] if test_problem_class == ConstrainedHartmann: outcome_names = outcome_names + ["constraint"] - runner = BotorchTestProblemRunner( + runner = runner_cls( + # pyre-fixme[6]: Incompatible parameter type: In call + # `BotorchTestProblemRunner.__init__`, for argument + # `test_problem_class`, expected `Type[BaseTestProblem]` but got + # `Union[Type[ConstrainedHartmann], Type[Hartmann], + # Type[TestParamBasedTestProblem]]`. test_problem_class=test_problem_class, test_problem_kwargs=test_problem_kwargs, outcome_names=outcome_names, @@ -51,11 +95,9 @@ def test_botorch_test_problem_runner(self) -> None: with self.subTest(f"Test basic construction, {test_description}"): self.assertIsInstance(runner.test_problem, test_problem_class) - self.assertEqual(runner.test_problem.dim, test_problem_kwargs["dim"]) - self.assertEqual(runner.test_problem.bounds.dtype, torch.double) self.assertEqual( runner._is_constrained, - isinstance(test_problem, ConstrainedBaseTestProblem), + issubclass(test_problem_class, ConstrainedBaseTestProblem), ) self.assertEqual(runner._modified_bounds, modified_bounds) if noise_std is not None: @@ -66,6 +108,17 @@ def test_botorch_test_problem_runner(self) -> None: # check equality with different class self.assertNotEqual(runner, Hartmann(dim=6)) self.assertEqual(runner, runner) + self.assertEqual(runner._is_moo, num_objectives > 1) + if issubclass(test_problem_class, BaseTestProblem): + self.assertEqual( + runner.test_problem.dim, test_problem_kwargs["dim"] + ) + self.assertEqual( + assert_is_instance( + runner.test_problem, BaseTestProblem + ).bounds.dtype, + torch.double, + ) with self.subTest(f"test `get_Y_true()`, {test_description}"): X = torch.rand(1, 6, dtype=torch.double) @@ -80,16 +133,22 @@ def test_botorch_test_problem_runner(self) -> None: ) else: X_tf = X - obj = test_problem.evaluate_true(X_tf) - if test_problem.negate: - obj = -obj - if runner._is_constrained: - expected_Y = torch.cat( - [obj.view(-1), test_problem.evaluate_slack(X_tf).view(-1)], - dim=-1, - ) + test_problem = runner.test_problem + if issubclass(test_problem_class, BaseTestProblem): + obj = test_problem.evaluate_true(X_tf) + if test_problem.negate: + obj = -obj + if runner._is_constrained: + expected_Y = torch.cat( + [obj.view(-1), test_problem.evaluate_slack(X_tf).view(-1)], + dim=-1, + ) + else: + expected_Y = obj else: - expected_Y = obj + expected_Y = torch.full( + torch.Size([2]), X.pow(2).sum().item(), dtype=torch.double + ) self.assertTrue(torch.allclose(Y, expected_Y)) with self.subTest(f"test `run()`, {test_description}"): @@ -116,21 +175,19 @@ def test_botorch_test_problem_runner(self) -> None: ) with self.subTest(f"test `serialize_init_args()`, {test_description}"): - serialize_init_args = BotorchTestProblemRunner.serialize_init_args( - obj=runner - ) + serialize_init_args = runner_cls.serialize_init_args(obj=runner) self.assertEqual( serialize_init_args, { "test_problem_module": runner._test_problem_class.__module__, "test_problem_class_name": runner._test_problem_class.__name__, "test_problem_kwargs": runner._test_problem_kwargs, - "outcome_names": runner._outcome_names, + "outcome_names": runner.outcome_names, "modified_bounds": runner._modified_bounds, }, ) # test deserialize args - deserialize_init_args = BotorchTestProblemRunner.deserialize_init_args( + deserialize_init_args = runner_cls.deserialize_init_args( serialize_init_args ) self.assertEqual( diff --git a/ax/benchmark/tests/stubs.py b/ax/benchmark/tests/stubs.py new file mode 100644 index 00000000000..0d01d980a50 --- /dev/null +++ b/ax/benchmark/tests/stubs.py @@ -0,0 +1,26 @@ +# 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 Optional, Union + +import torch +from ax.benchmark.runners.botorch_test import ParamBasedTestProblem + + +class TestParamBasedTestProblem(ParamBasedTestProblem): + optimal_value: float = 0.0 + + def __init__( + self, num_objectives: int, noise_std: Optional[Union[float, list[float]]] + ) -> None: + self.num_objectives = num_objectives + self.noise_std = noise_std + + # pyre-fixme[14]: Inconsistent override, as dict[str, float] is not a + # `TParameterization` + def evaluate_true(self, params: dict[str, float]) -> torch.Tensor: + value = sum(elt**2 for elt in params.values()) + return value * torch.ones(self.num_objectives, dtype=torch.double) diff --git a/ax/benchmark/tests/test_benchmark.py b/ax/benchmark/tests/test_benchmark.py index 3d0ae2eeda3..366c5611122 100644 --- a/ax/benchmark/tests/test_benchmark.py +++ b/ax/benchmark/tests/test_benchmark.py @@ -278,16 +278,22 @@ def test_create_benchmark_experiment(self) -> None: def test_replication_sobol_synthetic(self) -> None: method = get_sobol_benchmark_method() - problem = get_single_objective_benchmark_problem() - res = benchmark_replication(problem=problem, method=method, seed=0) + problems = [ + get_single_objective_benchmark_problem(), + get_problem("jenatton", num_trials=6), + ] + for problem in problems: + res = benchmark_replication(problem=problem, method=method, seed=0) - self.assertEqual( - min(problem.num_trials, not_none(method.scheduler_options.total_trials)), - len(not_none(res.experiment).trials), - ) + self.assertEqual( + min( + problem.num_trials, not_none(method.scheduler_options.total_trials) + ), + len(not_none(res.experiment).trials), + ) - self.assertTrue(np.isfinite(res.score_trace).all()) - self.assertTrue(np.all(res.score_trace <= 100)) + self.assertTrue(np.isfinite(res.score_trace).all()) + self.assertTrue(np.all(res.score_trace <= 100)) def test_replication_sobol_surrogate(self) -> None: method = get_sobol_benchmark_method() diff --git a/ax/storage/json_store/registry.py b/ax/storage/json_store/registry.py index 46ee42b8d6c..12d2ebde923 100644 --- a/ax/storage/json_store/registry.py +++ b/ax/storage/json_store/registry.py @@ -17,13 +17,15 @@ ) from ax.benchmark.benchmark_result import AggregatedBenchmarkResult, BenchmarkResult from ax.benchmark.metrics.benchmark import BenchmarkMetric, GroundTruthBenchmarkMetric -from ax.benchmark.metrics.jenatton import JenattonMetric from ax.benchmark.problems.hpo.pytorch_cnn import PyTorchCNNMetric from ax.benchmark.problems.hpo.torchvision import ( PyTorchCNNTorchvisionBenchmarkProblem, PyTorchCNNTorchvisionRunner, ) -from ax.benchmark.runners.botorch_test import BotorchTestProblemRunner +from ax.benchmark.runners.botorch_test import ( + BotorchTestProblemRunner, + ParamBasedTestProblemRunner, +) from ax.benchmark.runners.surrogate import SurrogateRunner from ax.core import Experiment, ObservationFeatures from ax.core.arm import Arm @@ -210,7 +212,6 @@ Hartmann6Metric: metric_to_dict, ImprovementGlobalStoppingStrategy: improvement_global_stopping_strategy_to_dict, Interval: botorch_component_to_dict, - JenattonMetric: metric_to_dict, L2NormMetric: metric_to_dict, LogNormalPrior: botorch_component_to_dict, MapData: map_data_to_dict, @@ -238,6 +239,7 @@ OrEarlyStoppingStrategy: logical_early_stopping_strategy_to_dict, OrderConstraint: order_parameter_constraint_to_dict, OutcomeConstraint: outcome_constraint_to_dict, + ParamBasedTestProblemRunner: runner_to_dict, ParameterConstraint: parameter_constraint_to_dict, ParameterDistribution: parameter_distribution_to_dict, pathlib.Path: pathlib_to_dict, @@ -333,7 +335,6 @@ "HierarchicalSearchSpace": HierarchicalSearchSpace, "ImprovementGlobalStoppingStrategy": ImprovementGlobalStoppingStrategy, "Interval": Interval, - "JenattonMetric": JenattonMetric, "LifecycleStage": LifecycleStage, "ListSurrogate": Surrogate, # For backwards compatibility "L2NormMetric": L2NormMetric, @@ -363,6 +364,7 @@ "OrEarlyStoppingStrategy": OrEarlyStoppingStrategy, "OrderConstraint": OrderConstraint, "OutcomeConstraint": OutcomeConstraint, + "ParamBasedTestProblemRunner": ParamBasedTestProblemRunner, "ParameterConstraint": ParameterConstraint, "ParameterConstraintType": ParameterConstraintType, "ParameterDistribution": ParameterDistribution, diff --git a/ax/storage/json_store/tests/test_json_store.py b/ax/storage/json_store/tests/test_json_store.py index 02e84bc9326..b53139ea3e5 100644 --- a/ax/storage/json_store/tests/test_json_store.py +++ b/ax/storage/json_store/tests/test_json_store.py @@ -13,7 +13,7 @@ import numpy as np import torch -from ax.benchmark.metrics.jenatton import JenattonMetric +from ax.benchmark.problems.synthetic.hss.jenatton import get_jenatton_benchmark_problem from ax.core.metric import Metric from ax.core.objective import Objective from ax.core.runner import Runner @@ -192,7 +192,6 @@ ("HierarchicalSearchSpace", get_hierarchical_search_space), ("ImprovementGlobalStoppingStrategy", get_improvement_global_stopping_strategy), ("Interval", get_interval), - ("JenattonMetric", JenattonMetric), ("MapData", get_map_data), ("MapData", get_map_data), ("MapKeyInfo", get_map_key_info), @@ -209,6 +208,7 @@ ("OrderConstraint", get_order_constraint), ("OutcomeConstraint", get_outcome_constraint), ("Path", get_pathlib_path), + ("Jenatton", get_jenatton_benchmark_problem), ("PercentileEarlyStoppingStrategy", get_percentile_early_stopping_strategy), ( "PercentileEarlyStoppingStrategy", diff --git a/ax/utils/testing/benchmark_stubs.py b/ax/utils/testing/benchmark_stubs.py index ecfb9de88ec..80d3ac9daa8 100644 --- a/ax/utils/testing/benchmark_stubs.py +++ b/ax/utils/testing/benchmark_stubs.py @@ -6,9 +6,10 @@ # pyre-strict -from typing import Any, Optional +from typing import Any, Optional, Union import numpy as np +import torch from ax.benchmark.benchmark_method import BenchmarkMethod from ax.benchmark.benchmark_problem import ( BenchmarkProblem, @@ -21,6 +22,7 @@ MOOSurrogateBenchmarkProblem, SOOSurrogateBenchmarkProblem, ) +from ax.benchmark.runners.botorch_test import ParamBasedTestProblem from ax.benchmark.runners.surrogate import SurrogateRunner from ax.core.experiment import Experiment from ax.core.optimization_config import ( @@ -218,3 +220,23 @@ def get_benchmark_result() -> BenchmarkResult: def get_aggregated_benchmark_result() -> AggregatedBenchmarkResult: result = get_benchmark_result() return AggregatedBenchmarkResult.from_benchmark_results([result, result]) + + +class TestParamBasedTestProblem(ParamBasedTestProblem): + optimal_value: float = 0.0 + + def __init__( + self, + num_objectives: int, + noise_std: Optional[Union[float, list[float]]] = None, + dim: int = 6, + ) -> None: + self.num_objectives = num_objectives + self.noise_std = noise_std + self.dim = dim + + # pyre-fixme[14]: Inconsistent override, as dict[str, float] is not a + # `TParameterization` + def evaluate_true(self, params: dict[str, float]) -> torch.Tensor: + value = sum(elt**2 for elt in params.values()) + return value * torch.ones(self.num_objectives, dtype=torch.double)