From a3bc1c36eaf774c17730700d2f30c6e38a82b394 Mon Sep 17 00:00:00 2001 From: CooperBigFoot Date: Sat, 27 Jun 2026 23:20:40 +0200 Subject: [PATCH] feat: optional evaluate_batch hook on ga()/nsga2() for whole-population evaluation --- pyproject.toml | 2 +- src/ctrl_freak/algorithms/ga.py | 37 +++++- src/ctrl_freak/algorithms/nsga2.py | 37 +++++- tests/test_evaluate_batch.py | 197 +++++++++++++++++++++++++++++ uv.lock | 2 +- 5 files changed, 270 insertions(+), 5 deletions(-) create mode 100644 tests/test_evaluate_batch.py diff --git a/pyproject.toml b/pyproject.toml index dc7dc94..ca91526 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "ctrl-freak" -version = "0.2.0" +version = "0.2.1" description = "Pure-numpy genetic algorithm framework for single-objective (GA) and multi-objective (NSGA-II) optimization" readme = "README.md" requires-python = ">=3.11" diff --git a/src/ctrl_freak/algorithms/ga.py b/src/ctrl_freak/algorithms/ga.py index 4ef04cc..6dc2e6d 100644 --- a/src/ctrl_freak/algorithms/ga.py +++ b/src/ctrl_freak/algorithms/ga.py @@ -47,6 +47,7 @@ def ga( select: str | ParentSelector = "tournament", survive: str | SurvivorSelector = "elitist", n_workers: int = 1, + evaluate_batch: Callable[[np.ndarray], np.ndarray] | None = None, ) -> GAResult: """Run a single-objective genetic algorithm. @@ -76,6 +77,14 @@ def ga( n_workers Number of workers for objective evaluation. Parallel evaluation is deterministic only when ``evaluate`` is pure. + evaluate_batch + Optional whole-population objective. When provided, it receives the + entire ``(pop_size, n_params)`` population matrix in a single call and + returns the objective for every individual with shape ``(pop_size,)`` or + ``(pop_size, 1)``. This bypasses the per-individual ``evaluate`` / ``lift`` + loop entirely, so ``evaluate`` is not called. When ``None`` (default), + evaluation falls back to the per-individual ``evaluate`` path and the + result is byte-for-byte identical to prior releases. Returns ------- @@ -109,6 +118,21 @@ def ga( ... ) >>> result.generations 2 + >>> # Whole-population evaluation via evaluate_batch (bypasses the per-individual loop): + >>> def evaluate_batch(pop): + ... return np.sum(pop**2, axis=1) + >>> batched = ga( + ... init=init, + ... evaluate=evaluate, + ... crossover=lambda p1, p2: (p1 + p2) / 2, + ... mutate=lambda x: x.copy(), + ... pop_size=10, + ... n_generations=2, + ... seed=1, + ... evaluate_batch=evaluate_batch, + ... ) + >>> batched.population.x.shape + (10, 2) """ # Validate inputs if pop_size <= 0: @@ -141,7 +165,18 @@ def evaluate_array(x: np.ndarray) -> np.ndarray: return np.asarray(evaluate(x)) # Shared lifted evaluation path. Parallel determinism assumes evaluate is pure. - lifted_evaluate = lift_parallel(evaluate_array, n_workers) if n_workers != 1 else lift(evaluate_array) + # When evaluate_batch is supplied it receives the whole (n, n_params) population + # matrix in a single call and returns (n,) or (n, 1), bypassing the per-individual + # lift / lift_parallel loop and the evaluate_array wrapper. When None, the + # per-individual path is preserved byte-for-byte. + if evaluate_batch is not None: + batch_fn = evaluate_batch + + def lifted_evaluate(x: np.ndarray) -> np.ndarray: + return np.asarray(batch_fn(x)) + + else: + lifted_evaluate = lift_parallel(evaluate_array, n_workers) if n_workers != 1 else lift(evaluate_array) # Initialize population init_x = np.stack([init(rng) for _ in range(pop_size)]) diff --git a/src/ctrl_freak/algorithms/nsga2.py b/src/ctrl_freak/algorithms/nsga2.py index 6e3297a..e239728 100644 --- a/src/ctrl_freak/algorithms/nsga2.py +++ b/src/ctrl_freak/algorithms/nsga2.py @@ -47,6 +47,7 @@ def nsga2( select: str | ParentSelector = "crowded", survive: str | SurvivorSelector = "nsga2", n_workers: int = 1, + evaluate_batch: Callable[[np.ndarray], np.ndarray] | None = None, ) -> NSGA2Result: """Run NSGA-II multi-objective optimization. @@ -75,6 +76,13 @@ def nsga2( n_workers Number of workers for objective evaluation. Parallel evaluation is deterministic only when ``evaluate`` is pure. + evaluate_batch + Optional whole-population objective. When provided, it receives the + entire ``(pop_size, n_params)`` population matrix in a single call and + returns the objective matrix of shape ``(pop_size, n_obj)``. This bypasses + the per-individual ``evaluate`` / ``lift`` loop entirely, so ``evaluate`` + is not called. When ``None`` (default), evaluation falls back to the + per-individual ``evaluate`` path and the result is unchanged. Returns ------- @@ -108,6 +116,21 @@ def nsga2( ... ) >>> result.generations 2 + >>> # Whole-population evaluation via evaluate_batch (returns (pop_size, n_obj)): + >>> def evaluate_batch(pop): + ... return np.stack([pop.sum(axis=1), (1.0 - pop).sum(axis=1)], axis=1) + >>> batched = nsga2( + ... init=init, + ... evaluate=evaluate, + ... crossover=lambda p1, p2: (p1 + p2) / 2, + ... mutate=lambda x: x.copy(), + ... pop_size=10, + ... n_generations=2, + ... seed=1, + ... evaluate_batch=evaluate_batch, + ... ) + >>> batched.population.x.shape + (10, 2) """ # Validate inputs if pop_size <= 0: @@ -124,8 +147,18 @@ def nsga2( survivor_selector = SurvivalRegistry.get(survive) if isinstance(survive, str) else survive - # Create evaluator (parallel or sequential) - lifted_evaluate = lift_parallel(evaluate, n_workers) if n_workers != 1 else lift(evaluate) + # Create evaluator. When evaluate_batch is supplied it receives the whole + # (n, n_params) population matrix in a single call and returns (n, n_obj), + # bypassing the per-individual lift / lift_parallel loop. When None, the + # per-individual path is preserved byte-for-byte. + if evaluate_batch is not None: + batch_fn = evaluate_batch + + def lifted_evaluate(x: np.ndarray) -> np.ndarray: + return np.asarray(batch_fn(x)) + + else: + lifted_evaluate = lift_parallel(evaluate, n_workers) if n_workers != 1 else lift(evaluate) # Derive independent per-phase RNG streams from the single master seed so one seed # reproduces init + parent selection + crossover + mutation bit-identically. diff --git a/tests/test_evaluate_batch.py b/tests/test_evaluate_batch.py new file mode 100644 index 0000000..1f12236 --- /dev/null +++ b/tests/test_evaluate_batch.py @@ -0,0 +1,197 @@ +"""Tests for the optional evaluate_batch hook on ga() and nsga2(). + +For BOTH algorithms these tests prove: + +1. Back-compat: evaluate_batch=None reproduces the default per-individual path. +2. Equivalence: supplying evaluate_batch (a vectorized form of the same + per-individual evaluate) yields results IDENTICAL to the per-individual path + on the same seed / pop_size. +3. Wiring: the batch callable provably receives the full (n, n_params) matrix + and the per-individual lift loop is NOT entered (the per-individual evaluate + is never called). +""" + +from typing import Any + +import numpy as np + +from ctrl_freak.algorithms.ga import ga +from ctrl_freak.algorithms.nsga2 import nsga2 + + +def _init2(rng: np.random.Generator) -> np.ndarray: + return rng.uniform(0.0, 1.0, size=2) + + +def _evaluate_ga(x: np.ndarray) -> float: + return float(np.sum(x**2)) + + +def _evaluate_nsga2(x: np.ndarray) -> np.ndarray: + return np.array([np.sum(x), np.sum((1.0 - x) ** 2)]) + + +def _crossover(p1: np.ndarray, p2: np.ndarray) -> np.ndarray: + return (p1 + p2) / 2.0 + + +def _mutate(x: np.ndarray) -> np.ndarray: + return x.copy() + + +# --------------------------------------------------------------------------- +# GA +# --------------------------------------------------------------------------- +def test_ga_evaluate_batch_matches_per_individual(): + """evaluate_batch yields ga() results IDENTICAL to the per-individual path.""" + + def evaluate_batch(pop: np.ndarray) -> np.ndarray: + # Row-wise application of the SAME per-individual evaluate. + return np.array([_evaluate_ga(row) for row in pop]) + + common: dict[str, Any] = { + "init": _init2, + "crossover": _crossover, + "mutate": _mutate, + "pop_size": 10, + "n_generations": 5, + "seed": 123, + } + reference = ga(evaluate=_evaluate_ga, **common) + batched = ga(evaluate=_evaluate_ga, evaluate_batch=evaluate_batch, **common) + + np.testing.assert_array_equal(batched.population.x, reference.population.x) + np.testing.assert_array_equal(batched.population.objectives, reference.population.objectives) + np.testing.assert_array_equal(batched.fitness, reference.fitness) + assert batched.best_idx == reference.best_idx + assert batched.generations == reference.generations + assert batched.evaluations == reference.evaluations + + +def test_ga_evaluate_batch_receives_full_matrix_and_bypasses_loop(): + """evaluate_batch sees the (pop_size, n_params) matrix; per-individual evaluate is never called.""" + pop_size = 8 + n_params = 2 + seen_shapes: list[tuple[int, ...]] = [] + + def evaluate_batch(pop: np.ndarray) -> np.ndarray: + seen_shapes.append(pop.shape) + return np.sum(pop**2, axis=1) + + def forbidden_evaluate(x: np.ndarray) -> float: + raise AssertionError("per-individual evaluate must not be called when evaluate_batch is supplied") + + result = ga( + init=_init2, + evaluate=forbidden_evaluate, + evaluate_batch=evaluate_batch, + crossover=_crossover, + mutate=_mutate, + pop_size=pop_size, + n_generations=3, + seed=7, + ) + + assert result.population.x.shape == (pop_size, n_params) + assert seen_shapes # at least the initial-population evaluation happened + for shape in seen_shapes: + assert shape == (pop_size, n_params) + + +def test_ga_evaluate_batch_none_is_unchanged(): + """evaluate_batch=None reproduces the default per-individual ga() exactly (back-compat).""" + common: dict[str, Any] = { + "init": _init2, + "evaluate": _evaluate_ga, + "crossover": _crossover, + "mutate": _mutate, + "pop_size": 10, + "n_generations": 4, + "seed": 99, + } + default = ga(**common) + explicit_none = ga(evaluate_batch=None, **common) + + np.testing.assert_array_equal(default.population.x, explicit_none.population.x) + np.testing.assert_array_equal(default.population.objectives, explicit_none.population.objectives) + np.testing.assert_array_equal(default.fitness, explicit_none.fitness) + + +# --------------------------------------------------------------------------- +# NSGA-II +# --------------------------------------------------------------------------- +def test_nsga2_evaluate_batch_matches_per_individual(): + """evaluate_batch yields nsga2() results IDENTICAL to the per-individual path.""" + + def evaluate_batch(pop: np.ndarray) -> np.ndarray: + return np.stack([_evaluate_nsga2(row) for row in pop]) + + common: dict[str, Any] = { + "init": _init2, + "crossover": _crossover, + "mutate": _mutate, + "pop_size": 10, + "n_generations": 5, + "seed": 321, + } + reference = nsga2(evaluate=_evaluate_nsga2, **common) + batched = nsga2(evaluate=_evaluate_nsga2, evaluate_batch=evaluate_batch, **common) + + np.testing.assert_array_equal(batched.population.x, reference.population.x) + np.testing.assert_array_equal(batched.population.objectives, reference.population.objectives) + np.testing.assert_array_equal(batched.rank, reference.rank) + np.testing.assert_array_equal(batched.crowding_distance, reference.crowding_distance) + assert batched.generations == reference.generations + assert batched.evaluations == reference.evaluations + + +def test_nsga2_evaluate_batch_receives_full_matrix_and_bypasses_loop(): + """evaluate_batch sees the (pop_size, n_params) matrix; per-individual evaluate is never called.""" + pop_size = 8 + n_params = 2 + seen_shapes: list[tuple[int, ...]] = [] + + def evaluate_batch(pop: np.ndarray) -> np.ndarray: + seen_shapes.append(pop.shape) + return np.stack([pop.sum(axis=1), (1.0 - pop).sum(axis=1)], axis=1) + + def forbidden_evaluate(x: np.ndarray) -> np.ndarray: + raise AssertionError("per-individual evaluate must not be called when evaluate_batch is supplied") + + result = nsga2( + init=_init2, + evaluate=forbidden_evaluate, + evaluate_batch=evaluate_batch, + crossover=_crossover, + mutate=_mutate, + pop_size=pop_size, + n_generations=3, + seed=7, + ) + + assert result.population.x.shape == (pop_size, n_params) + assert result.population.objectives is not None + assert result.population.objectives.shape == (pop_size, 2) + assert seen_shapes + for shape in seen_shapes: + assert shape == (pop_size, n_params) + + +def test_nsga2_evaluate_batch_none_is_unchanged(): + """evaluate_batch=None reproduces the default per-individual nsga2() exactly (back-compat).""" + common: dict[str, Any] = { + "init": _init2, + "evaluate": _evaluate_nsga2, + "crossover": _crossover, + "mutate": _mutate, + "pop_size": 10, + "n_generations": 4, + "seed": 99, + } + default = nsga2(**common) + explicit_none = nsga2(evaluate_batch=None, **common) + + np.testing.assert_array_equal(default.population.x, explicit_none.population.x) + np.testing.assert_array_equal(default.population.objectives, explicit_none.population.objectives) + np.testing.assert_array_equal(default.rank, explicit_none.rank) + np.testing.assert_array_equal(default.crowding_distance, explicit_none.crowding_distance) diff --git a/uv.lock b/uv.lock index d02b066..24c4468 100644 --- a/uv.lock +++ b/uv.lock @@ -435,7 +435,7 @@ toml = [ [[package]] name = "ctrl-freak" -version = "0.2.0" +version = "0.2.1" source = { editable = "." } dependencies = [ { name = "joblib" },