Remove AxClient.get_optimization_trace (#5168)

Summary:
Pull Request resolved: #5168

Cleaning up some unused methods from AxClient. This functionality has been superseeded by `UtilityProgressionAnalysis`.

Reviewed By: mpolson64, lena-kashtelyan

Differential Revision: D99308562

fbshipit-source-id: b3d660f789e2e000f80d0ba7ed999db765058fe9

Author: saitcakmak

ax/service/ax_client.py

@@ -16,14 +16,12 @@
 from typing import Any, TypeVar
 
 import ax.service.utils.early_stopping as early_stopping_utils
-import numpy as np
 import pandas as pd
 import torch
 from ax.adapter.prediction_utils import predict_by_features
 from ax.api.configs import ChoiceParameterConfig, RangeParameterConfig
 from ax.api.utils.instantiation.from_config import parameter_from_config
 from ax.core.arm import Arm
-from ax.core.base_trial import BaseTrial
 from ax.core.evaluations_to_data import raw_evaluations_to_data
 from ax.core.experiment import Experiment
 from ax.core.generator_run import GeneratorRun
@@ -58,12 +56,9 @@
 from ax.generation_strategy.generation_strategy import GenerationStrategy
 from ax.generation_strategy.transition_criterion import MaxGenerationParallelism
 from ax.global_stopping.strategies.base import BaseGlobalStoppingStrategy
-from ax.global_stopping.strategies.improvement import constraint_satisfaction
 from ax.plot.base import AxPlotConfig
 from ax.plot.contour import plot_contour
 from ax.plot.feature_importances import plot_feature_importance_by_feature
-from ax.plot.helper import _format_dict
-from ax.plot.trace import optimization_trace_single_method
 from ax.service.utils.analysis_base import AnalysisBase
 from ax.service.utils.best_point_mixin import BestPointMixin
 from ax.service.utils.instantiation import (
@@ -1087,77 +1082,6 @@ def get_max_concurrency(self) -> list[tuple[int, int]]:
     def get_max_parallelism(self) -> list[tuple[int, int]]:
         raise NotImplementedError("Use `get_max_concurrency` instead.")
 
-    def get_optimization_trace(
-        self, objective_optimum: float | None = None
-    ) -> AxPlotConfig:
-        """Retrieves the plot configuration for optimization trace, which shows
-        the evolution of the objective mean over iterations.
-
-        Args:
-            objective_optimum: Optimal objective, if known, for display in the
-                visualization.
-        """
-        if not self.experiment.trials:
-            raise ValueError("Cannot generate plot as there are no trials.")
-
-        objective = self.objective
-        if objective.is_multi_objective:
-            raise UnsupportedError(
-                "`get_optimization_trace` is not supported "
-                "for multi-objective experiments"
-            )
-
-        if objective.is_scalarized_objective:
-            raise UnsupportedError(
-                "`get_optimization_trace` is not supported for scalarized "
-                "objectives. `trial.objective_mean` returns the value of "
-                "the first metric, not the weighted combination."
-            )
-
-        opt_config = none_throws(self.experiment.optimization_config)
-        if any(len(oc.metric_names) > 1 for oc in opt_config.outcome_constraints):
-            raise UnsupportedError(
-                "`get_optimization_trace` is not supported for scalarized "
-                "outcome constraints."
-            )
-
-        minimize: bool = objective.minimize
-
-        # Setting the objective values of infeasible points to be infinitely
-        # bad prevents them from increasing or decreasing the
-        # optimization trace.
-        def _constrained_trial_objective_mean(trial: BaseTrial) -> float:
-            if constraint_satisfaction(trial):
-                return assert_is_instance(trial, Trial).objective_mean
-            return float("inf") if minimize else float("-inf")
-
-        objective_name = self.objective_name
-        best_objectives = np.array(
-            [
-                [
-                    _constrained_trial_objective_mean(trial)
-                    for trial in self.experiment.trials.values()
-                    if trial.status.is_completed
-                ]
-            ]
-        )
-        hover_labels = [
-            _format_dict(none_throws(assert_is_instance(trial, Trial).arm).parameters)
-            for trial in self.experiment.trials.values()
-            if trial.status.is_completed
-        ]
-        return optimization_trace_single_method(
-            y=(
-                np.minimum.accumulate(best_objectives, axis=1)
-                if minimize
-                else np.maximum.accumulate(best_objectives, axis=1)
-            ),
-            optimum=objective_optimum,
-            title="Best objective found vs. # of iterations",
-            ylabel=objective_name.capitalize(),
-            hover_labels=hover_labels,
-        )
-
     def get_contour_plot(
         self,
         param_x: str | None = None,

ax/service/tests/test_ax_client.py

@@ -26,11 +26,8 @@
 from ax.core.generator_run import GeneratorRun
 from ax.core.metric import Metric
 from ax.core.multi_type_experiment import MultiTypeExperiment
-from ax.core.objective import MultiObjective, Objective
-from ax.core.optimization_config import (
-    MultiObjectiveOptimizationConfig,
-    OptimizationConfig,
-)
+from ax.core.objective import MultiObjective
+from ax.core.optimization_config import MultiObjectiveOptimizationConfig
 from ax.core.outcome_constraint import ObjectiveThreshold, OutcomeConstraint
 from ax.core.parameter import (
     ChoiceParameter,
@@ -493,8 +490,6 @@ def test_default_generation_strategy_continuous(self) -> None:
             ],
             [Generators.SOBOL, Generators.BOTORCH_MODULAR],
         )
-        with self.assertRaisesRegex(ValueError, ".* no trials"):
-            ax_client.get_optimization_trace(objective_optimum=branin.fmin)
         for i in range(6):
             gen_limit, opt_complete = ax_client.get_current_trial_generation_limit()
             self.assertFalse(opt_complete)
@@ -522,7 +517,6 @@ def test_default_generation_strategy_continuous(self) -> None:
         self.assertEqual(
             none_throws(ax_client.generation_strategy.adapter)._generator_key, "BoTorch"
         )
-        ax_client.get_optimization_trace(objective_optimum=branin.fmin)
         ax_client.get_contour_plot()
         trials_df = ax_client.get_trials_data_frame()
         self.assertIn("x", trials_df)
@@ -690,8 +684,6 @@ def test_default_generation_strategy_continuous_for_moo(self) -> None:
             ],
             [Generators.SOBOL, Generators.BOTORCH_MODULAR],
         )
-        with self.assertRaisesRegex(ValueError, ".* no trials"):
-            ax_client.get_optimization_trace(objective_optimum=branin.fmin)
         for i in range(6):
             with mock.patch("ax.service.ax_client.logger.info") as mock_log:
                 parameterization, trial_index = ax_client.get_next_trial()
@@ -739,10 +731,6 @@ def test_default_generation_strategy_continuous_for_moo(self) -> None:
         self.assertIn("b", trials_df)
         self.assertEqual(len(trials_df), 6)
 
-        with self.subTest("it raises UnsupportedError for get_optimization_trace"):
-            with self.assertRaises(UnsupportedError):
-                ax_client.get_optimization_trace(objective_optimum=branin.fmin)
-
         with self.subTest(
             "it raises UnsupportedError for get_contour_plot without metric"
         ):
@@ -3641,49 +3629,6 @@ def test_with_node_based_gs(self) -> None:
             )._generator_key,
             "Sobol",
         )
-        with mock.patch(
-            "ax.service.ax_client.optimization_trace_single_method"
-        ) as mock_plot:
-            ax_client.get_optimization_trace()
-        mock_plot.assert_called_once()
-
-    def test_get_optimization_trace_scalarized(self) -> None:
-        """get_optimization_trace raises UnsupportedError for scalarized."""
-        ax_client = get_branin_optimization()
-        params, idx = ax_client.get_next_trial()
-        ax_client.complete_trial(trial_index=idx, raw_data={"branin": (1.0, 0.0)})
-        ax_client.experiment.add_tracking_metric(Metric(name="other_metric"))
-        ax_client.experiment._optimization_config = OptimizationConfig(
-            objective=Objective(
-                expression="2*branin + -1*other_metric",
-                metric_name_to_signature={
-                    "branin": "branin",
-                    "other_metric": "other_metric",
-                },
-            ),
-        )
-        with self.assertRaisesRegex(UnsupportedError, "not supported for scalarized"):
-            ax_client.get_optimization_trace()
-
-    def test_get_optimization_trace_scalarized_outcome_constraint(self) -> None:
-        """get_optimization_trace raises UnsupportedError for scalarized
-        outcome constraints."""
-        ax_client = get_branin_optimization()
-        params, idx = ax_client.get_next_trial()
-        ax_client.complete_trial(trial_index=idx, raw_data={"branin": (1.0, 0.0)})
-        ax_client.experiment.add_tracking_metric(Metric(name="m1"))
-        ax_client.experiment.add_tracking_metric(Metric(name="m2"))
-        ax_client.experiment._optimization_config = OptimizationConfig(
-            objective=Objective(metric=Metric(name="branin"), minimize=True),
-            outcome_constraints=[
-                OutcomeConstraint(
-                    expression="2*m1 + 3*m2 <= 10",
-                    metric_name_to_signature={"m1": "m1", "m2": "m2"},
-                ),
-            ],
-        )
-        with self.assertRaisesRegex(UnsupportedError, "not supported for scalarized"):
-            ax_client.get_optimization_trace()
 
 
 # Utility functions for testing get_model_predictions without calling