Repository: optuna/optuna
Pull request: Unify `plot_contour`s
<!-- Thank you for creating a pull request! In general, we merge your pull request after it gets two or more approvals. To proceed to the review process by the maintainers, please make sure that the PR meets the following conditions: (1) it passes all CI checks, and (2) it is neither in the draft nor WIP state. If you wish to discuss the PR in the draft state or need any other help, please mention the Optuna development team in the PR. -->

Partly, resolve https://github.com/optuna/optuna/issues/3586, and an alternative to #3612.

## Motivation
<!-- Describe your motivation why you will submit this PR. This is useful for reviewers to understand the context of PR. -->
This PR suggests to unify the implementation of `plot_contour` for two different backends, plotly and matplotlib. The design of the unification can be summarized as follows.

- 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.

Note that, in matplotlib, we need to transform the information of `_ContourInfo` to make the plot. In detail, we need to construct the int representation of the categorical variables and need to compute the exact grid information to make a resemble plot of plotly. As a result, in my eyes, the benefits of implementing a commonization process are not that great. I would like to discuss this point often.

## Description of the changes
<!-- Describe the changes in this PR. -->
- 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: Follow review comments

Patch to review:
diff --git a/optuna/visualization/_contour.py b/optuna/visualization/_contour.py
index 148e5bb2cdf..23e3c17bdf9 100644
--- a/optuna/visualization/_contour.py
+++ b/optuna/visualization/_contour.py
@@ -2,10 +2,10 @@
 from typing import Callable
 from typing import Dict
 from typing import List
+from typing import NamedTuple
 from typing import Optional
 from typing import Tuple
-
-from packaging import version
+from typing import Union
 
 from optuna.logging import get_logger
 from optuna.study import Study
@@ -14,7 +14,6 @@
 from optuna.visualization._plotly_imports import _imports
 from optuna.visualization._utils import _check_plot_args
 from optuna.visualization._utils import _filter_nonfinite
-from optuna.visualization._utils import _get_param_values
 from optuna.visualization._utils import _is_log_scale
 from optuna.visualization._utils import _is_numerical
 from optuna.visualization._utils import _is_reverse_scale
@@ -24,13 +23,35 @@
     from optuna.visualization._plotly_imports import Contour
     from optuna.visualization._plotly_imports import go
     from optuna.visualization._plotly_imports import make_subplots
-    from optuna.visualization._plotly_imports import plotly
     from optuna.visualization._plotly_imports import Scatter
     from optuna.visualization._utils import COLOR_SCALE
 
 _logger = get_logger(__name__)
 
 
+PADDING_RATIO = 0.05
+
+
+class _AxisInfo(NamedTuple):
+    name: str
+    range: Tuple[float, float]
+    is_log: bool
+    is_cat: bool
+    indices: List[Union[str, int, float]]
+    values: List[Union[str, float, None]]
+
+
+class _SubContourInfo(NamedTuple):
+    xaxis: _AxisInfo
+    yaxis: _AxisInfo
+    z_values: Dict[Tuple[int, int], float]
+
+
+class _ContourInfo(NamedTuple):
+    sorted_params: List[str]
+    sub_plot_infos: List[List[_SubContourInfo]]
+
+
 def plot_contour(
     study: Study,
     params: Optional[List[str]] = None,
@@ -96,78 +117,33 @@ def _get_contour_plot(
 
     layout = go.Layout(title="Contour Plot")
 
-    trials = _filter_nonfinite(
-        study.get_trials(deepcopy=False, states=(TrialState.COMPLETE,)), target=target
-    )
+    info = _get_contour_info(study, params, target)
+    sorted_params = info.sorted_params
+    sub_plot_infos = info.sub_plot_infos
 
-    if len(trials) == 0:
-        _logger.warning("Your study does not have any completed trials.")
+    if len(sorted_params) <= 1:
         return go.Figure(data=[], layout=layout)
 
-    all_params = {p_name for t in trials for p_name in t.params.keys()}
-    if params is None:
-        sorted_params = sorted(all_params)
-    elif len(params) <= 1:
-        _logger.warning("The length of params must be greater than 1.")
-        return go.Figure(data=[], layout=layout)
-    else:
-        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))
-        sorted_params = sorted(set(params))
-
-    padding_ratio = 0.05
-    param_values_range = {}
-    for p_name in sorted_params:
-        values = _get_param_values(trials, p_name)
-
-        min_value = min(values)
-        max_value = max(values)
-
-        if _is_log_scale(trials, p_name):
-            padding = (math.log10(max_value) - math.log10(min_value)) * padding_ratio
-            min_value = math.pow(10, math.log10(min_value) - padding)
-            max_value = math.pow(10, math.log10(max_value) + padding)
-
-        elif _is_numerical(trials, p_name):
-            padding = (max_value - min_value) * padding_ratio
-            min_value = min_value - padding
-            max_value = max_value + padding
-
-        else:
-            # Plotly>=4.12.0 draws contours using the indices of categorical variables instead of
-            # raw values and the range should be updated based on the cardinality of categorical
-            # variables. See https://github.com/optuna/optuna/issues/1967.
-            if version.parse(plotly.__version__) >= version.parse("4.12.0"):
-                span = len(set(values)) - 1
-                padding = span * padding_ratio
-                min_value = -padding
-                max_value = span + padding
-
-        param_values_range[p_name] = (min_value, max_value)
-
     reverse_scale = _is_reverse_scale(study, target)
 
     if len(sorted_params) == 2:
         x_param = sorted_params[0]
         y_param = sorted_params[1]
-        sub_plots = _generate_contour_subplot(
-            trials, x_param, y_param, reverse_scale, param_values_range, target, target_name
-        )
+        sub_plots = _generate_contour_subplot(sub_plot_infos[0][0], reverse_scale, target_name)
         figure = go.Figure(data=sub_plots, layout=layout)
-        figure.update_xaxes(title_text=x_param, range=param_values_range[x_param])
-        figure.update_yaxes(title_text=y_param, range=param_values_range[y_param])
+        figure.update_xaxes(title_text=x_param, range=sub_plot_infos[0][0].xaxis.range)
+        figure.update_yaxes(title_text=y_param, range=sub_plot_infos[0][0].yaxis.range)
 
-        if not _is_numerical(trials, x_param):
+        if sub_plot_infos[0][0].xaxis.is_cat:
             figure.update_xaxes(type="category")
-        if not _is_numerical(trials, y_param):
+        if sub_plot_infos[0][0].yaxis.is_cat:
             figure.update_yaxes(type="category")
 
-        if _is_log_scale(trials, x_param):
-            log_range = [math.log10(p) for p in param_values_range[x_param]]
+        if sub_plot_infos[0][0].xaxis.is_log:
+            log_range = [math.log10(p) for p in sub_plot_infos[0][0].xaxis.range]
             figure.update_xaxes(range=log_range, type="log")
-        if _is_log_scale(trials, y_param):
-            log_range = [math.log10(p) for p in param_values_range[y_param]]
+        if sub_plot_infos[0][0].yaxis.is_log:
+            log_range = [math.log10(p) for p in sub_plot_infos[0][0].yaxis.range]
             figure.update_yaxes(range=log_range, type="log")
     else:
         figure = make_subplots(
@@ -181,13 +157,7 @@ def _get_contour_plot(
                     figure.add_trace(go.Scatter(), row=y_i + 1, col=x_i + 1)
                 else:
                     sub_plots = _generate_contour_subplot(
-                        trials,
-                        x_param,
-                        y_param,
-                        reverse_scale,
-                        param_values_range,
-                        target,
-                        target_name,
+                        sub_plot_infos[y_i][x_i], reverse_scale, target_name
                     )
                     contour = sub_plots[0]
                     scatter = sub_plots[1]
@@ -197,19 +167,21 @@ def _get_contour_plot(
                     figure.add_trace(contour, row=y_i + 1, col=x_i + 1)
                     figure.add_trace(scatter, row=y_i + 1, col=x_i + 1)
 
-                figure.update_xaxes(range=param_values_range[x_param], row=y_i + 1, col=x_i + 1)
-                figure.update_yaxes(range=param_values_range[y_param], row=y_i + 1, col=x_i + 1)
+                xaxis = sub_plot_infos[y_i][x_i].xaxis
+                yaxis = sub_plot_infos[y_i][x_i].yaxis
+                figure.update_xaxes(range=xaxis.range, row=y_i + 1, col=x_i + 1)
+                figure.update_yaxes(range=yaxis.range, row=y_i + 1, col=x_i + 1)
 
-                if not _is_numerical(trials, x_param):
+                if xaxis.is_cat:
                     figure.update_xaxes(type="category", row=y_i + 1, col=x_i + 1)
-                if not _is_numerical(trials, y_param):
+                if yaxis.is_cat:
                     figure.update_yaxes(type="category", row=y_i + 1, col=x_i + 1)
 
-                if _is_log_scale(trials, x_param):
-                    log_range = [math.log10(p) for p in param_values_range[x_param]]
+                if xaxis.is_log:
+                    log_range = [math.log10(p) for p in xaxis.range]
                     figure.update_xaxes(range=log_range, type="log", row=y_i + 1, col=x_i + 1)
-                if _is_log_scale(trials, y_param):
-                    log_range = [math.log10(p) for p in param_values_range[y_param]]
+                if yaxis.is_log:
+                    log_range = [math.log10(p) for p in yaxis.range]
                     figure.update_yaxes(range=log_range, type="log", row=y_i + 1, col=x_i + 1)
 
                 if x_i == 0:
@@ -221,71 +193,35 @@ def _get_contour_plot(
 
 
 def _generate_contour_subplot(
-    trials: List[FrozenTrial],
-    x_param: str,
-    y_param: str,
+    info: _SubContourInfo,
     reverse_scale: bool,
-    param_values_range: Optional[Dict[str, Tuple[float, float]]] = None,
-    target: Optional[Callable[[FrozenTrial], float]] = None,
     target_name: str = "Objective Value",
 ) -> Tuple["Contour", "Scatter"]:
 
-    if param_values_range is None:
-        param_values_range = {}
+    x_indices = info.xaxis.indices
+    y_indices = info.yaxis.indices
+    x_values = []
+    y_values = []
+    for x_value, y_value in zip(info.xaxis.values, info.yaxis.values):
+        if x_value is not None and y_value is not None:
+            x_values.append(x_value)
+            y_values.append(y_value)
+    z_values = [
+        [float("nan") for _ in range(len(info.xaxis.indices))]
+        for _ in range(len(info.yaxis.indices))
+    ]
+    for (x_i, y_i), z_value in info.z_values.items():
+        z_values[y_i][x_i] = z_value
 
-    x_indices = sorted(set(_get_param_values(trials, x_param)))
-    y_indices = sorted(set(_get_param_values(trials, y_param)))
     if len(x_indices) < 2:
-        _logger.warning("Param {} unique value length is less than 2.".format(x_param))
         return go.Contour(), go.Scatter()
     if len(y_indices) < 2:
-        _logger.warning("Param {} unique value length is less than 2.".format(y_param))
         return go.Contour(), go.Scatter()
 
-    # Padding to the plot for non-categorical params.
-    x_range = param_values_range[x_param]
-    if _is_numerical(trials, x_param):
-        x_indices = [x_range[0]] + x_indices + [x_range[1]]
-
-    y_range = param_values_range[y_param]
-    if _is_numerical(trials, y_param):
-        y_indices = [y_range[0]] + y_indices + [y_range[1]]
-
-    z = [[float("nan") for _ in range(len(x_indices))] for _ in range(len(y_indices))]
-
-    x_values = []
-    y_values = []
-    for trial in trials:
-        if x_param not in trial.params or y_param not in trial.params:
-            continue
-        x_value = trial.params[x_param]
-        y_value = trial.params[y_param]
-        if not _is_numerical(trials, x_param):
-            x_value = str(x_value)
-        if not _is_numerical(trials, y_param):
-            y_value = str(y_value)
-        x_values.append(x_value)
-        y_values.append(y_value)
-        x_i = x_indices.index(x_value)
-        y_i = y_indices.index(y_value)
-
-        if target is None:
-            value = trial.value
-        else:
-            value = target(trial)
-
-        if isinstance(value, int):
-            value = float(value)
-        elif not isinstance(value, float):
-            raise ValueError(
-                f"Trial{trial.number} has COMPLETE state, but its target value is non-numeric."
-            )
-        z[y_i][x_i] = value
-
     contour = go.Contour(
         x=x_indices,
         y=y_indices,
-        z=z,
+        z=z_values,
         colorbar={"title": target_name},
         colorscale=COLOR_SCALE,
         connectgaps=True,
@@ -303,4 +239,157 @@ def _generate_contour_subplot(
         showlegend=False,
     )
 
-    return (contour, scatter)
+    return contour, scatter
+
+
+def _get_contour_info(
+    study: Study,
+    params: Optional[List[str]] = None,
+    target: Optional[Callable[[FrozenTrial], float]] = None,
+) -> _ContourInfo:
+
+    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 len(trials) == 0:
+        _logger.warning("Your study does not have any completed trials.")
+        sorted_params = []
+    elif params is None:
+        sorted_params = sorted(all_params)
+    elif len(params) <= 1:
+        _logger.warning("The length of params must be greater than 1.")
+        sorted_params = list(params)
+    else:
+        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))
+        sorted_params = sorted(set(params))
+
+    sub_plot_infos: List[List[_SubContourInfo]]
+    if len(sorted_params) == 2:
+        x_param = sorted_params[0]
+        y_param = sorted_params[1]
+        sub_plot_info = _generate_contour_subplot_info(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 = _generate_contour_subplot_info(trials, x_param, y_param, target)
+                sub_plot_infos[i].append(sub_plot_info)
+
+    return _ContourInfo(sorted_params=sorted_params, sub_plot_infos=sub_plot_infos)
+
+
+def _generate_contour_subplot_info(
+    trials: List[FrozenTrial],
+    x_param: str,
+    y_param: str,
+    target: Optional[Callable[[FrozenTrial], float]],
+) -> _SubContourInfo:
+
+    xaxis = _generate_axis_info(trials, x_param)
+    yaxis = _generate_axis_info(trials, y_param)
+
+    if x_param == y_param:
+        return _SubContourInfo(xaxis=xaxis, yaxis=yaxis, z_values={})
+
+    if len(xaxis.indices) < 2:
+        _logger.warning("Param {} unique value length is less than 2.".format(x_param))
+        return _SubContourInfo(xaxis=xaxis, yaxis=yaxis, z_values={})
+    if len(yaxis.indices) < 2:
+        _logger.warning("Param {} unique value length is less than 2.".format(y_param))
+        return _SubContourInfo(xaxis=xaxis, yaxis=yaxis, z_values={})
+
+    z_values = {}
+    for i, trial in enumerate(trials):
+        if x_param not in trial.params or y_param not in trial.params:
+            continue
+        x_value = xaxis.values[i]
+        y_value = yaxis.values[i]
+        assert x_value is not None
+        assert y_value is not None
+        x_i = xaxis.indices.index(x_value)
+        y_i = yaxis.indices.index(y_value)
+
+        if target is None:
+            value = trial.value
+        else:
+            value = target(trial)
+
+        if isinstance(value, int):
+            value = float(value)
+        elif not isinstance(value, float):
+            raise ValueError(
+                f"Trial{trial.number} has COMPLETE state, but its target value is non-numeric."
+            )
+        z_values[(x_i, y_i)] = value
+
+    return _SubContourInfo(xaxis=xaxis, yaxis=yaxis, z_values=z_values)
+
+
+def _generate_axis_info(trials: List[FrozenTrial], param_name: str) -> _AxisInfo:
+    values: List[Union[str, float, None]]
+    if _is_numerical(trials, param_name):
+        values = [t.params.get(param_name) for t in trials]
+    else:
+        values = [
+            str(t.params.get(param_name)) if param_name in t.params else None for t in trials
+        ]
+
+    min_value = min([v for v in values if v is not None])
+    max_value = max([v for v in values if v is not None])
+
+    if _is_log_scale(trials, param_name):
+        min_value = float(min_value)
+        max_value = float(max_value)
+        padding = (math.log10(max_value) - math.log10(min_value)) * PADDING_RATIO
+        min_value = math.pow(10, math.log10(min_value) - padding)
+        max_value = math.pow(10, math.log10(max_value) + padding)
+        is_log = True
+        is_cat = False
+
+    elif _is_numerical(trials, param_name):
+        min_value = float(min_value)
+        max_value = float(max_value)
+        padding = (max_value - min_value) * PADDING_RATIO
+        min_value = min_value - padding
+        max_value = max_value + padding
+        is_log = False
+        is_cat = False
+
+    else:
+        span = len(set(values)) - 1
+        padding = span * PADDING_RATIO
+        min_value = -padding
+        max_value = span + padding
+        is_log = False
+        is_cat = True
+
+    indices = sorted(set([v for v in values if v is not None]))
+
+    if len(indices) < 2:
+        return _AxisInfo(
+            name=param_name,
+            range=(min_value, max_value),
+            is_log=is_log,
+            is_cat=is_cat,
+            indices=indices,
+            values=values,
+        )
+
+    if _is_numerical(trials, param_name):
+        indices.insert(0, min_value)
+        indices.append(max_value)
+
+    return _AxisInfo(
+        name=param_name,
+        range=(min_value, max_value),
+        is_log=is_log,
+        is_cat=is_cat,
+        indices=indices,
+        values=values,
+    )
diff --git a/optuna/visualization/matplotlib/_contour.py b/optuna/visualization/matplotlib/_contour.py
index 2b75bd397ea..7500d6d4117 100644
--- a/optuna/visualization/matplotlib/_contour.py
+++ b/optuna/visualization/matplotlib/_contour.py
@@ -13,13 +13,11 @@
 from optuna.study import Study
 from optuna.study import StudyDirection
 from optuna.trial import FrozenTrial
-from optuna.trial import TrialState
+from optuna.visualization._contour import _AxisInfo
+from optuna.visualization._contour import _get_contour_info
+from optuna.visualization._contour import _SubContourInfo
 from optuna.visualization._utils import _check_plot_args
-from optuna.visualization._utils import _filter_nonfinite
-from optuna.visualization._utils import _get_param_values
 from optuna.visualization.matplotlib._matplotlib_imports import _imports
-from optuna.visualization.matplotlib._utils import _is_log_scale
-from optuna.visualization.matplotlib._utils import _is_numerical
 
 
 if _imports.is_successful():
@@ -31,7 +29,7 @@
 _logger = get_logger(__name__)
 
 
-AXES_PADDING_RATIO = 5e-2
+CONTOUR_POINT_NUM = 100
 
 
 @experimental_func("2.2.0")
@@ -109,29 +107,13 @@ def _get_contour_plot(
     target_name: str = "Objective Value",
 ) -> "Axes":
 
-    # Calculate basic numbers for plotting.
-    trials = _filter_nonfinite(
-        study.get_trials(deepcopy=False, states=(TrialState.COMPLETE,)), target=target
-    )
-
-    if len(trials) == 0:
-        _logger.warning("Your study does not have any completed trials.")
-        _, ax = plt.subplots()
-        return ax
-
-    all_params = {p_name for t in trials for p_name in t.params.keys()}
+    info = _get_contour_info(study, params, target)
+    sorted_params = info.sorted_params
+    sub_plot_infos = info.sub_plot_infos
 
-    if params is None:
-        sorted_params = sorted(all_params)
-    elif len(params) <= 1:
-        _logger.warning("The length of params must be greater than 1.")
+    if len(sorted_params) <= 1:
         _, ax = plt.subplots()
         return ax
-    else:
-        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))
-        sorted_params = sorted(set(params))
     n_params = len(sorted_params)
 
     plt.style.use("ggplot")  # Use ggplot style sheet for similar outputs to plotly.
@@ -140,18 +122,8 @@ def _get_contour_plot(
         fig, axs = plt.subplots()
         axs.set_title("Contour Plot")
         cmap = _set_cmap(study, target)
-        contour_point_num = 100
 
-        # Prepare data and draw contour plots.
-        if params:
-            x_param = params[0]
-            y_param = params[1]
-        else:
-            x_param = sorted_params[0]
-            y_param = sorted_params[1]
-        cs = _generate_contour_subplot(
-            trials, x_param, y_param, axs, cmap, contour_point_num, target
-        )
+        cs = _generate_contour_subplot(sub_plot_infos[0][0], axs, cmap)
         if isinstance(cs, ContourSet):
             axcb = fig.colorbar(cs)
             axcb.set_label(target_name)
@@ -160,16 +132,13 @@ def _get_contour_plot(
         fig, axs = plt.subplots(n_params, n_params)
         fig.suptitle("Contour Plot")
         cmap = _set_cmap(study, target)
-        contour_point_num = 100
 
         # Prepare data and draw contour plots.
         cs_list = []
         for x_i, x_param in enumerate(sorted_params):
             for y_i, y_param in enumerate(sorted_params):
                 ax = axs[y_i, x_i]
-                cs = _generate_contour_subplot(
-                    trials, x_param, y_param, ax, cmap, contour_point_num, target
-                )
+                cs = _generate_contour_subplot(sub_plot_infos[y_i][x_i], ax, cmap)
                 if isinstance(cs, ContourSet):
                     cs_list.append(cs)
         if cs_list:
@@ -192,290 +161,168 @@ def fit(self, labels: List[str]) -> "_LabelEncoder":
         self.labels = sorted(set(labels))
         return self
 
-    def transform(self, labels: List[str]) -> List[int]:
+    def transform(self, labels: List[str]) -> List[Union[str, float]]:
         return [self.labels.index(label) for label in labels]
 
-    def fit_transform(self, labels: List[str]) -> List[int]:
+    def fit_transform(self, labels: List[str]) -> List[Union[str, float]]:
         return self.fit(labels).transform(labels)
 
     def get_labels(self) -> List[str]:
         return self.labels
 
-    def get_indices(self) -> List[int]:
+    def get_indices(self) -> List[Union[str, int, float]]:
         return list(range(len(self.labels)))
 
 
 def _calculate_griddata(
-    trials: List[FrozenTrial],
-    x_param: str,
-    x_indices: List[Union[str, int, float]],
-    y_param: str,
-    y_indices: List[Union[str, int, float]],
-    contour_point_num: int,
-    target: Optional[Callable[[FrozenTrial], float]],
+    xaxis: _AxisInfo,
+    yaxis: _AxisInfo,
+    z_values: Dict[Tuple[int, int], float],
 ) -> Tuple[
     np.ndarray,
     np.ndarray,
     np.ndarray,
-    List[Union[int, float]],
-    List[Union[int, float]],
-    List[Union[int, float]],
-    List[Union[int, float]],
-    List[int],
+    List[Union[str, int, float]],
     List[str],
-    List[int],
+    List[Union[str, int, float]],
     List[str],
-    int,
-    int,
+    List[Union[str, float]],
+    List[Union[str, float]],
 ]:
 
-    # Extract values for x, y, z axes from each trail.
     x_values = []
     y_values = []
-    z_values = []
-    x_range_values = []
-    y_range_values = []
-    for trial in trials:
-        contains_x_param = x_param in trial.params
-        if contains_x_param:
-            x_range_values.append(trial.params[x_param])
-
-        contains_y_param = y_param in trial.params
-        if contains_y_param:
-            y_range_values.append(trial.params[y_param])
-
-        if not contains_x_param or not contains_y_param:
-            continue
-        x_values.append(trial.params[x_param])
-        y_values.append(trial.params[y_param])
-
-        if target is None:
-            value = trial.value
-        else:
-            value = target(trial)
-
-        if isinstance(value, int):
-            value = float(value)
-        elif not isinstance(value, float):
-            raise ValueError(
-                "Trial{} has COMPLETE state, but its target value is non-numeric.".format(
-                    trial.number
-                )
-            )
-        z_values.append(value)
+    for x_value, y_value in zip(xaxis.values, yaxis.values):
+        if x_value is not None and y_value is not None:
+            x_values.append(x_value)
+            y_values.append(y_value)
 
     # Return empty values when x or y has no value.
     if len(x_values) == 0 or len(y_values) == 0:
-        return (
-            np.array([]),
-            np.array([]),
-            np.array([]),
-            x_values,
-            y_values,
-            [],
-            [],
-            [],
-            [],
-            [],
-            [],
-            0,
-            0,
-        )
+        return np.array([]), np.array([]), np.array([]), [], [], [], [], [], []
 
     # Add dummy values for grid data calculation when a parameter has one unique value.
-    x_values_dummy = []
-    y_values_dummy = []
     if len(set(x_values)) == 1:
-        x_values_dummy = [x for x in x_indices if x not in x_values]
+        x_values_dummy = [x for x in xaxis.indices if x not in x_values]
         x_values = x_values + x_values_dummy * len(x_values)
         y_values = y_values + (y_values * len(x_values_dummy))
-        z_values = z_values + (z_values * len(x_values_dummy))
     if len(set(y_values)) == 1:
-        y_values_dummy = [y for y in y_indices if y not in y_values]
+        y_values_dummy = [y for y in yaxis.indices if y not in y_values]
         y_values = y_values + y_values_dummy * len(y_values)
         x_values = x_values + (x_values * len(y_values_dummy))
-        z_values = z_values + (z_values * len(y_values_dummy))
-
-    # Convert categorical values to int.
-    cat_param_labels_x = []  # type: List[str]
-    cat_param_pos_x = []  # type: List[int]
-    cat_param_labels_y = []  # type: List[str]
-    cat_param_pos_y = []  # type: List[int]
-    if not _is_numerical(trials, x_param):
-        enc = _LabelEncoder()
-        x_range_values = enc.fit_transform(list(map(str, x_range_values)))
-        x_values = enc.transform(list(map(str, x_values)))
-        cat_param_labels_x = enc.get_labels()
-        cat_param_pos_x = enc.get_indices()
-    if not _is_numerical(trials, y_param):
-        enc = _LabelEncoder()
-        y_range_values = enc.fit_transform(list(map(str, y_range_values)))
-        y_values = enc.transform(list(map(str, y_values)))
-        cat_param_labels_y = enc.get_labels()
-        cat_param_pos_y = enc.get_indices()
-
-    # Calculate min and max of x and y.
-    x_values_min = min(x_range_values)
-    x_values_max = max(x_range_values)
-    y_values_min = min(y_range_values)
-    y_values_max = max(y_range_values)
 
-    # Calculate grid data points.
-    # For x and y, create 1-D array of evenly spaced coordinates on linear or log scale.
-    xi: np.ndarray = np.array([])
-    yi: np.ndarray = np.array([])
-    zi: np.ndarray = np.array([])
+    def _calculate_axis_data(
+        axis: _AxisInfo,
+        values: List[Union[str, float]],
+    ) -> Tuple[
+        np.ndarray,
+        List[str],
+        List[Union[str, int, float]],
+        List[Union[str, int, float]],
+        List[Union[str, float]],
+    ]:
+
+        # Convert categorical values to int.
+        cat_param_labels = []  # type: List[str]
+        cat_param_pos = []  # type: List[Union[str, int, float]]
+        indices = axis.indices
+        if axis.is_cat:
+            enc = _LabelEncoder()
+            values = enc.fit_transform(list(map(str, values)))
+            cat_param_labels = enc.get_labels()
+            cat_param_pos = enc.get_indices()
+            indices = cat_param_pos
+
+        # For x and y, create 1-D array of evenly spaced coordinates on linear or log scale.
+        if axis.is_log:
+            ci = np.logspace(np.log10(axis.range[0]), np.log10(axis.range[1]), CONTOUR_POINT_NUM)
+        else:
+            ci = np.linspace(axis.range[0], axis.range[1], CONTOUR_POINT_NUM)
 
-    if _is_log_scale(trials, x_param):
-        padding_x = (np.log10(x_values_max) - np.log10(x_values_min)) * AXES_PADDING_RATIO
-        x_values_min = np.power(10, np.log10(x_values_min) - padding_x)
-        x_values_max = np.power(10, np.log10(x_values_max) + padding_x)
-        xi = np.logspace(np.log10(x_values_min), np.log10(x_values_max), contour_point_num)
-    else:
-        padding_x = (x_values_max - x_values_min) * AXES_PADDING_RATIO
-        x_values_min -= padding_x
-        x_values_max += padding_x
-        xi = np.linspace(x_values_min, x_values_max, contour_point_num)
-
-    if _is_log_scale(trials, y_param):
-        padding_y = (np.log10(y_values_max) - np.log10(y_values_min)) * AXES_PADDING_RATIO
-        y_values_min = np.power(10, np.log10(y_values_min) - padding_y)
-        y_values_max = np.power(10, np.log10(y_values_max) + padding_y)
-        yi = np.logspace(np.log10(y_values_min), np.log10(y_values_max), contour_point_num)
-    else:
-        padding_y = (y_values_max - y_values_min) * AXES_PADDING_RATIO
-        y_values_min -= padding_y
-        y_values_max += padding_y
-        yi = np.linspace(y_values_min, y_values_max, contour_point_num)
+        return ci, cat_param_labels, cat_param_pos, indices, values
+
+    xi, cat_param_labels_x, cat_param_pos_x, x_indices, x_values = _calculate_axis_data(
+        xaxis,
+        x_values,
+    )
+    yi, cat_param_labels_y, cat_param_pos_y, y_indices, y_values = _calculate_axis_data(
+        yaxis,
+        y_values,
+    )
 
+    # Calculate grid data points.
+    zi: np.ndarray = np.array([])
     # create irregularly spaced map of trial values
     # and interpolate it with Plotly's interpolation formulation
-    if x_param != y_param:
-        zmap = _create_zmap(x_values, y_values, z_values, xi, yi)
-        zi = _interpolate_zmap(zmap, contour_point_num)
+    if xaxis.name != yaxis.name:
+        zi = _interpolate_zmap(z_values, CONTOUR_POINT_NUM)
 
     return (
         xi,
         yi,
         zi,
-        x_values,
-        y_values,
-        [x_values_min, x_values_max],
-        [y_values_min, y_values_max],
         cat_param_pos_x,
         cat_param_labels_x,
         cat_param_pos_y,
         cat_param_labels_y,
-        len(x_values_dummy),
-        len(y_values_dummy),
+        x_values,
+        y_values,
     )
 
 
-def _generate_contour_subplot(
-    trials: List[FrozenTrial],
-    x_param: str,
-    y_param: str,
-    ax: "Axes",
-    cmap: "Colormap",
-    contour_point_num: int,
-    target: Optional[Callable[[FrozenTrial], float]],
-) -> "ContourSet":
-
-    x_indices = sorted(set(_get_param_values(trials, x_param)))
-    y_indices = sorted(set(_get_param_values(trials, y_param)))
-    if len(x_indices) < 2:
-        _logger.warning("Param {} unique value length is less than 2.".format(x_param))
+def _generate_contour_subplot(info: _SubContourInfo, ax: "Axes", cmap: "Colormap") -> "ContourSet":
+
+    if len(info.xaxis.indices) < 2:
+        ax.label_outer()
         return ax
-    if len(y_indices) < 2:
-        _logger.warning("Param {} unique value length is less than 2.".format(y_param))
+    if len(info.yaxis.indices) < 2:
+        ax.label_outer()
         return ax
 
     (
         xi,
         yi,
         zi,
-        x_values,
-        y_values,
-        x_values_range,
-        y_values_range,
         x_cat_param_pos,
         x_cat_param_label,
         y_cat_param_pos,
         y_cat_param_label,
-        x_values_dummy_count,
-        y_values_dummy_count,
-    ) = _calculate_griddata(
-        trials, x_param, x_indices, y_param, y_indices, contour_point_num, target
-    )
+        x_values,
+        y_values,
+    ) = _calculate_griddata(info.xaxis, info.yaxis, info.z_values)
     cs = None
-    ax.set(xlabel=x_param, ylabel=y_param)
-    ax.set_xlim(x_values_range[0], x_values_range[1])
-    ax.set_ylim(y_values_range[0], y_values_range[1])
+    ax.set(xlabel=info.xaxis.name, ylabel=info.yaxis.name)
+    ax.set_xlim(info.xaxis.range[0], info.xaxis.range[1])
+    ax.set_ylim(info.yaxis.range[0], info.yaxis.range[1])
     if len(zi) > 0:
-        if _is_log_scale(trials, x_param):
+        if info.xaxis.is_log:
             ax.set_xscale("log")
-        if _is_log_scale(trials, y_param):
+        if info.yaxis.is_log:
             ax.set_yscale("log")
-        if x_param != y_param:
+        if info.xaxis.name != info.yaxis.name:
             # Contour the gridded data.
             ax.contour(xi, yi, zi, 15, linewidths=0.5, colors="k")
             cs = ax.contourf(xi, yi, zi, 15, cmap=cmap.reversed())
             # Plot data points.
-            if x_values_dummy_count > 0:
-                x_org_len = int(len(x_values) / (x_values_dummy_count + 1))
-                y_org_len = int(len(y_values) / (x_values_dummy_count + 1))
-            elif y_values_dummy_count > 0:
-                x_org_len = int(len(x_values) / (y_values_dummy_count + 1))
-                y_org_len = int(len(y_values) / (y_values_dummy_count + 1))
-            else:
-                x_org_len = len(x_values)
-                y_org_len = len(x_values)
             ax.scatter(
-                x_values[:x_org_len],
-                y_values[:y_org_len],
+                x_values,
+                y_values,
                 marker="o",
                 c="black",
                 s=20,
                 edgecolors="grey",
                 linewidth=2.0,
             )
-    if x_cat_param_pos:
+    if info.xaxis.is_cat:
         ax.set_xticks(x_cat_param_pos)
         ax.set_xticklabels(x_cat_param_label)
-    if y_cat_param_pos:
+    if info.yaxis.is_cat:
         ax.set_yticks(y_cat_param_pos)
         ax.set_yticklabels(y_cat_param_label)
     ax.label_outer()
     return cs
 
 
-def _create_zmap(
-    x_values: List[Union[int, float]],
-    y_values: List[Union[int, float]],
-    z_values: List[float],
-    xi: np.ndarray,
-    yi: np.ndarray,
-) -> Dict[Tuple[int, int], float]:
-
-    # creates z-map from trial values and params.
-    # z-map is represented by hashmap of coordinate and trial value pairs
-    #
-    # coordinates are represented by tuple of integers, where the first item
-    # indicates x-axis index and the second item indicates y-axis index
-    # and refer to a position of trial value on irregular param grid
-    #
-    # since params were resampled either with linspace or logspace
-    # original params might not be on the x and y axes anymore
-    # so we are going with close approximations of trial value positions
-    zmap = dict()
-    for x, y, z in zip(x_values, y_values, z_values):
-        xindex = int(np.argmin(np.abs(xi - x)))
-        yindex = int(np.argmin(np.abs(yi - y)))
-        zmap[(xindex, yindex)] = z
-
-    return zmap
-
-
 def _interpolate_zmap(zmap: Dict[Tuple[int, int], float], contour_plot_num: int) -> np.ndarray:
 
     # implements interpolation formulation used in Plotly
diff --git a/tests/visualization_tests/matplotlib_tests/test_contour.py b/tests/visualization_tests/matplotlib_tests/test_contour.py
index ac53500ab65..2cff02497b3 100644
--- a/tests/visualization_tests/matplotlib_tests/test_contour.py
+++ b/tests/visualization_tests/matplotlib_tests/test_contour.py
@@ -16,30 +16,9 @@
 from optuna.testing.visualization import prepare_study_with_trials
 from optuna.trial import create_trial
 from optuna.trial import Trial
+from optuna.visualization._contour import PADDING_RATIO
 from optuna.visualization.matplotlib import plot_contour
-from optuna.visualization.matplotlib._contour import _create_zmap
 from optuna.visualization.matplotlib._contour import _interpolate_zmap
-from optuna.visualization.matplotlib._contour import AXES_PADDING_RATIO
-
-
-def test_create_zmap() -> None:
-
-    x_values = np.arange(10)
-    y_values = np.arange(10)
-    z_values = list(np.random.rand(10))
-
-    # we are testing for exact placement of z_values
-    # so also passing x_values and y_values as xi and yi
-    zmap = _create_zmap(x_values.tolist(), y_values.tolist(), z_values, x_values, y_values)
-
-    assert len(zmap) == len(z_values)
-    for coord, value in zmap.items():
-        # test if value under coordinate
-        # still refers to original trial value
-        xidx = coord[0]
-        yidx = coord[1]
-        assert xidx == yidx
-        assert z_values[xidx] == value
 
 
 def test_interpolate_zmap() -> None:
@@ -292,7 +271,7 @@ def test_contour_subplots_have_correct_axis_labels_and_ranges() -> None:
     }
     for index, (param_name, param_range) in enumerate(param_ranges.items()):
         minimum, maximum = param_range
-        padding = (maximum - minimum) * AXES_PADDING_RATIO
+        padding = (maximum - minimum) * PADDING_RATIO
         param_range_with_padding = (minimum - padding, maximum + padding)
         assert subplots[index, 0].get_ylabel() == param_name
         assert subplots[-1, index].get_xlabel() == param_name
diff --git a/tests/visualization_tests/test_contour.py b/tests/visualization_tests/test_contour.py
index 4bf3a5e062a..282fbefdf9f 100644
--- a/tests/visualization_tests/test_contour.py
+++ b/tests/visualization_tests/test_contour.py
@@ -16,6 +16,7 @@
 from optuna.trial import Trial
 from optuna.visualization import plot_contour
 from optuna.visualization._contour import _generate_contour_subplot
+from optuna.visualization._contour import _generate_contour_subplot_info
 from optuna.visualization._utils import _is_reverse_scale
 from optuna.visualization._utils import COLOR_SCALE
 
@@ -142,12 +143,14 @@ def test_generate_contour_plot_for_few_observations() -> None:
 
     # `x_axis` has one observation.
     params = ["param_a", "param_b"]
-    contour, scatter = _generate_contour_subplot(trials, params[0], params[1], reverse_scale)
+    info = _generate_contour_subplot_info(trials, params[0], params[1], None)
+    contour, scatter = _generate_contour_subplot(info, reverse_scale)
     assert contour.x is None and contour.y is None and scatter.x is None and scatter.y is None
 
     # `y_axis` has one observation.
     params = ["param_b", "param_a"]
-    contour, scatter = _generate_contour_subplot(trials, params[0], params[1], reverse_scale)
+    info = _generate_contour_subplot_info(trials, params[0], params[1], None)
+    contour, scatter = _generate_contour_subplot(info, reverse_scale)
     assert contour.x is None and contour.y is None and scatter.x is None and scatter.y is None
 
 


Code review comments:
diff_hunk: @@ -303,4 +239,157 @@ def _generate_contour_subplot(
         showlegend=False,
     )
 
-    return (contour, scatter)
+    return contour, scatter
+
+
+def _get_contour_info(
+    study: Study,
+    params: Optional[List[str]] = None,
+    target: Optional[Callable[[FrozenTrial], float]] = None,
+) -> _ContourInfo:
+
+    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 len(trials) == 0:
+        _logger.warning("Your study does not have any completed trials.")
+        sorted_params = []
+    elif params is None:
+        sorted_params = sorted(all_params)
+    elif len(params) <= 1:
+        _logger.warning("The length of params must be greater than 1.")
+        sorted_params = list(params)
+    else:
+        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))
+        sorted_params = sorted(set(params))
+
+    sub_plot_infos: List[List[_SubContourInfo]]
+    if len(sorted_params) == 2:
+        x_param = sorted_params[0]
+        y_param = sorted_params[1]
+        sub_plot_info = _generate_contour_subplot_info(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 = _generate_contour_subplot_info(trials, x_param, y_param, target)
+                sub_plot_infos[i].append(sub_plot_info)
+
+    return _ContourInfo(sorted_params=sorted_params, sub_plot_infos=sub_plot_infos)
+
+
+def _generate_contour_subplot_info(
path: optuna/visualization/_contour.py
text: @user1:
[nit] can we unify the naming rule to obtain the `_*Info`? Currently, we have two ways: `_get_*` and `_generate_*` as follows:

- `_get_contour_info` 
- `_generate_contour_subplot_info`
- `_generate_axis_info`

@author:
Thanks! I will use `_get_*_info`.

diff_hunk: @@ -192,290 +161,168 @@ def fit(self, labels: List[str]) -> "_LabelEncoder":
         self.labels = sorted(set(labels))
         return self
 
-    def transform(self, labels: List[str]) -> List[int]:
+    def transform(self, labels: List[str]) -> List[Union[str, float]]:
         return [self.labels.index(label) for label in labels]
 
-    def fit_transform(self, labels: List[str]) -> List[int]:
+    def fit_transform(self, labels: List[str]) -> List[Union[str, float]]:
         return self.fit(labels).transform(labels)
 
     def get_labels(self) -> List[str]:
         return self.labels
 
-    def get_indices(self) -> List[int]:
+    def get_indices(self) -> List[Union[str, int, float]]:
         return list(range(len(self.labels)))
 
 
 def _calculate_griddata(
-    trials: List[FrozenTrial],
-    x_param: str,
-    x_indices: List[Union[str, int, float]],
-    y_param: str,
-    y_indices: List[Union[str, int, float]],
-    contour_point_num: int,
-    target: Optional[Callable[[FrozenTrial], float]],
+    xaxis: _AxisInfo,
+    yaxis: _AxisInfo,
+    z_values: Dict[Tuple[int, int], float],
 ) -> Tuple[
     np.ndarray,
     np.ndarray,
     np.ndarray,
-    List[Union[int, float]],
-    List[Union[int, float]],
-    List[Union[int, float]],
-    List[Union[int, float]],
-    List[int],
+    List[Union[str, int, float]],
     List[str],
-    List[int],
+    List[Union[str, int, float]],
     List[str],
-    int,
-    int,
+    List[Union[str, float]],
+    List[Union[str, float]],
 ]:
 
-    # Extract values for x, y, z axes from each trail.
     x_values = []
     y_values = []
-    z_values = []
-    x_range_values = []
-    y_range_values = []
-    for trial in trials:
-        contains_x_param = x_param in trial.params
-        if contains_x_param:
-            x_range_values.append(trial.params[x_param])
-
-        contains_y_param = y_param in trial.params
-        if contains_y_param:
-            y_range_values.append(trial.params[y_param])
-
-        if not contains_x_param or not contains_y_param:
-            continue
-        x_values.append(trial.params[x_param])
-        y_values.append(trial.params[y_param])
-
-        if target is None:
-            value = trial.value
-        else:
-            value = target(trial)
-
-        if isinstance(value, int):
-            value = float(value)
-        elif not isinstance(value, float):
-            raise ValueError(
-                "Trial{} has COMPLETE state, but its target value is non-numeric.".format(
-                    trial.number
-                )
-            )
-        z_values.append(value)
+    for x_value, y_value in zip(xaxis.values, yaxis.values):
+        if x_value is not None and y_value is not None:
+            x_values.append(x_value)
+            y_values.append(y_value)
 
     # Return empty values when x or y has no value.
     if len(x_values) == 0 or len(y_values) == 0:
-        return (
-            np.array([]),
-            np.array([]),
-            np.array([]),
-            x_values,
-            y_values,
-            [],
-            [],
-            [],
-            [],
-            [],
-            [],
-            0,
-            0,
-        )
+        return np.array([]), np.array([]), np.array([]), [], [], [], [], [], []
 
     # Add dummy values for grid data calculation when a parameter has one unique value.
-    x_values_dummy = []
-    y_values_dummy = []
     if len(set(x_values)) == 1:
-        x_values_dummy = [x for x in x_indices if x not in x_values]
+        x_values_dummy = [x for x in xaxis.indices if x not in x_values]
         x_values = x_values + x_values_dummy * len(x_values)
         y_values = y_values + (y_values * len(x_values_dummy))
-        z_values = z_values + (z_values * len(x_values_dummy))
     if len(set(y_values)) == 1:
-        y_values_dummy = [y for y in y_indices if y not in y_values]
+        y_values_dummy = [y for y in yaxis.indices if y not in y_values]
         y_values = y_values + y_values_dummy * len(y_values)
         x_values = x_values + (x_values * len(y_values_dummy))
-        z_values = z_values + (z_values * len(y_values_dummy))
-
-    # Convert categorical values to int.
-    cat_param_labels_x = []  # type: List[str]
-    cat_param_pos_x = []  # type: List[int]
-    cat_param_labels_y = []  # type: List[str]
-    cat_param_pos_y = []  # type: List[int]
-    if not _is_numerical(trials, x_param):
-        enc = _LabelEncoder()
-        x_range_values = enc.fit_transform(list(map(str, x_range_values)))
-        x_values = enc.transform(list(map(str, x_values)))
-        cat_param_labels_x = enc.get_labels()
-        cat_param_pos_x = enc.get_indices()
-    if not _is_numerical(trials, y_param):
-        enc = _LabelEncoder()
-        y_range_values = enc.fit_transform(list(map(str, y_range_values)))
-        y_values = enc.transform(list(map(str, y_values)))
-        cat_param_labels_y = enc.get_labels()
-        cat_param_pos_y = enc.get_indices()
-
-    # Calculate min and max of x and y.
-    x_values_min = min(x_range_values)
-    x_values_max = max(x_range_values)
-    y_values_min = min(y_range_values)
-    y_values_max = max(y_range_values)
 
-    # Calculate grid data points.
-    # For x and y, create 1-D array of evenly spaced coordinates on linear or log scale.
-    xi: np.ndarray = np.array([])
-    yi: np.ndarray = np.array([])
-    zi: np.ndarray = np.array([])
+    def _calculate_axis_data(
+        axis: _AxisInfo,
+        values: List[Union[str, float]],
+    ) -> Tuple[
+        np.ndarray,
+        List[str],
+        List[Union[str, int, float]],
+        List[Union[str, int, float]],
+        List[Union[str, float]],
+    ]:
+
+        # Convert categorical values to int.
+        cat_param_labels = []  # type: List[str]
+        cat_param_pos = []  # type: List[Union[str, int, float]]
+        indices = axis.indices
+        if axis.is_cat:
+            enc = _LabelEncoder()
+            values = enc.fit_transform(list(map(str, values)))
+            cat_param_labels = enc.get_labels()
+            cat_param_pos = enc.get_indices()
+            indices = cat_param_pos
+
+        # For x and y, create 1-D array of evenly spaced coordinates on linear or log scale.
+        if axis.is_log:
+            ci = np.logspace(np.log10(axis.range[0]), np.log10(axis.range[1]), CONTOUR_POINT_NUM)
+        else:
+            ci = np.linspace(axis.range[0], axis.range[1], CONTOUR_POINT_NUM)
 
-    if _is_log_scale(trials, x_param):
-        padding_x = (np.log10(x_values_max) - np.log10(x_values_min)) * AXES_PADDING_RATIO
-        x_values_min = np.power(10, np.log10(x_values_min) - padding_x)
-        x_values_max = np.power(10, np.log10(x_values_max) + padding_x)
-        xi = np.logspace(np.log10(x_values_min), np.log10(x_values_max), contour_point_num)
-    else:
-        padding_x = (x_values_max - x_values_min) * AXES_PADDING_RATIO
-        x_values_min -= padding_x
-        x_values_max += padding_x
-        xi = np.linspace(x_values_min, x_values_max, contour_point_num)
-
-    if _is_log_scale(trials, y_param):
-        padding_y = (np.log10(y_values_max) - np.log10(y_values_min)) * AXES_PADDING_RATIO
-        y_values_min = np.power(10, np.log10(y_values_min) - padding_y)
-        y_values_max = np.power(10, np.log10(y_values_max) + padding_y)
-        yi = np.logspace(np.log10(y_values_min), np.log10(y_values_max), contour_point_num)
-    else:
-        padding_y = (y_values_max - y_values_min) * AXES_PADDING_RATIO
-        y_values_min -= padding_y
-        y_values_max += padding_y
-        yi = np.linspace(y_values_min, y_values_max, contour_point_num)
+        return ci, cat_param_labels, cat_param_pos, indices, values
+
+    xi, cat_param_labels_x, cat_param_pos_x, x_indices, x_values = _calculate_axis_data(
path: optuna/visualization/matplotlib/_contour.py
text: @user1:
`x_indices` is unused in the following part. Can we remove it from the returned value?

diff_hunk: @@ -192,290 +161,168 @@ def fit(self, labels: List[str]) -> "_LabelEncoder":
         self.labels = sorted(set(labels))
         return self
 
-    def transform(self, labels: List[str]) -> List[int]:
+    def transform(self, labels: List[str]) -> List[Union[str, float]]:
         return [self.labels.index(label) for label in labels]
 
-    def fit_transform(self, labels: List[str]) -> List[int]:
+    def fit_transform(self, labels: List[str]) -> List[Union[str, float]]:
         return self.fit(labels).transform(labels)
 
     def get_labels(self) -> List[str]:
         return self.labels
 
-    def get_indices(self) -> List[int]:
+    def get_indices(self) -> List[Union[str, int, float]]:
         return list(range(len(self.labels)))
 
 
 def _calculate_griddata(
-    trials: List[FrozenTrial],
-    x_param: str,
-    x_indices: List[Union[str, int, float]],
-    y_param: str,
-    y_indices: List[Union[str, int, float]],
-    contour_point_num: int,
-    target: Optional[Callable[[FrozenTrial], float]],
+    xaxis: _AxisInfo,
+    yaxis: _AxisInfo,
+    z_values: Dict[Tuple[int, int], float],
 ) -> Tuple[
     np.ndarray,
     np.ndarray,
     np.ndarray,
-    List[Union[int, float]],
-    List[Union[int, float]],
-    List[Union[int, float]],
-    List[Union[int, float]],
-    List[int],
+    List[Union[str, int, float]],
     List[str],
-    List[int],
+    List[Union[str, int, float]],
     List[str],
-    int,
-    int,
+    List[Union[str, float]],
+    List[Union[str, float]],
 ]:
 
-    # Extract values for x, y, z axes from each trail.
     x_values = []
     y_values = []
-    z_values = []
-    x_range_values = []
-    y_range_values = []
-    for trial in trials:
-        contains_x_param = x_param in trial.params
-        if contains_x_param:
-            x_range_values.append(trial.params[x_param])
-
-        contains_y_param = y_param in trial.params
-        if contains_y_param:
-            y_range_values.append(trial.params[y_param])
-
-        if not contains_x_param or not contains_y_param:
-            continue
-        x_values.append(trial.params[x_param])
-        y_values.append(trial.params[y_param])
-
-        if target is None:
-            value = trial.value
-        else:
-            value = target(trial)
-
-        if isinstance(value, int):
-            value = float(value)
-        elif not isinstance(value, float):
-            raise ValueError(
-                "Trial{} has COMPLETE state, but its target value is non-numeric.".format(
-                    trial.number
-                )
-            )
-        z_values.append(value)
+    for x_value, y_value in zip(xaxis.values, yaxis.values):
+        if x_value is not None and y_value is not None:
+            x_values.append(x_value)
+            y_values.append(y_value)
 
     # Return empty values when x or y has no value.
     if len(x_values) == 0 or len(y_values) == 0:
-        return (
-            np.array([]),
-            np.array([]),
-            np.array([]),
-            x_values,
-            y_values,
-            [],
-            [],
-            [],
-            [],
-            [],
-            [],
-            0,
-            0,
-        )
+        return np.array([]), np.array([]), np.array([]), [], [], [], [], [], []
 
     # Add dummy values for grid data calculation when a parameter has one unique value.
-    x_values_dummy = []
-    y_values_dummy = []
     if len(set(x_values)) == 1:
-        x_values_dummy = [x for x in x_indices if x not in x_values]
+        x_values_dummy = [x for x in xaxis.indices if x not in x_values]
         x_values = x_values + x_values_dummy * len(x_values)
         y_values = y_values + (y_values * len(x_values_dummy))
-        z_values = z_values + (z_values * len(x_values_dummy))
     if len(set(y_values)) == 1:
-        y_values_dummy = [y for y in y_indices if y not in y_values]
+        y_values_dummy = [y for y in yaxis.indices if y not in y_values]
         y_values = y_values + y_values_dummy * len(y_values)
         x_values = x_values + (x_values * len(y_values_dummy))
-        z_values = z_values + (z_values * len(y_values_dummy))
-
-    # Convert categorical values to int.
-    cat_param_labels_x = []  # type: List[str]
-    cat_param_pos_x = []  # type: List[int]
-    cat_param_labels_y = []  # type: List[str]
-    cat_param_pos_y = []  # type: List[int]
-    if not _is_numerical(trials, x_param):
-        enc = _LabelEncoder()
-        x_range_values = enc.fit_transform(list(map(str, x_range_values)))
-        x_values = enc.transform(list(map(str, x_values)))
-        cat_param_labels_x = enc.get_labels()
-        cat_param_pos_x = enc.get_indices()
-    if not _is_numerical(trials, y_param):
-        enc = _LabelEncoder()
-        y_range_values = enc.fit_transform(list(map(str, y_range_values)))
-        y_values = enc.transform(list(map(str, y_values)))
-        cat_param_labels_y = enc.get_labels()
-        cat_param_pos_y = enc.get_indices()
-
-    # Calculate min and max of x and y.
-    x_values_min = min(x_range_values)
-    x_values_max = max(x_range_values)
-    y_values_min = min(y_range_values)
-    y_values_max = max(y_range_values)
 
-    # Calculate grid data points.
-    # For x and y, create 1-D array of evenly spaced coordinates on linear or log scale.
-    xi: np.ndarray = np.array([])
-    yi: np.ndarray = np.array([])
-    zi: np.ndarray = np.array([])
+    def _calculate_axis_data(
+        axis: _AxisInfo,
+        values: List[Union[str, float]],
+    ) -> Tuple[
+        np.ndarray,
+        List[str],
+        List[Union[str, int, float]],
+        List[Union[str, int, float]],
+        List[Union[str, float]],
+    ]:
+
+        # Convert categorical values to int.
+        cat_param_labels = []  # type: List[str]
+        cat_param_pos = []  # type: List[Union[str, int, float]]
+        indices = axis.indices
+        if axis.is_cat:
+            enc = _LabelEncoder()
+            values = enc.fit_transform(list(map(str, values)))
+            cat_param_labels = enc.get_labels()
+            cat_param_pos = enc.get_indices()
+            indices = cat_param_pos
+
+        # For x and y, create 1-D array of evenly spaced coordinates on linear or log scale.
+        if axis.is_log:
+            ci = np.logspace(np.log10(axis.range[0]), np.log10(axis.range[1]), CONTOUR_POINT_NUM)
+        else:
+            ci = np.linspace(axis.range[0], axis.range[1], CONTOUR_POINT_NUM)
 
-    if _is_log_scale(trials, x_param):
-        padding_x = (np.log10(x_values_max) - np.log10(x_values_min)) * AXES_PADDING_RATIO
-        x_values_min = np.power(10, np.log10(x_values_min) - padding_x)
-        x_values_max = np.power(10, np.log10(x_values_max) + padding_x)
-        xi = np.logspace(np.log10(x_values_min), np.log10(x_values_max), contour_point_num)
-    else:
-        padding_x = (x_values_max - x_values_min) * AXES_PADDING_RATIO
-        x_values_min -= padding_x
-        x_values_max += padding_x
-        xi = np.linspace(x_values_min, x_values_max, contour_point_num)
-
-    if _is_log_scale(trials, y_param):
-        padding_y = (np.log10(y_values_max) - np.log10(y_values_min)) * AXES_PADDING_RATIO
-        y_values_min = np.power(10, np.log10(y_values_min) - padding_y)
-        y_values_max = np.power(10, np.log10(y_values_max) + padding_y)
-        yi = np.logspace(np.log10(y_values_min), np.log10(y_values_max), contour_point_num)
-    else:
-        padding_y = (y_values_max - y_values_min) * AXES_PADDING_RATIO
-        y_values_min -= padding_y
-        y_values_max += padding_y
-        yi = np.linspace(y_values_min, y_values_max, contour_point_num)
+        return ci, cat_param_labels, cat_param_pos, indices, values
+
+    xi, cat_param_labels_x, cat_param_pos_x, x_indices, x_values = _calculate_axis_data(
+        xaxis,
+        x_values,
+    )
+    yi, cat_param_labels_y, cat_param_pos_y, y_indices, y_values = _calculate_axis_data(
+        yaxis,
+        y_values,
+    )
 
+    # Calculate grid data points.
+    zi: np.ndarray = np.array([])
     # create irregularly spaced map of trial values
     # and interpolate it with Plotly's interpolation formulation
-    if x_param != y_param:
-        zmap = _create_zmap(x_values, y_values, z_values, xi, yi)
-        zi = _interpolate_zmap(zmap, contour_point_num)
+    if xaxis.name != yaxis.name:
+        zi = _interpolate_zmap(z_values, CONTOUR_POINT_NUM)
 
     return (
         xi,
         yi,
         zi,
-        x_values,
-        y_values,
-        [x_values_min, x_values_max],
-        [y_values_min, y_values_max],
         cat_param_pos_x,
         cat_param_labels_x,
         cat_param_pos_y,
         cat_param_labels_y,
-        len(x_values_dummy),
-        len(y_values_dummy),
+        x_values,
+        y_values,
     )
 
 
-def _generate_contour_subplot(
-    trials: List[FrozenTrial],
-    x_param: str,
-    y_param: str,
-    ax: "Axes",
-    cmap: "Colormap",
-    contour_point_num: int,
-    target: Optional[Callable[[FrozenTrial], float]],
-) -> "ContourSet":
-
-    x_indices = sorted(set(_get_param_values(trials, x_param)))
-    y_indices = sorted(set(_get_param_values(trials, y_param)))
-    if len(x_indices) < 2:
-        _logger.warning("Param {} unique value length is less than 2.".format(x_param))
+def _generate_contour_subplot(info: _SubContourInfo, ax: "Axes", cmap: "Colormap") -> "ContourSet":
+
+    if len(info.xaxis.indices) < 2:
path: optuna/visualization/matplotlib/_contour.py
text: @user1:
The inside of if block looks identical. How about adding the condition as follows and removing the second if block?
```suggestion
    if len(info.xaxis.indices) < 2 or len(info.yaxis.indices) < 2:
```

diff_hunk: @@ -96,78 +117,33 @@ def _get_contour_plot(
 
     layout = go.Layout(title="Contour Plot")
 
-    trials = _filter_nonfinite(
-        study.get_trials(deepcopy=False, states=(TrialState.COMPLETE,)), target=target
-    )
+    info = _get_contour_info(study, params, target)
+    sorted_params = info.sorted_params
+    sub_plot_infos = info.sub_plot_infos
 
-    if len(trials) == 0:
-        _logger.warning("Your study does not have any completed trials.")
+    if len(sorted_params) <= 1:
         return go.Figure(data=[], layout=layout)
 
-    all_params = {p_name for t in trials for p_name in t.params.keys()}
-    if params is None:
-        sorted_params = sorted(all_params)
-    elif len(params) <= 1:
-        _logger.warning("The length of params must be greater than 1.")
-        return go.Figure(data=[], layout=layout)
-    else:
-        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))
-        sorted_params = sorted(set(params))
-
-    padding_ratio = 0.05
-    param_values_range = {}
-    for p_name in sorted_params:
-        values = _get_param_values(trials, p_name)
-
-        min_value = min(values)
-        max_value = max(values)
-
-        if _is_log_scale(trials, p_name):
-            padding = (math.log10(max_value) - math.log10(min_value)) * padding_ratio
-            min_value = math.pow(10, math.log10(min_value) - padding)
-            max_value = math.pow(10, math.log10(max_value) + padding)
-
-        elif _is_numerical(trials, p_name):
-            padding = (max_value - min_value) * padding_ratio
-            min_value = min_value - padding
-            max_value = max_value + padding
-
-        else:
-            # Plotly>=4.12.0 draws contours using the indices of categorical variables instead of
-            # raw values and the range should be updated based on the cardinality of categorical
-            # variables. See https://github.com/optuna/optuna/issues/1967.
-            if version.parse(plotly.__version__) >= version.parse("4.12.0"):
-                span = len(set(values)) - 1
-                padding = span * padding_ratio
-                min_value = -padding
-                max_value = span + padding
-
-        param_values_range[p_name] = (min_value, max_value)
-
     reverse_scale = _is_reverse_scale(study, target)
 
     if len(sorted_params) == 2:
         x_param = sorted_params[0]
         y_param = sorted_params[1]
-        sub_plots = _generate_contour_subplot(
-            trials, x_param, y_param, reverse_scale, param_values_range, target, target_name
-        )
+        sub_plots = _generate_contour_subplot(sub_plot_infos[0][0], reverse_scale, target_name)
path: optuna/visualization/_contour.py
text: @user1:
Accessing `sub_plot_infos[0][0]` several times looks a little bit wordy. How about storing a variable like 
```suggestion
        sub_plot_info = sub_plot_infos[0][0]
        sub_plots = _generate_contour_subplot(sub_plot_info, reverse_scale, target_name)
```

diff_hunk: @@ -192,290 +161,168 @@ def fit(self, labels: List[str]) -> "_LabelEncoder":
         self.labels = sorted(set(labels))
         return self
 
-    def transform(self, labels: List[str]) -> List[int]:
+    def transform(self, labels: List[str]) -> List[Union[str, float]]:
path: optuna/visualization/matplotlib/_contour.py
text: @user1:
I might miss something, but can you explain why we need to change the returned type? Since the `index` method returns an int, the returned value is supposed to be a list of int.

@author:
Thanks. I revert the type change. This simplify the annotations not only `_LabelEncoder` but also `_calculate_griddata`.