From 25edf5ccbe73d1be74cd07478df529660159c446 Mon Sep 17 00:00:00 2001 From: Naoto Mizuno Date: Mon, 25 Sep 2023 19:37:41 +0900 Subject: [PATCH 1/3] Add _split_trials instead of _get_observation_pairs and _split_observation_pairs --- optuna/samplers/_tpe/sampler.py | 286 +++++++--------- .../tpe_tests/test_multi_objective_sampler.py | 107 ++---- .../samplers_tests/tpe_tests/test_sampler.py | 324 +++++++----------- 3 files changed, 271 insertions(+), 446 deletions(-) diff --git a/optuna/samplers/_tpe/sampler.py b/optuna/samplers/_tpe/sampler.py index a725176c8d..43e478e1e9 100644 --- a/optuna/samplers/_tpe/sampler.py +++ b/optuna/samplers/_tpe/sampler.py @@ -3,11 +3,10 @@ import math from typing import Any from typing import Callable +from typing import cast from typing import Dict -from typing import List from typing import Optional from typing import Sequence -from typing import Tuple from typing import Union import warnings @@ -444,19 +443,15 @@ def _sample( use_cache = not self._constant_liar trials = study._get_trials(deepcopy=False, states=states, use_cache=use_cache) - scores, violations = _get_observation_pairs( + # We divide data into below and above. + n = sum(trial.state != TrialState.RUNNING for trial in trials) # Ignore running trials. + below_trials, above_trials = _split_trials( study, trials, + self._gamma(n), self._constraints_func is not None, ) - n = sum(s < float("inf") for s, v in scores) # Ignore running trials. - - # We divide data into below and above. - indices_below, indices_above = _split_observation_pairs(scores, self._gamma(n), violations) - - below_trials = np.asarray(trials, dtype=object)[indices_below].tolist() - above_trials = np.asarray(trials, dtype=object)[indices_above].tolist() below = self._get_internal_repr(below_trials, search_space) above = self._get_internal_repr(above_trials, search_space) @@ -584,165 +579,146 @@ def _calculate_nondomination_rank(loss_vals: np.ndarray) -> np.ndarray: return ranks -def _get_observation_pairs( +def _split_trials( study: Study, trials: list[FrozenTrial], - constraints_enabled: bool = False, -) -> tuple[list[tuple[float, list[float]]], list[float] | None]: - """Get observation pairs from the study. - - This function collects observation pairs from the complete or pruned trials of the study. - In addition, if ``constant_liar`` is :obj:`True`, the running trials are considered. - The values for trials that don't contain the parameter in the ``param_names`` are skipped. - - An observation pair fundamentally consists of a parameter value and an objective value. - However, due to the pruning mechanism of Optuna, final objective values are not always - available. Therefore, this function uses intermediate values in addition to the final - ones, and reports the value with its step count as ``(-step, value)``. - Consequently, the structure of the observation pair is as follows: - ``(param_value, (-step, value))``. - - The second element of an observation pair is used to rank observations in - ``_split_observation_pairs`` method (i.e., observations are sorted lexicographically by - ``(-step, value)``). - - When ``constraints_enabled`` is :obj:`True`, 1-dimensional violation values are returned - as the third element (:obj:`None` otherwise). Each value is a float of 0 or greater and a - trial is feasible if and only if its violation score is 0. - """ - - signs = [] - for d in study.directions: - if d == StudyDirection.MINIMIZE: - signs.append(1) - else: - signs.append(-1) + n_below: int, + constraints_enabled: bool, +) -> tuple[list[FrozenTrial], list[FrozenTrial]]: + complete_trials = [] + pruned_trials = [] + infeasible_trials = [] - scores = [] - violations: Optional[List[float]] = [] if constraints_enabled else None for trial in trials: - # We extract score from the trial. - if trial.state is TrialState.COMPLETE: - assert trial.values is not None - score = (-float("inf"), [sign * v for sign, v in zip(signs, trial.values)]) - elif trial.state is TrialState.PRUNED: - assert not study._is_multi_objective() - - if len(trial.intermediate_values) > 0: - step, intermediate_value = max(trial.intermediate_values.items()) - if math.isnan(intermediate_value): - score = (-step, [float("inf")]) - else: - score = (-step, [signs[0] * intermediate_value]) - else: - score = (1, [0.0]) - elif trial.state is TrialState.RUNNING: - assert not study._is_multi_objective() - score = (float("inf"), [signs[0] * float("inf")]) + if constraints_enabled and _get_infeasible_trial_score(trial) > 0: + infeasible_trials.append(trial) + elif trial.state == TrialState.COMPLETE: + complete_trials.append(trial) + elif trial.state == TrialState.PRUNED: + pruned_trials.append(trial) else: - assert False - scores.append(score) - - if constraints_enabled: - assert violations is not None - if trial.state != TrialState.RUNNING: - constraint = trial.system_attrs.get(_CONSTRAINTS_KEY) - if constraint is None: - warnings.warn( - f"Trial {trial.number} does not have constraint values." - " It will be treated as a lower priority than other trials." - ) - violation = float("inf") - else: - # Violation values of infeasible dimensions are summed up. - violation = sum(v for v in constraint if v > 0) - violations.append(violation) - else: - violations.append(float("inf")) + assert trial.state == TrialState.RUNNING + + # We divide data into below and above. + if len(complete_trials) >= n_below: + below_trials = _split_complete_trials(complete_trials, study, n_below) + elif len(complete_trials) + len(pruned_trials) >= n_below: + below_pruned_trials = _split_pruned_trials( + pruned_trials, study, n_below - len(complete_trials) + ) + below_trials = complete_trials + below_pruned_trials + else: + below_infeasible_trials = _split_infeasible_trials( + infeasible_trials, n_below - len(complete_trials) - len(pruned_trials) + ) + below_trials = complete_trials + pruned_trials + below_infeasible_trials + + below_trials.sort(key=lambda trial: trial.number) + below_trial_numbers = set(trial.number for trial in below_trials) + above_trials = [trial for trial in trials if trial.number not in below_trial_numbers] + + return below_trials, above_trials - return scores, violations + +def _split_complete_trials( + trials: Sequence[FrozenTrial], study: Study, n_below: int +) -> list[FrozenTrial]: + if len(study.directions) <= 1: + return _split_complete_trials_single_objective(trials, study, n_below) + else: + return _split_complete_trials_multi_objective(trials, study, n_below) -def _split_observation_pairs( - loss_vals: List[Tuple[float, List[float]]], +def _split_complete_trials_single_objective( + trials: Sequence[FrozenTrial], + study: Study, n_below: int, - violations: Optional[List[float]], -) -> Tuple[np.ndarray, np.ndarray]: - # When constrains is not None, trials are split into below and above - # according to the following rules. - # 1. Feasible trials are better than infeasible trials. - # 2. Infeasible trials are sorted by sum of how much they violate each constraint. - # 3. Feasible trials are sorted by loss_vals. - if violations is not None: - violation_1d = np.array(violations, dtype=float) - idx = violation_1d.argsort(kind="stable") - if n_below >= len(idx) or violation_1d[idx[n_below]] > 0: - # Below is filled by all feasible trials and trials with smaller violation values. - indices_below = idx[:n_below] - indices_above = idx[n_below:] +) -> list[FrozenTrial]: + if study.direction == StudyDirection.MINIMIZE: + return sorted(trials, key=lambda trial: cast(float, trial.value))[:n_below] + else: + return sorted(trials, key=lambda trial: cast(float, trial.value), reverse=True)[:n_below] + + +def _split_complete_trials_multi_objective( + trials: Sequence[FrozenTrial], + study: Study, + n_below: int, +) -> list[FrozenTrial]: + if n_below == 0: + return [] + + lvals = np.asarray([trial.values for trial in trials]) + for i, direction in enumerate(study.directions): + if direction == StudyDirection.MAXIMIZE: + lvals[:, i] *= -1 + + # Solving HSSP for variables number of times is a waste of time. + nondomination_ranks = _calculate_nondomination_rank(lvals) + assert 0 <= n_below <= len(lvals) + + indices = np.array(range(len(lvals))) + indices_below = np.empty(n_below, dtype=int) + + # Nondomination rank-based selection + i = 0 + last_idx = 0 + while last_idx < n_below and last_idx + sum(nondomination_ranks == i) <= n_below: + length = indices[nondomination_ranks == i].shape[0] + indices_below[last_idx : last_idx + length] = indices[nondomination_ranks == i] + last_idx += length + i += 1 + + # Hypervolume subset selection problem (HSSP)-based selection + subset_size = n_below - last_idx + if subset_size > 0: + rank_i_lvals = lvals[nondomination_ranks == i] + rank_i_indices = indices[nondomination_ranks == i] + worst_point = np.max(rank_i_lvals, axis=0) + reference_point = np.maximum(1.1 * worst_point, 0.9 * worst_point) + reference_point[reference_point == 0] = EPS + selected_indices = _solve_hssp(rank_i_lvals, rank_i_indices, subset_size, reference_point) + indices_below[last_idx:] = selected_indices + + return [trials[index] for index in indices_below] + + +def _get_pruned_trial_score(trial: FrozenTrial, study: Study) -> tuple[float, float]: + if len(trial.intermediate_values) > 0: + step, intermediate_value = max(trial.intermediate_values.items()) + if math.isnan(intermediate_value): + return -step, float("inf") + elif study.direction == StudyDirection.MINIMIZE: + return -step, intermediate_value else: - # All trials in below are feasible. - # Feasible trials with smaller loss_vals are selected. - (feasible_idx,) = (violation_1d == 0).nonzero() - (infeasible_idx,) = (violation_1d > 0).nonzero() - assert len(feasible_idx) >= n_below - feasible_below, feasible_above = _split_observation_pairs( - [loss_vals[i] for i in feasible_idx], n_below, None - ) - indices_below = feasible_idx[feasible_below] - indices_above = np.concatenate([feasible_idx[feasible_above], infeasible_idx]) - # `np.sort` is used to keep chronological order. - return np.sort(indices_below), np.sort(indices_above) - - n_objectives = 1 - if len(loss_vals) > 0: - n_objectives = len(loss_vals[0][1]) - - if n_objectives <= 1: - loss_values = np.asarray( - [(s, v[0]) for s, v in loss_vals], dtype=[("step", float), ("score", float)] - ) + return -step, -intermediate_value + else: + return 1, 0.0 + - index_loss_ascending = np.argsort(loss_values, kind="stable") - # `np.sort` is used to keep chronological order. - indices_below = np.sort(index_loss_ascending[:n_below]) - indices_above = np.sort(index_loss_ascending[n_below:]) +def _split_pruned_trials( + trials: Sequence[FrozenTrial], + study: Study, + n_below: int, +) -> list[FrozenTrial]: + return sorted(trials, key=lambda trial: _get_pruned_trial_score(trial, study))[:n_below] + + +def _get_infeasible_trial_score(trial: FrozenTrial) -> float: + constraint = trial.system_attrs.get(_CONSTRAINTS_KEY) + if constraint is None: + warnings.warn( + f"Trial {trial.number} does not have constraint values." + " It will be treated as a lower priority than other trials." + ) + return float("inf") else: - # Multi-objective TPE does not support pruning, so it ignores the ``step``. - lvals = np.asarray([v for _, v in loss_vals]) - - # Solving HSSP for variables number of times is a waste of time. - nondomination_ranks = _calculate_nondomination_rank(lvals) - assert 0 <= n_below <= len(lvals) - - indices = np.array(range(len(lvals))) - indices_below = np.empty(n_below, dtype=int) - - # Nondomination rank-based selection - i = 0 - last_idx = 0 - while last_idx < n_below and last_idx + sum(nondomination_ranks == i) <= n_below: - length = indices[nondomination_ranks == i].shape[0] - indices_below[last_idx : last_idx + length] = indices[nondomination_ranks == i] - last_idx += length - i += 1 - - # Hypervolume subset selection problem (HSSP)-based selection - subset_size = n_below - last_idx - if subset_size > 0: - rank_i_lvals = lvals[nondomination_ranks == i] - rank_i_indices = indices[nondomination_ranks == i] - worst_point = np.max(rank_i_lvals, axis=0) - reference_point = np.maximum(1.1 * worst_point, 0.9 * worst_point) - reference_point[reference_point == 0] = EPS - selected_indices = _solve_hssp( - rank_i_lvals, rank_i_indices, subset_size, reference_point - ) - indices_below[last_idx:] = selected_indices + # Violation values of infeasible dimensions are summed up. + return sum(v for v in constraint if v > 0) - indices_above = np.setdiff1d(indices, indices_below) - return indices_below, indices_above +def _split_infeasible_trials(trials: Sequence[FrozenTrial], n_below: int) -> list[FrozenTrial]: + return sorted(trials, key=_get_infeasible_trial_score)[:n_below] def _calculate_weights_below_for_multi_objective( diff --git a/tests/samplers_tests/tpe_tests/test_multi_objective_sampler.py b/tests/samplers_tests/tpe_tests/test_multi_objective_sampler.py index 12db0d8aca..0d3ea7679a 100644 --- a/tests/samplers_tests/tpe_tests/test_multi_objective_sampler.py +++ b/tests/samplers_tests/tpe_tests/test_multi_objective_sampler.py @@ -3,7 +3,6 @@ from typing import Dict from typing import List from typing import Optional -from typing import Tuple from typing import Union from unittest.mock import patch from unittest.mock import PropertyMock @@ -305,93 +304,37 @@ def test_multi_objective_sample_independent_ignored_states() -> None: assert len(set(suggestions)) == 1 -@pytest.mark.parametrize("int_value", [-5, 5, 0]) -@pytest.mark.parametrize( - "categorical_value", [1, 0.0, "A", None, True, float("inf"), float("nan")] -) -@pytest.mark.parametrize("objective_value", [-5.0, 5.0, 0.0, -float("inf"), float("inf")]) -@pytest.mark.parametrize("multivariate", [True, False]) -@pytest.mark.parametrize("constant_liar", [True, False]) -@pytest.mark.filterwarnings("ignore::optuna.exceptions.ExperimentalWarning") -def test_multi_objective_get_observation_pairs( - int_value: int, - categorical_value: optuna.distributions.CategoricalChoiceType, - objective_value: float, - multivariate: bool, - constant_liar: bool, -) -> None: - def objective(trial: optuna.trial.Trial) -> Tuple[float, float]: - trial.suggest_int("x", int_value, int_value) - trial.suggest_categorical("y", [categorical_value]) - return objective_value, objective_value - - sampler = TPESampler(seed=0, multivariate=multivariate, constant_liar=constant_liar) - study = optuna.create_study(directions=["minimize", "maximize"], sampler=sampler) - study.optimize(objective, n_trials=2) - study.add_trial( - optuna.create_trial( - state=optuna.trial.TrialState.RUNNING, - params={"x": int_value, "y": categorical_value}, - distributions={ - "x": optuna.distributions.IntDistribution(int_value, int_value), - "y": optuna.distributions.CategoricalDistribution([categorical_value]), - }, +@pytest.mark.parametrize("direction0", ["minimize", "maximize"]) +@pytest.mark.parametrize("direction1", ["minimize", "maximize"]) +def test_split_complete_trials_multi_objective(direction0: str, direction1: str) -> None: + study = optuna.create_study(directions=(direction0, direction1)) + + for values in ([-2.0, -1.0], [3.0, 3.0], [0.0, 1.0], [-1.0, 0.0]): + value0, value1 = values + if direction0 == "maximize": + value0 = -value0 + if direction1 == "maximize": + value1 = -value1 + study.add_trial( + optuna.create_trial( + state=optuna.trial.TrialState.COMPLETE, + values=(value0, value1), + params={"x": 0}, + distributions={"x": optuna.distributions.FloatDistribution(-1.0, 1.0)}, + ) ) - ) - - states = [optuna.trial.TrialState.COMPLETE, optuna.trial.TrialState.PRUNED] - assert _tpe.sampler._get_observation_pairs(study, study.get_trials(states=states)) == ( - [(-float("inf"), [objective_value, -objective_value]) for _ in range(2)], - None, - ) - -@pytest.mark.parametrize("constraint_value", [-2, 2]) -def test_multi_objective_get_observation_pairs_constrained(constraint_value: int) -> None: - def objective(trial: optuna.trial.Trial) -> Tuple[float, float]: - trial.suggest_int("x", 5, 5) - trial.set_user_attr("constraint", (constraint_value, -1)) - return 5.0, 5.0 - - sampler = TPESampler(constraints_func=lambda trial: trial.user_attrs["constraint"], seed=0) - study = optuna.create_study(directions=["minimize", "maximize"], sampler=sampler) - study.optimize(objective, n_trials=5) - - violations = [max(0, constraint_value) for _ in range(5)] - states = (optuna.trial.TrialState.COMPLETE, optuna.trial.TrialState.PRUNED) - assert _tpe.sampler._get_observation_pairs( - study, study.get_trials(states=states), constraints_enabled=True - ) == ( - [(-float("inf"), [5.0, -5.0]) for _ in range(5)], - violations, - ) - - -def test_multi_objective_split_observation_pairs() -> None: - indices_below, indices_above = _tpe.sampler._split_observation_pairs( - [ - (-float("inf"), [-2.0, -1.0]), - (-float("inf"), [3.0, 3.0]), - (-float("inf"), [0.0, 1.0]), - (-float("inf"), [-1.0, 0.0]), - ], + below_trials = _tpe.sampler._split_complete_trials_multi_objective( + study.trials, + study, 2, - None, ) - assert list(indices_below) == [0, 3] - assert list(indices_above) == [1, 2] + assert [trial.number for trial in below_trials] == [0, 3] -def test_multi_objective_split_observation_pairs_with_all_indices_below() -> None: - indices_below, indices_above = _tpe.sampler._split_observation_pairs( - [ - (-float("inf"), [1.0, 1.0]), - ], - 1, - None, - ) - assert list(indices_below) == [0] - assert list(indices_above) == [] +def test_split_complete_trials_multi_objective_empty() -> None: + study = optuna.create_study(directions=("minimize", "minimize")) + _tpe.sampler._split_complete_trials_multi_objective([], study, 0) == [] def test_calculate_nondomination_rank() -> None: diff --git a/tests/samplers_tests/tpe_tests/test_sampler.py b/tests/samplers_tests/tpe_tests/test_sampler.py index d530b9bb4c..76d72d0c4a 100644 --- a/tests/samplers_tests/tpe_tests/test_sampler.py +++ b/tests/samplers_tests/tpe_tests/test_sampler.py @@ -1,9 +1,7 @@ import random from typing import Callable from typing import Dict -from typing import List from typing import Optional -from typing import Sequence from typing import Union from unittest.mock import Mock from unittest.mock import patch @@ -15,7 +13,6 @@ import optuna from optuna import distributions -from optuna import TrialPruned from optuna.samplers import _tpe from optuna.samplers import TPESampler from optuna.samplers._base import _CONSTRAINTS_KEY @@ -722,168 +719,10 @@ def test_constrained_sample_independent_zero_startup() -> None: sampler.sample_independent(study, trial, "param-a", dist) -@pytest.mark.parametrize("direction", ["minimize", "maximize"]) -@pytest.mark.parametrize( - "constraints_enabled, constraints_func, expected_violations", - [ - (False, None, None), - (True, lambda trial: [(-1, -1), (0, -1), (1, -1), (2, -1)][trial.number], [0, 0, 1, 2]), - ], -) -def test_get_observation_pairs( - direction: str, - constraints_enabled: bool, - constraints_func: Optional[Callable[[optuna.trial.FrozenTrial], Sequence[float]]], - expected_violations: List[float], -) -> None: - def objective(trial: Trial) -> float: - x = trial.suggest_int("x", 5, 5) - z = trial.suggest_categorical("z", [None]) - if trial.number == 0: - return x * int(z is None) - elif trial.number == 1: - trial.report(1, 4) - trial.report(2, 7) - raise TrialPruned() - elif trial.number == 2: - trial.report(float("nan"), 3) - raise TrialPruned() - elif trial.number == 3: - raise TrialPruned() - else: - raise RuntimeError() - - sampler = TPESampler(constraints_func=constraints_func) - study = optuna.create_study(direction=direction, sampler=sampler) - study.optimize(objective, n_trials=5, catch=(RuntimeError,)) - - sign = 1 if direction == "minimize" else -1 - scores = [ - (-float("inf"), [sign * 5.0]), # COMPLETE - (-7, [sign * 2]), # PRUNED (with intermediate values) - (-3, [float("inf")]), # PRUNED (with a NaN intermediate value; it's treated as infinity) - (1, [sign * 0.0]), # PRUNED (without intermediate values) - ] - states = (optuna.trial.TrialState.COMPLETE, optuna.trial.TrialState.PRUNED) - assert _tpe.sampler._get_observation_pairs( - study, study.get_trials(states=states), constraints_enabled=constraints_enabled - ) == ( - scores, - expected_violations, - ) - - -@pytest.mark.parametrize("direction", ["minimize", "maximize"]) -@pytest.mark.parametrize( - "constraints_enabled, constraints_func, expected_violations", - [ - (False, None, None), - (True, lambda trial: [(-1, -1), (0, -1), (1, -1), (2, -1)][trial.number], [0, 0, 1, 2]), - ], -) -def test_get_observation_pairs_multi( - direction: str, - constraints_enabled: bool, - constraints_func: Optional[Callable[[optuna.trial.FrozenTrial], Sequence[float]]], - expected_violations: List[float], -) -> None: - def objective(trial: Trial) -> float: - x = trial.suggest_int("x", 5, 5) - y = trial.suggest_int("y", 6, 6) - if trial.number == 0: - return x + y - elif trial.number == 1: - trial.report(1, 4) - trial.report(2, 7) - raise TrialPruned() - elif trial.number == 2: - trial.report(float("nan"), 3) - raise TrialPruned() - elif trial.number == 3: - raise TrialPruned() - else: - raise RuntimeError() - - sampler = TPESampler(constraints_func=constraints_func) - study = optuna.create_study(direction=direction, sampler=sampler) - study.optimize(objective, n_trials=5, catch=(RuntimeError,)) - - sign = 1 if direction == "minimize" else -1 - states = (optuna.trial.TrialState.COMPLETE, optuna.trial.TrialState.PRUNED) - assert _tpe.sampler._get_observation_pairs( - study, study.get_trials(states=states), constraints_enabled=constraints_enabled - ) == ( - [ - (-float("inf"), [sign * 11.0]), # COMPLETE - (-7, [sign * 2]), # PRUNED (with intermediate values) - ( - -3, - [float("inf")], - ), # PRUNED (with a NaN intermediate value; it's treated as infinity) - (1, [sign * 0.0]), # PRUNED (without intermediate values) - ], - expected_violations, - ) - - -def test_split_observation_pairs() -> None: - indices_below, indices_above = _tpe.sampler._split_observation_pairs( - [ - (-7, [-2]), # PRUNED (with intermediate values) - (float("inf"), [0.0]), # PRUNED (without intermediate values) - ( - -3, - [float("inf")], - ), # PRUNED (with a NaN intermediate value; it's treated as infinity) - (-float("inf"), [-5.0]), # COMPLETE - ], - 2, - None, - ) - assert list(indices_below) == [0, 3] - assert list(indices_above) == [1, 2] - - -def test_split_observation_pairs_with_constraints_below_all_feasible() -> None: - indices_below, indices_above = _tpe.sampler._split_observation_pairs( - [ - (-7, [-2]), # PRUNED (with intermediate values) - (float("inf"), [0.0]), # PRUNED (without intermediate values) - ( - -3, - [float("inf")], - ), # PRUNED (with a NaN intermediate value; it's treated as infinity) - (-float("inf"), [-5.0]), # COMPLETE - ], - 1, - [1, 0, 0, 2], - ) - assert list(indices_below) == [2] - assert list(indices_above) == [0, 1, 3] - - -def test_split_observation_pairs_with_constraints_below_include_infeasible() -> None: - indices_below, indices_above = _tpe.sampler._split_observation_pairs( - [ - (-7, [-2]), # PRUNED (with intermediate values) - (float("inf"), [0.0]), # PRUNED (without intermediate values) - ( - -3, - [float("inf")], - ), # PRUNED (with a NaN intermediate value; it's treated as infinity) - (-float("inf"), [-5.0]), # COMPLETE - ], - 3, - [1, 0, 0, 2], - ) - assert list(indices_below) == [0, 1, 2] - assert list(indices_above) == [3] - - @pytest.mark.parametrize("direction", ["minimize", "maximize"]) @pytest.mark.parametrize("constant_liar", [True, False]) @pytest.mark.parametrize("constraints", [True, False]) -def test_split_order(direction: str, constant_liar: bool, constraints: bool) -> None: +def test_split_trials(direction: str, constant_liar: bool, constraints: bool) -> None: study = optuna.create_study(direction=direction) for value in [-float("inf"), 0, 1, float("inf")]: @@ -919,7 +758,7 @@ def test_split_order(direction: str, constant_liar: bool, constraints: bool) -> ) if constraints: - for value in [1, 2]: + for value in [1, 2, float("inf")]: study.add_trial( optuna.create_trial( state=optuna.trial.TrialState.COMPLETE, @@ -939,6 +778,18 @@ def test_split_order(direction: str, constant_liar: bool, constraints: bool) -> ) ) + study.add_trial( + optuna.create_trial( + state=optuna.trial.TrialState.FAIL, + ) + ) + + study.add_trial( + optuna.create_trial( + state=optuna.trial.TrialState.WAITING, + ) + ) + if constant_liar: states = [ optuna.trial.TrialState.COMPLETE, @@ -948,24 +799,111 @@ def test_split_order(direction: str, constant_liar: bool, constraints: bool) -> else: states = [optuna.trial.TrialState.COMPLETE, optuna.trial.TrialState.PRUNED] - scores, violations = _tpe.sampler._get_observation_pairs( - study, - study.get_trials(states=states), - constraints, + trials = study.get_trials(states=states) + finished_trials = study.get_trials( + states=(optuna.trial.TrialState.COMPLETE, optuna.trial.TrialState.PRUNED) ) + for n_below in range(len(finished_trials) + 1): + below_trials, above_trials = _tpe.sampler._split_trials( + study, + trials, + n_below, + constraints, + ) + + below_trial_numbers = [trial.number for trial in below_trials] + assert below_trial_numbers == list(range(n_below)) + above_trial_numbers = [trial.number for trial in above_trials] + assert above_trial_numbers == list(range(n_below, len(trials))) - assert len(scores) == len(study.get_trials(states=states)) - if constraints: - assert violations is not None - assert len(violations) == len(study.get_trials(states=states)) - else: - assert violations is None - for gamma in range(1, len(scores)): - indices_below, indices_above = _tpe.sampler._split_observation_pairs( - scores, gamma, violations +@pytest.mark.parametrize("direction", ["minimize", "maximize"]) +def test_split_complete_trials_single_objective(direction: str) -> None: + study = optuna.create_study(direction=direction) + + for value in [-float("inf"), 0, 1, float("inf")]: + study.add_trial( + optuna.create_trial( + state=optuna.trial.TrialState.COMPLETE, + value=(value if direction == "minimize" else -value), + params={"x": 0}, + distributions={"x": optuna.distributions.FloatDistribution(-1.0, 1.0)}, + ) ) - assert np.max(indices_below) < np.min(indices_above) + + for n_below in range(len(study.trials) + 1): + below_trials = _tpe.sampler._split_complete_trials_single_objective( + study.trials, + study, + n_below, + ) + assert [trial.number for trial in below_trials] == list(range(n_below)) + + +def test_split_complete_trials_single_objective_empty() -> None: + study = optuna.create_study() + _tpe.sampler._split_complete_trials_single_objective([], study, 0) == [] + + +@pytest.mark.parametrize("direction", ["minimize", "maximize"]) +def test_split_pruned_trials(direction: str) -> None: + study = optuna.create_study(direction=direction) + + for step in [2, 1]: + for value in [-float("inf"), 0, 1, float("inf"), float("nan")]: + study.add_trial( + optuna.create_trial( + state=optuna.trial.TrialState.PRUNED, + params={"x": 0}, + distributions={"x": optuna.distributions.FloatDistribution(-1.0, 1.0)}, + intermediate_values={step: (value if direction == "minimize" else -value)}, + ) + ) + + study.add_trial( + optuna.create_trial( + state=optuna.trial.TrialState.PRUNED, + params={"x": 0}, + distributions={"x": optuna.distributions.FloatDistribution(-1.0, 1.0)}, + ) + ) + + for n_below in range(len(study.trials) + 1): + below_trials = _tpe.sampler._split_pruned_trials( + study.trials, + study, + n_below, + ) + assert [trial.number for trial in below_trials] == list(range(n_below)) + + +def test_split_pruned_trials_empty() -> None: + study = optuna.create_study() + _tpe.sampler._split_pruned_trials([], study, 0) == [] + + +@pytest.mark.parametrize("direction", ["minimize", "maximize"]) +def test_split_infeasible_trials(direction: str) -> None: + study = optuna.create_study(direction=direction) + + for value in [1, 2, float("inf")]: + study.add_trial( + optuna.create_trial( + state=optuna.trial.TrialState.COMPLETE, + value=0, + params={"x": 0}, + distributions={"x": optuna.distributions.FloatDistribution(-1.0, 1.0)}, + system_attrs={_CONSTRAINTS_KEY: [value]}, + ) + ) + + for n_below in range(len(study.trials) + 1): + below_trials = _tpe.sampler._split_infeasible_trials(study.trials, n_below) + assert [trial.number for trial in below_trials] == list(range(n_below)) + + +def test_split_infeasible_trials_empty() -> None: + _tpe.sampler._split_infeasible_trials([], 0) == [] def frozen_trial_factory( @@ -1105,38 +1043,6 @@ def test_group_experimental_warning() -> None: _ = TPESampler(multivariate=True, group=True) -@pytest.mark.parametrize("direction", ["minimize", "maximize"]) -def test_constant_liar_observation_pairs(direction: str) -> None: - with warnings.catch_warnings(): - warnings.simplefilter("ignore", optuna.exceptions.ExperimentalWarning) - sampler = TPESampler(constant_liar=True) - - study = optuna.create_study(sampler=sampler, direction=direction) - - trial = study.ask() - trial.suggest_int("x", 2, 2) - - assert ( - len(study.trials) == 1 and study.trials[0].state == optuna.trial.TrialState.RUNNING - ), "Precondition" - - # The value of the constant liar should be penalizing, i.e. `float("inf")` during minimization - # and `-float("inf")` during maximization. - expected_values = [(float("inf"), [float("inf") * (-1 if direction == "maximize" else 1)])] - - states = [optuna.trial.TrialState.COMPLETE, optuna.trial.TrialState.PRUNED] - assert _tpe.sampler._get_observation_pairs(study, study.get_trials(states=states)) == ( - [], - None, - ) - - states.append(optuna.trial.TrialState.RUNNING) - assert _tpe.sampler._get_observation_pairs(study, study.get_trials(states=states)) == ( - expected_values, - None, - ) - - def test_constant_liar_experimental_warning() -> None: with pytest.warns(optuna.exceptions.ExperimentalWarning): _ = TPESampler(constant_liar=True) From ce93d40b5ecaa919dec021af4d24fb2f3959b5f2 Mon Sep 17 00:00:00 2001 From: Naoto Mizuno Date: Fri, 29 Sep 2023 17:17:14 +0900 Subject: [PATCH 2/3] Return below and above in each split functions --- optuna/samplers/_tpe/sampler.py | 66 +++++++++++-------- .../tpe_tests/test_multi_objective_sampler.py | 5 +- .../samplers_tests/tpe_tests/test_sampler.py | 15 +++-- 3 files changed, 51 insertions(+), 35 deletions(-) diff --git a/optuna/samplers/_tpe/sampler.py b/optuna/samplers/_tpe/sampler.py index 43e478e1e9..fbb27ac39c 100644 --- a/optuna/samplers/_tpe/sampler.py +++ b/optuna/samplers/_tpe/sampler.py @@ -587,6 +587,7 @@ def _split_trials( ) -> tuple[list[FrozenTrial], list[FrozenTrial]]: complete_trials = [] pruned_trials = [] + running_trials = [] infeasible_trials = [] for trial in trials: @@ -596,33 +597,30 @@ def _split_trials( complete_trials.append(trial) elif trial.state == TrialState.PRUNED: pruned_trials.append(trial) + elif trial.state == TrialState.RUNNING: + running_trials.append(trial) else: - assert trial.state == TrialState.RUNNING + assert False # We divide data into below and above. - if len(complete_trials) >= n_below: - below_trials = _split_complete_trials(complete_trials, study, n_below) - elif len(complete_trials) + len(pruned_trials) >= n_below: - below_pruned_trials = _split_pruned_trials( - pruned_trials, study, n_below - len(complete_trials) - ) - below_trials = complete_trials + below_pruned_trials - else: - below_infeasible_trials = _split_infeasible_trials( - infeasible_trials, n_below - len(complete_trials) - len(pruned_trials) - ) - below_trials = complete_trials + pruned_trials + below_infeasible_trials - + below_complete, above_complete = _split_complete_trials(complete_trials, study, n_below) + n_below -= len(below_complete) + below_pruned, above_pruned = _split_pruned_trials(pruned_trials, study, n_below) + n_below -= len(below_pruned) + below_infeasible, above_infeasible = _split_infeasible_trials(infeasible_trials, n_below) + + below_trials = below_complete + below_pruned + below_infeasible + above_trials = above_complete + above_pruned + above_infeasible + running_trials below_trials.sort(key=lambda trial: trial.number) - below_trial_numbers = set(trial.number for trial in below_trials) - above_trials = [trial for trial in trials if trial.number not in below_trial_numbers] + above_trials.sort(key=lambda trial: trial.number) return below_trials, above_trials def _split_complete_trials( trials: Sequence[FrozenTrial], study: Study, n_below: int -) -> list[FrozenTrial]: +) -> tuple[list[FrozenTrial], list[FrozenTrial]]: + n_below = min(n_below, len(trials)) if len(study.directions) <= 1: return _split_complete_trials_single_objective(trials, study, n_below) else: @@ -633,20 +631,21 @@ def _split_complete_trials_single_objective( trials: Sequence[FrozenTrial], study: Study, n_below: int, -) -> list[FrozenTrial]: +) -> tuple[list[FrozenTrial], list[FrozenTrial]]: if study.direction == StudyDirection.MINIMIZE: - return sorted(trials, key=lambda trial: cast(float, trial.value))[:n_below] + sorted_trials = sorted(trials, key=lambda trial: cast(float, trial.value)) else: - return sorted(trials, key=lambda trial: cast(float, trial.value), reverse=True)[:n_below] + sorted_trials = sorted(trials, key=lambda trial: cast(float, trial.value), reverse=True) + return sorted_trials[:n_below], sorted_trials[n_below:] def _split_complete_trials_multi_objective( trials: Sequence[FrozenTrial], study: Study, n_below: int, -) -> list[FrozenTrial]: +) -> tuple[list[FrozenTrial], list[FrozenTrial]]: if n_below == 0: - return [] + return [], [] lvals = np.asarray([trial.values for trial in trials]) for i, direction in enumerate(study.directions): @@ -680,7 +679,14 @@ def _split_complete_trials_multi_objective( selected_indices = _solve_hssp(rank_i_lvals, rank_i_indices, subset_size, reference_point) indices_below[last_idx:] = selected_indices - return [trials[index] for index in indices_below] + below_trials = [] + above_trials = [] + for index in range(len(trials)): + if index in indices_below: + below_trials.append(trials[index]) + else: + above_trials.append(trials[index]) + return below_trials, above_trials def _get_pruned_trial_score(trial: FrozenTrial, study: Study) -> tuple[float, float]: @@ -700,8 +706,10 @@ def _split_pruned_trials( trials: Sequence[FrozenTrial], study: Study, n_below: int, -) -> list[FrozenTrial]: - return sorted(trials, key=lambda trial: _get_pruned_trial_score(trial, study))[:n_below] +) -> tuple[list[FrozenTrial], list[FrozenTrial]]: + n_below = min(n_below, len(trials)) + sorted_trials = sorted(trials, key=lambda trial: _get_pruned_trial_score(trial, study)) + return sorted_trials[:n_below], sorted_trials[n_below:] def _get_infeasible_trial_score(trial: FrozenTrial) -> float: @@ -717,8 +725,12 @@ def _get_infeasible_trial_score(trial: FrozenTrial) -> float: return sum(v for v in constraint if v > 0) -def _split_infeasible_trials(trials: Sequence[FrozenTrial], n_below: int) -> list[FrozenTrial]: - return sorted(trials, key=_get_infeasible_trial_score)[:n_below] +def _split_infeasible_trials( + trials: Sequence[FrozenTrial], n_below: int +) -> tuple[list[FrozenTrial], list[FrozenTrial]]: + n_below = min(n_below, len(trials)) + sorted_trials = sorted(trials, key=_get_infeasible_trial_score) + return sorted_trials[:n_below], sorted_trials[n_below:] def _calculate_weights_below_for_multi_objective( diff --git a/tests/samplers_tests/tpe_tests/test_multi_objective_sampler.py b/tests/samplers_tests/tpe_tests/test_multi_objective_sampler.py index 0d3ea7679a..5c8d97fb34 100644 --- a/tests/samplers_tests/tpe_tests/test_multi_objective_sampler.py +++ b/tests/samplers_tests/tpe_tests/test_multi_objective_sampler.py @@ -324,17 +324,18 @@ def test_split_complete_trials_multi_objective(direction0: str, direction1: str) ) ) - below_trials = _tpe.sampler._split_complete_trials_multi_objective( + below_trials, above_trials = _tpe.sampler._split_complete_trials_multi_objective( study.trials, study, 2, ) assert [trial.number for trial in below_trials] == [0, 3] + assert [trial.number for trial in above_trials] == [1, 2] def test_split_complete_trials_multi_objective_empty() -> None: study = optuna.create_study(directions=("minimize", "minimize")) - _tpe.sampler._split_complete_trials_multi_objective([], study, 0) == [] + _tpe.sampler._split_complete_trials_multi_objective([], study, 0) == ([], []) def test_calculate_nondomination_rank() -> None: diff --git a/tests/samplers_tests/tpe_tests/test_sampler.py b/tests/samplers_tests/tpe_tests/test_sampler.py index 76d72d0c4a..5973f22721 100644 --- a/tests/samplers_tests/tpe_tests/test_sampler.py +++ b/tests/samplers_tests/tpe_tests/test_sampler.py @@ -832,17 +832,18 @@ def test_split_complete_trials_single_objective(direction: str) -> None: ) for n_below in range(len(study.trials) + 1): - below_trials = _tpe.sampler._split_complete_trials_single_objective( + below_trials, above_trials = _tpe.sampler._split_complete_trials_single_objective( study.trials, study, n_below, ) assert [trial.number for trial in below_trials] == list(range(n_below)) + assert [trial.number for trial in above_trials] == list(range(n_below, len(study.trials))) def test_split_complete_trials_single_objective_empty() -> None: study = optuna.create_study() - _tpe.sampler._split_complete_trials_single_objective([], study, 0) == [] + _tpe.sampler._split_complete_trials_single_objective([], study, 0) == ([], []) @pytest.mark.parametrize("direction", ["minimize", "maximize"]) @@ -869,17 +870,18 @@ def test_split_pruned_trials(direction: str) -> None: ) for n_below in range(len(study.trials) + 1): - below_trials = _tpe.sampler._split_pruned_trials( + below_trials, above_trials = _tpe.sampler._split_pruned_trials( study.trials, study, n_below, ) assert [trial.number for trial in below_trials] == list(range(n_below)) + assert [trial.number for trial in above_trials] == list(range(n_below, len(study.trials))) def test_split_pruned_trials_empty() -> None: study = optuna.create_study() - _tpe.sampler._split_pruned_trials([], study, 0) == [] + _tpe.sampler._split_pruned_trials([], study, 0) == ([], []) @pytest.mark.parametrize("direction", ["minimize", "maximize"]) @@ -898,12 +900,13 @@ def test_split_infeasible_trials(direction: str) -> None: ) for n_below in range(len(study.trials) + 1): - below_trials = _tpe.sampler._split_infeasible_trials(study.trials, n_below) + below_trials, above_trials = _tpe.sampler._split_infeasible_trials(study.trials, n_below) assert [trial.number for trial in below_trials] == list(range(n_below)) + assert [trial.number for trial in above_trials] == list(range(n_below, len(study.trials))) def test_split_infeasible_trials_empty() -> None: - _tpe.sampler._split_infeasible_trials([], 0) == [] + _tpe.sampler._split_infeasible_trials([], 0) == ([], []) def frozen_trial_factory( From 58c0ad2a302539c315b37c2dab687566de9f30f8 Mon Sep 17 00:00:00 2001 From: Naoto Mizuno Date: Tue, 3 Oct 2023 16:29:49 +0900 Subject: [PATCH 3/3] Fix missing assert --- .../tpe_tests/test_multi_objective_sampler.py | 2 +- tests/samplers_tests/tpe_tests/test_sampler.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/samplers_tests/tpe_tests/test_multi_objective_sampler.py b/tests/samplers_tests/tpe_tests/test_multi_objective_sampler.py index 5c8d97fb34..4a58eee10b 100644 --- a/tests/samplers_tests/tpe_tests/test_multi_objective_sampler.py +++ b/tests/samplers_tests/tpe_tests/test_multi_objective_sampler.py @@ -335,7 +335,7 @@ def test_split_complete_trials_multi_objective(direction0: str, direction1: str) def test_split_complete_trials_multi_objective_empty() -> None: study = optuna.create_study(directions=("minimize", "minimize")) - _tpe.sampler._split_complete_trials_multi_objective([], study, 0) == ([], []) + assert _tpe.sampler._split_complete_trials_multi_objective([], study, 0) == ([], []) def test_calculate_nondomination_rank() -> None: diff --git a/tests/samplers_tests/tpe_tests/test_sampler.py b/tests/samplers_tests/tpe_tests/test_sampler.py index 5973f22721..d54491c79d 100644 --- a/tests/samplers_tests/tpe_tests/test_sampler.py +++ b/tests/samplers_tests/tpe_tests/test_sampler.py @@ -843,7 +843,7 @@ def test_split_complete_trials_single_objective(direction: str) -> None: def test_split_complete_trials_single_objective_empty() -> None: study = optuna.create_study() - _tpe.sampler._split_complete_trials_single_objective([], study, 0) == ([], []) + assert _tpe.sampler._split_complete_trials_single_objective([], study, 0) == ([], []) @pytest.mark.parametrize("direction", ["minimize", "maximize"]) @@ -881,7 +881,7 @@ def test_split_pruned_trials(direction: str) -> None: def test_split_pruned_trials_empty() -> None: study = optuna.create_study() - _tpe.sampler._split_pruned_trials([], study, 0) == ([], []) + assert _tpe.sampler._split_pruned_trials([], study, 0) == ([], []) @pytest.mark.parametrize("direction", ["minimize", "maximize"]) @@ -906,7 +906,7 @@ def test_split_infeasible_trials(direction: str) -> None: def test_split_infeasible_trials_empty() -> None: - _tpe.sampler._split_infeasible_trials([], 0) == ([], []) + assert _tpe.sampler._split_infeasible_trials([], 0) == ([], []) def frozen_trial_factory(