Add summarize method (#3506)

Summary:
Pull Request resolved: #3506

Add a convenience function for accessing a DF explaning the experiment state. Note that we do not call this "to_df", which IMO implies some sort of “lossless”-ness and makes it seem more analogous to “to_json” etc

The DataFrame computed will contain one row per arm and the following columns:
            - trial_index: The trial index of the arm
            - arm_name: The name of the arm
            - trial_status: The status of the trial (e.g. RUNNING, SUCCEDED, FAILED)
            - failure_reason: The reason for the failure, if applicable
            - generation_method: The model_key of the model that generated the arm
            - generation_node: The name of the ``GenerationNode`` that generated the arm
            - **METADATA: Any metadata associated with the trial, as specified by the
                Experiment's runner.run_metadata_report_keys field
            - **METRIC_NAME: The observed mean of the metric specified, for each metric
            - **PARAMETER_NAME: The parameter value for the arm, for each parameter

Reviewed By: lena-kashtelyan

Differential Revision: D70923695

Author: mpolson64

ax/preview/api/client.py

@@ -11,6 +11,7 @@
 from typing import Any
 
 import numpy as np
+import pandas as pd
 
 from ax.analysis.analysis import (  # Used as a return type
     Analysis,
@@ -20,6 +21,7 @@
 from ax.analysis.markdown.markdown_analysis import (
     markdown_analysis_card_from_analysis_e,
 )
+from ax.analysis.summary import Summary
 from ax.analysis.utils import choose_analyses
 from ax.core.experiment import Experiment
 from ax.core.metric import Metric
@@ -662,6 +664,36 @@ def compute_analyses(
 
         return cards
 
+    def summarize(self) -> pd.DataFrame:
+        """
+        Special convenience method for producing the DataFrame produced by the Summary
+        Analysis. This method is a convenient way to inspect the state of the
+        experiment, but because the shape of the resultant DataFrame can change based
+        on the experiment state both users and Ax developers should prefer to use other
+        methods for extracting information from the experiment to consume downstream.
+
+        The DataFrame computed will contain one row per arm and the following columns
+        (though empty columns are omitted):
+            - trial_index: The trial index of the arm
+            - arm_name: The name of the arm
+            - trial_status: The status of the trial (e.g. RUNNING, SUCCEDED, FAILED)
+            - failure_reason: The reason for the failure, if applicable
+            - generation_node: The name of the ``GenerationNode`` that generated the arm
+            - **METADATA: Any metadata associated with the trial, as specified by the
+                Experiment's runner.run_metadata_report_keys field
+            - **METRIC_NAME: The observed mean of the metric specified, for each metric
+            - **PARAMETER_NAME: The parameter value for the arm, for each parameter
+        """
+
+        return (
+            Summary(omit_empty_columns=True)
+            .compute(
+                experiment=self._experiment,
+                generation_strategy=self._generation_strategy,
+            )
+            .df
+        )
+
     def get_best_parameterization(
         self, use_model_predictions: bool = True
     ) -> tuple[TParameterization, TOutcome, int, str]:

ax/preview/api/tests/test_client.py

@@ -834,6 +834,75 @@ def test_get_next_trials_then_run_trials(self) -> None:
             5,
         )
 
+    def test_summarize(self) -> None:
+        client = Client()
+
+        client.configure_experiment(
+            experiment_config=ExperimentConfig(
+                name="test_experiment",
+                parameters=[
+                    RangeParameterConfig(
+                        name="x1",
+                        parameter_type=ParameterType.FLOAT,
+                        bounds=(0, 1),
+                    ),
+                    RangeParameterConfig(
+                        name="x2",
+                        parameter_type=ParameterType.FLOAT,
+                        bounds=(0, 1),
+                    ),
+                ],
+            )
+        )
+        client.configure_optimization(objective="foo, bar")
+
+        # Get two trials and fail one, giving us a ragged structure
+        client.get_next_trials(maximum_trials=2)
+        client.complete_trial(trial_index=0, raw_data={"foo": 1.0, "bar": 2.0})
+        client.mark_trial_failed(trial_index=1)
+
+        summary_df = client.summarize()
+
+        self.assertEqual(
+            {*summary_df.columns},
+            {
+                "trial_index",
+                "arm_name",
+                "trial_status",
+                "generation_node",
+                "foo",
+                "bar",
+                "x1",
+                "x2",
+            },
+        )
+
+        trial_0_parameters = none_throws(
+            assert_is_instance(client._experiment.trials[0], Trial).arm
+        ).parameters
+        trial_1_parameters = none_throws(
+            assert_is_instance(client._experiment.trials[1], Trial).arm
+        ).parameters
+        expected = pd.DataFrame(
+            {
+                "trial_index": {0: 0, 1: 1},
+                "arm_name": {0: "0_0", 1: "1_0"},
+                "trial_status": {0: "COMPLETED", 1: "FAILED"},
+                "generation_node": {0: "Sobol", 1: "Sobol"},
+                "foo": {0: 1.0, 1: np.nan},  # NaN because trial 1 failed
+                "bar": {0: 2.0, 1: np.nan},
+                "x1": {
+                    0: trial_0_parameters["x1"],
+                    1: trial_1_parameters["x1"],
+                },
+                "x2": {
+                    0: trial_0_parameters["x2"],
+                    1: trial_1_parameters["x2"],
+                },
+            }
+        )
+        pd.testing.assert_frame_equal(summary_df, expected)
+
     def test_compute_analyses(self) -> None:
         client = Client()