Genetic Algorithms & Stochastic Search

"""A genetic algorithm on 0/1 knapsack - every operator written out.

Population, fitness, tournament selection, one-point crossover, bit-flip
mutation, elitism. The only thing the search ever asks of the problem is
fitness(genome) -> number: no gradient, no structure, just a ranking.
"""

import random


def fitness(genome, values, weights, capacity):
    total_value = 0
    total_weight = 0
    for j in range(len(genome)):
        if genome[j] == 1:
            total_value += values[j]
            total_weight += weights[j]

    if total_weight > capacity:
        # infeasible: rank by how far over the limit, so any feasible genome wins
        return -(total_weight - capacity)
    return total_value


def hamming(a, b):
    return sum(1 for j in range(len(a)) if a[j] != b[j])


def diversity(population):
    """Mean pairwise Hamming distance: it collapsing is premature convergence."""
    total = 0
    pairs = 0
    for i in range(len(population)):
        for k in range(i + 1, len(population)):
            total += hamming(population[i], population[k])
            pairs += 1
    return total / pairs


def tournament(population, scores, size):
    best = random.randrange(len(population))
    for _ in range(size - 1):
        challenger = random.randrange(len(population))
        if scores[challenger] > scores[best]:
            best = challenger
    return population[best]


def evolve(values, weights, capacity, pop_size=60, generations=200,
           tournament_size=3, crossover_rate=0.9, elite=2):
    length = len(values)
    mutation_rate = 1.0 / length          # about one flipped bit per child
    population = [[random.randint(0, 1) for _ in range(length)]
                  for _ in range(pop_size)]
    history = []

    for generation in range(generations):
        scores = [fitness(g, values, weights, capacity) for g in population]
        history.append((max(scores), diversity(population)))

        # elitism: the best genomes pass to the next generation unchanged
        ranked = sorted(range(pop_size), key=lambda i: scores[i], reverse=True)
        next_population = [population[i][:] for i in ranked[:elite]]

        while len(next_population) < pop_size:
            mother = tournament(population, scores, tournament_size)
            father = tournament(population, scores, tournament_size)

            if random.random() < crossover_rate:
                cut = random.randint(1, length - 1)      # one-point crossover
                child = mother[:cut] + father[cut:]
            else:
                child = mother[:]

            for j in range(length):
                if random.random() < mutation_rate:
                    child[j] = 1 - child[j]              # bit flip
            next_population.append(child)

        population = next_population

    scores = [fitness(g, values, weights, capacity) for g in population]
    best = max(range(pop_size), key=lambda i: scores[i])
    return population[best], scores[best], history

"""Genetic algorithm - typed, seeded, validated, bit-identical on rerun."""

from dataclasses import dataclass
from typing import Callable

import numpy as np
from numpy.typing import NDArray

Genome = NDArray[np.bool_]           # shape (L,)
Population = NDArray[np.bool_]       # shape (N, L)
Fitness = Callable[[Genome], float]  # the only thing the problem supplies


@dataclass(frozen=True)
class GAConfig:
    pop_size: int = 100
    generations: int = 200
    tournament_size: int = 3
    crossover_rate: float = 0.9
    mutation_rate: float | None = None   # None means 1/L
    elite: int = 2
    seed: int = 0

    def __post_init__(self) -> None:
        if self.pop_size < 2:
            raise ValueError(f"pop_size must be at least 2, got {self.pop_size}")
        if not 0 <= self.elite < self.pop_size:
            raise ValueError(f"elite must be in [0, pop_size), got {self.elite}")
        if self.tournament_size < 1:
            raise ValueError("tournament_size must be at least 1")
        for name in ("crossover_rate", "mutation_rate"):
            rate = getattr(self, name)
            if rate is not None and not 0.0 <= rate <= 1.0:
                raise ValueError(f"{name} must be in [0, 1], got {rate}")


@dataclass(frozen=True)
class GAResult:
    best_genome: Genome
    best_fitness: float
    best_history: NDArray[np.float64]
    diversity_history: NDArray[np.float64]
    evaluations: int


def mean_pairwise_hamming(population: Population) -> float:
    """Mean Hamming distance over all pairs in O(N*L), not O(N^2 * L).

    Per locus, k ones and (N - k) zeros contribute k * (N - k) differing pairs.
    """
    n = population.shape[0]
    ones = population.sum(axis=0, dtype=np.int64)
    return float((ones * (n - ones)).sum()) / (n * (n - 1) / 2)


def _tournament(scores: NDArray[np.float64], size: int, count: int,
                rng: np.random.Generator) -> NDArray[np.intp]:
    contenders = rng.integers(0, scores.size, size=(count, size))
    best_slot = scores[contenders].argmax(axis=1)
    return contenders[np.arange(count), best_slot]


def _crossover(mothers: Population, fathers: Population, rate: float,
               rng: np.random.Generator) -> Population:
    count, length = mothers.shape
    cut = rng.integers(1, length, size=(count, 1))
    children = np.where(np.arange(length) < cut, mothers, fathers)
    keep_mother = rng.random(count) >= rate
    children[keep_mother] = mothers[keep_mother]
    return children


def evolve(fitness: Fitness, genome_length: int, config: GAConfig) -> GAResult:
    """Maximise fitness over bit strings. Same config and seed, same answer."""
    if genome_length < 2:
        raise ValueError(f"genome_length must be at least 2, got {genome_length}")

    rng = np.random.default_rng(config.seed)   # explicit stream, never global
    rate = config.mutation_rate if config.mutation_rate is not None else 1.0 / genome_length
    n_children = config.pop_size - config.elite

    population: Population = rng.random((config.pop_size, genome_length)) < 0.5
    scores = np.array([fitness(genome) for genome in population], dtype=np.float64)
    evaluations = config.pop_size
    best_history = np.empty(config.generations)
    diversity_history = np.empty(config.generations)

    for generation in range(config.generations):
        best_history[generation] = scores.max()
        diversity_history[generation] = mean_pairwise_hamming(population)

        elite_rows = np.argsort(scores)[::-1][: config.elite]
        mothers = population[_tournament(scores, config.tournament_size, n_children, rng)]
        fathers = population[_tournament(scores, config.tournament_size, n_children, rng)]
        children = _crossover(mothers, fathers, config.crossover_rate, rng)
        children ^= rng.random(children.shape) < rate           # bit-flip mutation

        # Elite scores are carried over rather than recomputed. That is only
        # sound for a deterministic fitness; re-score elites if it is noisy.
        child_scores = np.array([fitness(genome) for genome in children])
        evaluations += n_children
        population = np.concatenate([population[elite_rows], children])
        scores = np.concatenate([scores[elite_rows], child_scores])

    best = int(scores.argmax())
    return GAResult(population[best].copy(), float(scores[best]),
                    best_history, diversity_history, evaluations)

"""Genetic algorithm - whole-generation array ops, preallocated, one pass per score."""

import numpy as np
from numpy.typing import NDArray

Matrix = NDArray[np.float32]


def knapsack_scores(population: Matrix, value_weight: Matrix, capacity: float) -> NDArray[np.float32]:
    """Score every genome in one matrix product.

    value_weight is (L, 2): column 0 values, column 1 weights. One gemm reads
    the population once and yields both totals for every genome.
    """
    totals = population @ value_weight
    overweight = np.maximum(totals[:, 1] - capacity, 0.0)
    return np.where(overweight > 0.0, -overweight, totals[:, 0])


def evolve_fast(values: Matrix, weights: Matrix, capacity: float, *, pop_size: int,
                generations: int, tournament_size: int, crossover_rate: float,
                elite: int, seed: int) -> tuple[NDArray[np.float32], float]:
    rng = np.random.default_rng(seed)
    length = values.size
    n_children = pop_size - elite
    rate = 1.0 / length
    value_weight = np.ascontiguousarray(np.stack([values, weights], axis=1), dtype=np.float32)
    columns = np.arange(length)[None, :]

    # float32 0/1 genomes: the matmul needs no cast copy on every generation.
    current = (rng.random((pop_size, length), dtype=np.float32) < 0.5).astype(np.float32)
    following = np.empty_like(current)
    scores = knapsack_scores(current, value_weight, capacity)

    for _ in range(generations):
        top = np.argpartition(scores, -elite)[-elite:]            # O(N), not a full sort
        elite_scores = scores[top].copy()
        following[:elite] = current[top]

        # one RNG call draws every contender for both parents of every child
        contenders = rng.integers(0, pop_size, size=(2, n_children, tournament_size))
        winners = np.take_along_axis(
            contenders, scores[contenders].argmax(axis=2)[..., None], axis=2
        )[..., 0]

        children = following[elite:]                              # a view, not a copy
        children[...] = current[winners[1]]                       # father everywhere
        cut = rng.integers(1, length, size=(n_children, 1))
        np.copyto(children, current[winners[0]], where=columns < cut)           # mother's prefix
        np.copyto(children, current[winners[0]],
                  where=(rng.random(n_children) >= crossover_rate)[:, None])    # no crossover
        flips = rng.random(children.shape, dtype=np.float32) < rate
        np.subtract(1.0, children, out=children, where=flips)     # in-place bit flip

        scores[elite:] = knapsack_scores(children, value_weight, capacity)
        scores[:elite] = elite_scores
        current, following = following, current                   # swap buffers, no allocation

    best = int(scores.argmax())
    return current[best].astype(np.uint8), float(scores[best])