finish ask part of the algorithm;

use jax.lax.while_loop in graph algorithms and forward function;
fix "enabled not care" bug in forward

neat/genome/mutate.py

@@ -1,155 +1,64 @@
-from typing import Tuple
+"""
+Mutate a genome.
+The calculation method is the same as the mutation operation in NEAT-python.
+See https://neat-python.readthedocs.io/en/latest/_modules/genome.html#DefaultGenome.mutate
+"""
+from typing import Tuple, Dict
 from functools import partial
 
 import jax
-import numpy as np
 from jax import numpy as jnp
-from jax import jit, vmap, Array
+from jax import jit, Array
 
-from .utils import fetch_random, fetch_first, I_INT, unflatten_connections
+from .utils import fetch_random, fetch_first, I_INT
 from .genome import add_node, delete_node_by_idx, delete_connection_by_idx, add_connection
 from .graph import check_cycles
 
 
-# TODO: Temporally delete single_structural_mutation, for i need to run it as soon as possible.
 @jit
-def mutate(rand_key: Array,
-           nodes: Array,
-           connections: Array,
-           new_node_key: int,
-           input_idx: Array,
-           output_idx: Array,
-           bias_mean: float = 0,
-           bias_std: float = 1,
-           bias_mutate_strength: float = 0.5,
-           bias_mutate_rate: float = 0.7,
-           bias_replace_rate: float = 0.1,
-           response_mean: float = 1.,
-           response_std: float = 0.,
-           response_mutate_strength: float = 0.,
-           response_mutate_rate: float = 0.,
-           response_replace_rate: float = 0.,
-           weight_mean: float = 0.,
-           weight_std: float = 1.,
-           weight_mutate_strength: float = 0.5,
-           weight_mutate_rate: float = 0.7,
-           weight_replace_rate: float = 0.1,
-           act_default: int = 0,
-           act_list: Array = None,
-           act_replace_rate: float = 0.1,
-           agg_default: int = 0,
-           agg_list: Array = None,
-           agg_replace_rate: float = 0.1,
-           enabled_reverse_rate: float = 0.1,
-           add_node_rate: float = 0.2,
-           delete_node_rate: float = 0.2,
-           add_connection_rate: float = 0.4,
-           delete_connection_rate: float = 0.4,
-           ):
+def mutate(rand_key: Array, nodes: Array, connections: Array, new_node_key: int, jit_config: Dict):
     """
-    :param output_idx:
-    :param input_idx:
-    :param agg_default:
-    :param act_default:
     :param rand_key:
     :param nodes: (N, 5)
     :param connections: (2, N, N)
     :param new_node_key:
-    :param bias_mean:
-    :param bias_std:
-    :param bias_mutate_strength:
-    :param bias_mutate_rate:
-    :param bias_replace_rate:
-    :param response_mean:
-    :param response_std:
-    :param response_mutate_strength:
-    :param response_mutate_rate:
-    :param response_replace_rate:
-    :param weight_mean:
-    :param weight_std:
-    :param weight_mutate_strength:
-    :param weight_mutate_rate:
-    :param weight_replace_rate:
-    :param act_list:
-    :param act_replace_rate:
-    :param agg_list:
-    :param agg_replace_rate:
-    :param enabled_reverse_rate:
-    :param add_node_rate:
-    :param delete_node_rate:
-    :param add_connection_rate:
-    :param delete_connection_rate:
+    :param jit_config:
     :return:
     """
-
-    def m_add_node(rk, n, c):
-        return mutate_add_node(rk, n, c, new_node_key, bias_mean, response_mean, act_default, agg_default)
-
-    def m_add_connection(rk, n, c):
-        return mutate_add_connection(rk, n, c, input_idx, output_idx)
-
-    def m_delete_node(rk, n, c):
-        return mutate_delete_node(rk, n, c, input_idx, output_idx)
-
-    def m_delete_connection(rk, n, c):
-        return mutate_delete_connection(rk, n, c)
-
     r1, r2, r3, r4, rand_key = jax.random.split(rand_key, 5)
 
+    # structural mutations
     # mutate add node
-    aux_nodes, aux_connections = m_add_node(r1, nodes, connections)
-    nodes = jnp.where(rand(r1) < add_node_rate, aux_nodes, nodes)
-    connections = jnp.where(rand(r1) < add_node_rate, aux_connections, connections)
+    r = rand(r1)
+    aux_nodes, aux_connections = mutate_add_node(r1, nodes, connections, new_node_key, jit_config)
+    nodes = jnp.where(r < jit_config['node_add_prob'], aux_nodes, nodes)
+    connections = jnp.where(r < jit_config['node_add_prob'], aux_connections, connections)
 
     # mutate add connection
-    aux_nodes, aux_connections = m_add_connection(r3, nodes, connections)
-    nodes = jnp.where(rand(r3) < add_connection_rate, aux_nodes, nodes)
-    connections = jnp.where(rand(r3) < add_connection_rate, aux_connections, connections)
+    r = rand(r2)
+    aux_nodes, aux_connections = mutate_add_connection(r3, nodes, connections, jit_config)
+    nodes = jnp.where(r < jit_config['conn_add_prob'], aux_nodes, nodes)
+    connections = jnp.where(r < jit_config['conn_add_prob'], aux_connections, connections)
 
     # mutate delete node
-    aux_nodes, aux_connections = m_delete_node(r2, nodes, connections)
-    nodes = jnp.where(rand(r2) < delete_node_rate, aux_nodes, nodes)
-    connections = jnp.where(rand(r2) < delete_node_rate, aux_connections, connections)
+    r = rand(r3)
+    aux_nodes, aux_connections = mutate_delete_node(r2, nodes, connections, jit_config)
+    nodes = jnp.where(r < jit_config['node_delete_prob'], aux_nodes, nodes)
+    connections = jnp.where(r < jit_config['node_delete_prob'], aux_connections, connections)
 
     # mutate delete connection
-    aux_nodes, aux_connections = m_delete_connection(r4, nodes, connections)
-    nodes = jnp.where(rand(r4) < delete_connection_rate, aux_nodes, nodes)
-    connections = jnp.where(rand(r4) < delete_connection_rate, aux_connections, connections)
+    r = rand(r4)
+    aux_nodes, aux_connections = mutate_delete_connection(r4, nodes, connections)
+    nodes = jnp.where(r < jit_config['conn_delete_prob'], aux_nodes, nodes)
+    connections = jnp.where(r < jit_config['conn_delete_prob'], aux_connections, connections)
 
-    nodes, connections = mutate_values(rand_key, nodes, connections, bias_mean, bias_std, bias_mutate_strength,
-                                       bias_mutate_rate, bias_replace_rate, response_mean, response_std,
-                                       response_mutate_strength, response_mutate_rate, response_replace_rate,
-                                       weight_mean, weight_std, weight_mutate_strength,
-                                       weight_mutate_rate, weight_replace_rate, act_list, act_replace_rate, agg_list,
-                                       agg_replace_rate, enabled_reverse_rate)
+    # value mutations
+    nodes, connections = mutate_values(rand_key, nodes, connections, jit_config)
 
     return nodes, connections
 
 
-@jit
-def mutate_values(rand_key: Array,
-                  nodes: Array,
-                  cons: Array,
-                  bias_mean: float = 0,
-                  bias_std: float = 1,
-                  bias_mutate_strength: float = 0.5,
-                  bias_mutate_rate: float = 0.7,
-                  bias_replace_rate: float = 0.1,
-                  response_mean: float = 1.,
-                  response_std: float = 0.,
-                  response_mutate_strength: float = 0.,
-                  response_mutate_rate: float = 0.,
-                  response_replace_rate: float = 0.,
-                  weight_mean: float = 0.,
-                  weight_std: float = 1.,
-                  weight_mutate_strength: float = 0.5,
-                  weight_mutate_rate: float = 0.7,
-                  weight_replace_rate: float = 0.1,
-                  act_list: Array = None,
-                  act_replace_rate: float = 0.1,
-                  agg_list: Array = None,
-                  agg_replace_rate: float = 0.1,
-                  enabled_reverse_rate: float = 0.1) -> Tuple[Array, Array]:
+def mutate_values(rand_key: Array, nodes: Array, cons: Array, jit_config: Dict) -> Tuple[Array, Array]:
     """
     Mutate values of nodes and connections.
 
@@ -157,56 +66,48 @@ def mutate_values(rand_key: Array,
         rand_key: A random key for generating random values.
         nodes: A 2D array representing nodes.
         cons: A 3D array representing connections.
-        bias_mean: Mean of the bias values.
-        bias_std: Standard deviation of the bias values.
-        bias_mutate_strength: Strength of the bias mutation.
-        bias_mutate_rate: Rate of the bias mutation.
-        bias_replace_rate: Rate of the bias replacement.
-        response_mean: Mean of the response values.
-        response_std: Standard deviation of the response values.
-        response_mutate_strength: Strength of the response mutation.
-        response_mutate_rate: Rate of the response mutation.
-        response_replace_rate: Rate of the response replacement.
-        weight_mean: Mean of the weight values.
-        weight_std: Standard deviation of the weight values.
-        weight_mutate_strength: Strength of the weight mutation.
-        weight_mutate_rate: Rate of the weight mutation.
-        weight_replace_rate: Rate of the weight replacement.
-        act_list: List of the activation function values.
-        act_replace_rate: Rate of the activation function replacement.
-        agg_list: List of the aggregation function values.
-        agg_replace_rate: Rate of the aggregation function replacement.
-        enabled_reverse_rate: Rate of reversing enabled state of connections.
+        jit_config: A dict containing configuration for jit-able functions.
 
     Returns:
         A tuple containing mutated nodes and connections.
     """
 
     k1, k2, k3, k4, k5, rand_key = jax.random.split(rand_key, num=6)
-    bias_new = mutate_float_values(k1, nodes[:, 1], bias_mean, bias_std,
-                                   bias_mutate_strength, bias_mutate_rate, bias_replace_rate)
-    response_new = mutate_float_values(k2, nodes[:, 2], response_mean, response_std,
-                                       response_mutate_strength, response_mutate_rate, response_replace_rate)
-    weight_new = mutate_float_values(k3, cons[:, 2], weight_mean, weight_std,
-                                     weight_mutate_strength, weight_mutate_rate, weight_replace_rate)
-    act_new = mutate_int_values(k4, nodes[:, 3], act_list, act_replace_rate)
-    agg_new = mutate_int_values(k5, nodes[:, 4], agg_list, agg_replace_rate)
 
-    # mutate enabled
+    # bias
+    bias_new = mutate_float_values(k1, nodes[:, 1], jit_config['bias_init_mean'], jit_config['bias_init_std'],
+                                   jit_config['bias_mutate_power'], jit_config['bias_mutate_rate'],
+                                   jit_config['bias_replace_rate'])
+
+    # response
+    response_new = mutate_float_values(k2, nodes[:, 2], jit_config['response_init_mean'],
+                                       jit_config['response_init_std'], jit_config['response_mutate_power'],
+                                       jit_config['response_mutate_rate'], jit_config['response_replace_rate'])
+
+    # weight
+    weight_new = mutate_float_values(k3, cons[:, 2], jit_config['weight_init_mean'], jit_config['weight_init_std'],
+                                     jit_config['weight_mutate_power'], jit_config['weight_mutate_rate'],
+                                     jit_config['weight_replace_rate'])
+
+    # activation
+    act_new = mutate_int_values(k4, nodes[:, 3], jit_config['activation_options'],
+                                jit_config['activation_replace_rate'])
+
+    # aggregation
+    agg_new = mutate_int_values(k5, nodes[:, 4], jit_config['aggregation_options'],
+                                jit_config['aggregation_replace_rate'])
+
+    # enabled
     r = jax.random.uniform(rand_key, cons[:, 3].shape)
-    enabled_new = jnp.where(r < enabled_reverse_rate, 1 - cons[:, 3], cons[:, 3])
-    enabled_new = jnp.where(~jnp.isnan(cons[:, 3]), enabled_new, jnp.nan)
+    enabled_new = jnp.where(r < jit_config['enable_mutate_rate'], 1 - cons[:, 3], cons[:, 3])
+
+    # merge
+    nodes = jnp.column_stack([nodes[:, 0], bias_new, response_new, act_new, agg_new])
+    cons = jnp.column_stack([cons[:, 0], cons[:, 1], weight_new, enabled_new])
 
-    nodes = nodes.at[:, 1].set(bias_new)
-    nodes = nodes.at[:, 2].set(response_new)
-    nodes = nodes.at[:, 3].set(act_new)
-    nodes = nodes.at[:, 4].set(agg_new)
-    cons = cons.at[:, 2].set(weight_new)
-    cons = cons.at[:, 3].set(enabled_new)
     return nodes, cons
 
 
-@jit
 def mutate_float_values(rand_key: Array, old_vals: Array, mean: float, std: float,
                         mutate_strength: float, mutate_rate: float, replace_rate: float) -> Array:
     """
@@ -227,19 +128,26 @@ def mutate_float_values(rand_key: Array, old_vals: Array, mean: float, std: floa
     k1, k2, k3, rand_key = jax.random.split(rand_key, num=4)
     noise = jax.random.normal(k1, old_vals.shape) * mutate_strength
     replace = jax.random.normal(k2, old_vals.shape) * std + mean
+
     r = jax.random.uniform(k3, old_vals.shape)
+
+    # default
     new_vals = old_vals
+
+    # r in [0, mutate_rate), mutate
     new_vals = jnp.where(r < mutate_rate, new_vals + noise, new_vals)
+
+    # r in [mutate_rate, mutate_rate + replace_rate), replace
     new_vals = jnp.where(
-        jnp.logical_and(mutate_rate < r, r < mutate_rate + replace_rate),
-        replace,
+        (mutate_rate < r) & (r < mutate_rate + replace_rate),
+        replace + new_vals * 0.0,  # in case of nan replace to values
         new_vals
     )
+
     new_vals = jnp.where(~jnp.isnan(old_vals), new_vals, jnp.nan)
     return new_vals
 
 
-@jit
 def mutate_int_values(rand_key: Array, old_vals: Array, val_list: Array, replace_rate: float) -> Array:
     """
     Mutate integer values (act, agg) of a given array.
@@ -256,26 +164,20 @@ def mutate_int_values(rand_key: Array, old_vals: Array, val_list: Array, replace
     k1, k2, rand_key = jax.random.split(rand_key, num=3)
     replace_val = jax.random.choice(k1, val_list, old_vals.shape)
     r = jax.random.uniform(k2, old_vals.shape)
-    new_vals = old_vals
-    new_vals = jnp.where(r < replace_rate, replace_val, new_vals)
-    new_vals = jnp.where(~jnp.isnan(old_vals), new_vals, jnp.nan)
+    new_vals = jnp.where(r < replace_rate, replace_val + old_vals * 0.0, old_vals)  # in case of nan replace to values
+
     return new_vals
 
 
-@jit
 def mutate_add_node(rand_key: Array, nodes: Array, cons: Array, new_node_key: int,
-                    default_bias: float = 0, default_response: float = 1,
-                    default_act: int = 0, default_agg: int = 0) -> Tuple[Array, Array]:
+                    jit_config: Dict) -> Tuple[Array, Array]:
     """
     Randomly add a new node from splitting a connection.
     :param rand_key:
     :param new_node_key:
     :param nodes:
     :param cons:
-    :param default_bias:
-    :param default_response:
-    :param default_act:
-    :param default_agg:
+    :param jit_config:
     :return:
     """
     # randomly choose a connection
@@ -287,12 +189,13 @@ def mutate_add_node(rand_key: Array, nodes: Array, cons: Array, new_node_key: in
     def successful_add_node():
         # disable the connection
         new_nodes, new_cons = nodes, cons
+
+        # set enable to false
         new_cons = new_cons.at[idx, 3].set(False)
 
         # add a new node
-        new_nodes, new_cons = \
-            add_node(new_nodes, new_cons, new_node_key,
-                     bias=default_bias, response=default_response, act=default_act, agg=default_agg)
+        new_nodes, new_cons = add_node(new_nodes, new_cons, new_node_key, bias=0, response=1,
+                                       act=jit_config['activation_default'], agg=jit_config['aggregation_default'])
 
         # add two new connections
         w = new_cons[idx, 2]
@@ -306,59 +209,53 @@ def mutate_add_node(rand_key: Array, nodes: Array, cons: Array, new_node_key: in
     return nodes, cons
 
 
-# TODO: Need we really need to delete a node?
-@jit
-def mutate_delete_node(rand_key: Array, nodes: Array, cons: Array,
-                       input_keys: Array, output_keys: Array) -> Tuple[Array, Array]:
+# TODO: Do we really need to delete a node?
+def mutate_delete_node(rand_key: Array, nodes: Array, cons: Array, jit_config: Dict) -> Tuple[Array, Array]:
     """
     Randomly delete a node. Input and output nodes are not allowed to be deleted.
     :param rand_key:
     :param nodes:
     :param cons:
-    :param input_keys:
-    :param output_keys:
+    :param jit_config:
     :return:
     """
     # randomly choose a node
-    node_key, node_idx = choice_node_key(rand_key, nodes, input_keys, output_keys,
-                                         allow_input_keys=False, allow_output_keys=False)
+    key, idx = choice_node_key(rand_key, nodes, jit_config['input_idx'], jit_config['output_idx'],
+                               allow_input_keys=False, allow_output_keys=False)
 
     def nothing():
         return nodes, cons
 
     def successful_delete_node():
         # delete the node
-        aux_nodes, aux_cons = delete_node_by_idx(nodes, cons, node_idx)
+        aux_nodes, aux_cons = delete_node_by_idx(nodes, cons, idx)
 
         # delete all connections
-        aux_cons = jnp.where(((aux_cons[:, 0] == node_key) | (aux_cons[:, 1] == node_key))[:, jnp.newaxis],
+        aux_cons = jnp.where(((aux_cons[:, 0] == key) | (aux_cons[:, 1] == key))[:, None],
                              jnp.nan, aux_cons)
 
         return aux_nodes, aux_cons
 
-    nodes, cons = jax.lax.cond(node_idx == I_INT, nothing, successful_delete_node)
+    nodes, cons = jax.lax.cond(idx == I_INT, nothing, successful_delete_node)
 
     return nodes, cons
 
 
-@jit
-def mutate_add_connection(rand_key: Array, nodes: Array, cons: Array,
-                          input_keys: Array, output_keys: Array) -> Tuple[Array, Array]:
+def mutate_add_connection(rand_key: Array, nodes: Array, cons: Array, jit_config: Dict) -> Tuple[Array, Array]:
     """
     Randomly add a new connection. The output node is not allowed to be an input node. If in feedforward networks,
     cycles are not allowed.
     :param rand_key:
     :param nodes:
     :param cons:
-    :param input_keys:
-    :param output_keys:
+    :param jit_config:
     :return:
     """
     # randomly choose two nodes
     k1, k2 = jax.random.split(rand_key, num=2)
-    i_key, from_idx = choice_node_key(k1, nodes, input_keys, output_keys,
+    i_key, from_idx = choice_node_key(k1, nodes, jit_config['input_idx'], jit_config['output_idx'],
                                       allow_input_keys=True, allow_output_keys=True)
-    o_key, to_idx = choice_node_key(k2, nodes, input_keys, output_keys,
+    o_key, to_idx = choice_node_key(k2, nodes, jit_config['input_idx'], jit_config['output_idx'],
                                     allow_input_keys=False, allow_output_keys=True)
 
     con_idx = fetch_first((cons[:, 0] == i_key) & (cons[:, 1] == o_key))
@@ -375,15 +272,14 @@ def mutate_add_connection(rand_key: Array, nodes: Array, cons: Array,
         return nodes, cons
 
     is_already_exist = con_idx != I_INT
-    unflattened = unflatten_connections(nodes, cons)
-    is_cycle = check_cycles(nodes, unflattened, from_idx, to_idx)
+
+    is_cycle = check_cycles(nodes, cons, from_idx, to_idx)
 
     choice = jnp.where(is_already_exist, 0, jnp.where(is_cycle, 1, 2))
     nodes, cons = jax.lax.switch(choice, [already_exist, cycle, successful])
     return nodes, cons
 
 
-@jit
 def mutate_delete_connection(rand_key: Array, nodes: Array, cons: Array):
     """
     Randomly delete a connection.
@@ -406,7 +302,6 @@ def mutate_delete_connection(rand_key: Array, nodes: Array, cons: Array):
     return nodes, cons
 
 
-@partial(jit, static_argnames=('allow_input_keys', 'allow_output_keys'))
 def choice_node_key(rand_key: Array, nodes: Array,
                     input_keys: Array, output_keys: Array,
                     allow_input_keys: bool = False, allow_output_keys: bool = False) -> Tuple[Array, Array]:
@@ -435,7 +330,6 @@ def choice_node_key(rand_key: Array, nodes: Array,
     return key, idx
 
 
-@jit
 def choice_connection_key(rand_key: Array, nodes: Array, cons: Array) -> Tuple[Array, Array, Array]:
     """
     Randomly choose a connection key from the given connections.
@@ -452,6 +346,5 @@ def choice_connection_key(rand_key: Array, nodes: Array, cons: Array) -> Tuple[A
     return i_key, o_key, idx
 
 
-@jit
 def rand(rand_key):
     return jax.random.uniform(rand_key, ())