diff --git a/src/evomo/algorithms/ts_nsga2.py b/src/evomo/algorithms/ts_nsga2.py index b464d7d..e70a32d 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,122 @@ 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): + """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. 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) + 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) + positions = torch.arange(n, device=fit.device) + front = 0 + while int(torch.sum(front_no < unsorted_value)) < min(n_sort, n): + front += 1 + 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): + """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 ===