Add Progression Plots for MapMetric experiments to ResultsAnalysis (facebook#4705)

Summary:

This diff adds learning curve visualization (progression plots) to ResultsAnalysis for experiments with MapData and MapMetrics.

Differential Revision:
D89776181

Privacy Context Container: L1307644

Author: shrutipatel31

ax/analysis/plotly/progression.py

@@ -26,6 +26,15 @@
 from plotly import graph_objects as go
 from pyre_extensions import none_throws, override
 
+PROGRESSION_CARDGROUP_TITLE = "Learning Curves: Metric progression over trials"
+PROGRESSION_CARDGROUP_SUBTITLE = (
+    "These plots show curve metrics (learning curves) that track the evolution of "
+    "each metric over the course of the experiment. The plots display how metrics "
+    "change during trial execution, both by progression (e.g., epochs or steps) "
+    "and by wallclock time. This is useful for monitoring optimization progress and "
+    "informing early stopping decisions."
+)
+
 
 @final
 class ProgressionPlot(Analysis):

ax/analysis/results.py

@@ -13,6 +13,11 @@
 from ax.analysis.best_trials import BestTrials
 from ax.analysis.plotly.arm_effects import ArmEffectsPlot
 from ax.analysis.plotly.bandit_rollout import BanditRollout
+from ax.analysis.plotly.progression import (
+    PROGRESSION_CARDGROUP_SUBTITLE,
+    PROGRESSION_CARDGROUP_TITLE,
+    ProgressionPlot,
+)
 from ax.analysis.plotly.scatter import (
     SCATTER_CARDGROUP_SUBTITLE,
     SCATTER_CARDGROUP_TITLE,
@@ -25,6 +30,8 @@
 from ax.core.arm import Arm
 from ax.core.batch_trial import BatchTrial
 from ax.core.experiment import Experiment
+from ax.core.map_data import MapData
+from ax.core.map_metric import MapMetric
 from ax.core.outcome_constraint import ScalarizedOutcomeConstraint
 from ax.core.trial_status import TrialStatus
 from ax.core.utils import is_bandit_experiment
@@ -240,6 +247,33 @@ def compute(
             adapter=adapter,
         )
 
+        # Compute progression plots for MapMetrics (learning curves)
+        progression_group = None
+        data = experiment.lookup_data()
+        has_map_data = isinstance(data, MapData)
+        metrics = experiment.metrics.values()
+        map_metrics = [m for m in metrics if isinstance(m, MapMetric)]
+        if has_map_data and len(map_metrics) > 0:
+            map_metric_names = [m.name for m in map_metrics]
+            progression_cards = [
+                ProgressionPlot(
+                    metric_name=metric_name, by_wallclock_time=by_wallclock_time
+                ).compute_or_error_card(
+                    experiment=experiment,
+                    generation_strategy=generation_strategy,
+                    adapter=adapter,
+                )
+                for metric_name in map_metric_names
+                for by_wallclock_time in (False, True)
+            ]
+            if progression_cards:
+                progression_group = AnalysisCardGroup(
+                    name="ProgressionAnalysis",
+                    title=PROGRESSION_CARDGROUP_TITLE,
+                    subtitle=PROGRESSION_CARDGROUP_SUBTITLE,
+                    children=progression_cards,
+                )
+
         return self._create_analysis_card_group(
             title=RESULTS_CARDGROUP_TITLE,
             subtitle=RESULTS_CARDGROUP_SUBTITLE,
@@ -252,6 +286,7 @@ def compute(
                     bandit_rollout_card,
                     best_trials_card,
                     utility_progression_card,
+                    progression_group,
                     summary,
                 )
                 if child is not None

ax/analysis/tests/test_results.py

@@ -35,6 +35,7 @@
     get_experiment_with_scalarized_objective_and_outcome_constraint,
     get_offline_experiments,
     get_online_experiments,
+    get_test_map_data_experiment,
 )
 from ax.utils.testing.mock import mock_botorch_optimize
 from ax.utils.testing.modeling_stubs import get_default_generation_strategy_at_MBM_node
@@ -499,6 +500,39 @@ def test_offline_experiments(self) -> None:
             self.assertIsNotNone(card_group)
             self.assertGreater(len(card_group.children), 0)
 
+    @mock_botorch_optimize
+    def test_compute_with_map_data_includes_progression_plots(self) -> None:
+        # Setup: Create experiment with MapData and MapMetrics
+        experiment = get_test_map_data_experiment(
+            num_trials=3, num_fetches=2, num_complete=2
+        )
+        generation_strategy = get_default_generation_strategy_at_MBM_node(
+            experiment=experiment
+        )
+
+        # Execute: Compute ResultsAnalysis
+        card_group = ResultsAnalysis().compute(
+            experiment=experiment,
+            generation_strategy=generation_strategy,
+        )
+
+        # Assert: ProgressionAnalysis group exists with children
+        progression_group = None
+        for child in card_group.children:
+            if child.name == "ProgressionAnalysis":
+                progression_group = child
+                break
+
+        self.assertIsNotNone(
+            progression_group,
+            "ProgressionAnalysis group should be present for MapMetric experiments",
+        )
+        self.assertGreater(
+            len(assert_is_instance(progression_group, AnalysisCardGroup).children),
+            0,
+            "ProgressionAnalysis group should have at least one progression plot",
+        )
+
 
 class TestArmEffectsPair(TestCase):
     @mock_botorch_optimize