"""Typed GP loop for engine_v2. Init via ramped half-and-half. Tournament selection. Elitism. Subtree crossover with strict type matching. Subtree + point mutation. A single seed governs all RNG (init, selection, crossover, mutation, FeatureSet sampling). The fitness cache is keyed by ``program.repr_typed()`` so elites that carry across generations don't recompute. """ from __future__ import annotations import random from typing import Callable, Sequence import numpy as np import pandas as pd from engine_v2.fitness import V2Objective, fitness_fn from engine_v2.nodes import ExecContext, Node from engine_v2.synthesize import ( crossover, mutate, ramped_population, ) def _assign_id(program: Node, generation: int, idx: int) -> str: return f"g{generation}c{idx}" def _tournament_select( rng: random.Random, population: list[Node], fitnesses: list[float], k: int, ) -> int: contenders = rng.sample(range(len(population)), k) return max(contenders, key=lambda i: fitnesses[i]) def run_gp_v2( ctx_train: ExecContext, y_train: np.ndarray | None, pool: Sequence[str], *, objective: V2Objective, population_size: int = 150, n_generations: int = 30, tournament_k: int = 3, elitism: int = 5, p_mutate: float = 0.7, lambda_size: float = 0.005, cv_folds: int = 5, seed: int = 42, max_depth: int = 4, max_genes_per_set: int = 8, max_nodes: int = 64, coherence_weight: float = 0.0, immigrant_fraction: float = 0.0, rates_override: dict | None = None, scalar_share_override: float | None = None, on_generation: Callable[[dict], None] | None = None, ) -> tuple[list[dict], Node, float]: """Returns ``(generation_log, winner, winner_fitness)``. Each ``generation_log`` entry is a complete snapshot of the population: every candidate with its id, program_repr, fitness, survived flag, n_nodes, and depth. Top-12 are also exposed inline via ``top_candidates`` for the SSE stream's bandwidth budget. """ py_rng = random.Random(seed) # Per-objective synthesis overrides (e.g. UNSUP forbids Scalar roots # and Associate / Effect / FitApply — Vector-only programs). The # caller may also pass `rates_override` (e.g. to gate Search off # per-run); those keys win over the objective's overrides which # win over DEFAULT_RATES inside synthesize. ``scalar_share`` lives # on its own param (it's a single float, not a per-operator rate); # an explicit `scalar_share_override` wins over the objective's # value. overrides = objective.synthesis_overrides() if hasattr(objective, "synthesis_overrides") else {} objective_rates = overrides.get("rates") or {} objective_scalar_share = overrides.get("scalar_share") effective_scalar_share = ( scalar_share_override if scalar_share_override is not None else objective_scalar_share ) merged_rates_override: dict | None if rates_override is None and not objective_rates: merged_rates_override = None else: merged_rates_override = { **(objective_rates or {}), **(rates_override or {}), } population: list[Node] = ramped_population( py_rng, pool, n=population_size, objective_target=objective.target, max_depth=max_depth, max_genes_per_set=max_genes_per_set, **( {"rates": merged_rates_override} if merged_rates_override is not None else {} ), **( {"scalar_share": effective_scalar_share} if effective_scalar_share is not None else {} ), ) ids: list[str] = [_assign_id(p, 0, i) for i, p in enumerate(population)] parents: list[list[str]] = [[] for _ in population] fitness_cache: dict[str, float] = {} def evaluate(pop: list[Node]) -> list[float]: out = [] for prog in pop: sig = prog.repr_typed() if sig not in fitness_cache: fitness_cache[sig] = fitness_fn( prog, ctx_train, y_train, objective=objective, lambda_size=lambda_size, n_folds=cv_folds, random_state=seed, coherence_weight=coherence_weight, ) out.append(fitness_cache[sig]) return out worst = objective.worst_score() def _finite(x: float) -> float: return float(x) if np.isfinite(x) else worst log: list[dict] = [] best_overall: Node | None = None best_overall_fit = -np.inf for gen in range(n_generations): fitnesses = evaluate(population) ranked = sorted( range(len(population)), key=lambda i: fitnesses[i], reverse=True, ) elite_ix = set(ranked[:elitism]) candidates_payload = [ { "id": ids[i], "program_repr": population[i].repr_typed(), "fitness": _finite(fitnesses[i]), "n_nodes": int(population[i].node_count()), "depth": int(population[i].depth()), "gene_ids": list(population[i].feature_ids()), "parents": list(parents[i]), "survived": i in elite_ix, } for i in ranked ] top_payload = candidates_payload[:12] # Median computed over the finite values only, so a generation # where many programs are degenerate doesn't collapse to -inf. finite_vals = [f for f in fitnesses if np.isfinite(f)] median_fit = float(np.median(finite_vals)) if finite_vals else worst entry = { "generation": gen, "best_fitness": _finite(fitnesses[ranked[0]]), "median_fitness": median_fit, "elitism": elitism, "population_size": len(population), "top_candidates": top_payload, "candidates": candidates_payload, } log.append(entry) if on_generation is not None: on_generation(entry) if fitnesses[ranked[0]] > best_overall_fit: best_overall_fit = fitnesses[ranked[0]] best_overall = population[ranked[0]] if gen == n_generations - 1: break # ----- Build next generation ----- new_pop: list[Node] = [] new_ids: list[str] = [] new_parents: list[list[str]] = [] # Elites carry over unchanged. for ei, idx in enumerate(ranked[:elitism]): new_pop.append(population[idx]) # share the node — it's evaluated. new_ids.append(_assign_id(population[idx], gen + 1, ei)) new_parents.append([ids[idx]]) i_offset = elitism # Random immigrants — fresh programs drawn from ramped_population # using the same grammar / objective / depth constraints as init. # They displace offspring slots (never elites), counter premature # convergence, and are gated to immigrant_fraction > 0 so existing # runs are byte-for-byte unchanged. Rounding: round(frac * pop), # capped so it never pushes past the budget after elites. n_immigrants = ( int(round(immigrant_fraction * population_size)) if immigrant_fraction > 0.0 else 0 ) n_immigrants = max(0, min(n_immigrants, population_size - len(new_pop))) if n_immigrants > 0: fresh = ramped_population( py_rng, pool, n=n_immigrants, objective_target=objective.target, max_depth=max_depth, max_genes_per_set=max_genes_per_set, **( {"rates": merged_rates_override} if merged_rates_override is not None else {} ), **( {"scalar_share": effective_scalar_share} if effective_scalar_share is not None else {} ), ) for child in fresh: new_pop.append(child) new_ids.append(_assign_id(child, gen + 1, i_offset)) new_parents.append([]) # immigrant — no parents i_offset += 1 while len(new_pop) < population_size: i1 = _tournament_select(py_rng, population, fitnesses, tournament_k) i2 = _tournament_select(py_rng, population, fitnesses, tournament_k) child = crossover( py_rng, population[i1], population[i2], max_depth=max_depth, max_nodes=max_nodes, ) child = mutate( py_rng, child, pool, objective_target=objective.target, p_mut=p_mutate, max_depth=max_depth, max_genes_per_set=max_genes_per_set, max_nodes=max_nodes, **( {"rates": merged_rates_override} if merged_rates_override is not None else {} ), ) new_pop.append(child) new_ids.append(_assign_id(child, gen + 1, i_offset)) new_parents.append([ids[i1], ids[i2]]) i_offset += 1 population, ids, parents = new_pop, new_ids, new_parents assert best_overall is not None return log, best_overall, float(best_overall_fit)