Unify the implementations of non-dominated sort

[Alnusjaponica]

Motivation

Currently, the Optuna codebase has implementations of fast non-dominated sort in three different places:

• optuna.study._multi_objective._get_pareto_front_trials_2d() and optuna.study._multi_objective._get_pareto_front_trials_nd()
• optuna.samplers._tpe.sampler._calculate_nondomination_rank()
• optuna.samplers.nsgaii._elite_population_selection_strategy._fast_non_dominated_sort()

This issue aims to unify those implementation to make the codebase simpler.

Suggestion

One of the possible steps are

1. Unify _fast_non_dominated_sort() and _calculate_nondomination_rank() into optuna.study._multi_objective._fast_non_dominated_sort(). The new function will
   • take either of a trial list or a np.ndarray,
   • use numpy matrix operation to compare each trials,
   • equip with early termination when top $k$ individuals (for TPE) or enough ranks which includes $k$ individuals (for NSGA-II) are obtained,
   • and acccept _constrained_dominates.
2. (Optional) Merge _get_pareto_front_trials_nd() into optuna.study._multi_objective._fast_non_dominated_sort(). This might makes pareto front calculation slower, so it will be necessary to run some benchmarks.
3. (Optional) It is also possible to accelerate _get_pareto_front_trials_nd() by using Kung's algorithm or eliminating unnecessary comparison between new item and dominated items. Note that this will also needs to confirm benchmarks before merging.
4. Resolve https://github.com/optuna/optuna/pull/5160#discussion_r1491890749

Additional context (optional)

Non dominated sorts

• optuna/optuna/study/_multi_objective.py

  Lines 11 to 57 in e54ae1f

  	def _get_pareto_front_trials_2d(
  	trials: Sequence[FrozenTrial], directions: Sequence[StudyDirection]
  	) -> List[FrozenTrial]:
  	trials = [trial for trial in 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], directions[0]),
  	_normalize_value(trial.values[1], 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, 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(
  	trials: Sequence[FrozenTrial], directions: Sequence[StudyDirection]
  	) -> List[FrozenTrial]:
  	pareto_front = []
  	trials = [t for t in trials if t.state == TrialState.COMPLETE]
  	# TODO(vincent): Optimize (use the fast non dominated sort defined in the NSGA-II paper).
  	for trial in trials:
  	dominated = False
  	for other in trials:
  	if _dominates(other, trial, directions):
  	dominated = True
  	break
  	if not dominated:
  	pareto_front.append(trial)
  	return pareto_front

• optuna/optuna/samplers/_tpe/sampler.py

  Lines 564 to 576 in e54ae1f

  	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

• optuna/optuna/samplers/nsgaii/_elite_population_selection_strategy.py

  Lines 112 to 150 in e54ae1f

  	def _fast_non_dominated_sort(
  	population: list[FrozenTrial],
  	directions: list[optuna.study.StudyDirection],
  	dominates: Callable[[FrozenTrial, FrozenTrial, list[optuna.study.StudyDirection]], bool],
  	) -> 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)
  	return population_per_rank

The above functions have some difference:

• _get_pareto_front_trials_2d() and _get_pareto_front_trials_nd()
  • only need to calculate the pareto front, which can be obtained $O(n\log n)$ for 2 or 3 dimensional objectives and $O(n(\log n)^{d-2})$ for $d$-dimensional objectives by Kungs algorithm,
  • _get_pareto_front_trials_2d() is $O(n\log n)$, which is optimal, while _get_pareto_front_trials_nd() performs comparison of trials $n^2$ times in the worst case, which is constant times slower than fast non-dominated sort.
• _calculate_nondomination_rank()
  • only need to calculate top n_below trials, which enables to early return before the ranks of all trials are settled,
  • implemented using matrix operation on numpy, which should make comparisons faster than for loop implementation,
  • takes $O(n^3)$ time in the worst case
• _fast_non_dominated_sort()
  • implements fast non-dominated sort,
  • requires to consider _constrained_dominates,
  • destroys the given list of trials.

[Alnusjaponica]

Follow-ups

#5160 (comment)