diff --git a/optuna/testing/trials.py b/optuna/testing/trials.py new file mode 100644 index 0000000000..6f06019263 --- /dev/null +++ b/optuna/testing/trials.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +import optuna +from optuna.distributions import BaseDistribution +from optuna.samplers._base import _CONSTRAINTS_KEY +from optuna.trial import FrozenTrial +from optuna.trial import TrialState + + +def _create_frozen_trial( + number: int = 0, + values: Sequence[float] | None = None, + constraints: Sequence[float] | None = None, + params: dict[str, Any] | None = None, + param_distributions: dict[str, BaseDistribution] | None = None, + state: TrialState = TrialState.COMPLETE, +) -> optuna.trial.FrozenTrial: + return FrozenTrial( + number=number, + value=1.0 if values is None else None, + values=values, + state=state, + user_attrs={}, + system_attrs={} if constraints is None else {_CONSTRAINTS_KEY: list(constraints)}, + params=params or {}, + distributions=param_distributions or {}, + intermediate_values={}, + datetime_start=None, + datetime_complete=None, + trial_id=number, + ) diff --git a/tests/integration_tests/test_cma.py b/tests/integration_tests/test_cma.py index 4d819bc71c..90330b5309 100644 --- a/tests/integration_tests/test_cma.py +++ b/tests/integration_tests/test_cma.py @@ -17,7 +17,7 @@ from optuna.integration.cma import _Optimizer from optuna.study._study_direction import StudyDirection from optuna.testing.distributions import UnsupportedDistribution -from optuna.trial import FrozenTrial +from optuna.testing.trials import _create_frozen_trial from optuna.trial import Trial from optuna.trial import TrialState @@ -226,10 +226,13 @@ def test_tell( search_space, x0, 0.2, None, {"popsize": 3, "seed": 1} ) - trials = [_create_frozen_trial(x0, search_space)] + trials = [_create_frozen_trial(params=x0, param_distributions=search_space)] assert -1 == optimizer.tell(trials, direction) - trials = [_create_frozen_trial(x0, search_space, number=i) for i in range(3)] + trials = [ + _create_frozen_trial(params=x0, param_distributions=search_space, number=i) + for i in range(3) + ] assert 2 == optimizer.tell(trials, direction) @staticmethod @@ -241,8 +244,12 @@ def test_tell_filter_by_state( search_space, x0, 0.2, None, {"popsize": 2, "seed": 1} ) - trials = [_create_frozen_trial(x0, search_space)] - trials.append(_create_frozen_trial(x0, search_space, state, len(trials))) + trials = [_create_frozen_trial(params=x0, param_distributions=search_space)] + trials.append( + _create_frozen_trial( + params=x0, param_distributions=search_space, state=state, number=len(trials) + ) + ) assert -1 == optimizer.tell(trials, StudyDirection.MINIMIZE) @staticmethod @@ -253,15 +260,18 @@ def test_tell_filter_by_distribution( search_space, x0, 0.2, None, {"popsize": 2, "seed": 1} ) - trials = [_create_frozen_trial(x0, search_space)] + trials = [_create_frozen_trial(params=x0, param_distributions=search_space)] distributions = trials[0].distributions.copy() distributions["additional"] = FloatDistribution(0, 100) - trials.append(_create_frozen_trial(x0, distributions, number=1)) + trials.append(_create_frozen_trial(params=x0, param_distributions=search_space, number=1)) assert 1 == optimizer.tell(trials, StudyDirection.MINIMIZE) @staticmethod def test_ask(search_space: Dict[str, BaseDistribution], x0: Dict[str, Any]) -> None: - trials = [_create_frozen_trial(x0, search_space, number=i) for i in range(3)] + trials = [ + _create_frozen_trial(params=x0, param_distributions=search_space, number=i) + for i in range(3) + ] # Create 0-th individual. optimizer = _Optimizer(search_space, x0, 0.2, None, {"popsize": 3, "seed": 1}) @@ -273,7 +283,9 @@ def test_ask(search_space: Dict[str, BaseDistribution], x0: Dict[str, Any]) -> N last_told = optimizer.tell(trials, StudyDirection.MINIMIZE) distributions = trials[0].distributions.copy() distributions["additional"] = FloatDistribution(0, 100) - trials.append(_create_frozen_trial(x0, distributions, number=len(trials))) + trials.append( + _create_frozen_trial(params=x0, param_distributions=search_space, number=len(trials)) + ) params1 = optimizer.ask(trials, last_told) assert params0 != params1 @@ -281,7 +293,9 @@ def test_ask(search_space: Dict[str, BaseDistribution], x0: Dict[str, Any]) -> N # Create first individual. optimizer = _Optimizer(search_space, x0, 0.2, None, {"popsize": 3, "seed": 1}) - trials.append(_create_frozen_trial(x0, search_space, number=len(trials))) + trials.append( + _create_frozen_trial(params=x0, param_distributions=search_space, number=len(trials)) + ) last_told = optimizer.tell(trials, StudyDirection.MINIMIZE) params2 = optimizer.ask(trials, last_told) @@ -291,7 +305,11 @@ def test_ask(search_space: Dict[str, BaseDistribution], x0: Dict[str, Any]) -> N last_told = optimizer.tell(trials, StudyDirection.MINIMIZE) # Other worker adds three trials. for _ in range(3): - trials.append(_create_frozen_trial(x0, search_space, number=len(trials))) + trials.append( + _create_frozen_trial( + params=x0, param_distributions=search_space, number=len(trials) + ) + ) params3 = optimizer.ask(trials, last_told) assert params0 != params3 @@ -302,16 +320,19 @@ def test_is_compatible(search_space: Dict[str, BaseDistribution], x0: Dict[str, optimizer = optuna.integration.cma._Optimizer(search_space, x0, 0.1, None, {}) # Compatible. - trial = _create_frozen_trial(x0, search_space) + trial = _create_frozen_trial(params=x0, param_distributions=search_space) assert optimizer._is_compatible(trial) # Compatible. - trial = _create_frozen_trial(x0, dict(search_space, u=FloatDistribution(-10, 10))) + trial = _create_frozen_trial( + params=x0, param_distributions=dict(search_space, u=FloatDistribution(-10, 10)) + ) assert optimizer._is_compatible(trial) # Compatible. trial = _create_frozen_trial( - dict(x0, unknown=7), dict(search_space, unknown=FloatDistribution(0, 10)) + params=dict(x0, unknown=7), + param_distributions=dict(search_space, unknown=FloatDistribution(0, 10)), ) assert optimizer._is_compatible(trial) @@ -320,37 +341,19 @@ def test_is_compatible(search_space: Dict[str, BaseDistribution], x0: Dict[str, del param["u"] dist = dict(search_space) del dist["u"] - trial = _create_frozen_trial(param, dist) + trial = _create_frozen_trial(params=param, param_distributions=dist) assert not optimizer._is_compatible(trial) # Incompatible (the value of 'u' is out of range). trial = _create_frozen_trial( - dict(x0, u=20), dict(search_space, u=FloatDistribution(-100, 100)) + params=dict(x0, u=20), + param_distributions=dict(search_space, u=FloatDistribution(-100, 100)), ) assert not optimizer._is_compatible(trial) # Error (different distribution class). - trial = _create_frozen_trial(x0, dict(search_space, u=IntDistribution(-2, 2))) + trial = _create_frozen_trial( + params=x0, param_distributions=dict(search_space, u=IntDistribution(-2, 2)) + ) with pytest.raises(ValueError): optimizer._is_compatible(trial) - - -def _create_frozen_trial( - params: Dict[str, Any], - param_distributions: Dict[str, BaseDistribution], - state: TrialState = TrialState.COMPLETE, - number: int = 0, -) -> FrozenTrial: - return FrozenTrial( - number=number, - value=1.0, - state=state, - user_attrs={}, - system_attrs={}, - params=params, - distributions=param_distributions, - intermediate_values={}, - datetime_start=None, - datetime_complete=None, - trial_id=number, - ) diff --git a/tests/samplers_tests/test_nsgaii.py b/tests/samplers_tests/test_nsgaii.py index fc13863d30..652790f3c8 100644 --- a/tests/samplers_tests/test_nsgaii.py +++ b/tests/samplers_tests/test_nsgaii.py @@ -45,6 +45,7 @@ from optuna.samplers.nsgaii._sampler import _GENERATION_KEY from optuna.study._multi_objective import _dominates from optuna.study._study_direction import StudyDirection +from optuna.testing.trials import _create_frozen_trial from optuna.trial import FrozenTrial @@ -246,8 +247,8 @@ def test_constrained_dominates_feasible_vs_feasible( # The results of _constrained_dominates match _dominates in all feasible cases. for values1, constraints1 in values_constraints_list: for values2, constraints2 in values_constraints_list: - t1 = _create_frozen_trial(0, values1, constraints1) - t2 = _create_frozen_trial(1, values2, constraints2) + t1 = _create_frozen_trial(number=0, values=values1, constraints=constraints1) + t2 = _create_frozen_trial(number=1, values=values2, constraints=constraints2) assert _constrained_dominates(t1, t2, directions) == _dominates(t1, t2, directions) @@ -276,13 +277,13 @@ def test_constrained_dominates_feasible_vs_infeasible( # the infeasible trials. for constraints1 in constraints_list1: for constraints2 in constraints_list2: - t1 = _create_frozen_trial(0, [0], constraints1) - t2 = _create_frozen_trial(1, [1], constraints2) + t1 = _create_frozen_trial(number=0, values=[0], constraints=constraints1) + t2 = _create_frozen_trial(number=1, values=[1], constraints=constraints2) assert _constrained_dominates(t1, t2, directions) assert not _constrained_dominates(t2, t1, directions) - t1 = _create_frozen_trial(0, [1], constraints1) - t2 = _create_frozen_trial(1, [0], constraints2) + t1 = _create_frozen_trial(number=0, values=[1], constraints=constraints1) + t2 = _create_frozen_trial(number=1, values=[0], constraints=constraints2) assert _constrained_dominates(t1, t2, directions) assert not _constrained_dominates(t2, t1, directions) @@ -342,13 +343,13 @@ def test_constrained_dominates_infeasible_vs_infeasible(direction: StudyDirectio # every constraint in constraints_infeasible_sorted[j]. for constraints1 in constraints_infeasible_sorted[i]: for constraints2 in constraints_infeasible_sorted[j]: - t1 = _create_frozen_trial(0, [0], constraints1) - t2 = _create_frozen_trial(1, [1], constraints2) + t1 = _create_frozen_trial(number=0, values=[0], constraints=constraints1) + t2 = _create_frozen_trial(number=1, values=[1], constraints=constraints2) assert _constrained_dominates(t1, t2, directions) assert not _constrained_dominates(t2, t1, directions) - t1 = _create_frozen_trial(0, [1], constraints1) - t2 = _create_frozen_trial(1, [0], constraints2) + t1 = _create_frozen_trial(number=0, values=[1], constraints=constraints1) + t2 = _create_frozen_trial(number=1, values=[0], constraints=constraints2) assert _constrained_dominates(t1, t2, directions) assert not _constrained_dominates(t2, t1, directions) @@ -356,8 +357,8 @@ def test_constrained_dominates_infeasible_vs_infeasible(direction: StudyDirectio for constraints_with_same_violations in constraints_infeasible_sorted: for constraints1 in constraints_with_same_violations: for constraints2 in constraints_with_same_violations: - t1 = _create_frozen_trial(0, [0], constraints1) - t2 = _create_frozen_trial(1, [1], constraints2) + t1 = _create_frozen_trial(number=0, values=[0], constraints=constraints1) + t2 = _create_frozen_trial(number=1, values=[1], constraints=constraints2) assert not _constrained_dominates(t1, t2, directions) assert not _constrained_dominates(t2, t1, directions) @@ -397,7 +398,7 @@ def test_fast_non_dominated_sort_no_constraints( value_list = [10, 20, 20, 30, float("inf"), float("inf"), -float("inf")] values = [[v1, v2] for v1 in value_list for v2 in value_list] - trials = [_create_frozen_trial(i, v) for i, v in enumerate(values)] + trials = [_create_frozen_trial(number=i, values=v) for i, v in enumerate(values)] population_per_rank = _fast_non_dominated_sort(copy.copy(trials), directions, _dominates) _assert_population_per_rank(trials, directions, population_per_rank) @@ -410,7 +411,7 @@ def test_fast_non_dominated_sort_with_constraints() -> None: constraints = [[c1, c2] for c1 in constraint_list for c2 in constraint_list] trials = [ - _create_frozen_trial(i, v, c) + _create_frozen_trial(number=i, values=v, constraints=c) for i, (v, c) in enumerate(itertools.product(values, constraints)) ] directions = [StudyDirection.MINIMIZE, StudyDirection.MAXIMIZE] @@ -423,7 +424,7 @@ def test_fast_non_dominated_sort_with_constraints() -> None: def test_validate_constraints() -> None: with pytest.raises(ValueError): _validate_constraints( - [_create_frozen_trial(0, [1], [0, float("nan")])], + [_create_frozen_trial(number=0, values=[1], constraints=[0, float("nan")])], constraints_func=lambda _: [0], ) @@ -451,7 +452,10 @@ def test_fast_non_dominated_sort_missing_constraint_values( for directions in itertools.product( [StudyDirection.MINIMIZE, StudyDirection.MAXIMIZE], repeat=values_dim ): - trials = [_create_frozen_trial(i, v, c) for i, (v, c) in enumerate(values_and_constraints)] + trials = [ + _create_frozen_trial(number=i, values=v, constraints=c) + for i, (v, c) in enumerate(values_and_constraints) + ] with pytest.warns(UserWarning): population_per_rank = _fast_non_dominated_sort( @@ -499,7 +503,7 @@ def test_fast_non_dominated_sort_empty(n_dims: int) -> None: ], ) def test_calc_crowding_distance(values: list[list[float]], expected_dist: list[float]) -> None: - trials = [_create_frozen_trial(i, value) for i, value in enumerate(values)] + trials = [_create_frozen_trial(number=i, values=value) for i, value in enumerate(values)] crowding_dist = _calc_crowding_distance(trials) for i in range(len(trials)): assert _nan_equal(crowding_dist[i], expected_dist[i]), i @@ -517,7 +521,7 @@ def test_calc_crowding_distance(values: list[list[float]], expected_dist: list[f ) def test_crowding_distance_sort(values: list[list[float]]) -> None: """Checks that trials are sorted by the values of `_calc_crowding_distance`.""" - trials = [_create_frozen_trial(i, value) for i, value in enumerate(values)] + trials = [_create_frozen_trial(number=i, values=value) for i, value in enumerate(values)] crowding_dist = _calc_crowding_distance(trials) _crowding_distance_sort(trials) sorted_dist = [crowding_dist[t.number] for t in trials] @@ -575,20 +579,6 @@ def test_after_trial_strategy_experimental_warning() -> None: NSGAIISampler(after_trial_strategy=lambda study, trial, state, value: None) -# TODO(ohta): Consider to move this utility function to `optuna.testing` module. -def _create_frozen_trial( - number: int, values: Sequence[float], constraints: Sequence[float] | None = None -) -> optuna.trial.FrozenTrial: - trial = optuna.trial.create_trial( - state=optuna.trial.TrialState.COMPLETE, - values=list(values), - system_attrs={} if constraints is None else {_CONSTRAINTS_KEY: list(constraints)}, - ) - trial.number = number - trial._trial_id = number - return trial - - def test_elite_population_selection_strategy_invalid_value() -> None: with pytest.raises(ValueError): NSGAIIElitePopulationSelectionStrategy(population_size=1)