From 52465d97789350c4baca91aa459c6fd1fd07a7d5 Mon Sep 17 00:00:00 2001 From: nishiji <34439525+ryota717@users.noreply.github.com> Date: Tue, 26 Sep 2023 17:18:01 +0900 Subject: [PATCH] Support constraints plot_rank (#4899) * refactor: Passing color as normalized rgb ndarray. * feat: Support optimization with constraints in plot_rank. * fix: Make color ditribution same between tick and points. * refactor: Passing color with 0~255 RGB not 0~1. * test: Test color converter by assert_array_equal. * docs: Add constraints to rank_plot example. * fix: Delete debug message in test. Co-authored-by: Shinichi Hemmi <50256998+Alnusjaponica@users.noreply.github.com> --------- Co-authored-by: Shinichi Hemmi <50256998+Alnusjaponica@users.noreply.github.com> --- optuna/visualization/_rank.py | 51 +++++++--- optuna/visualization/matplotlib/_rank.py | 17 +++- tests/visualization_tests/test_rank.py | 118 +++++++++++++++++++---- 3 files changed, 148 insertions(+), 38 deletions(-) diff --git a/optuna/visualization/_rank.py b/optuna/visualization/_rank.py index 8969a64297..6198538ba8 100644 --- a/optuna/visualization/_rank.py +++ b/optuna/visualization/_rank.py @@ -10,6 +10,7 @@ from optuna._experimental import experimental_func from optuna.logging import get_logger +from optuna.samplers._base import _CONSTRAINTS_KEY from optuna.study import Study from optuna.trial import FrozenTrial from optuna.trial import TrialState @@ -45,7 +46,7 @@ class _RankSubplotInfo(NamedTuple): ys: list[Any] trials: list[FrozenTrial] zs: np.ndarray - color_idxs: np.ndarray + colors: np.ndarray class _RankPlotInfo(NamedTuple): @@ -53,7 +54,7 @@ class _RankPlotInfo(NamedTuple): sub_plot_infos: list[list[_RankSubplotInfo]] target_name: str zs: np.ndarray - color_idxs: np.ndarray + colors: np.ndarray has_custom_target: bool @@ -81,10 +82,18 @@ def plot_rank( def objective(trial): x = trial.suggest_float("x", -100, 100) y = trial.suggest_categorical("y", [-1, 0, 1]) + + c0 = 400 - (x + y)**2 + trial.set_user_attr("constraint", [c0]) + return x ** 2 + y - sampler = optuna.samplers.TPESampler(seed=10) + def constraints(trial): + return trial.user_attrs["constraint"] + + + sampler = optuna.samplers.TPESampler(seed=10, constraints_func=constraints) study = optuna.create_study(sampler=sampler) study.optimize(objective, n_trials=30) @@ -160,16 +169,18 @@ def target(trial: FrozenTrial) -> float: target_values = np.array([target(trial) for trial in trials]) raw_ranks = _get_order_with_same_order_averaging(target_values) color_idxs = raw_ranks / (len(trials) - 1) if len(trials) >= 2 else np.array([0.5]) + colors = _convert_color_idxs_to_scaled_rgb_colors(color_idxs) + sub_plot_infos: list[list[_RankSubplotInfo]] if len(params) == 2: x_param = params[0] y_param = params[1] - sub_plot_info = _get_rank_subplot_info(trials, target_values, color_idxs, x_param, y_param) + sub_plot_info = _get_rank_subplot_info(trials, target_values, colors, x_param, y_param) sub_plot_infos = [[sub_plot_info]] else: sub_plot_infos = [ [ - _get_rank_subplot_info(trials, target_values, color_idxs, x_param, y_param) + _get_rank_subplot_info(trials, target_values, colors, x_param, y_param) for x_param in params ] for y_param in params @@ -180,7 +191,7 @@ def target(trial: FrozenTrial) -> float: sub_plot_infos=sub_plot_infos, target_name=target_name, zs=target_values, - color_idxs=color_idxs, + colors=colors, has_custom_target=has_custom_target, ) @@ -188,13 +199,21 @@ def target(trial: FrozenTrial) -> float: def _get_rank_subplot_info( trials: list[FrozenTrial], target_values: np.ndarray, - color_idxs: np.ndarray, + colors: np.ndarray, x_param: str, y_param: str, ) -> _RankSubplotInfo: xaxis = _get_axis_info(trials, x_param) yaxis = _get_axis_info(trials, y_param) + infeasible_trial_ids = [] + for i in range(len(trials)): + constraints = trials[i].system_attrs.get(_CONSTRAINTS_KEY) + if constraints is not None and any([x > 0.0 for x in constraints]): + infeasible_trial_ids.append(i) + + colors[infeasible_trial_ids] = plotly.colors.hex_to_rgb("#cccccc") + filtered_ids = [ i for i in range(len(trials)) @@ -204,7 +223,7 @@ def _get_rank_subplot_info( xs = [trial.params[x_param] for trial in filtered_trials] ys = [trial.params[y_param] for trial in filtered_trials] zs = target_values[filtered_ids] - color_idxs = color_idxs[filtered_ids] + colors = colors[filtered_ids] return _RankSubplotInfo( xaxis=xaxis, yaxis=yaxis, @@ -212,7 +231,7 @@ def _get_rank_subplot_info( ys=ys, trials=filtered_trials, zs=np.array(zs), - color_idxs=np.array(color_idxs), + colors=colors, ) @@ -269,10 +288,6 @@ def _get_axis_info(trials: list[FrozenTrial], param_name: str) -> _AxisInfo: def _get_rank_subplot( info: _RankSubplotInfo, target_name: str, print_raw_objectives: bool ) -> "Scatter": - colormap = "RdYlBu_r" - # sample_colorscale requires plotly >= 5.0.0. - colors = plotly.colors.sample_colorscale(colormap, info.color_idxs) - def get_hover_text(trial: FrozenTrial, target_value: float) -> str: lines = [f"Trial #{trial.number}"] lines += [f"{k}: {v}" for k, v in trial.params.items()] @@ -285,7 +300,7 @@ def get_hover_text(trial: FrozenTrial, target_value: float) -> str: x=info.xs, y=info.ys, marker={ - "color": colors, + "color": list(map(plotly.colors.label_rgb, info.colors)), "line": {"width": 0.5, "color": "Grey"}, }, mode="markers", @@ -403,3 +418,11 @@ def _get_rank_plot( ) figure.add_trace(colorbar_trace) return figure + + +def _convert_color_idxs_to_scaled_rgb_colors(color_idxs: np.ndarray) -> np.ndarray: + colormap = "RdYlBu_r" + # sample_colorscale requires plotly >= 5.0.0. + labeled_colors = plotly.colors.sample_colorscale(colormap, color_idxs) + scaled_rgb_colors = np.array([plotly.colors.unlabel_rgb(cl) for cl in labeled_colors]) + return scaled_rgb_colors diff --git a/optuna/visualization/matplotlib/_rank.py b/optuna/visualization/matplotlib/_rank.py index ca379a3ae5..1ad6faedb9 100644 --- a/optuna/visualization/matplotlib/_rank.py +++ b/optuna/visualization/matplotlib/_rank.py @@ -54,10 +54,18 @@ def plot_rank( def objective(trial): x = trial.suggest_float("x", -100, 100) y = trial.suggest_categorical("y", [-1, 0, 1]) + + c0 = 400 - (x + y)**2 + trial.set_user_attr("constraint", [c0]) + return x ** 2 + y - sampler = optuna.samplers.TPESampler(seed=10) + def constraints(trial): + return trial.user_attrs["constraint"] + + + sampler = optuna.samplers.TPESampler(seed=10, constraints_func=constraints) study = optuna.create_study(sampler=sampler) study.optimize(objective, n_trials=30) @@ -126,7 +134,8 @@ def _get_rank_plot( tick_info = _get_tick_info(info.zs) - cbar = fig.colorbar(pc, ax=axs, ticks=tick_info.coloridxs, cmap=plt.get_cmap("RdYlBu_r")) + pc.set_cmap(plt.get_cmap("RdYlBu_r")) + cbar = fig.colorbar(pc, ax=axs, ticks=tick_info.coloridxs) cbar.ax.set_yticklabels(tick_info.text) cbar.outline.set_edgecolor("gray") return axs @@ -151,6 +160,4 @@ def _add_rank_subplot( if info.yaxis.is_log: ax.set_yscale("log") - return ax.scatter( - x=info.xs, y=info.ys, c=info.color_idxs, cmap=plt.get_cmap("RdYlBu_r"), edgecolors="grey" - ) + return ax.scatter(x=info.xs, y=info.ys, c=info.colors / 255, edgecolors="grey") diff --git a/tests/visualization_tests/test_rank.py b/tests/visualization_tests/test_rank.py index c403fa193b..9b1d54ea1b 100644 --- a/tests/visualization_tests/test_rank.py +++ b/tests/visualization_tests/test_rank.py @@ -11,6 +11,7 @@ from optuna.distributions import BaseDistribution from optuna.distributions import CategoricalDistribution from optuna.distributions import FloatDistribution +from optuna.samplers._base import _CONSTRAINTS_KEY from optuna.study import create_study from optuna.study import Study from optuna.testing.objectives import fail_objective @@ -19,6 +20,8 @@ from optuna.visualization import plot_rank as plotly_plot_rank from optuna.visualization._plotly_imports import go from optuna.visualization._rank import _AxisInfo +from optuna.visualization._rank import _convert_color_idxs_to_scaled_rgb_colors +from optuna.visualization._rank import _get_axis_info from optuna.visualization._rank import _get_order_with_same_order_averaging from optuna.visualization._rank import _get_rank_info from optuna.visualization._rank import _RankPlotInfo @@ -78,6 +81,31 @@ def _create_study_with_log_scale_and_str_category_3d() -> Study: return study +def _create_study_with_constraints() -> Study: + study = create_study() + distributions: dict[str, BaseDistribution] = { + "param_a": FloatDistribution(0.1, 0.2), + "param_b": FloatDistribution(0.3, 0.4), + } + study.add_trial( + create_trial( + value=0.0, + params={"param_a": 0.11, "param_b": 0.31}, + distributions=distributions, + system_attrs={_CONSTRAINTS_KEY: [-0.1, 0.0]}, + ) + ) + study.add_trial( + create_trial( + value=1.0, + params={"param_a": 0.19, "param_b": 0.34}, + distributions=distributions, + system_attrs={_CONSTRAINTS_KEY: [0.1, 0.0]}, + ) + ) + return study + + def _create_study_mixture_category_types() -> Study: study = create_study() distributions: dict[str, BaseDistribution] = { @@ -141,6 +169,7 @@ def _get_nested_list_shape(nested_list: list[list[Any]]) -> tuple[int, int]: [_create_study_with_log_scale_and_str_category_2d, None], [_create_study_with_log_scale_and_str_category_3d, None], [_create_study_mixture_category_types, None], + [_create_study_with_constraints, None], ], ) def test_plot_rank( @@ -226,13 +255,13 @@ def test_get_rank_info_2_params() -> None: ys=[2.0, 1.0], trials=[study.trials[0], study.trials[2]], zs=np.array([0.0, 1.0]), - color_idxs=np.array([0.0, 0.5]), + colors=_convert_color_idxs_to_scaled_rgb_colors(np.array([0.0, 0.5])), ) ] ], target_name="Objective Value", zs=np.array([0.0, 2.0, 1.0]), - color_idxs=np.array([0.0, 1.0, 0.5]), + colors=_convert_color_idxs_to_scaled_rgb_colors(np.array([0.0, 1.0, 0.5])), has_custom_target=False, ), ) @@ -326,13 +355,15 @@ def test_generate_rank_plot_for_no_plots(params: list[str]) -> None: ys=[], trials=[], zs=np.array([]), - color_idxs=np.array([]), + colors=_convert_color_idxs_to_scaled_rgb_colors(np.array([])).reshape( + -1, 3 + ), ) ] ], target_name="Objective Value", zs=np.array([0.0, 2.0]), - color_idxs=np.array([0.0, 1.0]), + colors=_convert_color_idxs_to_scaled_rgb_colors(np.array([0.0, 1.0])), has_custom_target=False, ), ) @@ -393,13 +424,13 @@ def test_generate_rank_plot_for_few_observations(params: list[str]) -> None: ys=[study.get_trials()[0].params[params[1]]], trials=[study.get_trials()[0]], zs=np.array([0.0]), - color_idxs=np.array([0.0]), + colors=_convert_color_idxs_to_scaled_rgb_colors(np.array([0.0])), ) ] ], target_name="Objective Value", zs=np.array([0.0, 2.0]), - color_idxs=np.array([0.0, 1.0]), + colors=_convert_color_idxs_to_scaled_rgb_colors(np.array([0.0, 1.0])), has_custom_target=False, ), ) @@ -432,13 +463,13 @@ def test_get_rank_info_log_scale_and_str_category_2_params() -> None: ys=["101", "100"], trials=[study.trials[0], study.trials[1]], zs=np.array([0.0, 1.0]), - color_idxs=np.array([0.0, 1.0]), + colors=_convert_color_idxs_to_scaled_rgb_colors(np.array([0.0, 1.0])), ) ] ], target_name="Objective Value", zs=np.array([0.0, 1.0]), - color_idxs=np.array([0.0, 1.0]), + colors=_convert_color_idxs_to_scaled_rgb_colors(np.array([0.0, 1.0])), has_custom_target=False, ), ) @@ -461,7 +492,7 @@ def test_get_rank_info_log_scale_and_str_category_more_than_2_params() -> None: param_values = {"param_a": [1e-6, 1e-5], "param_b": ["101", "100"], "param_c": ["one", "two"]} zs = np.array([0.0, 1.0]) - color_idxs = np.array([0.0, 1.0]) + colors = _convert_color_idxs_to_scaled_rgb_colors(np.array([0.0, 1.0])) def _check_axis(axis: _AxisInfo, name: str) -> None: assert axis.name == name @@ -481,11 +512,11 @@ def _check_axis(axis: _AxisInfo, name: str) -> None: assert info.sub_plot_infos[yi][xi].ys == param_values[y_param] assert info.sub_plot_infos[yi][xi].trials == study.trials assert np.all(info.sub_plot_infos[yi][xi].zs == zs) - assert np.all(info.sub_plot_infos[yi][xi].color_idxs == color_idxs) + assert np.all(info.sub_plot_infos[yi][xi].colors == colors) info.target_name == "Objective Value" assert np.all(info.zs == zs) - assert np.all(info.color_idxs == color_idxs) + assert np.all(info.colors == colors) assert not info.has_custom_target @@ -515,13 +546,13 @@ def test_get_rank_info_mixture_category_types() -> None: ys=[101, 102.0], trials=study.trials, zs=np.array([0.0, 0.5]), - color_idxs=np.array([0.0, 1.0]), + colors=_convert_color_idxs_to_scaled_rgb_colors(np.array([0.0, 1.0])), ) ] ], target_name="Objective Value", zs=np.array([0.0, 0.5]), - color_idxs=np.array([0.0, 1.0]), + colors=_convert_color_idxs_to_scaled_rgb_colors(np.array([0.0, 1.0])), has_custom_target=False, ), ) @@ -534,7 +565,11 @@ def test_get_rank_info_nonfinite(value: float) -> None: study, params=["param_b", "param_d"], target=None, target_name="Objective Value" ) - color_idxs = np.array([0.0, 1.0, 0.5]) if value == float("-inf") else np.array([1.0, 0.5, 0.0]) + colors = ( + _convert_color_idxs_to_scaled_rgb_colors(np.array([0.0, 1.0, 0.5])) + if value == float("-inf") + else _convert_color_idxs_to_scaled_rgb_colors(np.array([1.0, 0.5, 0.0])) + ) assert _named_tuple_equal( info, _RankPlotInfo( @@ -558,13 +593,13 @@ def test_get_rank_info_nonfinite(value: float) -> None: ys=[4.0, 4.0, 2.0], trials=study.trials, zs=np.array([value, 2.0, 1.0]), - color_idxs=color_idxs, + colors=colors, ) ] ], target_name="Objective Value", zs=np.array([value, 2.0, 1.0]), - color_idxs=color_idxs, + colors=colors, has_custom_target=False, ), ) @@ -580,7 +615,11 @@ def test_get_rank_info_nonfinite_multiobjective(objective: int, value: float) -> target=lambda t: t.values[objective], target_name="Target Name", ) - color_idxs = np.array([0.0, 1.0, 0.5]) if value == float("-inf") else np.array([1.0, 0.5, 0.0]) + colors = ( + _convert_color_idxs_to_scaled_rgb_colors(np.array([0.0, 1.0, 0.5])) + if value == float("-inf") + else _convert_color_idxs_to_scaled_rgb_colors(np.array([1.0, 0.5, 0.0])) + ) assert _named_tuple_equal( info, _RankPlotInfo( @@ -604,18 +643,59 @@ def test_get_rank_info_nonfinite_multiobjective(objective: int, value: float) -> ys=[4.0, 4.0, 2.0], trials=study.trials, zs=np.array([value, 2.0, 1.0]), - color_idxs=color_idxs, + colors=colors, ) ] ], target_name="Target Name", zs=np.array([value, 2.0, 1.0]), - color_idxs=color_idxs, + colors=colors, has_custom_target=True, ), ) +def test_generate_rank_info_with_constraints() -> None: + study = _create_study_with_constraints() + info = _get_rank_info(study, params=None, target=None, target_name="Objective Value") + expected_color = _convert_color_idxs_to_scaled_rgb_colors(np.array([0.0, 1.0])) + expected_color[1] = [204, 204, 204] + + assert _named_tuple_equal( + info, + _RankPlotInfo( + params=["param_a", "param_b"], + sub_plot_infos=[ + [ + _RankSubplotInfo( + xaxis=_get_axis_info(study.trials, "param_a"), + yaxis=_get_axis_info(study.trials, "param_b"), + xs=[0.11, 0.19], + ys=[0.31, 0.34], + trials=study.trials, + zs=np.array([0.0, 1.0]), + colors=expected_color, + ) + ] + ], + target_name="Objective Value", + zs=np.array([0.0, 1.0]), + colors=expected_color, + has_custom_target=False, + ), + ) + + def test_get_order_with_same_order_averaging() -> None: x = np.array([6.0, 2.0, 3.0, 1.0, 4.5, 4.5, 8.0, 8.0, 0.0, 8.0]) assert np.all(x == _get_order_with_same_order_averaging(x)) + + +def test_convert_color_idxs_to_scaled_rgb_colors() -> None: + x1 = np.array([0.1, 0.2]) + result1 = _convert_color_idxs_to_scaled_rgb_colors(x1) + np.testing.assert_array_equal(result1, [[69, 117, 180], [116, 173, 209]]) + + x2 = np.array([]) + result2 = _convert_color_idxs_to_scaled_rgb_colors(x2) + np.testing.assert_array_equal(result2, [])