Merge pull request #2503 from parsiad/log-linear-2d-pareto

Add log-linear algorithm for 2d Pareto front.

optuna/_multi_objective.py

@@ -8,7 +8,34 @@
 from optuna.trial import TrialState
 
 
-def _get_pareto_front_trials(study: "optuna.study.BaseStudy") -> List[FrozenTrial]:
+def _get_pareto_front_trials_2d(study: "optuna.study.BaseStudy") -> List[FrozenTrial]:
+    trials = [trial for trial in study.trials if trial.state == TrialState.COMPLETE]
+
+    n_trials = len(trials)
+    if n_trials == 0:
+        return []
+
+    trials.sort(
+        key=lambda trial: (
+            _normalize_value(trial.values[0], study.directions[0]),
+            _normalize_value(trial.values[1], study.directions[1]),
+        ),
+    )
+
+    last_nondominated_trial = trials[0]
+    pareto_front = [last_nondominated_trial]
+    for i in range(1, n_trials):
+        trial = trials[i]
+        if _dominates(last_nondominated_trial, trial, study.directions):
+            continue
+        pareto_front.append(trial)
+        last_nondominated_trial = trial
+
+    pareto_front.sort(key=lambda trial: trial.number)
+    return pareto_front
+
+
+def _get_pareto_front_trials_nd(study: "optuna.study.BaseStudy") -> List[FrozenTrial]:
     pareto_front = []
     trials = [t for t in study.trials if t.state == TrialState.COMPLETE]
 
@@ -26,6 +53,12 @@ def _get_pareto_front_trials(study: "optuna.study.BaseStudy") -> List[FrozenTria
     return pareto_front
 
 
+def _get_pareto_front_trials(study: "optuna.study.BaseStudy") -> List[FrozenTrial]:
+    if len(study.directions) == 2:
+        return _get_pareto_front_trials_2d(study)  # Log-linear in number of trials.
+    return _get_pareto_front_trials_nd(study)  # Quadratic in number of trials.
+
+
 def _dominates(
     trial0: FrozenTrial, trial1: FrozenTrial, directions: Sequence[StudyDirection]
 ) -> bool:

tests/study_tests/test_study.py

@@ -822,7 +822,7 @@ def objective(trial: Trial) -> List[float]:
         assert len(trial.values) == n_objectives
 
 
-def test_pareto_front() -> None:
+def test_pareto_front_2d() -> None:
     def _trial_to_values(t: FrozenTrial) -> Tuple[float, ...]:
         assert t.values is not None
         return tuple(t.values)
@@ -839,6 +839,9 @@ def _trial_to_values(t: FrozenTrial) -> Tuple[float, ...]:
     study.optimize(lambda t: [3, 1], n_trials=1)
     assert {_trial_to_values(t) for t in study.best_trials} == {(1, 1), (2, 2)}
 
+    study.optimize(lambda t: [3, 2], n_trials=1)
+    assert {_trial_to_values(t) for t in study.best_trials} == {(1, 1), (2, 2)}
+
     study.optimize(lambda t: [1, 3], n_trials=1)
     assert {_trial_to_values(t) for t in study.best_trials} == {(1, 3)}
     assert len(study.best_trials) == 1
@@ -848,6 +851,37 @@ def _trial_to_values(t: FrozenTrial) -> Tuple[float, ...]:
     assert len(study.best_trials) == 2
 
 
+def test_pareto_front_3d() -> None:
+    def _trial_to_values(t: FrozenTrial) -> Tuple[float, ...]:
+        assert t.values is not None
+        return tuple(t.values)
+
+    study = create_study(directions=["minimize", "maximize", "minimize"])
+    assert {_trial_to_values(t) for t in study.best_trials} == set()
+
+    study.optimize(lambda t: [2, 2, 2], n_trials=1)
+    assert {_trial_to_values(t) for t in study.best_trials} == {(2, 2, 2)}
+
+    study.optimize(lambda t: [1, 1, 1], n_trials=1)
+    assert {_trial_to_values(t) for t in study.best_trials} == {(1, 1, 1), (2, 2, 2)}
+
+    study.optimize(lambda t: [3, 1, 3], n_trials=1)
+    assert {_trial_to_values(t) for t in study.best_trials} == {(1, 1, 1), (2, 2, 2)}
+
+    study.optimize(lambda t: [3, 2, 3], n_trials=1)
+    assert {_trial_to_values(t) for t in study.best_trials} == {(1, 1, 1), (2, 2, 2)}
+
+    study.optimize(lambda t: [1, 3, 1], n_trials=1)
+    assert {_trial_to_values(t) for t in study.best_trials} == {(1, 3, 1)}
+    assert len(study.best_trials) == 1
+
+    study.optimize(
+        lambda t: [1, 3, 1], n_trials=1
+    )  # The trial result is the same as the above one.
+    assert {_trial_to_values(t) for t in study.best_trials} == {(1, 3, 1)}
+    assert len(study.best_trials) == 2
+
+
 def test_wrong_n_objectives() -> None:
     n_objectives = 2
     directions = ["minimize" for _ in range(n_objectives)]