Remove get_best_parameters_from_model_predictions (#3295)

Summary:
Pull Request resolved: #3295

This helper wraps `get_best_parameters_from_model_predictions_with_trial_index` and throws away the trial index. We have too many best point helpers that make this part of the codebase into a maintenance nightmare. Removing this & replacing is only use case with `get_best_parameters_from_model_predictions_with_trial_index`  before making BC-breaking changes to the function signature.

Reviewed By: esantorella

Differential Revision: D68908977

fbshipit-source-id: 40cbea15f5e073466eed5cdf8278d5f225cef223

ax/service/managed_loop.py

@@ -32,7 +32,7 @@
 from ax.modelbridge.generation_strategy import GenerationStrategy
 from ax.modelbridge.registry import Models
 from ax.service.utils.best_point import (
-    get_best_parameters_from_model_predictions,
+    get_best_parameters_from_model_predictions_with_trial_index,
     get_best_raw_objective_point,
 )
 from ax.service.utils.instantiation import (
@@ -250,11 +250,12 @@ def get_best_point(self) -> tuple[TParameterization, TModelPredictArm | None]:
         """Obtains the best point encountered in the course
         of this optimization."""
         # Find latest trial which has a generator_run attached and get its predictions
-        model_predictions = get_best_parameters_from_model_predictions(
+        best_point = get_best_parameters_from_model_predictions_with_trial_index(
             experiment=self.experiment, models_enum=Models
         )
-        if model_predictions is not None:
-            return model_predictions
+        if best_point is not None:
+            _, parameterizations, predictions = best_point
+            return parameterizations, predictions
 
         # Could not find through model, default to using raw objective.
         parameterization, values = get_best_raw_objective_point(

ax/service/tests/test_managed_loop.py

@@ -271,9 +271,10 @@ def test_optimize(self) -> None:
         self.assertIn("objective", vals[1]["objective"])
 
     @patch(
-        "ax.service.managed_loop.get_best_parameters_from_model_predictions",
+        "ax.service.managed_loop."
+        "get_best_parameters_from_model_predictions_with_trial_index",
         autospec=True,
-        return_value=({"x1": 2.0, "x2": 3.0}, ({"a": 9.0}, {"a": {"a": 3.0}})),
+        return_value=(0, {"x1": 2.0, "x2": 3.0}, ({"a": 9.0}, {"a": {"a": 3.0}})),
     )
     @mock_botorch_optimize
     def test_optimize_with_predictions(self, _) -> None:

ax/service/utils/best_point.py

@@ -307,44 +307,6 @@ def get_best_parameters_from_model_predictions_with_trial_index(
     return None
 
 
-def get_best_parameters_from_model_predictions(
-    experiment: Experiment,
-    models_enum: type[ModelRegistryBase],
-    trial_indices: Iterable[int] | None = None,
-) -> tuple[TParameterization, TModelPredictArm | None] | None:
-    """Given an experiment, returns the best predicted parameterization and
-    corresponding prediction based on the most recent Trial with predictions. If no
-    trials have predictions returns None.
-
-    Only some models return predictions. For instance GPEI does while Sobol does not.
-
-    TModelPredictArm is of the form:
-        ({metric_name: mean}, {metric_name_1: {metric_name_2: cov_1_2}})
-
-    Args:
-        experiment: Experiment, on which to identify best raw objective arm.
-        models_enum: Registry of all models that may be in the experiment's
-            generation strategy.
-        optimization_config: Optimization config to use in place of the one stored
-            on the experiment.
-        trial_indices: Indices of trials for which to retrieve data. If None will
-            retrieve data from all available trials.
-
-    Returns:
-        Tuple of parameterization and model predictions for it.
-    """
-    res = get_best_parameters_from_model_predictions_with_trial_index(
-        experiment=experiment, models_enum=models_enum, trial_indices=trial_indices
-    )
-
-    if res is None:
-        return None
-
-    _, parameterization, vals = res
-
-    return parameterization, vals
-
-
 def get_best_by_raw_objective_with_trial_index(
     experiment: Experiment,
     optimization_config: OptimizationConfig | None = None,