Unify implementation of fast non-dominated sort

optuna/samplers/nsgaii/_elite_population_selection_strategy.py

@@ -128,6 +128,8 @@
     domination_ranks = _fast_non_dominated_sort(objective_values, penalty=penalty)
     population_per_rank: list[list[FrozenTrial]] = [[] for _ in range(max(domination_ranks) + 1)]
     for trial, rank in zip(population, domination_ranks):
+        if rank == -1:
+            continue
         population_per_rank[rank].append(trial)
 
     return population_per_rank

optuna/study/_multi_objective.py

@@ -85,32 +85,37 @@
         return ranks
 
     if len(penalty) != len(objective_values):
         raise ValueError(
             "The length of penalty and objective_values must be same, but got "
             "len(penalty)={} and len(objective_values)={}.".format(
                 len(penalty), len(objective_values)
             )
         )
-    nondomination_rank = np.zeros(len(objective_values), dtype=int)
+    nondomination_rank = np.full(len(objective_values), -1)
     is_nan = np.isnan(penalty)
+    n_below = n_below or len(objective_values)
 
     # First, we calculate the domination rank for feasible trials.
     is_feasible = np.logical_and(~is_nan, penalty <= 0)
-    ranks, feasible_bottom_rank = _calculate_nondomination_rank(
+    ranks, bottom_rank = _calculate_nondomination_rank(
         objective_values[is_feasible], n_below=n_below
     )
-    nondomination_rank[is_feasible] += ranks
+    nondomination_rank[is_feasible] += 1 + ranks
+    n_below -= np.count_nonzero(is_feasible)
 
     # Second, we calculate the domination rank for infeasible trials.
     is_infeasible = np.logical_and(~is_nan, penalty > 0)
-    ranks, infeasible_bottom_rank = _calculate_nondomination_rank(
-        penalty[is_infeasible, np.newaxis], n_below=n_below
+    ranks, bottom_rank = _calculate_nondomination_rank(
+        penalty[is_infeasible, np.newaxis], n_below=n_below, base_rank=bottom_rank + 1
     )
-    nondomination_rank[is_infeasible] += ranks + (feasible_bottom_rank + 1)
+    nondomination_rank[is_infeasible] += 1 + ranks
+    n_below -= np.count_nonzero(is_infeasible)
 
     # Third, we calculate the domination rank for trials with no penalty information.
-    ranks, _ = _calculate_nondomination_rank(objective_values[is_nan], n_below=n_below)
-    nondomination_rank[is_nan] += ranks + (feasible_bottom_rank + 1) + (infeasible_bottom_rank + 1)
+    ranks, _ = _calculate_nondomination_rank(
+        objective_values[is_nan], n_below=n_below, base_rank=bottom_rank + 1
+    )
+    nondomination_rank[is_nan] += 1 + ranks
 
     return nondomination_rank
 
@@ -119,12 +124,15 @@
     objective_values: np.ndarray,
     *,
     n_below: int | None = None,
+    base_rank: int = 0,
 ) -> tuple[np.ndarray, int]:
-    # Calculate the domination matrix.
-    # The resulting matrix `domination_matrix` is a boolean matrix where
+    # Normalize n_below.
+    n_below = n_below or len(objective_values)
+    n_below = min(n_below, len(objective_values))
+
+    # The ndarray `domination_matrix` is a boolean 2d matrix where
     # `domination_matrix[i, j] == True` means that the j-th trial dominates the i-th trial in the
     # given multi objective minimization problem.
-
     domination_mat = np.all(
         objective_values[:, np.newaxis, :] >= objective_values[np.newaxis, :, :], axis=2
     ) & np.any(objective_values[:, np.newaxis, :] > objective_values[np.newaxis, :, :], axis=2)
@@ -137,9 +145,8 @@
     ranks = np.full(len(objective_values), -1)
     dominated_count = np.sum(domination_mat, axis=1)
 
-    rank = -1
+    rank = base_rank - 1
     ranked_idx_num = 0
-    n_below = n_below or len(objective_values)
     while ranked_idx_num < n_below:
         # Find the non-dominated trials and assign the rank.
         (non_dominated_idxs,) = np.nonzero(dominated_count == 0)

tests/study_tests/test_multi_objective.py

@@ -151,6 +151,6 @@ def test_dominates_complete_vs_incomplete(t1_state: TrialState) -> None:
         ),  # Three objectives with duplicate values are included.
     ],
 )
-def test_calculate_nondomination_rank(trial_values: list[float], trial_ranks: list[int]) -> None:
+def test_fast_non_dominated_sort(trial_values: list[float], trial_ranks: list[int]) -> None:
     ranks = list(_fast_non_dominated_sort(np.array(trial_values)))
     assert np.array_equal(ranks, trial_ranks)