File size: 4,762 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
"""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)