remove create_func....

algorithm/neat/ga/operation.py

@@ -0,0 +1,40 @@
+import jax
+from jax import numpy as jnp, vmap
+
+from config import NeatConfig
+from core import Genome, State, Gene
+from .mutate import mutate
+from .crossover import crossover
+
+
+def create_next_generation(config: NeatConfig, gene: Gene, state: State, randkey, winner, loser, elite_mask):
+    # prepare random keys
+    pop_size = state.idx2species.shape[0]
+    new_node_keys = jnp.arange(pop_size) + state.next_node_key
+
+    k1, k2 = jax.random.split(randkey, 2)
+    crossover_rand_keys = jax.random.split(k1, pop_size)
+    mutate_rand_keys = jax.random.split(k2, pop_size)
+
+    # batch crossover
+    wpn, wpc = state.pop_genomes.nodes[winner], state.pop_genomes.conns[winner]
+    lpn, lpc = state.pop_genomes.nodes[loser], state.pop_genomes.conns[loser]
+    n_genomes = vmap(crossover)(crossover_rand_keys, Genome(wpn, wpc), Genome(lpn, lpc))
+
+    # batch mutation
+    mutate_func = vmap(mutate, in_axes=(None, None, None, 0, 0, 0))
+    m_n_genomes = mutate_func(config, gene, state, mutate_rand_keys, n_genomes, new_node_keys)  # mutate_new_pop_nodes
+
+    # elitism don't mutate
+    pop_nodes = jnp.where(elite_mask[:, None, None], n_genomes.nodes, m_n_genomes.nodes)
+    pop_conns = jnp.where(elite_mask[:, None, None], n_genomes.conns, m_n_genomes.conns)
+
+    # update next node key
+    all_nodes_keys = pop_nodes[:, :, 0]
+    max_node_key = jnp.max(jnp.where(jnp.isnan(all_nodes_keys), -jnp.inf, all_nodes_keys))
+    next_node_key = max_node_key + 1
+
+    return state.update(
+        pop_genomes=Genome(pop_nodes, pop_conns),
+        next_node_key=next_node_key,
+    )