Repository: optuna/optuna
Pull request: Unify the logic of `plot_parallel_coordinate`s
Partly, resolve https://github.com/optuna/optuna/issues/3586.

## Motivation

This PR suggests to unify the implementation of `plot_parallel_coordinate` for two different backends, plotly and matplotlib. The design of the unification can be summarized as follows.

- Introduce `_ParallelCoordinateInfo` and `_DimensionInfo` to represents the enough information to draw the parallel coordinate plot.
- Define `_get_parallel_coordinate_info` to get the `_ParallelCoordinateInfo` from different backends.
- Use `_get_parallel_coordinate_info` to draw the parallel coordinate plot in plotly and matplotlib backends.

## Description of the changes

Introduce _ContourInfo and _SubContourInfo to represents the enough information to draw the contour plot.
Define _get_contour_info to get the _ContourInfo from different backends.
Use _get_contour_info to draw the contour plot in plotly and matplotlib backends.
Fix some test cases.


Problem statement:
Unify the visualization tests of different backends
### What is an issue?

## Motivation

Currently, visualization tests have two test file for each visualization function, e.g. `tests/visualization_tests/test_contour.py` for plotly and `tests/visualization_tests/matplotlib_tests/test_contour.py` for matplotlib. There is a lot of overlap between those test codes, and they are not well-maintained tests. This issue aims to unify both files to reduce code redundancy and improve maintainability.

## How to do that

First, I will prototype a PR to unify the `test_contour.py`. Let's use that as a springboard for discussion.

Hints from issue discussion:
This issue has not seen any recent activity.
Let me close this PR since we have finished the prototyping.

Commit message: Unify implementation of

Patch to review:
diff --git a/optuna/visualization/_parallel_coordinate.py b/optuna/visualization/_parallel_coordinate.py
index 2d711b23181..5e9ca241b3a 100644
--- a/optuna/visualization/_parallel_coordinate.py
+++ b/optuna/visualization/_parallel_coordinate.py
@@ -6,7 +6,9 @@
 from typing import DefaultDict
 from typing import Dict
 from typing import List
+from typing import NamedTuple
 from typing import Optional
+from typing import Tuple
 
 import numpy as np
 
@@ -31,6 +33,21 @@
 _logger = get_logger(__name__)
 
 
+class _DimensionInfo(NamedTuple):
+    label: str
+    values: Tuple[float, ...]
+    range: Tuple[float, float]
+    is_log: bool = False
+    is_cat: bool = False
+    tickvals: Optional[List[int]] = None
+    ticktext: Optional[List[str]] = None
+
+
+class _ParallelCoordinateInfo(NamedTuple):
+    dim_objective: _DimensionInfo
+    dims_params: List[_DimensionInfo]
+
+
 def plot_parallel_coordinate(
     study: Study,
     params: Optional[List[str]] = None,
@@ -97,14 +114,48 @@ def _get_parallel_coordinate_plot(
     layout = go.Layout(title="Parallel Coordinate Plot")
     reverse_scale = _is_reverse_scale(study, target)
 
-    trials = _filter_nonfinite(
-        study.get_trials(deepcopy=False, states=(TrialState.COMPLETE,)), target=target
-    )
+    info = _get_parallel_coordinate_info(study, params, target, target_name)
 
-    if len(trials) == 0:
+    if len(info.dims_params) == 0:
         _logger.warning("Your study does not have any completed trials.")
         return go.Figure(data=[], layout=layout)
 
+    if len(info.dim_objective.values) == 0:
+        _logger.warning("Your study has only completed trials with missing parameters.")
+        return go.Figure(data=[], layout=layout)
+
+    dims = _get_dims_from_info(info)
+
+    traces = [
+        go.Parcoords(
+            dimensions=dims,
+            labelangle=30,
+            labelside="bottom",
+            line={
+                "color": dims[0]["values"],
+                "colorscale": COLOR_SCALE,
+                "colorbar": {"title": target_name},
+                "showscale": True,
+                "reversescale": reverse_scale,
+            },
+        )
+    ]
+
+    figure = go.Figure(data=traces, layout=layout)
+
+    return figure
+
+
+def _get_parallel_coordinate_info(
+    study: Study,
+    params: Optional[List[str]] = None,
+    target: Optional[Callable[[FrozenTrial], float]] = None,
+    target_name: str = "Objective Value",
+) -> _ParallelCoordinateInfo:
+    trials = _filter_nonfinite(
+        study.get_trials(deepcopy=False, states=(TrialState.COMPLETE,)), target=target
+    )
+
     all_params = {p_name for t in trials for p_name in t.params.keys()}
     if params is not None:
         for input_p_name in params:
@@ -123,20 +174,15 @@ def _target(t: FrozenTrial) -> float:
     skipped_trial_ids = _get_skipped_trial_numbers(trials, sorted_params)
 
     objectives = tuple([target(t) for t in trials if t.number not in skipped_trial_ids])
+    # The value of (0, 0) is a dummy range. It is ignored when we plot.
+    objective_range = (min(objectives), max(objectives)) if len(objectives) > 0 else (0, 0)
+    dim_objective = _DimensionInfo(label=target_name, values=objectives, range=objective_range)
 
-    if len(objectives) == 0:
-        _logger.warning("Your study has only completed trials with missing parameters.")
-        return go.Figure(data=[], layout=layout)
-
-    dims: List[Dict[str, Any]] = [
-        {
-            "label": target_name,
-            "values": objectives,
-            "range": (min(objectives), max(objectives)),
-        }
-    ]
+    if len(trials) == 0 or len(objectives) == 0:
+        return _ParallelCoordinateInfo(dim_objective=dim_objective, dims_params=[])
 
     numeric_cat_params_indices: List[int] = []
+    dims = []
     for dim_index, p_name in enumerate(sorted_params, start=1):
         values = []
         for t in trials:
@@ -155,13 +201,14 @@ def _target(t: FrozenTrial) -> float:
                 tickvals = [min_value] + tickvals
             if max_value not in tickvals:
                 tickvals = tickvals + [max_value]
-            dim = {
-                "label": p_name if len(p_name) < 20 else "{}...".format(p_name[:17]),
-                "values": tuple(values),
-                "range": (min_value, max_value),
-                "tickvals": tickvals,
-                "ticktext": ["{:.3g}".format(math.pow(10, x)) for x in tickvals],
-            }
+            dim = _DimensionInfo(
+                label=p_name if len(p_name) < 20 else "{}...".format(p_name[:17]),
+                values=tuple(values),
+                range=(min_value, max_value),
+                is_log=True,
+                tickvals=tickvals,
+                ticktext=["{:.3g}".format(math.pow(10, x)) for x in tickvals],
+            )
         elif _is_categorical(trials, p_name):
             vocab: DefaultDict[str, int] = defaultdict(lambda: len(vocab))
 
@@ -173,47 +220,72 @@ def _target(t: FrozenTrial) -> float:
             else:
                 values = [vocab[v] for v in values]
                 ticktext = list(sorted(vocab.keys(), key=lambda x: vocab[x]))
-
-            dim = {
-                "label": p_name if len(p_name) < 20 else "{}...".format(p_name[:17]),
-                "values": tuple(values),
-                "range": (min(values), max(values)),
-                "tickvals": list(range(len(vocab))),
-                "ticktext": ticktext,
-            }
+            dim = _DimensionInfo(
+                label=p_name if len(p_name) < 20 else "{}...".format(p_name[:17]),
+                values=tuple(values),
+                range=(min(values), max(values)),
+                is_cat=True,
+                tickvals=list(range(len(vocab))),
+                ticktext=ticktext,
+            )
         else:
-            dim = {
-                "label": p_name if len(p_name) < 20 else "{}...".format(p_name[:17]),
-                "values": tuple(values),
-                "range": (min(values), max(values)),
-            }
+            dim = _DimensionInfo(
+                label=p_name if len(p_name) < 20 else "{}...".format(p_name[:17]),
+                values=tuple(values),
+                range=(min(values), max(values)),
+            )
 
         dims.append(dim)
 
     if numeric_cat_params_indices:
+        dims.insert(0, dim_objective)
         # np.lexsort consumes the sort keys the order from back to front.
         # So the values of parameters have to be reversed the order.
-        idx = np.lexsort([dims[index]["values"] for index in numeric_cat_params_indices][::-1])
+        idx = np.lexsort([dims[index].values for index in numeric_cat_params_indices][::-1])
+        updated_dims = []
         for dim in dims:
             # Since the values are mapped to other categories by the index,
             # the index will be swapped according to the sorted index of numeric params.
-            dim.update({"values": tuple(np.array(dim["values"])[idx])})
-
-    traces = [
-        go.Parcoords(
-            dimensions=dims,
-            labelangle=30,
-            labelside="bottom",
-            line={
-                "color": dims[0]["values"],
-                "colorscale": COLOR_SCALE,
-                "colorbar": {"title": target_name},
-                "showscale": True,
-                "reversescale": reverse_scale,
-            },
-        )
+            updated_dims.append(
+                _DimensionInfo(
+                    label=dim.label,
+                    values=tuple(np.array(dim.values)[idx]),
+                    range=dim.range,
+                    is_log=dim.is_log,
+                    is_cat=dim.is_cat,
+                    tickvals=dim.tickvals,
+                    ticktext=dim.ticktext,
+                )
+            )
+        dim_objective = updated_dims[0]
+        dims = updated_dims[1:]
+
+    return _ParallelCoordinateInfo(dim_objective=dim_objective, dims_params=dims)
+
+
+def _get_dims_from_info(info: _ParallelCoordinateInfo) -> List[Dict[str, Any]]:
+    dims = [
+        {
+            "label": info.dim_objective.label,
+            "values": info.dim_objective.values,
+            "range": info.dim_objective.range,
+        }
     ]
 
-    figure = go.Figure(data=traces, layout=layout)
+    for dim in info.dims_params:
+        if dim.is_log or dim.is_cat:
+            assert dim.tickvals is not None
+            assert dim.ticktext is not None
+            dims.append(
+                {
+                    "label": dim.label,
+                    "values": dim.values,
+                    "range": dim.range,
+                    "tickvals": dim.tickvals,
+                    "ticktext": dim.ticktext,
+                }
+            )
+        else:
+            dims.append({"label": dim.label, "values": dim.values, "range": dim.range})
 
-    return figure
+    return dims
diff --git a/optuna/visualization/matplotlib/_parallel_coordinate.py b/optuna/visualization/matplotlib/_parallel_coordinate.py
index a6f6583d110..ae6b8813fc1 100644
--- a/optuna/visualization/matplotlib/_parallel_coordinate.py
+++ b/optuna/visualization/matplotlib/_parallel_coordinate.py
@@ -1,7 +1,4 @@
-from collections import defaultdict
 from typing import Callable
-from typing import cast
-from typing import DefaultDict
 from typing import List
 from typing import Optional
 
@@ -12,14 +9,9 @@
 from optuna.study import Study
 from optuna.study._study_direction import StudyDirection
 from optuna.trial import FrozenTrial
-from optuna.trial import TrialState
+from optuna.visualization._parallel_coordinate import _get_parallel_coordinate_info
 from optuna.visualization._utils import _check_plot_args
-from optuna.visualization._utils import _filter_nonfinite
-from optuna.visualization._utils import _get_skipped_trial_numbers
 from optuna.visualization.matplotlib._matplotlib_imports import _imports
-from optuna.visualization.matplotlib._utils import _is_categorical
-from optuna.visualization.matplotlib._utils import _is_log_scale
-from optuna.visualization.matplotlib._utils import _is_numerical
 
 
 if _imports.is_successful():
@@ -96,15 +88,12 @@ def _get_parallel_coordinate_plot(
 ) -> "Axes":
 
     if target is None:
-
-        def _target(t: FrozenTrial) -> float:
-            return cast(float, t.value)
-
-        target = _target
         reversescale = study.direction == StudyDirection.MINIMIZE
     else:
         reversescale = True
 
+    info = _get_parallel_coordinate_info(study, params, target, target_name)
+
     # Set up the graph style.
     fig, ax = plt.subplots()
     cmap = plt.get_cmap("Blues_r" if reversescale else "Blues")
@@ -113,115 +102,59 @@ def _target(t: FrozenTrial) -> float:
     ax.spines["bottom"].set_visible(False)
 
     # Prepare data for plotting.
-    trials = _filter_nonfinite(
-        study.get_trials(deepcopy=False, states=(TrialState.COMPLETE,)), target=target
-    )
-    if len(trials) == 0:
+    if len(info.dims_params) == 0:
         _logger.warning("Your study does not have any completed trials.")
         return ax
 
-    all_params = {p_name for t in trials for p_name in t.params.keys()}
-    if params is not None:
-        for input_p_name in params:
-            if input_p_name not in all_params:
-                raise ValueError("Parameter {} does not exist in your study.".format(input_p_name))
-        all_params = set(params)
-    sorted_params = sorted(all_params)
-
-    skipped_trial_numbers = _get_skipped_trial_numbers(trials, sorted_params)
-
-    obj_org = [target(t) for t in trials if t.number not in skipped_trial_numbers]
-
-    if len(obj_org) == 0:
+    if len(info.dim_objective.values) == 0:
         _logger.warning("Your study has only completed trials with missing parameters.")
         return ax
 
-    obj_min = min(obj_org)
-    obj_max = max(obj_org)
+    obj_min = info.dim_objective.range[0]
+    obj_max = info.dim_objective.range[1]
     obj_w = obj_max - obj_min
-    dims_obj_base = [[o] for o in obj_org]
-
-    cat_param_names = []
-    cat_param_values = []
-    cat_param_ticks = []
-    param_values = []
-    var_names = [target_name]
-    numeric_cat_params_indices: List[int] = []
-
-    for param_index, p_name in enumerate(sorted_params):
-        values = [t.params[p_name] for t in trials if t.number not in skipped_trial_numbers]
-
-        if _is_categorical(trials, p_name):
-            vocab = defaultdict(lambda: len(vocab))  # type: DefaultDict[str, int]
-
-            if _is_numerical(trials, p_name):
-                _ = [vocab[v] for v in sorted(values)]
-                numeric_cat_params_indices.append(param_index)
-
-            values = [vocab[v] for v in values]
-
-            cat_param_names.append(p_name)
-            vocab_item_sorted = sorted(vocab.items(), key=lambda x: x[1])
-            cat_param_values.append([v[0] for v in vocab_item_sorted])
-            cat_param_ticks.append([v[1] for v in vocab_item_sorted])
-
-        if _is_log_scale(trials, p_name):
-            values_for_lc = [np.log10(v) for v in values]
-        else:
-            values_for_lc = values
-
-        p_min = min(values_for_lc)
-        p_max = max(values_for_lc)
+    dims_obj_base = [[o] for o in info.dim_objective.values]
+    for dim in info.dims_params:
+        p_min = dim.range[0]
+        p_max = dim.range[1]
         p_w = p_max - p_min
 
         if p_w == 0.0:
             center = obj_w / 2 + obj_min
-            for i in range(len(values)):
+            for i in range(len(dim.values)):
                 dims_obj_base[i].append(center)
         else:
-            for i, v in enumerate(values_for_lc):
+            for i, v in enumerate(dim.values):
                 dims_obj_base[i].append((v - p_min) / p_w * obj_w + obj_min)
 
-        var_names.append(p_name if len(p_name) < 20 else "{}...".format(p_name[:17]))
-        param_values.append(values)
-
-    if numeric_cat_params_indices:
-        # np.lexsort consumes the sort keys the order from back to front.
-        # So the values of parameters have to be reversed the order.
-        sorted_idx = np.lexsort(
-            [param_values[index] for index in numeric_cat_params_indices][::-1]
-        )
-        # Since the values are mapped to other categories by the index,
-        # the index will be swapped according to the sorted index of numeric params.
-        param_values = [list(np.array(v)[sorted_idx]) for v in param_values]
-
     # Draw multiple line plots and axes.
     # Ref: https://stackoverflow.com/a/50029441
-    ax.set_xlim(0, len(sorted_params))
-    ax.set_ylim(obj_min, obj_max)
-    xs = [range(len(sorted_params) + 1) for _ in range(len(dims_obj_base))]
+    n_params = len(info.dims_params)
+    ax.set_xlim(0, n_params)
+    ax.set_ylim(info.dim_objective.range[0], info.dim_objective.range[1])
+    xs = [range(n_params + 1) for _ in range(len(dims_obj_base))]
     segments = [np.column_stack([x, y]) for x, y in zip(xs, dims_obj_base)]
     lc = LineCollection(segments, cmap=cmap)
-    lc.set_array(np.asarray(obj_org))
+    lc.set_array(np.asarray(info.dim_objective.values))
     axcb = fig.colorbar(lc, pad=0.1)
     axcb.set_label(target_name)
-    plt.xticks(range(len(sorted_params) + 1), var_names, rotation=330)
+    var_names = [info.dim_objective.label] + [dim.label for dim in info.dims_params]
+    plt.xticks(range(n_params + 1), var_names, rotation=330)
 
-    for i, p_name in enumerate(sorted_params):
+    for i, dim in enumerate(info.dims_params):
         ax2 = ax.twinx()
-        ax2.set_ylim(min(param_values[i]), max(param_values[i]))
-        if _is_log_scale(trials, p_name):
+        if dim.is_log:
+            ax2.set_ylim(np.power(10, dim.range[0]), np.power(10, dim.range[1]))
             ax2.set_yscale("log")
+        else:
+            ax2.set_ylim(dim.range[0], dim.range[1])
         ax2.spines["top"].set_visible(False)
         ax2.spines["bottom"].set_visible(False)
         ax2.xaxis.set_visible(False)
-        ax2.spines["right"].set_position(("axes", (i + 1) / len(sorted_params)))
-        if p_name in cat_param_names:
-            idx = cat_param_names.index(p_name)
-            tick_pos = cat_param_ticks[idx]
-            tick_labels = cat_param_values[idx]
-            ax2.set_yticks(tick_pos)
-            ax2.set_yticklabels(tick_labels)
+        ax2.spines["right"].set_position(("axes", (i + 1) / n_params))
+        if dim.is_cat:
+            ax2.set_yticks(dim.tickvals)
+            ax2.set_yticklabels(dim.ticktext)
 
     ax.add_collection(lc)
 
diff --git a/tests/visualization_tests/matplotlib_tests/test_parallel_coordinate.py b/tests/visualization_tests/matplotlib_tests/test_parallel_coordinate.py
index 364b72cf952..92d56f32979 100644
--- a/tests/visualization_tests/matplotlib_tests/test_parallel_coordinate.py
+++ b/tests/visualization_tests/matplotlib_tests/test_parallel_coordinate.py
@@ -6,6 +6,7 @@
 from matplotlib.axes._axes import Axes
 from matplotlib.collections import LineCollection
 import matplotlib.pyplot as plt
+import numpy as np
 import pytest
 
 from optuna.distributions import BaseDistribution
@@ -211,9 +212,7 @@ def test_plot_parallel_coordinate_categorical_numeric_params() -> None:
     assert axes[1].get_ylabel() == "Objective Value"
     assert axes[1].get_ylim() == (0.0, 2.0)
     objectives = _fetch_objectives_from_figure(figure)
-    # Objective values are not sorted by the other parameters,
-    # unlike Plotly's parallel_coordinate.
-    assert objectives == [0.0, 1.0, 2.0]
+    assert objectives == [1.0, 2.0, 0.0]
     num_choices_category_a = 2
     assert axes[2].get_ylim() == (0, num_choices_category_a - 1)
     assert [int(label.get_text()) for label in axes[2].get_yticklabels()] == [1, 2]
@@ -268,7 +267,7 @@ def test_plot_parallel_coordinate_log_params() -> None:
     objectives = _fetch_objectives_from_figure(figure)
     assert objectives == [0.0, 1.0, 0.1]
     assert axes[2].get_ylim() == (1e-6, 1e-4)
-    assert axes[3].get_ylim() == (10.0, 200)
+    np.testing.assert_almost_equal(axes[3].get_ylim(), (10.0, 200))
     expected_labels = ["Objective Value", "param_a", "param_b"]
     _test_xtick_labels(axes, expected_labels)
     plt.savefig(BytesIO())
@@ -301,15 +300,13 @@ def test_plot_parallel_coordinate_unique_hyper_param() -> None:
     # Optuna's parallel coordinate uses 10% padding for color map.
     assert axes[1].get_ylim() == (-0.1, 0.1)
     objectives = _fetch_objectives_from_figure(figure)
-    # Objective values are not sorted by the other parameters,
-    # unlike Plotly's parallel_coordinate.
     assert objectives == [0.0]
     assert axes[2].get_ylim() == (-default_padding_fraction, default_padding_fraction)
     assert [line.get_text() for line in axes[2].get_yticklabels()] == ["preferred"]
     assert [label.get_position()[1] for label in axes[2].get_yticklabels()] == [0]
-    assert axes[3].get_ylim() == (
-        30 * (1.0 - default_padding_fraction),
-        30 * (1.0 + default_padding_fraction),
+    np.testing.assert_almost_equal(
+        axes[3].get_ylim(),
+        (30 * (1.0 - default_padding_fraction), 30 * (1.0 + default_padding_fraction)),
     )
     expected_labels = ["Objective Value", "category_a", "param_b"]
     _test_xtick_labels(axes, expected_labels)
@@ -331,13 +328,11 @@ def test_plot_parallel_coordinate_unique_hyper_param() -> None:
     assert axes[1].get_ylabel() == "Objective Value"
     assert axes[1].get_ylim() == (0.0, 2.0)
     objectives = _fetch_objectives_from_figure(figure)
-    # Objective values are not sorted by the other parameters,
-    # unlike Plotly's parallel_coordinate.
     assert objectives == [0.0, 2.0]
     assert axes[2].get_ylim() == (-default_padding_fraction, default_padding_fraction)
     assert [line.get_text() for line in axes[2].get_yticklabels()] == ["preferred"]
     assert [label.get_position()[1] for label in axes[2].get_yticklabels()] == [0]
-    assert axes[3].get_ylim() == (20, 30)
+    np.testing.assert_almost_equal(axes[3].get_ylim(), (20, 30))
     expected_labels = ["Objective Value", "category_a", "param_b"]
     _test_xtick_labels(axes, expected_labels)
     plt.savefig(BytesIO())
@@ -385,15 +380,14 @@ def test_plot_parallel_coordinate_with_categorical_numeric_params() -> None:
         )
     )
     figure = plot_parallel_coordinate(study)
+    plt.savefig("hoge.png")
     axes = figure.get_figure().axes
     assert len(axes) == 5 + 1
     assert axes[0].get_ylim() == (0.0, 3.0)
     assert axes[1].get_ylabel() == "Objective Value"
     assert axes[1].get_ylim() == (0.0, 3.0)
     objectives = _fetch_objectives_from_figure(figure)
-    # Objective values are not sorted by the other parameters,
-    # unlike Plotly's parallel_coordinate.
-    assert objectives == [0.0, 1.0, 2.0, 3.0]
+    assert objectives == [1.0, 3.0, 0.0, 2.0]
     num_choices_param_a = 2
     assert axes[2].get_ylim() == (0, num_choices_param_a - 1)
     assert [line.get_text() for line in axes[2].get_yticklabels()] == ["preferred", "opt"]
@@ -406,7 +400,7 @@ def test_plot_parallel_coordinate_with_categorical_numeric_params() -> None:
     assert [label.get_position()[1] for label in axes[3].get_yticklabels()] == list(
         range(num_choices_param_b)
     )
-    assert axes[4].get_ylim() == (10, 200)
+    np.testing.assert_almost_equal(axes[4].get_ylim(), (10, 200))
     num_choices_param_d = 3
     assert axes[5].get_ylim() == (0.0, num_choices_param_d - 1)
     assert [int(label.get_text()) for label in axes[5].get_yticklabels()] == [-1, 1, 2]


Code review comments:
diff_hunk: @@ -97,14 +114,48 @@ def _get_parallel_coordinate_plot(
     layout = go.Layout(title="Parallel Coordinate Plot")
     reverse_scale = _is_reverse_scale(study, target)
 
-    trials = _filter_nonfinite(
-        study.get_trials(deepcopy=False, states=(TrialState.COMPLETE,)), target=target
-    )
+    info = _get_parallel_coordinate_info(study, params, target, target_name)
 
-    if len(trials) == 0:
+    if len(info.dims_params) == 0:
         _logger.warning("Your study does not have any completed trials.")
path: optuna/visualization/_parallel_coordinate.py
text: @user1:
The warning is also a common logic between plotly and matplotlib implementations. How about also DRYing those?