File size: 5,005 Bytes
0fff343
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
"""GP main loop with evolution logging.

Seeded RNG, tournament selection, elitism, crossover, mutation. We log the
top-K candidates per generation (not the full population) so the persisted
artefact stays small and the Streamlit replay can scrub instantly. Fitness
results are cached by program signature so elites carrying 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.fitness import fitness_fn
from engine.objectives import BinaryAUROCObjective, Objective
from engine.program import Program, crossover, mutate, random_program


def _assign_id(program: Program, generation: int, idx: int) -> None:
    program.program_id = f"g{generation}c{idx}"


def _tournament_select(
    rng: random.Random,
    population: list[Program],
    fitnesses: list[float],
    k: int,
) -> Program:
    contenders = rng.sample(range(len(population)), k)
    winner = max(contenders, key=lambda i: fitnesses[i])
    return population[winner]


def run_gp(
    M_train: pd.DataFrame,
    y_train: np.ndarray,
    pool: Sequence[str],
    *,
    objective: Objective | None = None,
    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,
    log_top_k: int = 8,
    on_generation: Callable[[dict], None] | None = None,
) -> tuple[list[dict], Program, float]:
    """Run the GP loop. Returns (per-generation log, best program, best fitness).

    ``on_generation`` is invoked once per generation with the just-built log
    entry; the API layer uses it to stream generation events over SSE.
    """
    obj = objective or BinaryAUROCObjective()
    py_rng = random.Random(seed)

    population: list[Program] = [
        random_program(py_rng, pool) for _ in range(population_size)
    ]
    for i, p in enumerate(population):
        _assign_id(p, 0, i)
        p.born = True
        p.parents = []

    fitness_cache: dict[tuple, float] = {}

    def evaluate(pop: list[Program]) -> list[float]:
        out = []
        for prog in pop:
            sig = prog.signature
            if sig not in fitness_cache:
                fitness_cache[sig] = fitness_fn(
                    M_train, y_train, prog,
                    objective=obj,
                    lambda_size=lambda_size,
                    n_folds=cv_folds,
                    random_state=seed,
                )
            out.append(fitness_cache[sig])
        return out

    log: list[dict] = []
    best_overall: Program | None = None
    best_overall_fit = -np.inf

    for gen in range(n_generations):
        fitnesses = evaluate(population)

        ranked = sorted(
            zip(population, fitnesses),
            key=lambda kv: kv[1],
            reverse=True,
        )
        top = ranked[:log_top_k]
        gen_entry = {
            "generation": gen,
            "best_fitness": float(top[0][1]),
            "median_fitness": float(np.median(fitnesses)),
            "elitism": elitism,
            "candidates": [
                {
                    "id": prog.program_id,
                    "gene_ids": prog.gene_ids,
                    "feature_sets": [list(s) for s in prog.feature_sets],
                    "program_repr": prog.program_repr(),
                    "fitness": float(fit),
                    "n_genes": prog.n_genes,
                    "parents": list(prog.parents),
                    "born": bool(prog.born),
                    "survived": rank < elitism,
                }
                for rank, (prog, fit) in enumerate(top)
            ],
        }
        log.append(gen_entry)
        if on_generation is not None:
            on_generation(gen_entry)

        if top[0][1] > best_overall_fit:
            best_overall = top[0][0]
            best_overall_fit = top[0][1]

        if gen == n_generations - 1:
            break

        elites = [prog for prog, _ in ranked[:elitism]]
        new_pop: list[Program] = []
        for ei, e in enumerate(elites):
            carried = Program(
                feature_sets=[list(s) for s in e.feature_sets],
                parents=[e.program_id],
                born=False,
            )
            _assign_id(carried, gen + 1, ei)
            new_pop.append(carried)

        i = elitism
        while len(new_pop) < population_size:
            p1 = _tournament_select(py_rng, population, fitnesses, tournament_k)
            p2 = _tournament_select(py_rng, population, fitnesses, tournament_k)
            child = crossover(py_rng, p1, p2)
            child = mutate(py_rng, child, pool, p_mut=p_mutate)
            _assign_id(child, gen + 1, i)
            child.born = True
            new_pop.append(child)
            i += 1
        population = new_pop

    assert best_overall is not None
    return log, best_overall, float(best_overall_fit)