Add black "x"s at sampled x coordinates in slice plot

Summary: As titled. These are moderately useful here, but I really want to add them to the contour plot and I want the slice plot to match since theyre often displayed together.

Differential Revision: D69988707

Author: mpolson64

ax/analysis/plotly/interaction.py

@@ -353,4 +353,5 @@ def _prepare_surface_plot(
         log_x=is_axis_log_scale(
             parameter=experiment.search_space.parameters[feature_name]
         ),
+        display_sampled=True,
     )

ax/analysis/plotly/surface/slice.py

@@ -36,21 +36,26 @@ class SlicePlot(PlotlyAnalysis):
         - PARAMETER_NAME: The value of the parameter specified
         - METRIC_NAME_mean: The predected mean of the metric specified
         - METRIC_NAME_sem: The predected sem of the metric specified
+        - sampled: Whether the parameter value was sampled in at least one trial
     """
 
     def __init__(
         self,
         parameter_name: str,
         metric_name: str | None = None,
+        display_sampled: bool = True,
     ) -> None:
         """
         Args:
             parameter_name: The name of the parameter to plot on the x axis.
             metric_name: The name of the metric to plot on the y axis. If not
                 specified the objective will be used.
+            display_sampled: If True, plot "x"s at x coordinates which have been
+                sampled in at least one trial.
         """
         self.parameter_name = parameter_name
         self.metric_name = metric_name
+        self._display_sampled = display_sampled
 
     def compute(
         self,
@@ -83,6 +88,7 @@ def compute(
             log_x=is_axis_log_scale(
                 parameter=experiment.search_space.parameters[self.parameter_name]
             ),
+            display_sampled=self._display_sampled,
         )
 
         return self._create_plotly_analysis_card(
@@ -104,10 +110,16 @@ def _prepare_data(
     parameter_name: str,
     metric_name: str,
 ) -> pd.DataFrame:
+    sampled_xs = [
+        arm.parameters[parameter_name]
+        for trial in experiment.trials.values()
+        for arm in trial.arms
+    ]
     # Choose which parameter values to predict points for.
-    xs = get_parameter_values(
+    unsampled_xs = get_parameter_values(
         parameter=experiment.search_space.parameters[parameter_name]
     )
+    xs = [*sampled_xs, *unsampled_xs]
 
     # Construct observation features for each parameter value previously chosen by
     # fixing all other parameters to their status-quo value or mean.
@@ -127,27 +139,32 @@ def _prepare_data(
 
     predictions = model.predict(observation_features=features)
 
-    return pd.DataFrame.from_records(
-        [
-            {
-                parameter_name: xs[i],
-                f"{metric_name}_mean": predictions[0][metric_name][i],
-                f"{metric_name}_sem": predictions[1][metric_name][metric_name][i]
-                ** 0.5,  # Convert the variance to the SEM
-            }
-            for i in range(len(xs))
-        ]
+    return none_throws(
+        pd.DataFrame.from_records(
+            [
+                {
+                    parameter_name: xs[i],
+                    f"{metric_name}_mean": predictions[0][metric_name][i],
+                    f"{metric_name}_sem": predictions[1][metric_name][metric_name][i]
+                    ** 0.5,  # Convert the variance to the SEM
+                    "sampled": xs[i] in sampled_xs,
+                }
+                for i in range(len(xs))
+            ]
+        ).drop_duplicates()
     ).sort_values(by=parameter_name)
 
 
 def _prepare_plot(
     df: pd.DataFrame,
     parameter_name: str,
     metric_name: str,
-    log_x: bool = False,
+    log_x: bool,
+    display_sampled: bool,
 ) -> go.Figure:
     x = df[parameter_name].tolist()
     y = df[f"{metric_name}_mean"].tolist()
+
     # Convert the SEMs to 95% confidence intervals
     y_upper = (df[f"{metric_name}_mean"] + 1.96 * df[f"{metric_name}_sem"]).tolist()
     y_lower = (df[f"{metric_name}_mean"] - 1.96 * df[f"{metric_name}_sem"]).tolist()
@@ -184,6 +201,25 @@ def _prepare_plot(
         ),
     )
 
+    if display_sampled:
+        sampled = df[df["sampled"]]
+        x_sampled = sampled[parameter_name].tolist()
+        y_sampled = sampled[f"{metric_name}_mean"].tolist()
+
+        samples = go.Scatter(
+            x=x_sampled,
+            y=y_sampled,
+            mode="markers",
+            marker={
+                "symbol": "x",
+                "color": "black",
+            },
+            name=f"Sampled {parameter_name}",
+            showlegend=False,
+        )
+
+        fig.add_trace(samples)
+
     # Set the x-axis scale to log if relevant
     if log_x:
         fig.update_xaxes(

ax/analysis/plotly/surface/tests/test_slice.py

@@ -7,11 +7,14 @@
 
 from ax.analysis.analysis import AnalysisCardCategory, AnalysisCardLevel
 from ax.analysis.plotly.surface.slice import SlicePlot
+from ax.core.trial import Trial
 from ax.exceptions.core import UserInputError
 from ax.service.ax_client import AxClient, ObjectiveProperties
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.mock import mock_botorch_optimize
 
+from pyre_extensions import assert_is_instance, none_throws
+
 
 class TestSlicePlot(TestCase):
     @mock_botorch_optimize
@@ -64,11 +67,20 @@ def test_compute(self) -> None:
         self.assertEqual(card.category, AnalysisCardCategory.INSIGHT)
         self.assertEqual(
             {*card.df.columns},
-            {
-                "x",
-                "bar_mean",
-                "bar_sem",
-            },
+            {"x", "bar_mean", "bar_sem", "sampled"},
         )
         self.assertIsNotNone(card.blob)
         self.assertEqual(card.blob_annotation, "plotly")
+
+        # Assert that any row where sampled is True has a value of x that is
+        # sampled in at least one trial.
+        x_values_sampled = {
+            none_throws(assert_is_instance(trial, Trial).arm).parameters["x"]
+            for trial in self.client.experiment.trials.values()
+        }
+        self.assertTrue(
+            card.df.apply(
+                lambda row: row["x"] in x_values_sampled if row["sampled"] else True,
+                axis=1,
+            ).all()
+        )