zero2geospatial3d-sdk 0.1.0__py3-none-any.whl

zero2geospatial3d/nsga2.py

@@ -0,0 +1,315 @@
+"""NSGA-II multi-objective optimiser.
+
+Deb, Pratap, Agarwal and Meyarivan (2002), "A fast and elitist multiobjective
+genetic algorithm: NSGA-II", IEEE Trans. Evol. Comput. 6(2):182-197,
+doi:10.1109/4235.996017, with constraint domination, simulated binary
+crossover and polynomial mutation. All objectives are minimised. The run is
+seeded (mulberry32, bit-identical to the plugin's viewer) so it can be
+repeated exactly.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field, replace
+from typing import Any, Callable, Optional, Sequence, Union
+
+__all__ = [
+    "Individual",
+    "NSGA2Result",
+    "Problem",
+    "crowding_distance",
+    "dominates",
+    "make_rng",
+    "non_dominated_sort",
+    "nsga2",
+]
+
+_M32 = 0xFFFFFFFF
+
+
+def _imul(a: int, b: int) -> int:
+    return (a * b) & _M32
+
+
+def make_rng(seed: int = 1) -> Callable[[], float]:
+    """mulberry32: small, fast, seedable PRNG returning floats in ``[0, 1)``."""
+    state = [seed & _M32]
+
+    def rng() -> float:
+        state[0] = (state[0] + 0x6D2B79F5) & _M32
+        t = state[0]
+        t = _imul(t ^ (t >> 15), t | 1)
+        t = (t ^ ((t + _imul(t ^ (t >> 7), t | 61)) & _M32)) & _M32
+        return ((t ^ (t >> 14)) & _M32) / 4294967296
+
+    return rng
+
+
+@dataclass
+class Individual:
+    """One candidate solution."""
+
+    x: list[float] = field(default_factory=list)
+    objectives: list[float] = field(default_factory=list)
+    violation: float = 0.0
+    extra: Any = None
+    rank: int = 0
+    crowding: float = 0.0
+
+
+EvalResult = Union[dict, Sequence[float]]
+
+
+@dataclass
+class Problem:
+    """Optimisation problem.
+
+    ``evaluate(x)`` returns either a sequence of objectives or a dict with
+    ``objectives`` and optional ``violation`` (>= 0, 0 = feasible) and
+    ``extra``.
+    """
+
+    lower: Sequence[float]
+    upper: Sequence[float]
+    evaluate: Callable[[list[float]], EvalResult]
+    integer: bool = False
+    seeds: Optional[Sequence[Sequence[float]]] = None
+
+
+@dataclass
+class NSGA2Result:
+    """Final population, its first front, evaluation count and per-generation history."""
+
+    population: list[Individual]
+    front: list[Individual]
+    evaluations: int
+    history: list[dict]
+
+
+def dominates(a: Individual, b: Individual) -> bool:
+    """``a`` dominates ``b`` (constraint domination: feasibility first)."""
+    if a.violation > 0 or b.violation > 0:
+        return a.violation < b.violation
+    better = False
+    for ai, bi in zip(a.objectives, b.objectives):
+        if ai > bi:
+            return False
+        if ai < bi:
+            better = True
+    return better
+
+
+def non_dominated_sort(pop: Sequence[Individual]) -> list[list[int]]:
+    """Fast non-dominated sort: list of fronts (lists of indices); sets ``rank``."""
+    n = len(pop)
+    dominated_by: list[list[int]] = [[] for _ in range(n)]
+    count = [0] * n
+    fronts: list[list[int]] = [[]]
+    for p in range(n):
+        for q in range(n):
+            if p == q:
+                continue
+            if dominates(pop[p], pop[q]):
+                dominated_by[p].append(q)
+            elif dominates(pop[q], pop[p]):
+                count[p] += 1
+        if count[p] == 0:
+            pop[p].rank = 0
+            fronts[0].append(p)
+    i = 0
+    while fronts[i]:
+        nxt: list[int] = []
+        for p in fronts[i]:
+            for q in dominated_by[p]:
+                count[q] -= 1
+                if count[q] == 0:
+                    pop[q].rank = i + 1
+                    nxt.append(q)
+        i += 1
+        fronts.append(nxt)
+    fronts.pop()
+    return fronts
+
+
+def crowding_distance(pop: Sequence[Individual], front: Sequence[int]) -> None:
+    """Crowding distance within one front; sets ``crowding``."""
+    for i in front:
+        pop[i].crowding = 0.0
+    if len(front) <= 2:
+        for i in front:
+            pop[i].crowding = math.inf
+        return
+    m = len(pop[front[0]].objectives)
+    for k in range(m):
+        srt = sorted(front, key=lambda i: pop[i].objectives[k])  # noqa: B023
+        lo = pop[srt[0]].objectives[k]
+        hi = pop[srt[-1]].objectives[k]
+        pop[srt[0]].crowding = math.inf
+        pop[srt[-1]].crowding = math.inf
+        if hi - lo <= 0:
+            continue
+        for j in range(1, len(srt) - 1):
+            pop[srt[j]].crowding += (
+                pop[srt[j + 1]].objectives[k] - pop[srt[j - 1]].objectives[k]
+            ) / (hi - lo)
+
+
+def _better(a: Individual, b: Individual) -> bool:
+    if a.rank != b.rank:
+        return a.rank < b.rank
+    return a.crowding > b.crowding
+
+
+def _sbx(
+    rng: Callable[[], float], x1: float, x2: float, lo: float, hi: float, eta: float
+) -> tuple[float, float]:
+    if rng() > 0.5 or abs(x1 - x2) < 1e-12:
+        return x1, x2
+    u = rng()
+    if u <= 0.5:
+        beta = (2 * u) ** (1 / (eta + 1))
+    else:
+        beta = (1 / (2 * (1 - u))) ** (1 / (eta + 1))
+    c1 = 0.5 * ((1 + beta) * x1 + (1 - beta) * x2)
+    c2 = 0.5 * ((1 - beta) * x1 + (1 + beta) * x2)
+    return min(hi, max(lo, c1)), min(hi, max(lo, c2))
+
+
+def _poly_mutate(rng: Callable[[], float], x: float, lo: float, hi: float, eta: float) -> float:
+    span = hi - lo
+    if span <= 0:
+        return x
+    u = rng()
+    if u < 0.5:
+        delta = (2 * u) ** (1 / (eta + 1)) - 1
+    else:
+        delta = 1 - (2 * (1 - u)) ** (1 / (eta + 1))
+    return min(hi, max(lo, x + delta * span))
+
+
+def nsga2(
+    problem: Problem,
+    population_size: int = 24,
+    generations: int = 20,
+    seed: int = 1,
+    eta_c: float = 15,
+    eta_m: float = 20,
+    crossover_prob: float = 0.9,
+    on_generation: Optional[Callable[[dict], None]] = None,
+    should_stop: Optional[Callable[[], bool]] = None,
+) -> NSGA2Result:
+    """Run NSGA-II on ``problem`` and return the final population and front."""
+    rng = make_rng(seed)
+    lower = [float(v) for v in problem.lower]
+    upper = [float(v) for v in problem.upper]
+    n = len(lower)
+    pm = 1 / max(1, n)
+
+    def rnd(x: list[float]) -> list[float]:
+        if not problem.integer:
+            return x
+        return [min(upper[i], max(lower[i], math.floor(v + 0.5))) for i, v in enumerate(x)]
+
+    cache: dict[tuple, Individual] = {}
+    evaluations = 0
+
+    def evaluate(x: list[float]) -> Individual:
+        nonlocal evaluations
+        key = tuple(x)
+        if key in cache:
+            return replace(cache[key])
+        r = problem.evaluate(list(x))
+        evaluations += 1
+        if isinstance(r, dict):
+            objs = list(r["objectives"])
+            viol = max(0.0, float(r.get("violation") or 0.0))
+            extra = r.get("extra")
+        else:
+            objs, viol, extra = list(r), 0.0, None
+        ind = Individual(x=list(x), objectives=objs, violation=viol, extra=extra)
+        cache[key] = ind
+        return replace(ind)
+
+    pop: list[Individual] = []
+    for s in problem.seeds or []:
+        if len(pop) < population_size:
+            pop.append(evaluate(rnd([float(v) for v in s])))
+    while len(pop) < population_size:
+        x = [lo + rng() * (upper[i] - lo) for i, lo in enumerate(lower)]
+        pop.append(evaluate(rnd(x)))
+
+    def rank_all(p: list[Individual]) -> list[list[int]]:
+        fronts = non_dominated_sort(p)
+        for f in fronts:
+            crowding_distance(p, f)
+        return fronts
+
+    history: list[dict] = []
+    rank_all(pop)
+    for g in range(generations):
+        if should_stop is not None and should_stop():
+            break
+
+        def tournament(pool: list[Individual]) -> Individual:
+            a = pool[int(math.floor(rng() * len(pool)))]
+            b = pool[int(math.floor(rng() * len(pool)))]
+            return a if _better(a, b) else b
+
+        children: list[Individual] = []
+        while len(children) < population_size:
+            p1 = tournament(pop).x
+            p2 = tournament(pop).x
+            c1 = list(p1)
+            c2 = list(p2)
+            if rng() < crossover_prob:
+                for i in range(n):
+                    c1[i], c2[i] = _sbx(rng, p1[i], p2[i], lower[i], upper[i], eta_c)
+            for i in range(n):
+                if rng() < pm:
+                    c1[i] = _poly_mutate(rng, c1[i], lower[i], upper[i], eta_m)
+                if rng() < pm:
+                    c2[i] = _poly_mutate(rng, c2[i], lower[i], upper[i], eta_m)
+            children.append(evaluate(rnd(c1)))
+            if len(children) < population_size:
+                children.append(evaluate(rnd(c2)))
+        seen: set = set()
+        unique: list[Individual] = []
+        dupes: list[Individual] = []
+        for ind in pop + children:
+            key = tuple(ind.x)
+            if key in seen:
+                dupes.append(ind)
+            else:
+                seen.add(key)
+                unique.append(ind)
+        merged = (
+            unique
+            if len(unique) >= population_size
+            else unique + dupes[: population_size - len(unique)]
+        )
+        fronts = rank_all(merged)
+        nxt: list[Individual] = []
+        for f in fronts:
+            if len(nxt) + len(f) <= population_size:
+                nxt.extend(merged[i] for i in f)
+                continue
+            srt = sorted(f, key=lambda i: -merged[i].crowding)
+            nxt.extend(merged[i] for i in srt[: population_size - len(nxt)])
+            break
+        pop = nxt
+        rank_all(pop)
+        front = [p for p in pop if p.rank == 0 and p.violation == 0]
+        history.append({"generation": g + 1, "evaluations": evaluations, "front_size": len(front)})
+        if on_generation is not None:
+            on_generation(
+                {
+                    "generation": g + 1,
+                    "generations": generations,
+                    "population": pop,
+                    "front": front,
+                    "evaluations": evaluations,
+                }
+            )
+    front = [p for p in pop if p.rank == 0]
+    return NSGA2Result(population=pop, front=front, evaluations=evaluations, history=history)