Merge pull request #4848 from xadrianzetx/contour-best-values

Display best objective value in contour plot for a given param pair, not the value from the most recent trial

optuna/visualization/_contour.py

@@ -6,6 +6,7 @@
 
 from optuna.logging import get_logger
 from optuna.study import Study
+from optuna.study import StudyDirection
 from optuna.trial import FrozenTrial
 from optuna.trial import TrialState
 from optuna.visualization._plotly_imports import _imports
@@ -267,14 +268,14 @@ def _get_contour_info(
     if len(sorted_params) == 2:
         x_param = sorted_params[0]
         y_param = sorted_params[1]
-        sub_plot_info = _get_contour_subplot_info(trials, x_param, y_param, target)
+        sub_plot_info = _get_contour_subplot_info(study, trials, x_param, y_param, target)
         sub_plot_infos = [[sub_plot_info]]
     else:
         sub_plot_infos = []
         for i, y_param in enumerate(sorted_params):
             sub_plot_infos.append([])
             for x_param in sorted_params:
-                sub_plot_info = _get_contour_subplot_info(trials, x_param, y_param, target)
+                sub_plot_info = _get_contour_subplot_info(study, trials, x_param, y_param, target)
                 sub_plot_infos[i].append(sub_plot_info)
 
     reverse_scale = _is_reverse_scale(study, target)
@@ -288,6 +289,7 @@ def _get_contour_info(
 
 
 def _get_contour_subplot_info(
+    study: Study,
     trials: list[FrozenTrial],
     x_param: str,
     y_param: str,
@@ -306,7 +308,7 @@ def _get_contour_subplot_info(
         _logger.warning("Param {} unique value length is less than 2.".format(y_param))
         return _SubContourInfo(xaxis=xaxis, yaxis=yaxis, z_values={})
 
-    z_values = {}
+    z_values: dict[tuple[int, int], float] = {}
     for i, trial in enumerate(trials):
         if x_param not in trial.params or y_param not in trial.params:
             continue
@@ -323,7 +325,17 @@ def _get_contour_subplot_info(
             value = target(trial)
         assert value is not None
 
-        z_values[(x_i, y_i)] = value
+        existing = z_values.get((x_i, y_i))
+        if existing is None or target is not None:
+            # When target function is present, we can't be sure what the z-value
+            # represents and therefore we don't know how to select the best one.
+            z_values[(x_i, y_i)] = value
+        else:
+            z_values[(x_i, y_i)] = (
+                min(existing, value)
+                if study.direction is StudyDirection.MINIMIZE
+                else max(existing, value)
+            )
 
     return _SubContourInfo(xaxis=xaxis, yaxis=yaxis, z_values=z_values)
 

tests/visualization_tests/test_contour.py

@@ -102,6 +102,37 @@ def _create_study_mixture_category_types() -> Study:
     return study
 
 
+def _create_study_with_overlapping_params(direction: str) -> Study:
+    study = create_study(direction=direction)
+    distributions = {
+        "param_a": FloatDistribution(1.0, 2.0),
+        "param_b": CategoricalDistribution(["100", "101"]),
+        "param_c": CategoricalDistribution(["foo", "bar"]),
+    }
+    study.add_trial(
+        create_trial(
+            value=0.0,
+            params={"param_a": 1.0, "param_b": "101", "param_c": "foo"},
+            distributions=distributions,
+        )
+    )
+    study.add_trial(
+        create_trial(
+            value=1.0,
+            params={"param_a": 1.0, "param_b": "101", "param_c": "bar"},
+            distributions=distributions,
+        )
+    )
+    study.add_trial(
+        create_trial(
+            value=1.0,
+            params={"param_a": 2.0, "param_b": "100", "param_c": "foo"},
+            distributions=distributions,
+        )
+    )
+    return study
+
+
 @parametrize_plot_contour
 def test_plot_contour_customized_target_name(plot_contour: Callable[..., Any]) -> None:
     params = ["param_a", "param_b"]
@@ -517,6 +548,40 @@ def test_get_contour_info_nonfinite_multiobjective(objective: int, value: float)
     )
 
 
+@pytest.mark.parametrize("direction,expected", (("minimize", 0.0), ("maximize", 1.0)))
+def test_get_contour_info_overlapping_params(direction: str, expected: float) -> None:
+    study = _create_study_with_overlapping_params(direction)
+    info = _get_contour_info(study, params=["param_a", "param_b"])
+    assert info == _ContourInfo(
+        sorted_params=["param_a", "param_b"],
+        sub_plot_infos=[
+            [
+                _SubContourInfo(
+                    xaxis=_AxisInfo(
+                        name="param_a",
+                        range=(0.95, 2.05),
+                        is_log=False,
+                        is_cat=False,
+                        indices=[0.95, 1.0, 2.0, 2.05],
+                        values=[1.0, 1.0, 2.0],
+                    ),
+                    yaxis=_AxisInfo(
+                        name="param_b",
+                        range=(-0.05, 1.05),
+                        is_log=False,
+                        is_cat=True,
+                        indices=["100", "101"],
+                        values=["101", "101", "100"],
+                    ),
+                    z_values={(1, 1): expected, (2, 0): 1.0},
+                )
+            ]
+        ],
+        reverse_scale=False if direction == "maximize" else True,
+        target_name="Objective Value",
+    )
+
+
 @pytest.mark.parametrize("direction", ["minimize", "maximize"])
 def test_color_map(direction: str) -> None:
     study = create_study(direction=direction)