From 1dc69fa4b67c5df5149d7de46467abaa0d195c70 Mon Sep 17 00:00:00 2001 From: flyingdoubleG <53119724+flyingdoubleG@users.noreply.github.com> Date: Wed, 28 Oct 2026 23:27:15 +0800 Subject: [PATCH 1/2] Align TS-NSGA-II's environmental selection with the reference - Stage 2 chose among all 2N solutions, dominated ones included, so it had no convergence pressure. Restrict it to the first non-dominated fronts holding at least N solutions and, as the reference does, take for each reference vector in turn the remaining solution with the largest cosine, normalizing the remaining ones each time. - Stage 1 ranked by the number of SPD-dominators. Peel SPD fronts instead, zero d1 and d2 of the extreme non-dominated solutions, and fill from the last front by d2. - Switch stages at (0.8 + 0.2 * (1 - 3 / M)) * max_fe rather than max_fe / 2, and use the region density as the second mating key in stage 2. - Drop the level computation whose result was discarded; the reference's level sort does not narrow the set either. --- src/evomo/algorithms/ts_nsga2.py | 212 ++++++++++++++------------- unit_test/algorithms/test_tsnsga2.py | 41 ++++++ 2 files changed, 152 insertions(+), 101 deletions(-) create mode 100644 unit_test/algorithms/test_tsnsga2.py diff --git a/src/evomo/algorithms/ts_nsga2.py b/src/evomo/algorithms/ts_nsga2.py index b464d7d..9818de6 100644 --- a/src/evomo/algorithms/ts_nsga2.py +++ b/src/evomo/algorithms/ts_nsga2.py @@ -7,7 +7,6 @@ from evox.utils import clamp, lexsort from evomo.operators.selection import non_dominate_rank -from evomo.utils import unique_rows_sorted class TSNSGAII(Algorithm): @@ -26,9 +25,10 @@ def __init__(self, pop_size: int, n_objs: int, lb: torch.Tensor, ub: torch.Tenso :param ub: Required. Upper decision bounds of shape ``(D,)``. The decision dimension is inferred from the bounds, rather than passed separately. Match ``lb`` in shape, dtype and device. :type ub: torch.Tensor - :param max_fe: Default: ``10000``. Positive evaluation horizon setting the phase boundary at ``max_fe // 2``. - The evaluation counter includes the initial population. This parameter controls the phase switch and does - not stop subsequent calls to ``step()``. + :param max_fe: Default: ``10000``. Positive evaluation horizon setting the phase boundary at + ``(0.8 + 0.2 * (1 - 3 / n_objs)) * max_fe`` evaluations, as in the reference implementation. The evaluation + counter includes the initial population. This parameter controls the phase switch and does not stop + subsequent calls to ``step()``. :type max_fe: int .. note:: @@ -100,105 +100,115 @@ def step(self) -> None: self.pop, self.fit, self.front_no, self.d2 = self._environmental_selection(merged_pop, merged_fit) def _environmental_selection(self, pop, fit): - N = self.pop_size - M = self.n_objs - - # Normalization (Bug #12) - f_min = torch.min(fit, dim=0)[0] - f_max = torch.max(fit, dim=0)[0] - norm_fit = (fit - f_min) / (f_max - f_min + 1e-6) - - # Stage Transition Logic - if self.fe < (self.max_fe // 2): - # Stage 1: SPD-Sort - # Association - norm_val = torch.norm(norm_fit, dim=1, keepdim=True) - w_norm = torch.norm(self.w, dim=1, keepdim=True) - cosine = (norm_fit @ self.w.T) / (norm_val @ w_norm.T + 1e-6) - - # d1 and d2 - d1_all = norm_val * cosine - d2_all = norm_val * torch.sqrt(clamp(1 - cosine**2, 0.0, 1.0) + 1e-6) - - # Find closest weight for each individual - closest_w_idx = torch.argmin(d2_all, dim=1) - idx_range = torch.arange(pop.shape[0], device=pop.device) - - d1 = d1_all[idx_range, closest_w_idx] - d2 = d2_all[idx_range, closest_w_idx] - spd_score = d1 + 5 * d2 - - front_no = self._spd_sort(fit, spd_score, closest_w_idx) - - # Selection via Lexsort - rank_indices = lexsort(torch.stack([spd_score, front_no.float()])) - survivor_idx = rank_indices[:N] - - return pop[survivor_idx], fit[survivor_idx], front_no[survivor_idx], d2[survivor_idx] - - else: - # Stage 2: Level-Sort & Decomposition - # Interval-based levels - interval = (f_max - f_min) / (2 * M + 1e-6) - torch.max((fit - f_min) / (interval + 1e-6), dim=1)[0].floor() - - # Greedy Selection (Vectorized) - survivor_idx = self._level_sort_selection(norm_fit, self.w, N) + # Stage 1 (SPD sorting) until (0.8 + 0.2 * (1 - 3 / M)) of the evaluation budget, then stage 2. + if self.fe < (0.8 + 0.2 * (1 - 3 / self.n_objs)) * self.max_fe: + return self._spd_selection(pop, fit) + return self._decomposition_selection(pop, fit) - # Re-calculate front_no and d2 for the next mating step - new_fit = fit[survivor_idx] - new_norm = norm_fit[survivor_idx] - new_front = non_dominate_rank(new_fit) - - # Update d2 for mating - new_norm_val = torch.norm(new_norm, dim=1, keepdim=True) - new_cosine = (new_norm @ self.w.T) / (new_norm_val @ torch.norm(self.w, dim=1, keepdim=True).T + 1e-6) - new_d2 = torch.min(new_norm_val * torch.sqrt(clamp(1 - new_cosine**2, 0.0, 1.0) + 1e-6), dim=1)[0] - - return pop[survivor_idx], new_fit, new_front, new_d2 - - def _spd_sort(self, fit, spd_score, w_idx): - # Bug #24: Pareto Dominance - pareto_dom = (fit.unsqueeze(1) <= fit.unsqueeze(0)).all(-1) & (fit.unsqueeze(1) < fit.unsqueeze(0)).any(-1) - - # SPD Dominance: better score on the same weight - same_w = w_idx.unsqueeze(1) == w_idx.unsqueeze(0) - dist_dom = (spd_score.unsqueeze(1) < spd_score.unsqueeze(0)) & same_w - - total_dom = pareto_dom | dist_dom - - # Peeling logic (Vectorized Rank Assignment) - num_dominated = torch.sum(total_dom, dim=0) - return num_dominated.int() - - def _level_sort_selection(self, norm_fit, w, n_required): - # Vectorized Greedy Selection - # Calculate cosine similarity between all individuals and all weights + def _spd_selection(self, pop, fit): + """Stage 1: keep the first SPD fronts, then the lowest-d2 solutions of the last one.""" + N = self.pop_size + norm_fit = _normalize(fit) + cosine = _cosine(norm_fit, self.w) norm_val = torch.norm(norm_fit, dim=1, keepdim=True) - w_norm = torch.norm(w, dim=1, keepdim=True) - cosine = (norm_fit @ w.T) / (norm_val @ w_norm.T + 1e-6) # [2N, N] - - # For each weight, find the individual with max cosine similarity - best_match_for_w = torch.argmax(cosine, dim=0) # [N] - - # Handle potential duplicates if one individual is best for multiple weights - # But for TSNSGAII Stage 2, we typically take the best for each weight - # If N_weights == N_required, this naturally selects N individuals - u_idx, _ = unique_rows_sorted(best_match_for_w.unsqueeze(1)) - u_idx = u_idx.squeeze(1) - - # If not enough unique individuals, pad with remaining best by cosine - # This is a JIT-safe way to ensure we return exactly N - mask = torch.zeros(norm_fit.shape[0], dtype=torch.bool, device=norm_fit.device) - mask[u_idx] = True - - remaining_indices = torch.where(~mask)[0] - # Sort remaining by max similarity to any weight - max_sim_to_any = torch.max(cosine, dim=1)[0] - rem_sorted = remaining_indices[torch.argsort(max_sim_to_any[remaining_indices], descending=True)] - - combined = torch.cat([u_idx, rem_sorted], dim=0) - return combined[:n_required] + d1_all = norm_val * cosine + d2_all = norm_val * torch.sqrt(clamp(1 - cosine**2, 0.0, 1.0)) + d2, rp = torch.min(d2_all, dim=1) + d1 = d1_all.gather(1, rp.unsqueeze(1)).squeeze(1) + + # The extreme non-dominated solution of each objective gets d1 = d2 = 0. + nd = torch.where(non_dominate_rank(norm_fit) == 0)[0] + extreme = nd[torch.argmax(norm_fit[nd], dim=0)] + d1 = d1.index_fill(0, extreme, 0.0) + d2 = d2.index_fill(0, extreme, 0.0) + + front_no = self._spd_sort(norm_fit, d1 + 5 * d2, rp, N) + max_front = torch.max(torch.where(front_no < torch.iinfo(front_no.dtype).max, front_no, -1)) + # Earlier fronts first, then the last front by d2; unsorted solutions sort last. + survivor_idx = lexsort([d2, front_no])[:N] + return pop[survivor_idx], fit[survivor_idx], torch.clamp(front_no[survivor_idx], max=max_front), d2[survivor_idx] + + def _spd_sort(self, fit, spd_score, rp, n_sort): + """Peel SPD fronts until at least ``n_sort`` solutions are sorted. + + A solution SPD-dominates another when it Pareto-dominates it, or when neither Pareto-dominates the other, + they share a reference vector, and it has the smaller d1 + 5 * d2. Unsorted solutions keep the dtype maximum. + """ + n = fit.shape[0] + better_any = (fit.unsqueeze(1) < fit.unsqueeze(0)).any(-1) + pareto_dom = (fit.unsqueeze(1) <= fit.unsqueeze(0)).all(-1) & better_any + incomparable = ~pareto_dom & ~pareto_dom.T + same_rp = rp.unsqueeze(1) == rp.unsqueeze(0) + dominates = pareto_dom | (incomparable & same_rp & (spd_score.unsqueeze(1) < spd_score.unsqueeze(0))) + + unsorted_value = torch.iinfo(torch.int32).max + front_no = torch.full((n,), unsorted_value, dtype=torch.int32, device=fit.device) + front = 0 + while int(torch.sum(front_no < unsorted_value)) < min(n_sort, n): + front += 1 + remaining = front_no == unsorted_value + dominated = (dominates & remaining.unsqueeze(1)).any(0) + current = remaining & ~dominated + if not bool(current.any()): + current = remaining + front_no = torch.where(current, front, front_no) + return front_no + + def _decomposition_selection(self, pop, fit): + """Stage 2: among the first non-dominated fronts holding at least N solutions, take for each reference + vector in turn the remaining solution with the largest cosine to it, after normalizing the remaining ones. + + The reference's level sort does not narrow this set (its early exit compares ``size(..., 1)`` of a row + of solutions, which is always 1), so it is not applied. With fewer reference vectors than ``pop_size``, + the vectors are reused in order until ``pop_size`` solutions are chosen. + """ + N = self.pop_size + rank = non_dominate_rank(fit) + sorted_rank, _ = torch.sort(rank) + candidate = rank <= sorted_rank[N - 1] + + chosen = [] + remaining = candidate.clone() + n_w = self.w.shape[0] + for t in range(N): + f_min = _masked_min(fit, remaining) + f_max = _masked_max(fit, remaining) + norm_fit = (fit - f_min) / torch.clamp(f_max - f_min, min=torch.finfo(fit.dtype).tiny) + cosine = _cosine(norm_fit, self.w[t % n_w].unsqueeze(0)).squeeze(1) + best = torch.argmax(torch.where(remaining, cosine, -torch.inf)) + chosen.append(best) + remaining = remaining.clone() + remaining[best] = False + survivor_idx = torch.stack(chosen) + + new_fit = fit[survivor_idx] + new_front = non_dominate_rank(new_fit) + # Density of each solution's region (number of selected solutions closest in cosine to the same vector). + region = torch.argmax(_cosine(_normalize(new_fit), self.w), dim=1) + density = torch.bincount(region, minlength=n_w)[region].to(new_fit.dtype) + return pop[survivor_idx], new_fit, new_front, density + + +def _normalize(fit): + f_min = torch.min(fit, dim=0)[0] + f_max = torch.max(fit, dim=0)[0] + return (fit - f_min) / torch.clamp(f_max - f_min, min=torch.finfo(fit.dtype).tiny) + + +def _masked_min(fit, mask): + return torch.min(torch.where(mask.unsqueeze(1), fit, torch.inf), dim=0)[0] + + +def _masked_max(fit, mask): + return torch.max(torch.where(mask.unsqueeze(1), fit, -torch.inf), dim=0)[0] + + +def _cosine(norm_fit, w): + """Cosine between each normalized solution and each reference vector; a zero vector has cosine 0.""" + tiny = torch.finfo(norm_fit.dtype).tiny + norm_val = torch.clamp(torch.norm(norm_fit, dim=1, keepdim=True), min=tiny) + w_norm = torch.clamp(torch.norm(w, dim=1, keepdim=True), min=tiny) + return (norm_fit @ w.T) / (norm_val @ w_norm.T) # === FIXED DEMO BLOCK === diff --git a/unit_test/algorithms/test_tsnsga2.py b/unit_test/algorithms/test_tsnsga2.py new file mode 100644 index 0000000..83ef04d --- /dev/null +++ b/unit_test/algorithms/test_tsnsga2.py @@ -0,0 +1,41 @@ +import unittest + +import torch + +from evomo.algorithms import TSNSGAII +from evomo.operators.selection import non_dominate_rank + + +class TestTSNSGAIISelection(unittest.TestCase): + def setUp(self): + torch.manual_seed(0) + self.algo = TSNSGAII(pop_size=4, n_objs=2, lb=torch.zeros(3), ub=torch.ones(3)) + + def test_spd_sort_peels_fronts(self): + fit = torch.tensor([[0.0, 1.0], [1.0, 0.0], [1.0, 1.0], [2.0, 2.0]]) + rp = torch.tensor([0, 1, 2, 3]) + front_no = self.algo._spd_sort(fit, torch.zeros(4), rp, 4) + self.assertEqual(front_no.tolist(), [1, 1, 2, 3]) + + def test_spd_sort_breaks_ties_on_a_shared_reference_vector(self): + # Mutually non-dominated, same reference vector: the smaller d1 + 5 * d2 comes first. + fit = torch.tensor([[0.0, 1.0], [1.0, 0.0]]) + front_no = self.algo._spd_sort(fit, torch.tensor([0.3, 0.1]), torch.tensor([0, 0]), 2) + self.assertEqual(front_no.tolist(), [2, 1]) + + def test_decomposition_selection_keeps_non_dominated_solutions(self): + # Five non-dominated solutions and three dominated ones: stage 2 picks pop_size = 4 non-dominated ones. + t = torch.linspace(0, 1, 5) + fit = torch.cat([torch.stack([t, 1 - t], dim=1), torch.tensor([[1.0, 1.0], [2.0, 2.0], [0.9, 0.9]])]) + pop = torch.arange(8.0).unsqueeze(1).repeat(1, 3) + new_pop, new_fit, new_front, density = self.algo._decomposition_selection(pop, fit) + self.assertEqual(new_pop.shape[0], 4) + self.assertEqual(len(set(new_pop[:, 0].tolist())), 4) + self.assertTrue((new_pop[:, 0] < 5).all()) + self.assertTrue((non_dominate_rank(new_fit) == 0).all()) + self.assertTrue(torch.equal(new_front, non_dominate_rank(new_fit))) + self.assertEqual(density.shape, (4,)) + + +if __name__ == "__main__": + unittest.main() From e1363f6cc8d3b6e8338b0c58375de3c89374f906 Mon Sep 17 00:00:00 2001 From: Zhenyu2Liang <1370096263@qq.com> Date: Fri, 30 Oct 2026 02:36:46 +0800 Subject: [PATCH 2/2] Match reference SPD elimination order and remove PR tests --- src/evomo/algorithms/ts_nsga2.py | 23 ++++++++++------ unit_test/algorithms/test_tsnsga2.py | 41 ---------------------------- 2 files changed, 15 insertions(+), 49 deletions(-) delete mode 100644 unit_test/algorithms/test_tsnsga2.py diff --git a/src/evomo/algorithms/ts_nsga2.py b/src/evomo/algorithms/ts_nsga2.py index 9818de6..e70a32d 100644 --- a/src/evomo/algorithms/ts_nsga2.py +++ b/src/evomo/algorithms/ts_nsga2.py @@ -129,10 +129,12 @@ def _spd_selection(self, pop, fit): return pop[survivor_idx], fit[survivor_idx], torch.clamp(front_no[survivor_idx], max=max_front), d2[survivor_idx] def _spd_sort(self, fit, spd_score, rp, n_sort): - """Peel SPD fronts until at least ``n_sort`` solutions are sorted. + """Assign SPD fronts in the reference's sequential elimination order. A solution SPD-dominates another when it Pareto-dominates it, or when neither Pareto-dominates the other, - they share a reference vector, and it has the smaller d1 + 5 * d2. Unsorted solutions keep the dtype maximum. + they share a reference vector, and it has the smaller d1 + 5 * d2. This relation is not transitive: + candidates eliminated during a front's scan must not eliminate later candidates. Unsorted solutions + keep the dtype maximum. """ n = fit.shape[0] better_any = (fit.unsqueeze(1) < fit.unsqueeze(0)).any(-1) @@ -143,15 +145,20 @@ def _spd_sort(self, fit, spd_score, rp, n_sort): unsorted_value = torch.iinfo(torch.int32).max front_no = torch.full((n,), unsorted_value, dtype=torch.int32, device=fit.device) + positions = torch.arange(n, device=fit.device) front = 0 while int(torch.sum(front_no < unsorted_value)) < min(n_sort, n): front += 1 - remaining = front_no == unsorted_value - dominated = (dominates & remaining.unsqueeze(1)).any(0) - current = remaining & ~dominated - if not bool(current.any()): - current = remaining - front_no = torch.where(current, front, front_no) + deleted = front_no != unsorted_value + for i in range(n): + active = ~deleted[i] + later = (positions > i) & ~deleted + # Stop at the first remaining candidate that defeats i. Comparisons before that point + # may delete candidates, but comparisons after it must have no effect. + first_defeat = torch.where(later & dominates[:, i], positions, n).amin() + deleted = deleted | (active & later & (positions < first_defeat) & dominates[i]) + deleted[i] = deleted[i] | (active & (first_defeat < n)) + front_no[i] = torch.where(~deleted[i], front, front_no[i]) return front_no def _decomposition_selection(self, pop, fit): diff --git a/unit_test/algorithms/test_tsnsga2.py b/unit_test/algorithms/test_tsnsga2.py deleted file mode 100644 index 83ef04d..0000000 --- a/unit_test/algorithms/test_tsnsga2.py +++ /dev/null @@ -1,41 +0,0 @@ -import unittest - -import torch - -from evomo.algorithms import TSNSGAII -from evomo.operators.selection import non_dominate_rank - - -class TestTSNSGAIISelection(unittest.TestCase): - def setUp(self): - torch.manual_seed(0) - self.algo = TSNSGAII(pop_size=4, n_objs=2, lb=torch.zeros(3), ub=torch.ones(3)) - - def test_spd_sort_peels_fronts(self): - fit = torch.tensor([[0.0, 1.0], [1.0, 0.0], [1.0, 1.0], [2.0, 2.0]]) - rp = torch.tensor([0, 1, 2, 3]) - front_no = self.algo._spd_sort(fit, torch.zeros(4), rp, 4) - self.assertEqual(front_no.tolist(), [1, 1, 2, 3]) - - def test_spd_sort_breaks_ties_on_a_shared_reference_vector(self): - # Mutually non-dominated, same reference vector: the smaller d1 + 5 * d2 comes first. - fit = torch.tensor([[0.0, 1.0], [1.0, 0.0]]) - front_no = self.algo._spd_sort(fit, torch.tensor([0.3, 0.1]), torch.tensor([0, 0]), 2) - self.assertEqual(front_no.tolist(), [2, 1]) - - def test_decomposition_selection_keeps_non_dominated_solutions(self): - # Five non-dominated solutions and three dominated ones: stage 2 picks pop_size = 4 non-dominated ones. - t = torch.linspace(0, 1, 5) - fit = torch.cat([torch.stack([t, 1 - t], dim=1), torch.tensor([[1.0, 1.0], [2.0, 2.0], [0.9, 0.9]])]) - pop = torch.arange(8.0).unsqueeze(1).repeat(1, 3) - new_pop, new_fit, new_front, density = self.algo._decomposition_selection(pop, fit) - self.assertEqual(new_pop.shape[0], 4) - self.assertEqual(len(set(new_pop[:, 0].tolist())), 4) - self.assertTrue((new_pop[:, 0] < 5).all()) - self.assertTrue((non_dominate_rank(new_fit) == 0).all()) - self.assertTrue(torch.equal(new_front, non_dominate_rank(new_fit))) - self.assertEqual(density.shape, (4,)) - - -if __name__ == "__main__": - unittest.main()