Fix Bug in Slice/Contour Plot (#4969)

Summary:
Pull Request resolved: #4969

Fix a bug that casted BatchTrials to Trials for all experiments for slice/contour plots

Reviewed By: mgarrard

Differential Revision: D94983088

fbshipit-source-id: e8ea8811bde8a00fcc6bd6f926947ac0e377c183

Author: eonofrey

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

@@ -7,7 +7,10 @@
 
 from unittest.mock import patch
 
-from ax.analysis.plotly.surface.utils import get_fixed_values_for_slice_or_contour
+from ax.analysis.plotly.surface.utils import (
+    _get_best_trial_info,
+    get_fixed_values_for_slice_or_contour,
+)
 from ax.core.arm import Arm
 from ax.service.ax_client import AxClient, ObjectiveProperties
 from ax.utils.common.testutils import TestCase
@@ -121,3 +124,50 @@ def test_returns_center_when_best_point_returns_none(self) -> None:
 
         self.assertEqual(fixed_values, {"x": 0.0, "y": 0.0})
         self.assertEqual(description, "the center of the search space")
+
+    def test_get_best_trial_info_with_batch_trial(self) -> None:
+        """Test that _get_best_trial_info works with BatchTrial (PTS experiments),
+        both when a matching arm is found and when no arm matches."""
+        client = self._create_client()
+        experiment = client.experiment
+        trial = experiment.new_batch_trial()
+        trial.add_arms_and_weights(
+            arms=[
+                Arm(parameters={"x": 0.1, "y": 0.2}, name="0_0"),
+                Arm(parameters={"x": 0.3, "y": 0.4}, name="0_1"),
+                Arm(parameters={"x": 0.5, "y": 0.6}, name="0_2"),
+            ],
+        )
+
+        # Subtest 1: matching arm found
+        with self.subTest("matching_arm_found"):
+            with patch(
+                "ax.analysis.plotly.surface.utils."
+                "get_best_parameters_from_model_predictions_with_trial_index",
+                return_value=(0, {"x": 0.3, "y": 0.4}, None),
+            ):
+                result = _get_best_trial_info(
+                    experiment=experiment,
+                    generation_strategy=client.generation_strategy,
+                )
+
+            self.assertIsNotNone(result)
+            # pyre-ignore[16]: result is not None per assertion above
+            parameterization, trial_index, arm_name = result
+            self.assertEqual(parameterization, {"x": 0.3, "y": 0.4})
+            self.assertEqual(trial_index, 0)
+            self.assertEqual(arm_name, "0_1")
+
+        # Subtest 2: no matching arm returns None
+        with self.subTest("no_matching_arm_returns_none"):
+            with patch(
+                "ax.analysis.plotly.surface.utils."
+                "get_best_parameters_from_model_predictions_with_trial_index",
+                return_value=(0, {"x": 0.9, "y": 0.9}, None),
+            ):
+                result = _get_best_trial_info(
+                    experiment=experiment,
+                    generation_strategy=client.generation_strategy,
+                )
+
+            self.assertIsNone(result)

ax/analysis/plotly/surface/utils.py

@@ -10,6 +10,7 @@
 from typing import TYPE_CHECKING
 
 import numpy as np
+from ax.core.batch_trial import BatchTrial
 from ax.core.observation import ObservationFeatures
 from ax.core.parameter import (
     ChoiceParameter,
@@ -68,9 +69,21 @@ def _get_best_trial_info(
         return None
 
     trial_index, parameterization, _prediction = result
-    # Get the arm name from the trial
-    trial = assert_is_instance(experiment.trials[trial_index], Trial)
-    arm_name = none_throws(trial.arm).name
+    # Get the arm name from the trial. Handle both single-arm Trial and
+    # multi-arm BatchTrial (used by PTS experiments).
+    base_trial = experiment.trials[trial_index]
+    if isinstance(base_trial, Trial):
+        arm_name = none_throws(base_trial.arm).name
+    else:
+        # BatchTrial: match the best parameterization to the corresponding arm.
+        batch_trial = assert_is_instance(base_trial, BatchTrial)
+        arm_name = None
+        for arm in batch_trial.arms:
+            if arm.parameters == parameterization:
+                arm_name = arm.name
+                break
+        if arm_name is None:
+            return None
     return parameterization, trial_index, arm_name