Unify implementation of fast non-dominated sort

optuna/samplers/_nsgaiii/_elite_population_selection_strategy.py

@@ -9,11 +9,9 @@
 import numpy as np
 
 from optuna.samplers._lazy_random_state import LazyRandomState
-from optuna.samplers.nsgaii._dominates import _constrained_dominates
-from optuna.samplers.nsgaii._dominates import _validate_constraints
-from optuna.samplers.nsgaii._elite_population_selection_strategy import _fast_non_dominated_sort
+from optuna.samplers.nsgaii._constraints_evaluation import _validate_constraints
+from optuna.samplers.nsgaii._elite_population_selection_strategy import _rank_population
 from optuna.study import Study
-from optuna.study._multi_objective import _dominates
 from optuna.trial import FrozenTrial
 
 
@@ -52,10 +50,11 @@ def __call__(self, study: Study, population: list[FrozenTrial]) -> list[FrozenTr
         Returns:
             A list of trials that are selected as elite population.
         """
-        _validate_constraints(population, self._constraints_func)
+        _validate_constraints(population, is_constrained=self._constraints_func is not None)
+        population_per_rank = _rank_population(
+            population, study.directions, is_constrained=self._constraints_func is not None
+        )
 
-        dominates = _dominates if self._constraints_func is None else _constrained_dominates
-        population_per_rank = _fast_non_dominated_sort(population, study.directions, dominates)
         elite_population: list[FrozenTrial] = []
         for population in population_per_rank:
             if len(elite_population) + len(population) < self._population_size:

optuna/samplers/_tpe/sampler.py

@@ -28,6 +28,7 @@
 from optuna.search_space.group_decomposed import _GroupDecomposedSearchSpace
 from optuna.search_space.group_decomposed import _SearchSpaceGroup
 from optuna.study import Study
+from optuna.study._multi_objective import _fast_non_dominated_sort
 from optuna.study._study_direction import StudyDirection
 from optuna.trial import FrozenTrial
 from optuna.trial import TrialState
@@ -607,21 +608,6 @@ def after_trial(
         self._random_sampler.after_trial(study, trial, state, values)
 
 
-def _calculate_nondomination_rank(loss_vals: np.ndarray, n_below: int) -> np.ndarray:
-    ranks = np.full(len(loss_vals), -1)
-    num_ranked = 0
-    rank = 0
-    domination_mat = np.all(loss_vals[:, None, :] >= loss_vals[None, :, :], axis=2) & np.any(
-        loss_vals[:, None, :] > loss_vals[None, :, :], axis=2
-    )
-    while num_ranked < n_below:
-        counts = np.sum((ranks == -1)[None, :] & domination_mat, axis=1)
-        num_ranked += np.sum((counts == 0) & (ranks == -1))
-        ranks[(counts == 0) & (ranks == -1)] = rank
-        rank += 1
-    return ranks
-
-
 def _split_trials(
     study: Study,
     trials: list[FrozenTrial],
@@ -693,13 +679,11 @@ def _split_complete_trials_multi_objective(
         # The type of trials must be `list`, but not `Sequence`.
         return [], list(trials)
 
-    lvals = np.asarray([trial.values for trial in trials])
-    for i, direction in enumerate(study.directions):
-        if direction == StudyDirection.MAXIMIZE:
-            lvals[:, i] *= -1
+    lvals = np.array([trial.values for trial in trials])
+    lvals *= np.array([-1.0 if d == StudyDirection.MAXIMIZE else 1.0 for d in study.directions])
 
     # Solving HSSP for variables number of times is a waste of time.
-    nondomination_ranks = _calculate_nondomination_rank(lvals, n_below)
+    nondomination_ranks = _fast_non_dominated_sort(lvals, n_below=n_below)
     assert 0 <= n_below <= len(lvals)
 
     indices = np.array(range(len(lvals)))

optuna/samplers/nsgaii/_child_generation_strategy.py

@@ -6,9 +6,9 @@
 
 from optuna.distributions import BaseDistribution
 from optuna.samplers._lazy_random_state import LazyRandomState
+from optuna.samplers.nsgaii._constraints_evaluation import _constrained_dominates
 from optuna.samplers.nsgaii._crossover import perform_crossover
 from optuna.samplers.nsgaii._crossovers._base import BaseCrossover
-from optuna.samplers.nsgaii._dominates import _constrained_dominates
 from optuna.study import Study
 from optuna.study._multi_objective import _dominates
 from optuna.trial import FrozenTrial

optuna/samplers/nsgaii/_dominates.py → optuna/samplers/nsgaii/_constraints_evaluation.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-from collections.abc import Callable
 from collections.abc import Sequence
 import warnings
 
@@ -86,15 +85,43 @@ def _constrained_dominates(
     return violation0 < violation1
 
 
+def _evaluate_penalty(population: Sequence[FrozenTrial]) -> np.ndarray:
+    """Evaluate feasibility of trials in population.
+    Returns:
+        A list of feasibility status T/F/None of trials in population, where T/F means
+        feasible/infeasible and None means that the trial does not have constraint values.
+    """
+
+    penalty: list[float] = []
+    for trial in population:
+        constraints = trial.system_attrs.get(_CONSTRAINTS_KEY)
+        if constraints is None:
+            penalty.append(np.nan)
+        else:
+            penalty.append(sum(v for v in constraints if v > 0))
+    return np.array(penalty)
+
+
 def _validate_constraints(
     population: list[FrozenTrial],
-    constraints_func: Callable[[FrozenTrial], Sequence[float]] | None = None,
+    *,
+    is_constrained: bool = False,
 ) -> None:
-    if constraints_func is None:
+    if not is_constrained:
         return
+
+    num_constraints = max(
+        [len(t.system_attrs.get(_CONSTRAINTS_KEY, [])) for t in population], default=0
+    )
     for _trial in population:
         _constraints = _trial.system_attrs.get(_CONSTRAINTS_KEY)
         if _constraints is None:
+            warnings.warn(
+                f"Trial {_trial.number} does not have constraint values."
+                " It will be dominated by the other trials."
+            )
             continue
         if np.any(np.isnan(np.array(_constraints))):
             raise ValueError("NaN is not acceptable as constraint value.")
+        elif len(_constraints) != num_constraints:
+            raise ValueError("Trials with different numbers of constraints cannot be compared.")

optuna/samplers/nsgaii/_elite_population_selection_strategy.py

@@ -3,13 +3,14 @@
 from collections import defaultdict
 from collections.abc import Callable
 from collections.abc import Sequence
-import itertools
 
-import optuna
-from optuna.samplers.nsgaii._dominates import _constrained_dominates
-from optuna.samplers.nsgaii._dominates import _validate_constraints
+import numpy as np
+
+from optuna.samplers.nsgaii._constraints_evaluation import _evaluate_penalty
+from optuna.samplers.nsgaii._constraints_evaluation import _validate_constraints
 from optuna.study import Study
-from optuna.study._multi_objective import _dominates
+from optuna.study import StudyDirection
+from optuna.study._multi_objective import _fast_non_dominated_sort
 from optuna.trial import FrozenTrial
 
 
@@ -38,10 +39,10 @@
         Returns:
             A list of trials that are selected as elite population.
         """
-        _validate_constraints(population, self._constraints_func)
-        dominates = _dominates if self._constraints_func is None else _constrained_dominates
-        population_per_rank = _fast_non_dominated_sort(population, study.directions, dominates)
-
+        _validate_constraints(population, is_constrained=self._constraints_func is not None)
+        population_per_rank = _rank_population(
+            population, study.directions, is_constrained=self._constraints_func is not None
+        )
         elite_population: list[FrozenTrial] = []
         for individuals in population_per_rank:
             if len(elite_population) + len(individuals) < self._population_size:
@@ -109,42 +110,26 @@
     population.reverse()
 
 
-def _fast_non_dominated_sort(
+def _rank_population(
     population: list[FrozenTrial],
-    directions: list[optuna.study.StudyDirection],
-    dominates: Callable[[FrozenTrial, FrozenTrial, list[optuna.study.StudyDirection]], bool],
+    directions: Sequence[StudyDirection],
+    *,
+    is_constrained: bool = False,
 ) -> list[list[FrozenTrial]]:
-    dominated_count: defaultdict[int, int] = defaultdict(int)
-    dominates_list = defaultdict(list)
-
-    for p, q in itertools.combinations(population, 2):
-        if dominates(p, q, directions):
-            dominates_list[p.number].append(q.number)
-            dominated_count[q.number] += 1
-        elif dominates(q, p, directions):
-            dominates_list[q.number].append(p.number)
-            dominated_count[p.number] += 1
-
-    population_per_rank = []
-    while population:
-        non_dominated_population = []
-        i = 0
-        while i < len(population):
-            if dominated_count[population[i].number] == 0:
-                individual = population[i]
-                if i == len(population) - 1:
-                    population.pop()
-                else:
-                    population[i] = population.pop()
-                non_dominated_population.append(individual)
-            else:
-                i += 1
-
-        for x in non_dominated_population:
-            for y in dominates_list[x.number]:
-                dominated_count[y] -= 1
-
-        assert non_dominated_population
-        population_per_rank.append(non_dominated_population)
+    if len(population) == 0:
+        return []
+
+    objective_values = np.array([trial.values for trial in population], dtype=np.float64)
+    objective_values *= np.array(
+        [-1.0 if d == StudyDirection.MAXIMIZE else 1.0 for d in directions]
+    )
+    penalty = _evaluate_penalty(population) if is_constrained else None
+
+    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

@@ -1,7 +1,12 @@
+from __future__ import annotations
+
+from collections import defaultdict
 from typing import List
 from typing import Optional
 from typing import Sequence
 
+import numpy as np
+
 import optuna
 from optuna.study._study_direction import StudyDirection
 from optuna.trial import FrozenTrial
@@ -69,6 +74,131 @@ def _get_pareto_front_trials(study: "optuna.study.Study") -> List[FrozenTrial]:
     return _get_pareto_front_trials_by_trials(study.trials, study.directions)
 
 
+def _fast_non_dominated_sort(
+    objective_values: np.ndarray,
+    *,
+    penalty: np.ndarray | None = None,
+    n_below: int | None = None,
+) -> np.ndarray:
+    """Perform the fast non-dominated sort algorithm.
+
+    The fast non-dominated sort algorithm assigns a rank to each trial based on the dominance
+    relationship of the trials, determined by the objective values and the penalty values. The
+    algorithm is based on `the constrained NSGA-II algorithm
+    <https://doi.org/10.1109/4235.99601>`_, but the handling of the case when penalty
+    values are None is different. The algorithm assigns the rank according to the following
+    rules:
+
+    1. Feasible trials: First, the algorithm assigns the rank to feasible trials, whose penalty
+        values are less than or equal to 0, according to unconstrained version of fast non-
+        dominated sort.
+    2. Infeasible trials: Next, the algorithm assigns the rank from the minimum penalty value of to
+        the maximum penalty value.
+    3. Trials with no penalty information (constraints value is None): Finally, The algorithm
+        assigns the rank to trials with no penalty information according to unconstrained version
+        of fast non-dominated sort. Note that only this step is different from the original
+        constrained NSGA-II algorithm.
+    Plus, the algorithm terminates whenever the number of sorted trials reaches n_below.
+
+    Args:
+        objective_values:
+            Objective values of each trials.
+        penalty:
+            Constraints values of each trials. Defaults to None.
+        n_below: The minimum number of top trials required to be sorted. The algorithm will
+            terminate when the number of sorted trials reaches n_below. Defaults to None.
+
+    Returns:
+        An ndarray in the shape of (n_trials,), where each element is the non-dominated rank of
+        each trial. The rank is 0-indexed and rank -1 means that the algorithm terminated before
+        the trial was sorted.
+    """
+    if penalty is None:
+        ranks, _ = _calculate_nondomination_rank(objective_values, n_below=n_below)
+        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.full(len(objective_values), -1)
+    is_penalty_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_penalty_nan, penalty <= 0)
+    ranks, bottom_rank = _calculate_nondomination_rank(
+        objective_values[is_feasible], n_below=n_below
+    )
+    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_penalty_nan, penalty > 0)
+    num_infeasible_trials = np.count_nonzero(is_infeasible)
+    if num_infeasible_trials > 0:
+        _, ranks = np.unique(penalty[is_infeasible], return_inverse=True)
+        ranks += 1
+        nondomination_rank[is_infeasible] += 1 + bottom_rank + ranks
+        bottom_rank += np.max(ranks)
+        n_below -= num_infeasible_trials
+
+    # Third, we calculate the domination rank for trials with no penalty information.
+    ranks, _ = _calculate_nondomination_rank(
+        objective_values[is_penalty_nan], n_below=n_below, base_rank=bottom_rank + 1
+    )
+    nondomination_rank[is_penalty_nan] += 1 + ranks
+
+    return nondomination_rank
+
+
+def _calculate_nondomination_rank(
+    objective_values: np.ndarray,
+    *,
+    n_below: int | None = None,
+    base_rank: int = 0,
+) -> tuple[np.ndarray, int]:
+    if n_below is not None and n_below <= 0:
+        return np.full(len(objective_values), -1), base_rank
+    # Normalize n_below.
+    n_below = n_below or len(objective_values)
+    n_below = min(n_below, len(objective_values))
+
+    # The ndarray `domination_mat` is a boolean 2d matrix where
+    # `domination_mat[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)
+
+    domination_list = np.nonzero(domination_mat)
+    domination_map = defaultdict(list)
+    for dominated_idx, dominating_idx in zip(*domination_list):
+        domination_map[dominating_idx].append(dominated_idx)
+
+    ranks = np.full(len(objective_values), -1)
+    dominated_count = np.sum(domination_mat, axis=1)
+
+    rank = base_rank - 1
+    ranked_idx_num = 0
+    while ranked_idx_num < n_below:
+        # Find the non-dominated trials and assign the rank.
+        (non_dominated_idxs,) = np.nonzero(dominated_count == 0)
+        ranked_idx_num += len(non_dominated_idxs)
+        rank += 1
+        ranks[non_dominated_idxs] = rank
+
+        # Update the dominated count.
+        dominated_count[non_dominated_idxs] = -1
+        for non_dominated_idx in non_dominated_idxs:
+            dominated_count[domination_map[non_dominated_idx]] -= 1
+
+    return ranks, rank
+
+
 def _dominates(
     trial0: FrozenTrial, trial1: FrozenTrial, directions: Sequence[StudyDirection]
 ) -> bool: