add "update_by_batch" in gene;

add flatten_conns as an inverse function for unflatten_conns;
add "test_flatten.ipynb" as test for them.

tensorneat/algorithm/neat/gene/base.py

@@ -32,6 +32,9 @@ class BaseGene:
     def forward(self, state, attrs, inputs):
         raise NotImplementedError
 
+    def update_by_batch(self, state, attrs, batch_inputs):
+        raise NotImplementedError
+
     @property
     def length(self):
         return len(self.fixed_attrs) + len(self.custom_attrs)

tensorneat/algorithm/neat/gene/conn/base.py

@@ -29,3 +29,10 @@ class BaseConnGene(BaseGene):
 
     def forward(self, state, attrs, inputs):
         raise NotImplementedError
+
+    def update_by_batch(self, state, attrs, batch_inputs):
+        # default: do not update attrs, but to calculate batch_res
+        return (
+            jax.vmap(self.forward, in_axes=(None, None, 0))(state, attrs, batch_inputs),
+            attrs,
+        )

tensorneat/algorithm/neat/gene/node/base.py

@@ -18,3 +18,12 @@ class BaseNodeGene(BaseGene):
 
     def forward(self, state, attrs, inputs, is_output_node=False):
         raise NotImplementedError
+
+    def update_by_batch(self, state, attrs, batch_inputs, is_output_node=False):
+        # default: do not update attrs, but to calculate batch_res
+        return (
+            jax.vmap(self.forward, in_axes=(None, None, 0, None))(
+                state, attrs, batch_inputs, is_output_node
+            ),
+            attrs,
+        )

tensorneat/algorithm/neat/gene/node/default_without_response.py

@@ -76,11 +76,11 @@ class NodeGeneWithoutResponse(BaseNodeGene):
         )
 
         act = mutate_int(
-            k3, node[3], self.activation_indices, self.activation_replace_rate
+            k3, node[2], self.activation_indices, self.activation_replace_rate
         )
 
         agg = mutate_int(
-            k4, node[4], self.aggregation_indices, self.aggregation_replace_rate
+            k4, node[3], self.aggregation_indices, self.aggregation_replace_rate
         )
 
         return jnp.array([index, bias, act, agg])
@@ -88,8 +88,8 @@ class NodeGeneWithoutResponse(BaseNodeGene):
     def distance(self, state, node1, node2):
         return (
             jnp.abs(node1[1] - node2[1])  # bias
-            + (node1[3] != node2[3])  # activation
+            + (node1[2] != node2[2])  # activation
-            + (node1[4] != node2[4])  # aggregation
+            + (node1[3] != node2[3])  # aggregation
         )
 
     def forward(self, state, attrs, inputs, is_output_node=False):

tensorneat/algorithm/neat/gene/node/normalized.py

@@ -0,0 +1,204 @@
+from typing import Tuple
+
+import jax, jax.numpy as jnp
+
+from utils import Act, Agg, act, agg, mutate_int, mutate_float
+from . import BaseNodeGene
+
+
+class NormalizedNode(BaseNodeGene):
+    """
+    Node with normalization before activation.
+    """
+
+    # alpha and beta is used for normalization, just like BatchNorm
+    # norm: (data - mean) / (std + eps) * alpha + beta
+    custom_attrs = ["bias", "aggregation", "activation", "mean", "std", "alpha", "beta"]
+    eps = 1e-6
+
+    def __init__(
+        self,
+        bias_init_mean: float = 0.0,
+        bias_init_std: float = 1.0,
+        bias_mutate_power: float = 0.5,
+        bias_mutate_rate: float = 0.7,
+        bias_replace_rate: float = 0.1,
+        activation_default: callable = Act.sigmoid,
+        activation_options: Tuple = (Act.sigmoid,),
+        activation_replace_rate: float = 0.1,
+        aggregation_default: callable = Agg.sum,
+        aggregation_options: Tuple = (Agg.sum,),
+        aggregation_replace_rate: float = 0.1,
+        alpha_init_mean: float = 0.0,
+        alpha_init_std: float = 1.0,
+        alpha_mutate_power: float = 0.5,
+        alpha_mutate_rate: float = 0.7,
+        alpha_replace_rate: float = 0.1,
+        beta_init_mean: float = 1.0,
+        beta_init_std: float = 1.0,
+        beta_mutate_power: float = 0.5,
+        beta_mutate_rate: float = 0.7,
+        beta_replace_rate: float = 0.1,
+    ):
+        super().__init__()
+        self.bias_init_mean = bias_init_mean
+        self.bias_init_std = bias_init_std
+        self.bias_mutate_power = bias_mutate_power
+        self.bias_mutate_rate = bias_mutate_rate
+        self.bias_replace_rate = bias_replace_rate
+
+        self.activation_default = activation_options.index(activation_default)
+        self.activation_options = activation_options
+        self.activation_indices = jnp.arange(len(activation_options))
+        self.activation_replace_rate = activation_replace_rate
+
+        self.aggregation_default = aggregation_options.index(aggregation_default)
+        self.aggregation_options = aggregation_options
+        self.aggregation_indices = jnp.arange(len(aggregation_options))
+        self.aggregation_replace_rate = aggregation_replace_rate
+
+        self.alpha_init_mean = alpha_init_mean
+        self.alpha_init_std = alpha_init_std
+        self.alpha_mutate_power = alpha_mutate_power
+        self.alpha_mutate_rate = alpha_mutate_rate
+        self.alpha_replace_rate = alpha_replace_rate
+
+        self.beta_init_mean = beta_init_mean
+        self.beta_init_std = beta_init_std
+        self.beta_mutate_power = beta_mutate_power
+        self.beta_mutate_rate = beta_mutate_rate
+        self.beta_replace_rate = beta_replace_rate
+
+    def new_custom_attrs(self, state):
+        return jnp.array(
+            [
+                self.bias_init_mean,
+                self.activation_default,
+                self.aggregation_default,
+                0,  # mean
+                1,  # std
+                self.alpha_init_mean,  # alpha
+                self.beta_init_mean,  # beta
+            ]
+        )
+
+    def new_random_attrs(self, state, randkey):
+        k1, k2, k3, k4, k5, k6 = jax.random.split(randkey, num=6)
+        bias = jax.random.normal(k1, ()) * self.bias_init_std + self.bias_init_mean
+        act = jax.random.randint(k3, (), 0, len(self.activation_options))
+        agg = jax.random.randint(k4, (), 0, len(self.aggregation_options))
+        mean = 0
+        std = 1
+        alpha = jax.random.normal(k5, ()) * self.alpha_init_std + self.alpha_init_mean
+        beta = jax.random.normal(k6, ()) * self.beta_init_std + self.beta_init_mean
+
+        return jnp.array([bias, act, agg, 0, 1, alpha, beta])
+
+    def mutate(self, state, randkey, node):
+        k1, k2, k3, k4, k5, k6 = jax.random.split(state.randkey, num=6)
+        index = node[0]
+
+        bias = mutate_float(
+            k1,
+            node[1],
+            self.bias_init_mean,
+            self.bias_init_std,
+            self.bias_mutate_power,
+            self.bias_mutate_rate,
+            self.bias_replace_rate,
+        )
+
+        act = mutate_int(
+            k3, node[2], self.activation_indices, self.activation_replace_rate
+        )
+
+        agg = mutate_int(
+            k4, node[3], self.aggregation_indices, self.aggregation_replace_rate
+        )
+
+        mean = node[4]
+        std = node[5]
+
+        alpha = mutate_float(
+            k5,
+            node[6],
+            self.alpha_init_mean,
+            self.alpha_init_std,
+            self.alpha_mutate_power,
+            self.alpha_mutate_rate,
+            self.alpha_replace_rate,
+        )
+
+        beta = mutate_float(
+            k6,
+            node[7],
+            self.beta_init_mean,
+            self.beta_init_std,
+            self.beta_mutate_power,
+            self.beta_mutate_rate,
+            self.beta_replace_rate,
+        )
+
+        return jnp.array([index, bias, act, agg, mean, std, alpha, beta])
+
+    def distance(self, state, node1, node2):
+        return (
+            jnp.abs(node1[1] - node2[1])  # bias
+            + (node1[2] != node2[2])  # activation
+            + (node1[3] != node2[3])  # aggregation
+            + (node1[6] - node2[6])  # alpha
+            + (node1[7] - node2[7])  # beta
+        )
+
+    def forward(self, state, attrs, inputs, is_output_node=False):
+        """
+        post_act = (agg(inputs) + bias - mean) / std * alpha + beta
+        """
+        bias, act_idx, agg_idx, mean, std, alpha, beta = attrs
+
+        z = agg(agg_idx, inputs, self.aggregation_options)
+        z = bias + z
+        z = (z - mean) / (std + self.eps) * alpha + beta  # normalization
+
+        # the last output node should not be activated
+        z = jax.lax.cond(
+            is_output_node, lambda: z, lambda: act(act_idx, z, self.activation_options)
+        )
+
+        return z
+
+    def update_by_batch(self, state, attrs, batch_inputs, is_output_node=False):
+
+        bias, act_idx, agg_idx, mean, std, alpha, beta = attrs
+
+        batch_z = jax.vmap(agg, in_axes=(None, 0, None))(
+            agg_idx, batch_inputs, self.aggregation_options
+        )
+
+        batch_z = bias + batch_z
+
+        # calculate mean
+        valid_values_count = jnp.sum(~jnp.isnan(batch_inputs))
+        valid_values_sum = jnp.sum(jnp.where(jnp.isnan(batch_inputs), 0, batch_inputs))
+        mean = valid_values_sum / valid_values_count
+
+        # calculate std
+        std = jnp.sqrt(
+            jnp.sum(jnp.where(jnp.isnan(batch_inputs), 0, (batch_inputs - mean) ** 2))
+            / valid_values_count
+        )
+
+        batch_z = (batch_z - mean) / (std + self.eps) * alpha + beta  # normalization
+        batch_z = jax.lax.cond(
+            is_output_node,
+            lambda: batch_z,
+            lambda: jax.vmap(act, in_axes=(None, 0, None))(
+                act_idx, batch_z, self.activation_options
+            ),
+        )
+
+        # update mean and std to the attrs
+        attrs = attrs.at[3].set(mean)
+        attrs = attrs.at[4].set(std)
+
+        return batch_z, attrs

tensorneat/algorithm/neat/genome/base.py

@@ -120,3 +120,9 @@ class BaseGenome:
         conns = conns.at[: len(conn_keys), 3:].set(random_conn_attrs)
 
         return nodes, conns
+
+    def update_by_batch(self, state, batch_input, nodes, conns):
+        """
+        Update the genome by a batch of data.
+        """
+        raise NotImplementedError

tensorneat/algorithm/neat/genome/default.py

@@ -93,3 +93,6 @@ class DefaultGenome(BaseGenome):
             return vals[self.output_idx]
         else:
             return self.output_transform(vals[self.output_idx])
+
+    def update_by_batch(self, state, batch_input, nodes, conns):
+        pass

tensorneat/test/test_flatten.ipynb

@@ -0,0 +1,207 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "initial_id",
+   "metadata": {
+    "collapsed": true,
+    "ExecuteTime": {
+     "end_time": "2024-05-30T11:40:55.584592400Z",
+     "start_time": "2024-05-30T11:40:53.016051600Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "from algorithm.neat.genome import DefaultGenome\n",
+    "from utils.tools import flatten_conns, unflatten_conns\n",
+    "import jax, jax.numpy as jnp\n",
+    "from jax import vmap"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "outputs": [
+    {
+     "data": {
+      "text/plain": "((10, 5), (10, 4))"
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "genome = DefaultGenome(num_inputs=3, num_outputs=2, max_nodes=10, max_conns=10)\n",
+    "state = genome.setup()\n",
+    "key = jax.random.PRNGKey(0)\n",
+    "nodes, conns = genome.initialize(state, key)\n",
+    "nodes.shape, conns.shape"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2024-05-30T11:40:59.021858400Z",
+     "start_time": "2024-05-30T11:40:55.592593400Z"
+    }
+   },
+   "id": "89fb5cd0e77a028d"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "outputs": [
+    {
+     "data": {
+      "text/plain": "(2, 10, 10)"
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "unflatten = unflatten_conns(nodes, conns)\n",
+    "unflatten.shape"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2024-05-30T11:40:59.472701700Z",
+     "start_time": "2024-05-30T11:40:59.021858400Z"
+    }
+   },
+   "id": "aaa88227bbf29936"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "outputs": [
+    {
+     "data": {
+      "text/plain": "(Array([[ 0.        ,  5.        ,  1.        , -0.41923347],\n        [ 1.        ,  5.        ,  1.        , -3.1815007 ],\n        [ 2.        ,  5.        ,  1.        ,  0.5184341 ],\n        [ 5.        ,  3.        ,  1.        , -1.9784615 ],\n        [ 5.        ,  4.        ,  1.        ,  0.7132204 ],\n        [        nan,         nan,         nan,         nan],\n        [        nan,         nan,         nan,         nan],\n        [        nan,         nan,         nan,         nan],\n        [        nan,         nan,         nan,         nan],\n        [        nan,         nan,         nan,         nan]],      dtype=float32, weak_type=True),\n Array([[ 0.        ,  5.        ,  1.        , -0.41923347],\n        [ 1.        ,  5.        ,  1.        , -3.1815007 ],\n        [ 2.        ,  5.        ,  1.        ,  0.5184341 ],\n        [ 5.        ,  3.        ,  1.        , -1.9784615 ],\n        [ 5.        ,  4.        ,  1.        ,  0.7132204 ],\n        [        nan,         nan,         nan,         nan],\n        [        nan,         nan,         nan,         nan],\n        [        nan,         nan,         nan,         nan],\n        [        nan,         nan,         nan,         nan],\n        [        nan,         nan,         nan,         nan]],      dtype=float32))"
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# single flatten\n",
+    "flatten = flatten_conns(nodes, unflatten, C=10)\n",
+    "conns, flatten"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2024-05-30T11:41:00.308954100Z",
+     "start_time": "2024-05-30T11:40:59.469541500Z"
+    }
+   },
+   "id": "f2c65de38fdcff8f"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "outputs": [
+    {
+     "data": {
+      "text/plain": "((3, 10, 5), (3, 10, 4))"
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# batch_flatten\n",
+    "key = jax.random.PRNGKey(0)\n",
+    "keys = jax.random.split(key, 3)\n",
+    "pop_nodes, pop_conns = jax.vmap(genome.initialize, in_axes=(None, 0))(state, keys)\n",
+    "pop_nodes.shape, pop_conns.shape"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2024-05-30T11:43:09.287012800Z",
+     "start_time": "2024-05-30T11:43:09.230179800Z"
+    }
+   },
+   "id": "fe89b178b721d656"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "outputs": [
+    {
+     "data": {
+      "text/plain": "(3, 2, 10, 10)"
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "pop_unflatten = jax.vmap(unflatten_conns)(pop_nodes, pop_conns)\n",
+    "pop_unflatten.shape"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2024-05-30T11:43:10.004429100Z",
+     "start_time": "2024-05-30T11:43:09.404949800Z"
+    }
+   },
+   "id": "14bbb257e5ddeab"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "outputs": [
+    {
+     "data": {
+      "text/plain": "(3, 10, 4)"
+     },
+     "execution_count": 10,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "flatten = jax.vmap(flatten_conns, in_axes=(0, 0, None))(pop_nodes, pop_unflatten, 10)\n",
+    "flatten.shape"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2024-05-30T11:43:39.983690700Z",
+     "start_time": "2024-05-30T11:43:39.208549Z"
+    }
+   },
+   "id": "8e5cdf6140c81dc0"
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 2
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython2",
+   "version": "2.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

tensorneat/utils/tools.py

@@ -13,24 +13,55 @@ def unflatten_conns(nodes, conns):
     connection length, N means the number of nodes, C means the number of connections
     returns the un_flattened connections with shape (CL-2, N, N)
     """
-    N = nodes.shape[0]
+    N = nodes.shape[0]  # max_nodes
-    CL = conns.shape[1]
+    CL = conns.shape[1]  # connection length = (fix_attrs + custom_attrs)
     node_keys = nodes[:, 0]
     i_keys, o_keys = conns[:, 0], conns[:, 1]
+
+    def key_to_indices(key, keys):
+        return fetch_first(key == keys)
+
     i_idxs = vmap(key_to_indices, in_axes=(0, None))(i_keys, node_keys)
     o_idxs = vmap(key_to_indices, in_axes=(0, None))(o_keys, node_keys)
-    res = jnp.full((CL - 2, N, N), jnp.nan)
+    unflatten = jnp.full((CL - 2, N, N), jnp.nan)
 
     # Is interesting that jax use clip when attach data in array
     # however, it will do nothing set values in an array
     # put all attributes include enable in res
-    res = res.at[:, i_idxs, o_idxs].set(conns[:, 2:].T)
+    unflatten = unflatten.at[:, i_idxs, o_idxs].set(conns[:, 2:].T)
+    assert unflatten.shape == (CL - 2, N, N)
 
-    return res
+    return unflatten
 
 
-def key_to_indices(key, keys):
+def flatten_conns(nodes, unflatten, C):
-    return fetch_first(key == keys)
+    """
+    the inverse function of unflatten_conns
+    transform the unflatten conn (CL-2, N, N) to (C, CL)
+    """
+    N = nodes.shape[0]
+    CL = unflatten.shape[0] + 2
+    node_keys = nodes[:, 0]
+
+    def extract_conn(i, j):
+        return jnp.where(
+            jnp.isnan(unflatten[0, i, j]),
+            jnp.nan,
+            jnp.concatenate([jnp.array([node_keys[i], node_keys[j]]), unflatten[:, i, j]]),
+        )
+
+    x, y = jnp.meshgrid(jnp.arange(N), jnp.arange(N), indexing="ij")
+    conns = vmap(extract_conn)(x.flatten(), y.flatten())
+    assert conns.shape == (N * N, CL)
+
+    # put nan to the tail of the conns
+    sorted_idx = jnp.argsort(conns[:, 0])
+    sorted_conn = conns[sorted_idx]
+
+    # truncate the conns to the number of connections
+    conns = sorted_conn[:C]
+    assert conns.shape == (C, CL)
+    return conns
 
 
 @jit