feat: optional evaluate_batch hook on ga()/nsga2() for whole-population evaluation

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"

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)])

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.

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)