change a lot a lot a lot!!!!!!!

algorithm/neat/species/__init__.py

@@ -0,0 +1 @@
+from .operations import update_species, create_speciate

algorithm/neat/species/distance.py

@@ -0,0 +1,73 @@
+from typing import Type
+
+from jax import Array, numpy as jnp, vmap
+
+from core import Gene
+
+
+def create_distance(gene_type: Type[Gene]):
+    def node_distance(state, nodes1: Array, nodes2: Array):
+        """
+        Calculate the distance between nodes of two genomes.
+        """
+        # statistics nodes count of two genomes
+        node_cnt1 = jnp.sum(~jnp.isnan(nodes1[:, 0]))
+        node_cnt2 = jnp.sum(~jnp.isnan(nodes2[:, 0]))
+        max_cnt = jnp.maximum(node_cnt1, node_cnt2)
+
+        # align homologous nodes
+        # this process is similar to np.intersect1d.
+        nodes = jnp.concatenate((nodes1, nodes2), axis=0)
+        keys = nodes[:, 0]
+        sorted_indices = jnp.argsort(keys, axis=0)
+        nodes = nodes[sorted_indices]
+        nodes = jnp.concatenate([nodes, jnp.full((1, nodes.shape[1]), jnp.nan)], axis=0)  # add a nan row to the end
+        fr, sr = nodes[:-1], nodes[1:]  # first row, second row
+
+        # flag location of homologous nodes
+        intersect_mask = (fr[:, 0] == sr[:, 0]) & ~jnp.isnan(nodes[:-1, 0])
+
+        # calculate the count of non_homologous of two genomes
+        non_homologous_cnt = node_cnt1 + node_cnt2 - 2 * jnp.sum(intersect_mask)
+
+        # calculate the distance of homologous nodes
+        hnd = vmap(gene_type.distance_node, in_axes=(None, 0, 0))(state, fr, sr)
+        hnd = jnp.where(jnp.isnan(hnd), 0, hnd)
+        homologous_distance = jnp.sum(hnd * intersect_mask)
+
+        val = non_homologous_cnt * state.compatibility_disjoint + homologous_distance * state.compatibility_weight
+
+        return jnp.where(max_cnt == 0, 0, val / max_cnt)  # avoid zero division
+
+    def connection_distance(state, cons1: Array, cons2: Array):
+        """
+        Calculate the distance between connections of two genomes.
+        Similar process as node_distance.
+        """
+        con_cnt1 = jnp.sum(~jnp.isnan(cons1[:, 0]))
+        con_cnt2 = jnp.sum(~jnp.isnan(cons2[:, 0]))
+        max_cnt = jnp.maximum(con_cnt1, con_cnt2)
+
+        cons = jnp.concatenate((cons1, cons2), axis=0)
+        keys = cons[:, :2]
+        sorted_indices = jnp.lexsort(keys.T[::-1])
+        cons = cons[sorted_indices]
+        cons = jnp.concatenate([cons, jnp.full((1, cons.shape[1]), jnp.nan)], axis=0)  # add a nan row to the end
+        fr, sr = cons[:-1], cons[1:]  # first row, second row
+
+        # both genome has such connection
+        intersect_mask = jnp.all(fr[:, :2] == sr[:, :2], axis=1) & ~jnp.isnan(fr[:, 0])
+
+        non_homologous_cnt = con_cnt1 + con_cnt2 - 2 * jnp.sum(intersect_mask)
+        hcd = vmap(gene_type.distance_conn, in_axes=(None, 0, 0))(state, fr, sr)
+        hcd = jnp.where(jnp.isnan(hcd), 0, hcd)
+        homologous_distance = jnp.sum(hcd * intersect_mask)
+
+        val = non_homologous_cnt * state.compatibility_disjoint + homologous_distance * state.compatibility_weight
+
+        return jnp.where(max_cnt == 0, 0, val / max_cnt)
+
+    def distance(state, genome1, genome2):
+        return node_distance(state, genome1.nodes, genome2.nodes) + connection_distance(state, genome1.conns, genome2.conns)
+
+    return distance

algorithm/neat/species/operations.py

@@ -0,0 +1,334 @@
+from typing import Type
+
+import jax
+from jax import numpy as jnp, vmap
+
+from core import Gene, Genome
+from utils import rank_elements, fetch_first
+from .distance import create_distance
+
+
+def update_species(state, randkey, fitness):
+    # update the fitness of each species
+    species_fitness = update_species_fitness(state, fitness)
+
+    # stagnation species
+    state, species_fitness = stagnation(state, species_fitness)
+
+    # sort species_info by their fitness. (push nan to the end)
+    sort_indices = jnp.argsort(species_fitness)[::-1]
+
+    center_nodes = state.center_genomes.nodes[sort_indices]
+    center_conns = state.center_genomes.conns[sort_indices]
+
+    state = state.update(
+        species_keys=state.species_keys[sort_indices],
+        best_fitness=state.best_fitness[sort_indices],
+        last_improved=state.last_improved[sort_indices],
+        member_count=state.member_count[sort_indices],
+        center_genomes=Genome(center_nodes, center_conns),
+    )
+
+    # decide the number of members of each species by their fitness
+    spawn_number = cal_spawn_numbers(state)
+
+    # crossover info
+    winner, loser, elite_mask = create_crossover_pair(state, randkey, spawn_number, fitness)
+
+    return state, winner, loser, elite_mask
+
+
+def update_species_fitness(state, fitness):
+    """
+    obtain the fitness of the species by the fitness of each individual.
+    use max criterion.
+    """
+
+    def aux_func(idx):
+        s_fitness = jnp.where(state.idx2species == state.species_keys[idx], fitness, -jnp.inf)
+        f = jnp.max(s_fitness)
+        return f
+
+    return vmap(aux_func)(jnp.arange(state.species_keys.shape[0]))
+
+
+def stagnation(state, species_fitness):
+    """
+    stagnation species.
+    those species whose fitness is not better than the best fitness of the species for a long time will be stagnation.
+    elitism species never stagnation
+    """
+
+    def aux_func(idx):
+        s_fitness = species_fitness[idx]
+        sk, bf, li = state.species_keys[idx], state.best_fitness[idx], state.last_improved[idx]
+        st = (s_fitness <= bf) & (state.generation - li > state.max_stagnation)
+        li = jnp.where(s_fitness > bf, state.generation, li)
+        bf = jnp.where(s_fitness > bf, s_fitness, bf)
+
+        return st, sk, bf, li
+
+    spe_st, species_keys, best_fitness, last_improved = vmap(aux_func)(jnp.arange(species_fitness.shape[0]))
+
+    # elite species will not be stagnation
+    species_rank = rank_elements(species_fitness)
+    spe_st = jnp.where(species_rank < state.species_elitism, False, spe_st)  # elitism never stagnation
+
+    # set stagnation species to nan
+    species_keys = jnp.where(spe_st, jnp.nan, species_keys)
+    best_fitness = jnp.where(spe_st, jnp.nan, best_fitness)
+    last_improved = jnp.where(spe_st, jnp.nan, last_improved)
+    member_count = jnp.where(spe_st, jnp.nan, state.member_count)
+    species_fitness = jnp.where(spe_st, -jnp.inf, species_fitness)
+
+    center_nodes = jnp.where(spe_st[:, None, None], jnp.nan, state.center_genomes.nodes)
+    center_conns = jnp.where(spe_st[:, None, None], jnp.nan, state.center_genomes.conns)
+
+    state = state.update(
+        species_keys=species_keys,
+        best_fitness=best_fitness,
+        last_improved=last_improved,
+        member_count=member_count,
+        center_genomes=state.center_genomes.update(center_nodes, center_conns)
+    )
+
+    return state, species_fitness
+
+
+def cal_spawn_numbers(state):
+    """
+    decide the number of members of each species by their fitness rank.
+    the species with higher fitness will have more members
+    Linear ranking selection
+        e.g. N = 3, P=10 -> probability = [0.5, 0.33, 0.17], spawn_number = [5, 3, 2]
+    """
+
+    is_species_valid = ~jnp.isnan(state.species_keys)
+    valid_species_num = jnp.sum(is_species_valid)
+    denominator = (valid_species_num + 1) * valid_species_num / 2  # obtain 3 + 2 + 1 = 6
+
+    rank_score = valid_species_num - jnp.arange(state.species_keys.shape[0])  # obtain [3, 2, 1]
+    spawn_number_rate = rank_score / denominator  # obtain [0.5, 0.33, 0.17]
+    spawn_number_rate = jnp.where(is_species_valid, spawn_number_rate, 0)  # set invalid species to 0
+
+    target_spawn_number = jnp.floor(spawn_number_rate * state.P)  # calculate member
+
+    # Avoid too much variation of numbers in a species
+    previous_size = state.member_count
+    spawn_number = previous_size + (target_spawn_number - previous_size) * state.spawn_number_change_rate
+    # jax.debug.print("previous_size: {}, spawn_number: {}", previous_size, spawn_number)
+    spawn_number = spawn_number.astype(jnp.int32)
+
+    # must control the sum of spawn_number to be equal to pop_size
+    error = state.P - jnp.sum(spawn_number)
+    spawn_number = spawn_number.at[0].add(error)  # add error to the first species to control the sum of spawn_number
+
+    return spawn_number
+
+
+def create_crossover_pair(state, randkey, spawn_number, fitness):
+    species_size = state.species_keys.shape[0]
+    pop_size = fitness.shape[0]
+    s_idx = jnp.arange(species_size)
+    p_idx = jnp.arange(pop_size)
+
+    # def aux_func(key, idx):
+    def aux_func(key, idx):
+        members = state.idx2species == state.species_keys[idx]
+        members_num = jnp.sum(members)
+
+        members_fitness = jnp.where(members, fitness, -jnp.inf)
+        sorted_member_indices = jnp.argsort(members_fitness)[::-1]
+
+        elite_size = state.genome_elitism
+        survive_size = jnp.floor(state.survival_threshold * members_num).astype(jnp.int32)
+
+        select_pro = (p_idx < survive_size) / survive_size
+        fa, ma = jax.random.choice(key, sorted_member_indices, shape=(2, pop_size), replace=True, p=select_pro)
+
+        # elite
+        fa = jnp.where(p_idx < elite_size, sorted_member_indices, fa)
+        ma = jnp.where(p_idx < elite_size, sorted_member_indices, ma)
+        elite = jnp.where(p_idx < elite_size, True, False)
+        return fa, ma, elite
+
+    fas, mas, elites = vmap(aux_func)(jax.random.split(randkey, species_size), s_idx)
+
+    spawn_number_cum = jnp.cumsum(spawn_number)
+
+    def aux_func(idx):
+        loc = jnp.argmax(idx < spawn_number_cum)
+
+        # elite genomes are at the beginning of the species
+        idx_in_species = jnp.where(loc > 0, idx - spawn_number_cum[loc - 1], idx)
+        return fas[loc, idx_in_species], mas[loc, idx_in_species], elites[loc, idx_in_species]
+
+    part1, part2, elite_mask = vmap(aux_func)(p_idx)
+
+    is_part1_win = fitness[part1] >= fitness[part2]
+    winner = jnp.where(is_part1_win, part1, part2)
+    loser = jnp.where(is_part1_win, part2, part1)
+
+    return winner, loser, elite_mask
+
+
+def create_speciate(gene_type: Type[Gene]):
+    distance = create_distance(gene_type)
+
+    def speciate(state):
+        pop_size, species_size = state.idx2species.shape[0], state.species_keys.shape[0]
+
+        # prepare distance functions
+        o2p_distance_func = vmap(distance, in_axes=(None, None, 0))  # one to population
+
+        # idx to specie key
+        idx2species = jnp.full((pop_size,), jnp.nan)  # NaN means not assigned to any species
+
+        # the distance between genomes to its center genomes
+        o2c_distances = jnp.full((pop_size,), jnp.inf)
+
+        # step 1: find new centers
+        def cond_func(carry):
+            i, i2s, cgs, o2c = carry
+
+            return (i < species_size) & (~jnp.isnan(state.species_keys[i]))  # current species is existing
+
+        def body_func(carry):
+            i, i2s, cgs, o2c = carry
+
+            distances = o2p_distance_func(state, Genome(cgs.nodes[i], cgs.conns[i]), state.pop_genomes)
+
+            # find the closest one
+            closest_idx = argmin_with_mask(distances, mask=jnp.isnan(i2s))
+            # jax.debug.print("closest_idx: {}", closest_idx)
+
+            i2s = i2s.at[closest_idx].set(state.species_keys[i])
+            cn = cgs.nodes.at[i].set(state.pop_genomes.nodes[closest_idx])
+            cc = cgs.conns.at[i].set(state.pop_genomes.conns[closest_idx])
+
+            # the genome with closest_idx will become the new center, thus its distance to center is 0.
+            o2c = o2c.at[closest_idx].set(0)
+
+            return i + 1, i2s, Genome(cn, cc), o2c
+
+        _, idx2species, center_genomes, o2c_distances = \
+            jax.lax.while_loop(cond_func, body_func, (0, idx2species, state.center_genomes, o2c_distances))
+
+        state = state.update(
+            idx2species=idx2species,
+            center_genomes=center_genomes,
+        )
+
+        # part 2: assign members to each species
+        def cond_func(carry):
+            i, i2s, cgs, sk, o2c, nsk = carry
+
+            current_species_existed = ~jnp.isnan(sk[i])
+            not_all_assigned = jnp.any(jnp.isnan(i2s))
+            not_reach_species_upper_bounds = i < species_size
+            return not_reach_species_upper_bounds & (current_species_existed | not_all_assigned)
+
+        def body_func(carry):
+            i, i2s, cgs, sk, o2c, nsk = carry
+
+            _, i2s, cgs, sk, o2c, nsk = jax.lax.cond(
+                jnp.isnan(sk[i]),  # whether the current species is existing or not
+                create_new_species,  # if not existing, create a new specie
+                update_exist_specie,  # if existing, update the specie
+                (i, i2s, cgs, sk, o2c, nsk)
+            )
+
+            return i + 1, i2s, cgs, sk, o2c, nsk
+
+        def create_new_species(carry):
+            i, i2s, cgs, sk, o2c, nsk = carry
+
+            # pick the first one who has not been assigned to any species
+            idx = fetch_first(jnp.isnan(i2s))
+
+            # assign it to the new species
+            # [key, best score, last update generation, members_count]
+            sk = sk.at[i].set(nsk)
+            i2s = i2s.at[idx].set(nsk)
+            o2c = o2c.at[idx].set(0)
+
+            # update center genomes
+            cn = cgs.nodes.at[i].set(state.pop_genomes.nodes[idx])
+            cc = cgs.conns.at[i].set(state.pop_genomes.conns[idx])
+            cgs = Genome(cn, cc)
+
+            i2s, o2c = speciate_by_threshold(i, i2s, cgs, sk, o2c)
+
+            # when a new species is created, it needs to be updated, thus do not change i
+            return i + 1, i2s, cgs, sk, o2c, nsk + 1  # change to next new speciate key
+
+        def update_exist_specie(carry):
+            i, i2s, cgs, sk, o2c, nsk = carry
+
+            i2s, o2c = speciate_by_threshold(i, i2s, cgs, sk, o2c)
+
+            # turn to next species
+            return i + 1, i2s, cgs, sk, o2c, nsk
+
+        def speciate_by_threshold(i, i2s, cgs, sk, o2c):
+            # distance between such center genome and ppo genomes
+
+            center = Genome(cgs.nodes[i], cgs.conns[i])
+            o2p_distance = o2p_distance_func(state, center, state.pop_genomes)
+            close_enough_mask = o2p_distance < state.compatibility_threshold
+
+            # when a genome is not assigned or the distance between its current center is bigger than this center
+            cacheable_mask = jnp.isnan(i2s) | (o2p_distance < o2c)
+            # jax.debug.print("{}", o2p_distance)
+            mask = close_enough_mask & cacheable_mask
+
+            # update species info
+            i2s = jnp.where(mask, sk[i], i2s)
+
+            # update distance between centers
+            o2c = jnp.where(mask, o2p_distance, o2c)
+
+            return i2s, o2c
+
+        # update idx2species
+        _, idx2species, center_genomes, species_keys, _, next_species_key = jax.lax.while_loop(
+            cond_func,
+            body_func,
+            (0, state.idx2species, state.center_genomes, state.species_keys, o2c_distances, state.next_species_key)
+        )
+
+        # if there are still some pop genomes not assigned to any species, add them to the last genome
+        # this condition can only happen when the number of species is reached species upper bounds
+        idx2species = jnp.where(jnp.isnan(idx2species), species_keys[-1], idx2species)
+
+        # complete info of species which is created in this generation
+        new_created_mask = (~jnp.isnan(species_keys)) & jnp.isnan(state.best_fitness)
+        best_fitness = jnp.where(new_created_mask, -jnp.inf, state.best_fitness)
+        last_improved = jnp.where(new_created_mask, state.generation, state.last_improved)
+
+        # update members count
+        def count_members(idx):
+            key = species_keys[idx]
+            count = jnp.sum(idx2species == key)
+            count = jnp.where(jnp.isnan(key), jnp.nan, count)
+            return count
+
+        member_count = vmap(count_members)(jnp.arange(species_size))
+
+        return state.update(
+            species_keys=species_keys,
+            best_fitness=best_fitness,
+            last_improved=last_improved,
+            members_count=member_count,
+            idx2species=idx2species,
+            center_genomes=center_genomes,
+            next_species_key=next_species_key
+        )
+
+    return speciate
+
+
+def argmin_with_mask(arr, mask):
+    masked_arr = jnp.where(mask, arr, jnp.inf)
+    min_idx = jnp.argmin(masked_arr)
+    return min_idx