change repo structure; modify readme

tensorneat/algorithm/neat/ga/__init__.py

@@ -0,0 +1,2 @@
+from .crossover import BaseCrossover, DefaultCrossover
+from .mutation import BaseMutation, DefaultMutation

tensorneat/algorithm/neat/ga/crossover/__init__.py

@@ -0,0 +1,2 @@
+from .base import BaseCrossover
+from .default import DefaultCrossover

tensorneat/algorithm/neat/ga/crossover/base.py

@@ -0,0 +1,3 @@
+class BaseCrossover:
+    def __call__(self, randkey, genome, nodes1, nodes2, conns1, conns2):
+        raise NotImplementedError

tensorneat/algorithm/neat/ga/crossover/default.py

@@ -0,0 +1,67 @@
+import jax, jax.numpy as jnp
+
+from .base import BaseCrossover
+
+class DefaultCrossover(BaseCrossover):
+
+    def __call__(self, randkey, genome, nodes1, conns1, nodes2, conns2):
+        """
+        use genome1 and genome2 to generate a new genome
+        notice that genome1 should have higher fitness than genome2 (genome1 is winner!)
+        """
+        randkey_1, randkey_2, key = jax.random.split(randkey, 3)
+
+        # crossover nodes
+        keys1, keys2 = nodes1[:, 0], nodes2[:, 0]
+        # make homologous genes align in nodes2 align with nodes1
+        nodes2 = self.align_array(keys1, keys2, nodes2, False)
+
+        # For not homologous genes, use the value of nodes1(winner)
+        # For homologous genes, use the crossover result between nodes1 and nodes2
+        new_nodes = jnp.where(jnp.isnan(nodes1) | jnp.isnan(nodes2), nodes1, self.crossover_gene(randkey_1, nodes1, nodes2))
+
+        # crossover connections
+        con_keys1, con_keys2 = conns1[:, :2], conns2[:, :2]
+        conns2 = self.align_array(con_keys1, con_keys2, conns2, True)
+
+        new_conns = jnp.where(jnp.isnan(conns1) | jnp.isnan(conns2), conns1, self.crossover_gene(randkey_2, conns1, conns2))
+
+        return new_nodes, new_conns
+
+    def align_array(self, seq1, seq2, ar2, is_conn: bool):
+        """
+        After I review this code, I found that it is the most difficult part of the code. Please never change it!
+        make ar2 align with ar1.
+        :param seq1:
+        :param seq2:
+        :param ar2:
+        :param is_conn:
+        :return:
+        align means to intersect part of ar2 will be at the same position as ar1,
+        non-intersect part of ar2 will be set to Nan
+        """
+        seq1, seq2 = seq1[:, jnp.newaxis], seq2[jnp.newaxis, :]
+        mask = (seq1 == seq2) & (~jnp.isnan(seq1))
+
+        if is_conn:
+            mask = jnp.all(mask, axis=2)
+
+        intersect_mask = mask.any(axis=1)
+        idx = jnp.arange(0, len(seq1))
+        idx_fixed = jnp.dot(mask, idx)
+
+        refactor_ar2 = jnp.where(intersect_mask[:, jnp.newaxis], ar2[idx_fixed], jnp.nan)
+
+        return refactor_ar2
+
+    def crossover_gene(self, rand_key, g1, g2):
+        """
+        crossover two genes
+        :param rand_key:
+        :param g1:
+        :param g2:
+        :return:
+        only gene with the same key will be crossover, thus don't need to consider change key
+        """
+        r = jax.random.uniform(rand_key, shape=g1.shape)
+        return jnp.where(r > 0.5, g1, g2)

tensorneat/algorithm/neat/ga/mutation/__init__.py

@@ -0,0 +1,2 @@
+from .base import BaseMutation
+from .default import DefaultMutation

tensorneat/algorithm/neat/ga/mutation/base.py

@@ -0,0 +1,3 @@
+class BaseMutation:
+    def __call__(self, key, genome, nodes, conns, new_node_key):
+        raise NotImplementedError

tensorneat/algorithm/neat/ga/mutation/default.py

@@ -0,0 +1,202 @@
+import jax, jax.numpy as jnp
+from . import BaseMutation
+from utils import fetch_first, fetch_random, I_INT, unflatten_conns, check_cycles
+
+
+class DefaultMutation(BaseMutation):
+
+    def __init__(
+            self,
+            conn_add: float = 0.4,
+            conn_delete: float = 0,
+            node_add: float = 0.2,
+            node_delete: float = 0,
+    ):
+        self.conn_add = conn_add
+        self.conn_delete = conn_delete
+        self.node_add = node_add
+        self.node_delete = node_delete
+
+    def __call__(self, randkey, genome, nodes, conns, new_node_key):
+        k1, k2 = jax.random.split(randkey)
+
+        nodes, conns = self.mutate_structure(k1, genome, nodes, conns, new_node_key)
+        nodes, conns = self.mutate_values(k2, genome, nodes, conns)
+
+        return nodes, conns
+
+    def mutate_structure(self, randkey, genome, nodes, conns, new_node_key):
+        def mutate_add_node(key_, nodes_, conns_):
+            i_key, o_key, idx = self.choice_connection_key(key_, conns_)
+
+            def successful_add_node():
+                # disable the connection
+                new_conns = conns_.at[idx, 2].set(False)
+
+                # add a new node
+                new_nodes = genome.add_node(nodes_, new_node_key, genome.node_gene.new_custom_attrs())
+
+                # add two new connections
+                new_conns = genome.add_conn(new_conns, i_key, new_node_key, True, genome.conn_gene.new_custom_attrs())
+                new_conns = genome.add_conn(new_conns, new_node_key, o_key, True, genome.conn_gene.new_custom_attrs())
+
+                return new_nodes, new_conns
+
+            return jax.lax.cond(
+                idx == I_INT,
+                lambda: (nodes_, conns_),  # do nothing
+                successful_add_node
+            )
+
+        def mutate_delete_node(key_, nodes_, conns_):
+
+            # randomly choose a node
+            key, idx = self.choice_node_key(key_, nodes_, genome.input_idx, genome.output_idx,
+                                            allow_input_keys=False, allow_output_keys=False)
+
+            def successful_delete_node():
+                # delete the node
+                new_nodes = genome.delete_node_by_pos(nodes_, idx)
+
+                # delete all connections
+                new_conns = jnp.where(
+                    ((conns_[:, 0] == key) | (conns_[:, 1] == key))[:, None],
+                    jnp.nan,
+                    conns_
+                )
+
+                return new_nodes, new_conns
+
+            return jax.lax.cond(
+                idx == I_INT,
+                lambda: (nodes_, conns_),  # do nothing
+                successful_delete_node
+            )
+
+        def mutate_add_conn(key_, nodes_, conns_):
+            # randomly choose two nodes
+            k1_, k2_ = jax.random.split(key_, num=2)
+
+            # input node of the connection can be any node
+            i_key, from_idx = self.choice_node_key(k1_, nodes_, genome.input_idx, genome.output_idx,
+                                                   allow_input_keys=True, allow_output_keys=True)
+
+            # output node of the connection can be any node except input node
+            o_key, to_idx = self.choice_node_key(k2_, nodes_, genome.input_idx, genome.output_idx,
+                                                 allow_input_keys=False, allow_output_keys=True)
+
+            conn_pos = fetch_first((conns_[:, 0] == i_key) & (conns_[:, 1] == o_key))
+            is_already_exist = conn_pos != I_INT
+
+            def nothing():
+                return nodes_, conns_
+
+            def successful():
+                return nodes_, genome.add_conn(conns_, i_key, o_key, True, genome.conn_gene.new_custom_attrs())
+
+            def already_exist():
+                return nodes_, conns_.at[conn_pos, 2].set(True)
+
+            if genome.network_type == 'feedforward':
+                u_cons = unflatten_conns(nodes_, conns_)
+                cons_exist = ~jnp.isnan(u_cons[0, :, :])
+                is_cycle = check_cycles(nodes_, cons_exist, from_idx, to_idx)
+
+                return jax.lax.cond(
+                    is_already_exist,
+                    already_exist,
+                    lambda:
+                        jax.lax.cond(
+                            is_cycle,
+                            nothing,
+                            successful
+                        )
+                )
+
+            elif genome.network_type == 'recurrent':
+                return jax.lax.cond(
+                    is_already_exist,
+                    already_exist,
+                    successful
+                )
+
+            else:
+                raise ValueError(f"Invalid network type: {genome.network_type}")
+
+        def mutate_delete_conn(key_, nodes_, conns_):
+            # randomly choose a connection
+            i_key, o_key, idx = self.choice_connection_key(key_, conns_)
+
+            def successfully_delete_connection():
+                return nodes_, genome.delete_conn_by_pos(conns_, idx)
+
+            return jax.lax.cond(
+                idx == I_INT,
+                lambda: (nodes_, conns_),  # nothing
+                successfully_delete_connection
+            )
+
+        k1, k2, k3, k4 = jax.random.split(randkey, num=4)
+        r1, r2, r3, r4 = jax.random.uniform(k1, shape=(4,))
+
+        def no(key_, nodes_, conns_):
+            return nodes_, conns_
+
+        nodes, conns = jax.lax.cond(r1 < self.node_add, mutate_add_node, no, k1, nodes, conns)
+        nodes, conns = jax.lax.cond(r2 < self.node_delete, mutate_delete_node, no, k2, nodes, conns)
+        nodes, conns = jax.lax.cond(r3 < self.conn_add, mutate_add_conn, no, k3, nodes, conns)
+        nodes, conns = jax.lax.cond(r4 < self.conn_delete, mutate_delete_conn, no, k4, nodes, conns)
+
+        return nodes, conns
+
+    def mutate_values(self, randkey, genome, nodes, conns):
+        k1, k2 = jax.random.split(randkey, num=2)
+        nodes_keys = jax.random.split(k1, num=nodes.shape[0])
+        conns_keys = jax.random.split(k2, num=conns.shape[0])
+
+        new_nodes = jax.vmap(genome.node_gene.mutate, in_axes=(0, 0))(nodes_keys, nodes)
+        new_conns = jax.vmap(genome.conn_gene.mutate, in_axes=(0, 0))(conns_keys, conns)
+
+        # nan nodes not changed
+        new_nodes = jnp.where(jnp.isnan(nodes), jnp.nan, new_nodes)
+        new_conns = jnp.where(jnp.isnan(conns), jnp.nan, new_conns)
+
+        return new_nodes, new_conns
+
+    def choice_node_key(self, rand_key, nodes, input_idx, output_idx,
+                        allow_input_keys: bool = False, allow_output_keys: bool = False):
+        """
+        Randomly choose a node key from the given nodes. It guarantees that the chosen node not be the input or output node.
+        :param rand_key:
+        :param nodes:
+        :param input_idx:
+        :param output_idx:
+        :param allow_input_keys:
+        :param allow_output_keys:
+        :return: return its key and position(idx)
+        """
+
+        node_keys = nodes[:, 0]
+        mask = ~jnp.isnan(node_keys)
+
+        if not allow_input_keys:
+            mask = jnp.logical_and(mask, ~jnp.isin(node_keys, input_idx))
+
+        if not allow_output_keys:
+            mask = jnp.logical_and(mask, ~jnp.isin(node_keys, output_idx))
+
+        idx = fetch_random(rand_key, mask)
+        key = jnp.where(idx != I_INT, nodes[idx, 0], jnp.nan)
+        return key, idx
+
+    def choice_connection_key(self, rand_key, conns):
+        """
+        Randomly choose a connection key from the given connections.
+        :return: i_key, o_key, idx
+        """
+
+        idx = fetch_random(rand_key, ~jnp.isnan(conns[:, 0]))
+        i_key = jnp.where(idx != I_INT, conns[idx, 0], jnp.nan)
+        o_key = jnp.where(idx != I_INT, conns[idx, 1], jnp.nan)
+
+        return i_key, o_key, idx