Repository: optuna/optuna

Issue:
Fix the axes of `optuna.visualization.matplotlib.plot_contour`
I found that the axes of `optuna.visualization.matplotlib.plot_contour` has bug.

```python
import optuna
from optuna.testing.visualization import prepare_study_with_trials
import matplotlib.pyplot as plt

study = prepare_study_with_trials(more_than_three=True, n_objectives=2, with_c_d=True)
print(study.trials_dataframe())
figure = optuna.visualization.plot_contour(study, target=lambda t: t.values[0])

figure.write_image('plotly.png')
optuna.visualization.matplotlib.plot_contour(study, target=lambda t: t.values[0])
plt.savefig('matplotlib.png')
```

![plotly](https://user-images.githubusercontent.com/1255726/146722151-9516e355-f9a7-4ccc-95fd-3158e5ea500b.png)

![matplotlib](https://user-images.githubusercontent.com/1255726/146722133-f7a48ce3-1985-4537-b672-bd4150e5307c.png)

To plot a-b subplot, the max and the min of params are determined with trials whose params have `a` and `b`. But the maximum and minimum values should be determined independently for `a` and `b`.

_Originally posted by @not522 in https://github.com/optuna/optuna/issues/2959#issuecomment-997638266_


Discussion:
This issue has not seen any recent activity.

Resolving patch:
diff --git a/optuna/visualization/matplotlib/_contour.py b/optuna/visualization/matplotlib/_contour.py
--- a/optuna/visualization/matplotlib/_contour.py
+++ b/optuna/visualization/matplotlib/_contour.py
@@ -1,6 +1,4 @@
-from collections import defaultdict
 from typing import Callable
-from typing import DefaultDict
 from typing import Dict
 from typing import List
 from typing import Optional
@@ -187,15 +185,25 @@ def _set_cmap(study: Study, target: Optional[Callable[[FrozenTrial], float]]) ->
     return plt.get_cmap(cmap)
 
 
-def _convert_categorical2int(values: List[str]) -> Tuple[List[int], List[str], List[int]]:
-    vocab = defaultdict(lambda: len(vocab))  # type: DefaultDict[str, int]
-    [vocab[v] for v in sorted(values)]
-    values_converted = [vocab[v] for v in values]
-    vocab_item_sorted = sorted(vocab.items(), key=lambda x: x[1])
-    cat_param_labels = [v[0] for v in vocab_item_sorted]
-    cat_param_pos = [v[1] for v in vocab_item_sorted]
+class _LabelEncoder:
+    def __init__(self) -> None:
+        self.labels: List[str] = []
 
-    return values_converted, cat_param_labels, cat_param_pos
+    def fit(self, labels: List[str]) -> "_LabelEncoder":
+        self.labels = sorted(set(labels))
+        return self
+
+    def transform(self, labels: List[str]) -> List[int]:
+        return [self.labels.index(label) for label in labels]
+
+    def fit_transform(self, labels: List[str]) -> List[int]:
+        return self.fit(labels).transform(labels)
+
+    def get_labels(self) -> List[str]:
+        return self.labels
+
+    def get_indices(self) -> List[int]:
+        return list(range(len(self.labels)))
 
 
 def _calculate_griddata(
@@ -226,8 +234,18 @@ def _calculate_griddata(
     x_values = []
     y_values = []
     z_values = []
+    x_range_values = []
+    y_range_values = []
     for trial in trials:
-        if x_param not in trial.params or y_param not in trial.params:
+        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])
@@ -285,25 +303,23 @@ def _calculate_griddata(
     cat_param_labels_y = []  # type: List[str]
     cat_param_pos_y = []  # type: List[int]
     if not _is_numerical(trials, x_param):
-        x_values = [str(x) for x in x_values]
-        (
-            x_values,
-            cat_param_labels_x,
-            cat_param_pos_x,
-        ) = _convert_categorical2int(x_values)
+        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):
-        y_values = [str(y) for y in y_values]
-        (
-            y_values,
-            cat_param_labels_y,
-            cat_param_pos_y,
-        ) = _convert_categorical2int(y_values)
+        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_values)
-    x_values_max = max(x_values)
-    y_values_min = min(y_values)
-    y_values_max = max(y_values)
+    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.


Test patch:
diff --git a/tests/visualization_tests/matplotlib_tests/test_contour.py b/tests/visualization_tests/matplotlib_tests/test_contour.py
--- a/tests/visualization_tests/matplotlib_tests/test_contour.py
+++ b/tests/visualization_tests/matplotlib_tests/test_contour.py
@@ -1,6 +1,8 @@
 import itertools
+from typing import Iterable
 from typing import List
 from typing import Optional
+from typing import Tuple
 
 import numpy as np
 import pytest
@@ -14,6 +16,7 @@
 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:
@@ -235,3 +238,53 @@ def test_generate_contour_plot_for_few_observations(params: List[str]) -> None:
     study = prepare_study_with_trials(less_than_two=True)
     figure = plot_contour(study, params)
     assert not figure.has_data()
+
+
+def all_equal(iterable: Iterable) -> bool:
+    """Returns True if all the elements are equal to each other"""
+    return len(set(iterable)) == 1
+
+
+def range_covers(range1: Tuple[float, float], range2: Tuple[float, float]) -> bool:
+    """Returns True if `range1` covers `range2`"""
+    min1, max1 = sorted(range1)
+    min2, max2 = sorted(range2)
+    return min1 <= min2 and max1 >= max2
+
+
+def test_contour_subplots_have_correct_axis_labels_and_ranges() -> None:
+    study = prepare_study_with_trials()
+    params = ["param_a", "param_b", "param_c"]
+    subplots = plot_contour(study, params=params)
+    # `subplots` should look like this:
+    # param_a [[subplot 1, subplot 2, subplot 3],
+    # param_b  [subplot 4, subplot 4, subplot 6],
+    # param_c  [subplot 7, subplot 8, subplot 9]]
+    #           param_a    param_b    param_c
+    #
+    # The folowing block ensures:
+    # - The y-axis label of subplot 1 is "param_a"
+    # - The x-axis label of subplot 7 is "param_a"
+    # - Subplot 1, 2, and 3 have the same y-axis range that covers the search space for `param_a`
+    # - Subplot 1, 4, and 7 have the same x-axis range that covers the search space for `param_a`
+    # - The y-axis label of subplot 4 is "param_b"
+    # - ...
+    # - The y-axis label of subplot 7 is "param_c"
+    # - ...
+    param_ranges = {
+        "param_a": (0.0, 3.0),
+        "param_b": (0.0, 3.0),
+        "param_c": (2.0, 5.0),
+    }
+    for index, (param_name, param_range) in enumerate(param_ranges.items()):
+        minimum, maximum = param_range
+        padding = (maximum - minimum) * AXES_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
+        ylims = [ax.get_ylim() for ax in subplots[index, :]]
+        assert all_equal(ylims)
+        assert all(range_covers(param_range_with_padding, ylim) for ylim in ylims)
+        xlims = [ax.get_xlim() for ax in subplots[:, index]]
+        assert all_equal(xlims)
+        assert all(range_covers(param_range_with_padding, xlim) for xlim in xlims)