Merge pull request optuna#5157 from Alnusjaponica/move-_create_frozen…

…_trial

Unify the implementation of `_create_frozen_trial()` under `testing` module

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,
+    )

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,15 +283,19 @@ 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
         assert "additional" not in params1
 
         # 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,
-    )

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,22 +343,22 @@ 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)
 
     # Check that constraints with same violations are incomparable.
     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)