Repository: optuna/optuna

Issue:
Make `optuna.visualization.matplotlib.plot_contour` more beautiful
<!-- Please write a clear and concise description of the feature proposal. -->

## Motivation

`plot_contour` with Matplotlib backend looks different from one with Plotly backend. 

## Description

As shown in https://github.com/optuna/optuna/pull/2712#pullrequestreview-679629017, we can improve a plot generated by using `optuna.visualization.matplotlib.plot_contour` by following points:

- [x] smoother contour lines
- [ ] fill blank area



<!-- Please write a detailed description of the new feature. -->

## Alternatives (optional)

<!-- Please write a clear and concise description of any alternative solutions or features you've considered. -->

## Additional context (optional)

<!-- Please add any other context or screenshots about the feature request here. -->



Discussion:
Hi, I've found which algorithm is used in Plotly to interpolate `z` matrix.  This should cover

* blank areas
* smoother plotly-like contour lines (although it's only partially due to algorithm)

I will also tackle `xrange` and `yrange` problem mentioned in review.

I've ported this algorithm from JS to Python already, so please give me few days (hopefully) and i will open PR for this :^)
@xadrianzetx  Wow! Awesome! I'm looking forward to your pull request!

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,13 +1,13 @@
 from collections import defaultdict
 from typing import Callable
 from typing import DefaultDict
+from typing import Dict
 from typing import List
 from typing import Optional
 from typing import Tuple
 from typing import Union
 
 import numpy as np
-from scipy.interpolate import griddata
 
 from optuna._experimental import experimental
 from optuna.logging import get_logger
@@ -30,6 +30,12 @@
 _logger = get_logger(__name__)
 
 
+NEIGHBOR_OFFSETS = [1 + 0j, -1 + 0j, 0 + 1j, 0 - 1j]
+NUM_OPTIMIZATION_ITERATIONS = 100
+FRACTIONAL_DELTA_THRESHOLD = 1e-2
+AXES_PADDING_RATIO = 5e-2
+
+
 @experimental("2.2.0")
 def plot_contour(
     study: Study,
@@ -138,7 +144,7 @@ def _get_contour_plot(
         fig, axs = plt.subplots()
         axs.set_title("Contour Plot")
         cmap = _set_cmap(study, target)
-        contour_point_num = 1000
+        contour_point_num = 100
 
         # Prepare data and draw contour plots.
         if params:
@@ -307,22 +313,36 @@ def _calculate_griddata(
     zi = np.array([])
     if x_param != y_param:
         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(np.log10(x_values_min) - padding_x, 10)
+            x_values_max = np.power(np.log10(x_values_max) + padding_x, 10)
             xi = np.logspace(np.log10(x_values_min), np.log10(x_values_max), contour_point_num)
+            x_array = np.log10(x_values)
         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)
+            x_array = np.array(x_values)
+
         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(np.log10(y_values_min) - padding_y, 10)
+            y_values_max = np.power(np.log10(y_values_max) + padding_y, 10)
             yi = np.logspace(np.log10(y_values_min), np.log10(y_values_max), contour_point_num)
+            y_array = np.log10(y_values)
         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)
+            y_array = np.array(y_values)
 
-        # Interpolate z-axis data on a grid with linear interpolator.
-        # TODO(ytknzw): Implement Plotly-like interpolation algorithm.
-        zi = griddata(
-            np.column_stack((x_values, y_values)),
-            z_values,
-            (xi[None, :], yi[:, None]),
-            method="linear",
-        )
+        # create irregularly spaced map of trial values
+        # and interpolate it with Plotly algorithm
+        zmap = _create_zmap(x_array, y_array, z_values, xi, yi)
+        _interpolate_zmap(zmap, contour_point_num)
+        zi = _create_zmatrix_from_zmap(zmap, contour_point_num)
 
     return (
         xi,
@@ -417,3 +437,186 @@ def _generate_contour_subplot(
         ax.set_yticklabels(y_cat_param_label)
     ax.label_outer()
     return cs
+
+
+def _create_zmap(
+    x_values: np.ndarray,
+    y_values: np.ndarray,
+    z_values: List[float],
+    xi: np.ndarray,
+    yi: np.ndarray,
+) -> Dict[complex, 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 complex numbers, where real part
+    # indicates x-axis index and imaginary part 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 = np.argmin(np.abs(xi - x))
+        yindex = np.argmin(np.abs(yi - y))
+        zmap[complex(xindex, yindex)] = z  # type: ignore
+
+    return zmap
+
+
+def _create_zmatrix_from_zmap(zmap: Dict[complex, float], shape: int) -> np.ndarray:
+
+    # converts hashmap of coordinates to grid
+    zmatrix = np.zeros(shape=(shape, shape))
+    for coord, value in zmap.items():
+        zmatrix[int(coord.imag), int(coord.real)] += value
+
+    return zmatrix
+
+
+def _interpolate_zmap(zmap: Dict[complex, float], contour_plot_num: int) -> None:
+
+    # implements interpolation algorithm used in Plotly
+    # to interpolate heatmaps and contour plots
+    # https://github.com/plotly/plotly.js/blob/master/src/traces/heatmap/interp2d.js#L30
+    # citing their doc:
+    #
+    # > Fill in missing data from a 2D array using an iterative
+    # > poisson equation solver with zero-derivative BC at edges.
+    # > Amazingly, this just amounts to repeatedly averaging all the existing
+    # > nearest neighbors
+    #
+    # our goal here is to assign value to every coordinate that is a part
+    # of 100x100 param grid (so from 0 + 0j up to 100 + 100j) that is not already
+    # occupied by actual trial value
+    max_fractional_delta = 1.0
+    empties = _find_coordinates_where_empty(zmap, contour_plot_num)
+
+    # one pass to fill in a starting value for all the empties
+    _run_iteration(zmap, empties)
+
+    for _ in range(NUM_OPTIMIZATION_ITERATIONS):
+        if max_fractional_delta > FRACTIONAL_DELTA_THRESHOLD:
+            # correct for overshoot and run again
+            max_fractional_delta = 0.5 - 0.25 * min(1, max_fractional_delta * 0.5)
+            max_fractional_delta = _run_iteration(zmap, empties, max_fractional_delta)
+
+        else:
+            break
+
+
+def _find_coordinates_where_empty(
+    zmap: Dict[complex, float], contour_point_num: int
+) -> List[complex]:
+
+    # this function implements missing value discovery and sorting
+    # algorithm used in Plotly to interpolate heatmaps and contour plots
+    # https://github.com/plotly/plotly.js/blob/master/src/traces/heatmap/find_empties.js
+    # it works by repeatedly iterating  over coordinate map in search for patches of
+    # missing values with existing or previously discovered neighbors
+    # when discovered, such patches are added to the iteration queue (list of coordinates)
+    # sorted by number of neighbors, marking iteration order for interpolation algorithm
+    # search ends when all missing patches have been discovered
+    # it's like playing minesweeper in reverse
+
+    iter_queue: List[complex] = []
+    zcopy = zmap.copy()
+    discovered = 0
+    n_missing = (contour_point_num ** 2) - len(zmap)
+    coordinates = [
+        complex(xaxis, yaxis)
+        for yaxis in range(contour_point_num)
+        for xaxis in range(contour_point_num)
+    ]
+
+    while discovered != n_missing:
+        # patch will contain coordinates
+        # of missing values, which happen to have neighbors
+        # at this iteration, and were not previously covered
+        patchmap: Dict[complex, int] = {}
+
+        for coord in coordinates:
+            # we will always iterate over entire grid
+            # surprisingly, this is faster than removing
+            # discovered coordinates from list after each iteration
+            value = zcopy.get(coord, None)
+
+            if value is not None:
+                # trial value or already discovered
+                continue
+
+            n_neighbors = 0
+            for offset in NEIGHBOR_OFFSETS:
+                neighbor = zcopy.get(coord + offset, None)
+                if neighbor is not None:
+                    n_neighbors += 1
+
+            if n_neighbors > 0:
+                # missing values, which still have no neighbors
+                # are left for subsequent iterations
+                patchmap[coord] = n_neighbors
+
+        # newly discovered values will form
+        # neighbors for another missing value
+        # in the next iteration
+        zcopy.update(patchmap)
+
+        # we are sorting patches independently, since neighbor counts for patch `b`
+        # are only valid after patch `a` has been discovered
+        patch = [k for k, _ in sorted(patchmap.items(), key=lambda i: i[1], reverse=True)]
+        iter_queue.extend(patch)
+        discovered += len(patch)
+
+    return iter_queue
+
+
+def _run_iteration(
+    zmap: Dict[complex, float], coordinates: List[complex], overshoot: float = 0.0
+) -> float:
+
+    max_fractional_delta = 0.0
+
+    for coord in coordinates:
+        # we will iterate the grid in order
+        # established by `find_coordinates_where_empty`
+        # this will ensure that we are not trying to interpolate
+        # value with zero neighbors
+        current_val = zmap.get(coord, None)
+        max_neighbor = -np.inf
+        min_neighbor = np.inf
+        sum_neighbors = 0.0
+        n_neighbors = 0
+
+        for offset in NEIGHBOR_OFFSETS:
+            # if we represent current value as a center
+            # of 3x3 matrix, neighbors are all other
+            # values excluding those at diagonals
+            neighbor = zmap.get(coord + offset, None)
+
+            if neighbor is None:
+                # off the edge or not filled in
+                continue
+
+            sum_neighbors += neighbor
+            n_neighbors += 1
+
+            if current_val is not None:
+                max_neighbor = max(max_neighbor, neighbor)
+                min_neighbor = min(min_neighbor, neighbor)
+
+        # fill value is just mean of its neighbors
+        new_val = sum_neighbors / n_neighbors
+
+        if current_val is None:
+            zmap[coord] = new_val
+            max_fractional_delta = 1.0
+
+        else:
+            zmap[coord] = (1 + overshoot) * new_val - overshoot * current_val
+            if max_neighbor > min_neighbor:
+                fractional_delta = abs(new_val - current_val) / (max_neighbor - min_neighbor)
+                max_fractional_delta = max(overshoot, fractional_delta)
+
+    return max_fractional_delta


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,12 +1,92 @@
 from typing import List
 from typing import Optional
 
+import numpy as np
 import pytest
 
 from optuna.study import create_study
 from optuna.testing.visualization import prepare_study_with_trials
 from optuna.trial import Trial
 from optuna.visualization.matplotlib import plot_contour
+from optuna.visualization.matplotlib._contour import _create_zmap
+from optuna.visualization.matplotlib._contour import _create_zmatrix_from_zmap
+from optuna.visualization.matplotlib._contour import _find_coordinates_where_empty
+from optuna.visualization.matplotlib._contour import _interpolate_zmap
+from optuna.visualization.matplotlib._contour import _run_iteration
+
+
+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, y_values, 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 = int(coord.real)
+        yidx = int(coord.imag)
+        assert xidx == yidx
+        assert z_values[xidx] == value
+
+
+def test_create_zmatrix_from_zmap() -> None:
+
+    contour_point_num = 2
+    zmap = {0 + 0j: 1.0, 1 + 0j: 2.0, 0 + 1j: 3.0, 1 + 1j: 4.0}
+    expected = np.array([[1.0, 2.0], [3.0, 4.0]])
+    zmatrix = _create_zmatrix_from_zmap(zmap, contour_point_num)
+
+    assert zmatrix.shape == (contour_point_num, contour_point_num)
+    assert np.array_equal(zmatrix, expected)
+
+
+def test_find_coordinates_where_empty() -> None:
+
+    contour_point_num = 2
+    zmap = {0 + 0j: 1.0, 1 + 1j: 4.0}
+    n_missing = (contour_point_num ** 2) - len(zmap)
+    empties = _find_coordinates_where_empty(zmap, contour_point_num)
+
+    # test if all missing are found
+    assert len(empties) == n_missing
+
+    # test if iter queue follows C style iteration
+    assert empties[0] == 1 + 0j
+    assert empties[1] == 0 + 1j
+
+
+def test_interpolation_first_iteration() -> None:
+
+    zmap = {0 + 0j: 1.0, 1 + 1j: 4.0}
+    empties = [1 + 0j, 0 + 1j]
+    initial_zmap_len = len(zmap)
+
+    max_fractional_change = _run_iteration(zmap, empties)
+
+    # test loss after first iter
+    assert max_fractional_change == 1.0
+
+    # test if initial pass filled all values
+    assert len(zmap) == initial_zmap_len + len(empties)
+
+
+def test_interpolate_zmap() -> None:
+
+    contour_point_num = 2
+    zmap = {0 + 0j: 1.0, 1 + 1j: 4.0}
+    interpolated = {1 + 0j: 2.5, 1 + 1j: 4.0, 0 + 0j: 1.0, 0 + 1j: 2.5}
+
+    _interpolate_zmap(zmap, contour_point_num)
+
+    for coord, value in zmap.items():
+        expected_at_coord = interpolated.get(coord)
+        assert value == expected_at_coord
 
 
 def test_target_is_none_and_study_is_multi_obj() -> None: