Enable BoTorch-based BenchmarkProblems with MapMetric (facebook#4252)

Summary:
Pull Request resolved: facebook#4252

* Enables BoTorchTestFunction to produce data with `steps` and `MapMetric`. The data is simply repeated, so this is not interesting for modeling, but it is useful for measuring runtime.
* Makes outputs 2d in accordance with the docstring.

Reviewed By: saitcakmak

Differential Revision: D81689093

fbshipit-source-id: 4ca953647818cc3ece21f9c60df3808e613ce860

Author: esantorella

ax/benchmark/benchmark_test_functions/botorch_test.py

@@ -44,6 +44,10 @@ class BoTorchTestFunction(BenchmarkTestFunction):
             ``self.botorch_problem``.
             self.botorch_problem.dim + len(dummy_param_names) should equal the
             number of parameters in the ``params`` passed to ``evaluate_true``.
+        n_steps: Number of data points produced per metric and per evaluation. 1
+            if data is not time-series. If data is time-series, this will
+            eventually become the number of values on a `MapMetric` for
+            evaluations that run to completion.
     """
 
     outcome_names: Sequence[str]
@@ -82,7 +86,7 @@ def __post_init__(self) -> None:
             )
 
     def tensorize_params(self, params: Mapping[str, int | float]) -> torch.Tensor:
-        """Converts parameters to a tensor.
+        """Converts parameters to a 1d tensor.
 
         If modified bounds are provided, we normalize the parameters from the modified
         bounds to the unit cube, and then unnormalize to the original problem bounds.
@@ -114,8 +118,17 @@ def evaluate_true(self, params: Mapping[str, float | int]) -> torch.Tensor:
                 f"Expected {expected_n_dims} parameters, got {len(params)}."
             )
         x = self.tensorize_params(params=params)
-        objectives = self.botorch_problem(x).view(-1)
+        # self.botorch_problem(x) has shape [n_metrics] if n_metrics > 1,
+        # otherwise []. So `objectives` has shape [n_metrics]
+        objectives = torch.atleast_1d(self.botorch_problem(x))
         if isinstance(self.botorch_problem, ConstrainedBaseTestProblem):
-            constraints = self.botorch_problem.evaluate_slack_true(x).view(-1)
-            return torch.cat([objectives, constraints], dim=-1)
-        return objectives
+            constraints = torch.atleast_1d(self.botorch_problem.evaluate_slack_true(x))
+            metrics = torch.cat([objectives, constraints], dim=-1)
+        else:
+            metrics = objectives
+        # shape (n_metrics, 1)
+        metrics = metrics.unsqueeze(1)
+        if self.n_steps == 1:
+            return metrics
+        # shape (n_metrics, n_steps)
+        return metrics.repeat(1, self.n_steps)

ax/benchmark/problems/synthetic/from_botorch.py

@@ -98,6 +98,8 @@ def create_problem_from_botorch(
         dict[AuxiliaryExperimentPurpose, list[AuxiliaryExperiment]] | None
     ) = None,
     n_dummy_dimensions: int = 0,
+    use_map_metric: bool = False,
+    n_steps: int = 1,
 ) -> BenchmarkProblem:
     """
     Create a ``BenchmarkProblem`` from a BoTorch ``BaseTestProblem``.
@@ -154,6 +156,12 @@ def create_problem_from_botorch(
         n_dummy_dimensions: If >0, the search space will be augmented
             with extra dimensions. The corresponding parameters will have no
             effect on function values.
+        use_map_metric: Whether to use a ``BenchmarkMapMetric`` (rather than a
+            ``BenchmarkMetric``).
+        n_steps: Number of steps (progression values) in each evaluation. The
+            default of 1 reflects a normal synthetic function evaluation. A
+            higher number results in repeating the evaluation and getting the
+            same result ``n_steps`` times (before IID noise is added).
 
     Example:
         >>> from ax.benchmark.benchmark_problem import create_problem_from_botorch
@@ -213,6 +221,7 @@ def create_problem_from_botorch(
         dummy_param_names={
             n for n in search_space.parameters if "embedding_dummy_" in n
         },
+        n_steps=n_steps,
     )
 
     if isinstance(test_problem, MultiObjectiveTestProblem):
@@ -224,12 +233,14 @@ def create_problem_from_botorch(
             observe_noise_sd=observe_noise_sd,
             outcome_names=test_function.outcome_names,
             ref_point=test_problem._ref_point,
+            use_map_metric=use_map_metric,
         )
     else:
         optimization_config = get_soo_opt_config(
             outcome_names=test_function.outcome_names,
             lower_is_better=lower_is_better,
             observe_noise_sd=observe_noise_sd,
+            use_map_metric=use_map_metric,
         )
 
     optimal_value = (

ax/benchmark/tests/benchmark_test_functions/test_botorch_test_function.py

@@ -92,7 +92,7 @@ def test_tensor_shapes(self) -> None:
         for name, result in evaluate_true_results.items():
             with self.subTest(name=name):
                 self.assertEqual(result.dtype, torch.double)
-                self.assertEqual(result.shape, torch.Size([expected_len[name]]))
+                self.assertEqual(result.shape, torch.Size([expected_len[name], 1]))
 
     def test_input_dimensions(self) -> None:
         test_function = self.botorch_test_functions["base Hartmann"]
@@ -115,3 +115,27 @@ def test_dummy_dimensions(self) -> None:
             test_function.evaluate_true(params=params),
             embedded_test_function.evaluate_true(params=embedded_params),
         )
+
+    def test_with_steps(self) -> None:
+        ha_params = {f"x{i}": 0.5 for i in range(6)}
+        br_params = {"x0": 0.0, "x1": 0.0}
+        for name, botorch_problem, outcome_names, params in [
+            ("Unconstrained", Hartmann(dim=6), ["y"], ha_params),
+            ("Constrained", ConstrainedHartmann(dim=6), ["y", "c"], ha_params),
+            ("Moo", BraninCurrin(), ["y1", "y2"], br_params),
+        ]:
+            for n_steps in [1, 3]:
+                with self.subTest(name=name):
+                    test_function = BoTorchTestFunction(
+                        outcome_names=outcome_names,
+                        botorch_problem=botorch_problem,
+                        n_steps=n_steps,
+                    )
+                    self.assertEqual(test_function.n_steps, n_steps)
+                    result = test_function.evaluate_true(params=params)
+                    self.assertEqual(
+                        result.shape, torch.Size([len(outcome_names), n_steps])
+                    )
+                    if n_steps == 2:
+                        # data is simply repeated down the step dimension
+                        self.assertEqual(result[0, 0].item(), result[0, 1].item())

ax/benchmark/tests/problems/synthetic/test_from_botorch.py

@@ -8,7 +8,7 @@
 from itertools import product
 
 import torch
-from ax.benchmark.benchmark_metric import BenchmarkMetric
+from ax.benchmark.benchmark_metric import BenchmarkMapMetric, BenchmarkMetric
 from ax.benchmark.benchmark_problem import get_continuous_search_space
 from ax.benchmark.benchmark_test_functions.botorch_test import BoTorchTestFunction
 from ax.benchmark.problems.synthetic.from_botorch import (
@@ -342,3 +342,40 @@ def test_get_name(self) -> None:
                 test_problem=Branin(), n_dummy_dimensions=24, observe_noise_sd=True
             )
             self.assertEqual(name, "Branin_observed_noise_26d")
+
+    def test_with_map_metric(self) -> None:
+        with self.subTest("With default n_steps"):
+            problem = create_problem_from_botorch(
+                test_problem_class=Branin,
+                test_problem_kwargs={},
+                num_trials=1,
+                use_map_metric=True,
+            )
+            self.assertIsInstance(
+                problem.optimization_config.objective.metric, BenchmarkMapMetric
+            )
+            self.assertEqual(problem.test_function.n_steps, 1)
+
+        with self.subTest("With non-default n_steps"):
+            n_steps = 4
+            problem = create_problem_from_botorch(
+                test_problem_class=Branin,
+                test_problem_kwargs={},
+                num_trials=1,
+                use_map_metric=True,
+                n_steps=4,
+            )
+            self.assertIsInstance(
+                problem.optimization_config.objective.metric, BenchmarkMapMetric
+            )
+            self.assertEqual(problem.test_function.n_steps, n_steps)
+
+        with self.subTest("MOO"):
+            problem = create_problem_from_botorch(
+                test_problem_class=BraninCurrin,
+                test_problem_kwargs={},
+                num_trials=1,
+                use_map_metric=True,
+            )
+            metric = next(iter(problem.optimization_config.metrics.values()))
+            self.assertIsInstance(metric, BenchmarkMapMetric)