Improve Summary Analysis by relativizing the metric results if there is a status quo to relativize against (#4342)

Summary: Pull Request resolved: #4342

Reviewed By: mpolson64

Differential Revision: D82658357

Author: shrutipatel31

ax/analysis/summary.py

@@ -14,6 +14,7 @@
 from ax.analysis.analysis_card import AnalysisCard
 from ax.analysis.utils import validate_experiment
 from ax.core.experiment import Experiment
+from ax.core.map_data import MapData
 from ax.core.trial_status import NON_STALE_STATUSES, TrialStatus
 from ax.exceptions.core import UserInputError
 from ax.generation_strategy.generation_strategy import GenerationStrategy
@@ -78,15 +79,34 @@ def compute(
         if experiment is None:
             raise UserInputError("`Summary` analysis requires an `Experiment` input")
 
+        # Determine if we should relativize based on:
+        # (1) experiment has metrics and (2) experiment has status quo
+        # (3) experiment data is not MapData (MapData doesn't support relativization
+        # due to time-series step alignment complexities.)
+        data = experiment.lookup_data(trial_indices=self.trial_indices)
+        should_relativize = (
+            len(experiment.metrics) > 0
+            and experiment.status_quo is not None
+            and not isinstance(data, MapData)
+        )
+
         return self._create_analysis_card(
             title=(
                 "Summary for "
                 f"{experiment.name if experiment.has_name else 'Experiment'}"
             ),
-            subtitle="High-level summary of the `Trial`-s in this `Experiment`",
+            subtitle=(
+                "High-level summary of the `Trial`-s in this `Experiment`"
+                if not should_relativize
+                else (
+                    "High-level summary of the `Trial`-s in this `Experiment` "
+                    "Metric results are relativized against status quo."
+                )
+            ),
             df=experiment.to_df(
                 trial_indices=self.trial_indices,
                 omit_empty_columns=self.omit_empty_columns,
                 trial_statuses=self.trial_statuses,
+                relativize=should_relativize,
             ),
         )

ax/analysis/tests/test_summary.py

@@ -7,21 +7,28 @@
 
 import numpy as np
 import pandas as pd
+
 from ax.analysis.summary import Summary
 from ax.api.client import Client
 from ax.api.configs import RangeParameterConfig
 from ax.core.base_trial import TrialStatus
+from ax.core.map_data import MapData
 from ax.core.trial import Trial
 from ax.exceptions.core import UserInputError
 from ax.utils.common.testutils import TestCase
-from ax.utils.testing.core_stubs import get_offline_experiments, get_online_experiments
+from ax.utils.testing.core_stubs import (
+    get_branin_experiment_with_status_quo_trials,
+    get_offline_experiments,
+    get_online_experiments,
+)
 from pyre_extensions import assert_is_instance, none_throws
 
 
 class TestSummary(TestCase):
-    def test_compute(self) -> None:
-        client = Client()
-        client.configure_experiment(
+    def setUp(self) -> None:
+        super().setUp()
+        self.client = Client()
+        self.client.configure_experiment(
             name="test_experiment",
             parameters=[
                 RangeParameterConfig(
@@ -36,7 +43,10 @@ def test_compute(self) -> None:
                 ),
             ],
         )
-        client.configure_optimization(objective="foo, bar")
+        self.client.configure_optimization(objective="foo, bar")
+
+    def test_compute(self) -> None:
+        client = self.client
 
         # Get two trials and fail one, giving us a ragged structure
         client.get_next_trials(max_trials=2)
@@ -142,23 +152,7 @@ def test_offline(self) -> None:
 
     def test_trial_indices_filter(self) -> None:
         """Test that Client.summarize correctly uses Summary."""
-        client = Client()
-        client.configure_experiment(
-            name="test_experiment",
-            parameters=[
-                RangeParameterConfig(
-                    name="x1",
-                    parameter_type="float",
-                    bounds=(0, 1),
-                ),
-                RangeParameterConfig(
-                    name="x2",
-                    parameter_type="float",
-                    bounds=(0, 1),
-                ),
-            ],
-        )
-        client.configure_optimization(objective="foo")
+        client = self.client
 
         # Get a trial
         client.get_next_trials(max_trials=1)
@@ -228,19 +222,7 @@ def test_trial_status_filter(self) -> None:
 
     def test_default_excludes_stale_trials(self) -> None:
         """Test that Summary defaults to excluding STALE trials."""
-        # Set up experiment with basic configuration
-        client = Client()
-        client.configure_experiment(
-            name="test_experiment",
-            parameters=[
-                RangeParameterConfig(
-                    name="x1",
-                    parameter_type="float",
-                    bounds=(0, 1),
-                ),
-            ],
-        )
-        client.configure_optimization(objective="foo")
+        client = self.client
 
         # Create 3 trials with different statuses to test default filtering behavior
         client.get_next_trials(max_trials=3)
@@ -275,3 +257,72 @@ def test_default_excludes_stale_trials(self) -> None:
         # Verify that no trials in the output have STALE status
         stale_statuses = card.df[card.df["trial_status"] == "STALE"]
         self.assertEqual(len(stale_statuses), 0)
+
+    def test_metrics_relativized_with_status_quo(self) -> None:
+        """Test that Summary relativizes metrics by default when status
+        quos are present."""
+        # Use helper function that creates batch trials with status quo
+        experiment = get_branin_experiment_with_status_quo_trials(num_sobol_trials=2)
+
+        analysis = Summary()
+        card = analysis.compute(experiment=experiment)
+
+        with self.subTest("subtitle_indicates_relativization"):
+            self.assertIn("relativized", card.subtitle.lower())
+
+        with self.subTest("metric_values_formatted_as_percentages"):
+            metric_values = card.df["branin"].dropna()
+            self.assertGreater(len(metric_values), 0)
+            for val in metric_values:
+                self.assertIsInstance(val, str)
+                self.assertTrue(val.endswith("%"))
+
+        with self.subTest("relativization_calculation_correct"):
+            raw_data = experiment.lookup_data().df
+            sq_name = none_throws(experiment.status_quo).name
+            trial_0_data = raw_data[raw_data["trial_index"] == 0]
+            treatment_arm = [a for a in experiment.trials[0].arms if a.name != sq_name][
+                0
+            ]
+
+            sq_val = trial_0_data[trial_0_data["arm_name"] == sq_name]["mean"].values[0]
+            arm_val = trial_0_data[trial_0_data["arm_name"] == treatment_arm.name][
+                "mean"
+            ].values[0]
+            expected = ((arm_val - sq_val) / sq_val) * 100
+
+            actual = float(
+                card.df[card.df["arm_name"] == treatment_arm.name]["branin"]
+                .values[0]
+                .rstrip("%")
+            )
+            self.assertAlmostEqual(actual, expected, places=1)
+
+    def test_mapdata_not_relativized(self) -> None:
+        """Test that Summary does not attempt relativization when data is MapData,
+        even when status quo is present."""
+        # Create an experiment with MapData
+        experiment = get_branin_experiment_with_status_quo_trials(num_sobol_trials=2)
+
+        # Replace the experiment's data with MapData
+        map_data = MapData(
+            df=experiment.lookup_data().df.assign(step=1.0)  # Add step column
+        )
+        # Store the MapData in the experiment
+        for trial in experiment.trials.values():
+            trial_data = map_data.filter(trial_indices=[trial.index])
+            experiment.attach_data(trial_data, combine_with_last_data=False)
+
+        # Compute the summary
+        analysis = Summary()
+        card = analysis.compute(experiment=experiment)
+
+        with self.subTest("subtitle_does_not_indicate_relativization"):
+            self.assertNotIn("relativized", card.subtitle.lower())
+
+        with self.subTest("metric_values_not_formatted_as_percentages"):
+            metric_values = card.df["branin"].dropna()
+            self.assertGreater(len(metric_values), 0)
+            # Values should be raw floats, not percentage strings
+            for val in metric_values:
+                self.assertIsInstance(val, (float, np.floating))

ax/core/data.py

@@ -410,9 +410,21 @@ def relativize_dataframe(
     for grp in grouped_df.groups.keys():
         subgroup_df = grouped_df.get_group(grp)
         is_sq = subgroup_df["arm_name"] == status_quo_name
-        sq_mean, sq_sem = (
-            subgroup_df[is_sq][["mean", "sem"]].drop_duplicates().values.flatten()
-        )
+
+        # Check if status quo exists in this subgroup (trial)
+        sq_data = subgroup_df[is_sq][["mean", "sem"]].drop_duplicates().values
+        if len(sq_data) == 0:
+            # No status quo in this trial - skip relativization and include raw data
+            logger.debug(
+                "Status quo '%s' not found in trial group %s - "
+                "skipping relativization for this group",
+                status_quo_name,
+                grp,
+            )
+            dfs.append(subgroup_df)
+            continue
+
+        sq_mean, sq_sem = sq_data.flatten()
 
         # rm status quo from final df to relativize
         if not include_sq:

ax/core/experiment.py

@@ -2079,6 +2079,7 @@ def to_df(
         trial_indices: Iterable[int] | None = None,
         trial_statuses: Sequence[TrialStatus] | None = None,
         omit_empty_columns: bool = True,
+        relativize: bool = False,
     ) -> pd.DataFrame:
         """
         High-level summary of the Experiment with one row per arm. Any values missing at
@@ -2100,10 +2101,32 @@ def to_df(
             trial_indices: If specified, only include these trial indices.
             omit_empty_columns: If True, omit columns where every value is None.
             trial_status: If specified, only include trials with this status.
+            relativize: If True and:
+                * experiment has a status quo on all of its ``BatchTrial``-s
+                * OR a status quo trial among its ``Trial``-s,
+                , relativize metrics against the status quo.
         """
 
         records = []
-        data_df = self.lookup_data(trial_indices=trial_indices).df
+        data = self.lookup_data(trial_indices=trial_indices)
+
+        # Relativize metrics if requested
+        if relativize:
+            if self.status_quo is None:
+                raise UserInputError(
+                    "Attempting to relativize the experiment data, however, "
+                    "the experiment status quo is None. Please set the experiment "
+                    "status quo, or set `relativize` = False"
+                )
+
+            data_df = data.relativize(
+                status_quo_name=self.status_quo.name,
+                as_percent=True,
+                include_sq=True,
+            ).df
+        else:
+            data_df = data.df
+
         trials = (
             self.get_trials_by_indices(trial_indices=trial_indices)
             if trial_indices
@@ -2165,6 +2188,20 @@ def to_df(
         df = pd.DataFrame(records)
         if omit_empty_columns:
             df = df.loc[:, df.notnull().any()]
+
+        # Format metric columns as percentages with 4 significant figures when
+        # relativized
+        if relativize:
+            for metric_name in self.metrics.keys():
+                if metric_name in df.columns:
+                    df[metric_name] = df[metric_name].apply(
+                        lambda x: (
+                            f"{x:.4g}%"
+                            if pd.notna(x) and x != 0.0
+                            else ("0%" if pd.notna(x) else None)
+                        )
+                    )
+
         return df
 
     def add_auxiliary_experiment(

ax/core/tests/test_experiment.py

@@ -1679,6 +1679,72 @@ def test_to_df(self) -> None:
         )
         self.assertTrue(df_completed.equals(expected_completed_df))
 
+    def test_to_df_with_relativize(self) -> None:
+        """Test the relativize flag in to_df method with status quo."""
+        # Create an experiment with status quo and completed trials
+        experiment = get_branin_experiment(
+            with_status_quo=True, with_completed_batch=True
+        )
+
+        with self.subTest("without relativization"):
+            df_no_rel = experiment.to_df(relativize=False)
+
+            # Verify dataframe has expected structure
+            self.assertGreater(len(df_no_rel), 0)
+            self.assertIn("trial_index", df_no_rel.columns)
+            self.assertIn("arm_name", df_no_rel.columns)
+
+            # Branin experiment has a single metric named "branin"
+            metric_name = "branin"
+            self.assertIn(metric_name, df_no_rel.columns)
+
+            # Verify metric values are numeric, not percentage strings
+            values = df_no_rel[metric_name]
+            for val in values:
+                self.assertIsInstance(
+                    val, float, "Non-relativized values should be floats"
+                )
+
+        with self.subTest("with relativization"):
+            df_with_rel = experiment.to_df(relativize=True)
+            df_no_rel = experiment.to_df(relativize=False)
+
+            # Verify structure is preserved
+            self.assertEqual(len(df_with_rel), len(df_no_rel))
+            self.assertEqual(set(df_with_rel.columns), set(df_no_rel.columns))
+
+            # Branin experiment has a single metric named "branin"
+            metric_name = "branin"
+
+            # Verify relativization for the metric
+            self.assertIsNotNone(experiment.status_quo)
+            status_quo_name = experiment.status_quo.name
+
+            # Status quo should be 0% after relativization (using .4g format)
+            sq_rel_values = df_with_rel[df_with_rel["arm_name"] == status_quo_name][
+                metric_name
+            ]
+            for val in sq_rel_values:
+                self.assertEqual(val, "0%", "Status quo should be relativized to 0%")
+
+            # Non-status-quo arms should have percentage strings
+            non_sq_rel_values = df_with_rel[df_with_rel["arm_name"] != status_quo_name][
+                metric_name
+            ]
+            for val in non_sq_rel_values:
+                self.assertIsInstance(val, str, "Relativized values should be strings")
+                self.assertTrue(
+                    val.endswith("%"), "Relativized values should end with %"
+                )
+
+            # Verify at least one non-status-quo value is non-zero
+            has_nonzero = any(float(v.rstrip("%")) != 0.0 for v in non_sq_rel_values)
+            self.assertTrue(
+                has_nonzero,
+                "At least one non-status-quo arm should have non-zero "
+                "relativized value",
+            )
+
 
 class ExperimentWithMapDataTest(TestCase):
     def setUp(self) -> None: