engine/program.py

"""GP program — a small composition of Select / Reduce / Fit.

The program picks 1 or 2 feature sets (each up to MAX_SET_SIZE opaque IDs),
reduces each to a mean-score per patient, and lets the fitness function fit
a logistic regression on the resulting 1- or 2-dim state. The program
itself is just the choice of feature sets; everything else is mechanical.
"""

from __future__ import annotations

import random
from dataclasses import dataclass, field
from typing import Sequence


MAX_SETS = 2
MIN_SET_SIZE = 2
MAX_SET_SIZE = 8


@dataclass
class Program:
    feature_sets: list[list[str]]
    program_id: str = ""
    parents: list[str] = field(default_factory=list)
    born: bool = True

    @property
    def gene_ids(self) -> list[str]:
        return [g for fs in self.feature_sets for g in fs]

    @property
    def n_genes(self) -> int:
        return sum(len(fs) for fs in self.feature_sets)

    @property
    def signature(self) -> tuple:
        """Canonical signature for caching — order-invariant within / across sets."""
        return tuple(sorted(tuple(sorted(fs)) for fs in self.feature_sets))

    def program_repr(self) -> str:
        parts = [
            f"Reduce(Select(M, [{len(fs)} ids]), mean)" for fs in self.feature_sets
        ]
        return f"Fit({', '.join(parts)} -> target)"


def random_program(rng: random.Random, pool: Sequence[str]) -> Program:
    """A random program: 1 or 2 distinct-gene feature sets sampled from `pool`."""
    k = rng.choice([1, 2])
    feature_sets: list[list[str]] = []
    for _ in range(k):
        used = {g for fs in feature_sets for g in fs}
        available = [g for g in pool if g not in used]
        if len(available) < MIN_SET_SIZE:
            break
        size = rng.randint(MIN_SET_SIZE, min(MAX_SET_SIZE, len(available)))
        feature_sets.append(rng.sample(available, size))
    if not feature_sets:
        feature_sets.append(rng.sample(list(pool), MIN_SET_SIZE))
    return Program(feature_sets=feature_sets)


def _dedupe_across_sets(sets: list[list[str]]) -> list[list[str]]:
    seen: set[str] = set()
    out: list[list[str]] = []
    for s in sets:
        kept = []
        for g in s:
            if g not in seen:
                kept.append(g)
                seen.add(g)
        if kept:
            out.append(kept)
    return out


def crossover(rng: random.Random, p1: Program, p2: Program) -> Program:
    """Single-set crossover: child takes one set from each parent."""
    sets1 = [list(s) for s in p1.feature_sets]
    sets2 = [list(s) for s in p2.feature_sets]
    if rng.random() < 0.5:
        sets1, sets2 = sets2, sets1

    child: list[list[str]] = [rng.choice([sets1[0], sets2[0]])]
    if len(sets1) > 1 and len(sets2) > 1:
        child.append(rng.choice([sets1[1], sets2[1]]))
    elif len(sets1) > 1 or len(sets2) > 1:
        if rng.random() < 0.5:
            child.append(sets1[1] if len(sets1) > 1 else sets2[1])

    child = _dedupe_across_sets(child)
    if not child:
        return random_program(rng, p1.gene_ids + p2.gene_ids)
    return Program(
        feature_sets=child,
        parents=[p1.program_id, p2.program_id],
        born=True,
    )


def mutate(
    rng: random.Random,
    p: Program,
    pool: Sequence[str],
    *,
    p_mut: float = 0.7,
) -> Program:
    """One of: swap a gene, add a gene, drop a gene, add a set, drop a set."""
    if rng.random() > p_mut:
        return p
    sets = [list(s) for s in p.feature_sets]
    used = {g for fs in sets for g in fs}

    options = ["swap_gene", "add_gene", "drop_gene"]
    if len(sets) < MAX_SETS:
        options.append("add_set")
    if len(sets) > 1:
        options.append("drop_set")
    op = rng.choice(options)

    available = [g for g in pool if g not in used]

    if op == "swap_gene":
        si = rng.randrange(len(sets))
        gi = rng.randrange(len(sets[si]))
        if available:
            sets[si][gi] = rng.choice(available)
    elif op == "add_gene":
        si = rng.randrange(len(sets))
        if len(sets[si]) < MAX_SET_SIZE and available:
            sets[si].append(rng.choice(available))
    elif op == "drop_gene":
        si = rng.randrange(len(sets))
        if len(sets[si]) > MIN_SET_SIZE:
            gi = rng.randrange(len(sets[si]))
            sets[si].pop(gi)
    elif op == "add_set" and len(available) >= MIN_SET_SIZE:
        size = rng.randint(MIN_SET_SIZE, min(MAX_SET_SIZE, len(available)))
        sets.append(rng.sample(available, size))
    elif op == "drop_set" and len(sets) > 1:
        si = rng.randrange(len(sets))
        sets.pop(si)

    sets = _dedupe_across_sets(sets)
    if not sets:
        return random_program(rng, pool)
    return Program(feature_sets=sets, parents=p.parents, born=p.born)