odify genome for the official release
This commit is contained in:
root
@@ -1,5 +1,4 @@
|
||||
from .base import BaseGenome
|
||||
from .default import DefaultGenome
|
||||
from .recurrent import RecurrentGenome
|
||||
from .hidden import HiddenInitialize
|
||||
from .dense import DenseInitialize
|
||||
|
||||
|
||||
@@ -1,8 +1,16 @@
|
||||
from typing import Callable, Sequence
|
||||
|
||||
import numpy as np
|
||||
import jax, jax.numpy as jnp
|
||||
import jax
|
||||
from jax import vmap, numpy as jnp
|
||||
from ..gene import BaseNodeGene, BaseConnGene
|
||||
from ..ga import BaseMutation, BaseCrossover
|
||||
from utils import State, StatefulBaseClass, topological_sort_python, hash_array
|
||||
from .operations import BaseMutation, BaseCrossover, BaseDistance
|
||||
from tensorneat.common import (
|
||||
State,
|
||||
StatefulBaseClass,
|
||||
hash_array,
|
||||
)
|
||||
from .utils import valid_cnt
|
||||
|
||||
|
||||
class BaseGenome(StatefulBaseClass):
|
||||
@@ -18,120 +26,159 @@ class BaseGenome(StatefulBaseClass):
|
||||
conn_gene: BaseConnGene,
|
||||
mutation: BaseMutation,
|
||||
crossover: BaseCrossover,
|
||||
distance: BaseDistance,
|
||||
output_transform: Callable = None,
|
||||
input_transform: Callable = None,
|
||||
init_hidden_layers: Sequence[int] = (),
|
||||
):
|
||||
|
||||
# check transform functions
|
||||
if input_transform is not None:
|
||||
try:
|
||||
_ = input_transform(jnp.zeros(num_inputs))
|
||||
except Exception as e:
|
||||
raise ValueError(f"Output transform function failed: {e}")
|
||||
|
||||
if output_transform is not None:
|
||||
try:
|
||||
_ = output_transform(jnp.zeros(num_outputs))
|
||||
except Exception as e:
|
||||
raise ValueError(f"Output transform function failed: {e}")
|
||||
|
||||
# prepare for initialization
|
||||
all_layers = [num_inputs] + list(init_hidden_layers) + [num_outputs]
|
||||
layer_indices = []
|
||||
next_index = 0
|
||||
for layer in all_layers:
|
||||
layer_indices.append(list(range(next_index, next_index + layer)))
|
||||
next_index += layer
|
||||
|
||||
all_init_nodes = []
|
||||
all_init_conns_in_idx = []
|
||||
all_init_conns_out_idx = []
|
||||
for i in range(len(layer_indices) - 1):
|
||||
in_layer = layer_indices[i]
|
||||
out_layer = layer_indices[i + 1]
|
||||
for in_idx in in_layer:
|
||||
for out_idx in out_layer:
|
||||
all_init_conns_in_idx.append(in_idx)
|
||||
all_init_conns_out_idx.append(out_idx)
|
||||
all_init_nodes.extend(in_layer)
|
||||
|
||||
if max_nodes < len(all_init_nodes):
|
||||
raise ValueError(
|
||||
f"max_nodes={max_nodes} must be greater than or equal to the number of initial nodes={len(all_init_nodes)}"
|
||||
)
|
||||
|
||||
if max_conns < len(all_init_conns_in_idx):
|
||||
raise ValueError(
|
||||
f"max_conns={max_conns} must be greater than or equal to the number of initial connections={len(all_init_conns_in_idx)}"
|
||||
)
|
||||
|
||||
self.num_inputs = num_inputs
|
||||
self.num_outputs = num_outputs
|
||||
self.input_idx = np.arange(num_inputs)
|
||||
self.output_idx = np.arange(num_inputs, num_inputs + num_outputs)
|
||||
self.max_nodes = max_nodes
|
||||
self.max_conns = max_conns
|
||||
self.node_gene = node_gene
|
||||
self.conn_gene = conn_gene
|
||||
self.mutation = mutation
|
||||
self.crossover = crossover
|
||||
self.distance = distance
|
||||
self.output_transform = output_transform
|
||||
self.input_transform = input_transform
|
||||
|
||||
self.input_idx = np.array(layer_indices[0])
|
||||
self.output_idx = np.array(layer_indices[-1])
|
||||
self.all_init_nodes = np.array(all_init_nodes)
|
||||
self.all_init_conns = np.c_[all_init_conns_in_idx, all_init_conns_out_idx]
|
||||
|
||||
def setup(self, state=State()):
|
||||
state = self.node_gene.setup(state)
|
||||
state = self.conn_gene.setup(state)
|
||||
state = self.mutation.setup(state)
|
||||
state = self.crossover.setup(state)
|
||||
state = self.mutation.setup(state, self)
|
||||
state = self.crossover.setup(state, self)
|
||||
state = self.distance.setup(state, self)
|
||||
return state
|
||||
|
||||
def transform(self, state, nodes, conns):
|
||||
raise NotImplementedError
|
||||
|
||||
def restore(self, state, transformed):
|
||||
raise NotImplementedError
|
||||
|
||||
def forward(self, state, transformed, inputs):
|
||||
raise NotImplementedError
|
||||
|
||||
def sympy_func(self):
|
||||
raise NotImplementedError
|
||||
|
||||
def visualize(self):
|
||||
raise NotImplementedError
|
||||
|
||||
def execute_mutation(self, state, randkey, nodes, conns, new_node_key):
|
||||
return self.mutation(state, randkey, self, nodes, conns, new_node_key)
|
||||
return self.mutation(state, randkey, nodes, conns, new_node_key)
|
||||
|
||||
def execute_crossover(self, state, randkey, nodes1, conns1, nodes2, conns2):
|
||||
return self.crossover(state, randkey, self, nodes1, conns1, nodes2, conns2)
|
||||
return self.crossover(state, randkey, nodes1, conns1, nodes2, conns2)
|
||||
|
||||
def execute_distance(self, state, nodes1, conns1, nodes2, conns2):
|
||||
return self.distance(state, nodes1, conns1, nodes2, conns2)
|
||||
|
||||
def initialize(self, state, randkey):
|
||||
"""
|
||||
Default initialization method for the genome.
|
||||
Add an extra hidden node.
|
||||
Make all input nodes and output nodes connected to the hidden node.
|
||||
All attributes will be initialized randomly using gene.new_random_attrs method.
|
||||
|
||||
For example, a network with 2 inputs and 1 output, the structure will be:
|
||||
nodes:
|
||||
[
|
||||
[0, attrs0], # input node 0
|
||||
[1, attrs1], # input node 1
|
||||
[2, attrs2], # output node 0
|
||||
[3, attrs3], # hidden node
|
||||
[NaN, NaN], # empty node
|
||||
]
|
||||
conns:
|
||||
[
|
||||
[0, 3, attrs0], # input node 0 -> hidden node
|
||||
[1, 3, attrs1], # input node 1 -> hidden node
|
||||
[3, 2, attrs2], # hidden node -> output node 0
|
||||
[NaN, NaN],
|
||||
[NaN, NaN],
|
||||
]
|
||||
"""
|
||||
|
||||
k1, k2 = jax.random.split(randkey) # k1 for nodes, k2 for conns
|
||||
|
||||
all_nodes_cnt = len(self.all_init_nodes)
|
||||
all_conns_cnt = len(self.all_init_conns)
|
||||
|
||||
# initialize nodes
|
||||
new_node_key = (
|
||||
max([*self.input_idx, *self.output_idx]) + 1
|
||||
) # the key for the hidden node
|
||||
node_keys = jnp.concatenate(
|
||||
[self.input_idx, self.output_idx, jnp.array([new_node_key])]
|
||||
) # the list of all node keys
|
||||
|
||||
# initialize nodes and connections with NaN
|
||||
nodes = jnp.full((self.max_nodes, self.node_gene.length), jnp.nan)
|
||||
conns = jnp.full((self.max_conns, self.conn_gene.length), jnp.nan)
|
||||
# create node indices
|
||||
node_indices = self.all_init_nodes
|
||||
# create node attrs
|
||||
rand_keys_n = jax.random.split(k1, num=all_nodes_cnt)
|
||||
node_attr_func = vmap(self.node_gene.new_random_attrs, in_axes=(None, 0))
|
||||
node_attrs = node_attr_func(state, rand_keys_n)
|
||||
|
||||
# set keys for input nodes, output nodes and hidden node
|
||||
nodes = nodes.at[node_keys, 0].set(node_keys)
|
||||
|
||||
# generate random attributes for nodes
|
||||
node_keys = jax.random.split(k1, len(node_keys))
|
||||
random_node_attrs = jax.vmap(
|
||||
self.node_gene.new_random_attrs, in_axes=(None, 0)
|
||||
)(state, node_keys)
|
||||
nodes = nodes.at[: len(node_keys), 1:].set(random_node_attrs)
|
||||
nodes = nodes.at[:all_nodes_cnt, 0].set(node_indices) # set node indices
|
||||
nodes = nodes.at[:all_nodes_cnt, 1:].set(node_attrs) # set node attrs
|
||||
|
||||
# initialize conns
|
||||
# input-hidden connections
|
||||
input_conns = jnp.c_[
|
||||
self.input_idx, jnp.full_like(self.input_idx, new_node_key)
|
||||
]
|
||||
conns = conns.at[self.input_idx, :2].set(input_conns) # in-keys, out-keys
|
||||
conns = jnp.full((self.max_conns, self.conn_gene.length), jnp.nan)
|
||||
# create input and output indices
|
||||
conn_indices = self.all_init_conns
|
||||
# create conn attrs
|
||||
rand_keys_c = jax.random.split(k2, num=all_conns_cnt)
|
||||
conns_attr_func = jax.vmap(
|
||||
self.conn_gene.new_random_attrs,
|
||||
in_axes=(
|
||||
None,
|
||||
0,
|
||||
),
|
||||
)
|
||||
conns_attrs = conns_attr_func(state, rand_keys_c)
|
||||
|
||||
# output-hidden connections
|
||||
output_conns = jnp.c_[
|
||||
jnp.full_like(self.output_idx, new_node_key), self.output_idx
|
||||
]
|
||||
conns = conns.at[self.output_idx, :2].set(output_conns) # in-keys, out-keys
|
||||
|
||||
conn_keys = jax.random.split(k2, num=len(self.input_idx) + len(self.output_idx))
|
||||
# generate random attributes for conns
|
||||
random_conn_attrs = jax.vmap(
|
||||
self.conn_gene.new_random_attrs, in_axes=(None, 0)
|
||||
)(state, conn_keys)
|
||||
conns = conns.at[: len(conn_keys), 2:].set(random_conn_attrs)
|
||||
conns = conns.at[:all_conns_cnt, :2].set(conn_indices) # set conn indices
|
||||
conns = conns.at[:all_conns_cnt, 2:].set(conns_attrs) # set conn attrs
|
||||
|
||||
return nodes, conns
|
||||
|
||||
def update_by_batch(self, state, batch_input, transformed):
|
||||
"""
|
||||
Update the genome by a batch of data.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
def network_dict(self, state, nodes, conns):
|
||||
return {
|
||||
"nodes": self._get_node_dict(state, nodes),
|
||||
"conns": self._get_conn_dict(state, conns),
|
||||
}
|
||||
|
||||
def get_input_idx(self):
|
||||
return self.input_idx.tolist()
|
||||
|
||||
def get_output_idx(self):
|
||||
return self.output_idx.tolist()
|
||||
|
||||
def hash(self, nodes, conns):
|
||||
nodes_hashs = vmap(hash_array)(nodes)
|
||||
conns_hashs = vmap(hash_array)(conns)
|
||||
return hash_array(jnp.concatenate([nodes_hashs, conns_hashs]))
|
||||
|
||||
def repr(self, state, nodes, conns, precision=2):
|
||||
nodes, conns = jax.device_get([nodes, conns])
|
||||
nodes_cnt, conns_cnt = self.valid_cnt(nodes), self.valid_cnt(conns)
|
||||
nodes_cnt, conns_cnt = valid_cnt(nodes), valid_cnt(conns)
|
||||
s = f"{self.__class__.__name__}(nodes={nodes_cnt}, conns={conns_cnt}):\n"
|
||||
s += f"\tNodes:\n"
|
||||
for node in nodes:
|
||||
@@ -152,11 +199,7 @@ class BaseGenome(StatefulBaseClass):
|
||||
s += f"\t\t{self.conn_gene.repr(state, conn, precision=precision)}\n"
|
||||
return s
|
||||
|
||||
@classmethod
|
||||
def valid_cnt(cls, arr):
|
||||
return jnp.sum(~jnp.isnan(arr[:, 0]))
|
||||
|
||||
def get_conn_dict(self, state, conns):
|
||||
def _get_conn_dict(self, state, conns):
|
||||
conns = jax.device_get(conns)
|
||||
conn_dict = {}
|
||||
for conn in conns:
|
||||
@@ -167,7 +210,7 @@ class BaseGenome(StatefulBaseClass):
|
||||
conn_dict[(in_idx, out_idx)] = cd
|
||||
return conn_dict
|
||||
|
||||
def get_node_dict(self, state, nodes):
|
||||
def _get_node_dict(self, state, nodes):
|
||||
nodes = jax.device_get(nodes)
|
||||
node_dict = {}
|
||||
for node in nodes:
|
||||
@@ -177,92 +220,3 @@ class BaseGenome(StatefulBaseClass):
|
||||
idx = nd["idx"]
|
||||
node_dict[idx] = nd
|
||||
return node_dict
|
||||
|
||||
def network_dict(self, state, nodes, conns):
|
||||
return {
|
||||
"nodes": self.get_node_dict(state, nodes),
|
||||
"conns": self.get_conn_dict(state, conns),
|
||||
}
|
||||
|
||||
def get_input_idx(self):
|
||||
return self.input_idx.tolist()
|
||||
|
||||
def get_output_idx(self):
|
||||
return self.output_idx.tolist()
|
||||
|
||||
def sympy_func(self, state, network, sympy_output_transform=None):
|
||||
raise NotImplementedError
|
||||
|
||||
def visualize(
|
||||
self,
|
||||
network,
|
||||
rotate=0,
|
||||
reverse_node_order=False,
|
||||
size=(300, 300, 300),
|
||||
color=("blue", "blue", "blue"),
|
||||
save_path="network.svg",
|
||||
save_dpi=800,
|
||||
**kwargs,
|
||||
):
|
||||
import networkx as nx
|
||||
from matplotlib import pyplot as plt
|
||||
|
||||
nodes_list = list(network["nodes"])
|
||||
conns_list = list(network["conns"])
|
||||
input_idx = self.get_input_idx()
|
||||
output_idx = self.get_output_idx()
|
||||
topo_order, topo_layers = topological_sort_python(nodes_list, conns_list)
|
||||
node2layer = {
|
||||
node: layer for layer, nodes in enumerate(topo_layers) for node in nodes
|
||||
}
|
||||
if reverse_node_order:
|
||||
topo_order = topo_order[::-1]
|
||||
|
||||
G = nx.DiGraph()
|
||||
|
||||
if not isinstance(size, tuple):
|
||||
size = (size, size, size)
|
||||
if not isinstance(color, tuple):
|
||||
color = (color, color, color)
|
||||
|
||||
for node in topo_order:
|
||||
if node in input_idx:
|
||||
G.add_node(node, subset=node2layer[node], size=size[0], color=color[0])
|
||||
elif node in output_idx:
|
||||
G.add_node(node, subset=node2layer[node], size=size[2], color=color[2])
|
||||
else:
|
||||
G.add_node(node, subset=node2layer[node], size=size[1], color=color[1])
|
||||
|
||||
for conn in conns_list:
|
||||
G.add_edge(conn[0], conn[1])
|
||||
pos = nx.multipartite_layout(G)
|
||||
|
||||
def rotate_layout(pos, angle):
|
||||
angle_rad = np.deg2rad(angle)
|
||||
cos_angle, sin_angle = np.cos(angle_rad), np.sin(angle_rad)
|
||||
rotated_pos = {}
|
||||
for node, (x, y) in pos.items():
|
||||
rotated_pos[node] = (
|
||||
cos_angle * x - sin_angle * y,
|
||||
sin_angle * x + cos_angle * y,
|
||||
)
|
||||
return rotated_pos
|
||||
|
||||
rotated_pos = rotate_layout(pos, rotate)
|
||||
|
||||
node_sizes = [n["size"] for n in G.nodes.values()]
|
||||
node_colors = [n["color"] for n in G.nodes.values()]
|
||||
|
||||
nx.draw(
|
||||
G,
|
||||
pos=rotated_pos,
|
||||
node_size=node_sizes,
|
||||
node_color=node_colors,
|
||||
**kwargs,
|
||||
)
|
||||
plt.savefig(save_path, dpi=save_dpi)
|
||||
|
||||
def hash(self, nodes, conns):
|
||||
nodes_hashs = jax.vmap(hash_array)(nodes)
|
||||
conns_hashs = jax.vmap(hash_array)(conns)
|
||||
return hash_array(jnp.concatenate([nodes_hashs, conns_hashs]))
|
||||
|
||||
@@ -1,25 +1,23 @@
|
||||
import warnings
|
||||
from typing import Callable
|
||||
|
||||
import jax, jax.numpy as jnp
|
||||
import jax
|
||||
from jax import vmap, numpy as jnp
|
||||
import numpy as np
|
||||
import sympy as sp
|
||||
from utils import (
|
||||
unflatten_conns,
|
||||
|
||||
from . import BaseGenome
|
||||
from ..gene import DefaultNodeGene, DefaultConnGene
|
||||
from .operations import DefaultMutation, DefaultCrossover, DefaultDistance
|
||||
from .utils import unflatten_conns, extract_node_attrs, extract_conn_attrs
|
||||
|
||||
from tensorneat.common import (
|
||||
topological_sort,
|
||||
topological_sort_python,
|
||||
I_INF,
|
||||
extract_node_attrs,
|
||||
extract_conn_attrs,
|
||||
set_node_attrs,
|
||||
set_conn_attrs,
|
||||
attach_with_inf,
|
||||
SYMPY_FUNCS_MODULE_NP,
|
||||
SYMPY_FUNCS_MODULE_JNP,
|
||||
)
|
||||
from . import BaseGenome
|
||||
from ..gene import BaseNodeGene, BaseConnGene, DefaultNodeGene, DefaultConnGene
|
||||
from ..ga import BaseMutation, BaseCrossover, DefaultMutation, DefaultCrossover
|
||||
|
||||
|
||||
class DefaultGenome(BaseGenome):
|
||||
@@ -31,15 +29,18 @@ class DefaultGenome(BaseGenome):
|
||||
self,
|
||||
num_inputs: int,
|
||||
num_outputs: int,
|
||||
max_nodes=5,
|
||||
max_conns=4,
|
||||
node_gene: BaseNodeGene = DefaultNodeGene(),
|
||||
conn_gene: BaseConnGene = DefaultConnGene(),
|
||||
mutation: BaseMutation = DefaultMutation(),
|
||||
crossover: BaseCrossover = DefaultCrossover(),
|
||||
output_transform: Callable = None,
|
||||
input_transform: Callable = None,
|
||||
max_nodes=50,
|
||||
max_conns=100,
|
||||
node_gene=DefaultNodeGene(),
|
||||
conn_gene=DefaultConnGene(),
|
||||
mutation=DefaultMutation(),
|
||||
crossover=DefaultCrossover(),
|
||||
distance=DefaultDistance(),
|
||||
output_transform=None,
|
||||
input_transform=None,
|
||||
init_hidden_layers=(),
|
||||
):
|
||||
|
||||
super().__init__(
|
||||
num_inputs,
|
||||
num_outputs,
|
||||
@@ -49,22 +50,12 @@ class DefaultGenome(BaseGenome):
|
||||
conn_gene,
|
||||
mutation,
|
||||
crossover,
|
||||
distance,
|
||||
output_transform,
|
||||
input_transform,
|
||||
init_hidden_layers,
|
||||
)
|
||||
|
||||
if input_transform is not None:
|
||||
try:
|
||||
_ = input_transform(np.zeros(num_inputs))
|
||||
except Exception as e:
|
||||
raise ValueError(f"Output transform function failed: {e}")
|
||||
self.input_transform = input_transform
|
||||
|
||||
if output_transform is not None:
|
||||
try:
|
||||
_ = output_transform(np.zeros(num_outputs))
|
||||
except Exception as e:
|
||||
raise ValueError(f"Output transform function failed: {e}")
|
||||
self.output_transform = output_transform
|
||||
|
||||
def transform(self, state, nodes, conns):
|
||||
u_conns = unflatten_conns(nodes, conns)
|
||||
conn_exist = u_conns != I_INF
|
||||
@@ -73,10 +64,6 @@ class DefaultGenome(BaseGenome):
|
||||
|
||||
return seqs, nodes, conns, u_conns
|
||||
|
||||
def restore(self, state, transformed):
|
||||
seqs, nodes, conns, u_conns = transformed
|
||||
return nodes, conns
|
||||
|
||||
def forward(self, state, transformed, inputs):
|
||||
|
||||
if self.input_transform is not None:
|
||||
@@ -86,8 +73,8 @@ class DefaultGenome(BaseGenome):
|
||||
|
||||
ini_vals = jnp.full((self.max_nodes,), jnp.nan)
|
||||
ini_vals = ini_vals.at[self.input_idx].set(inputs)
|
||||
nodes_attrs = jax.vmap(extract_node_attrs)(nodes)
|
||||
conns_attrs = jax.vmap(extract_conn_attrs)(conns)
|
||||
nodes_attrs = vmap(extract_node_attrs)(nodes)
|
||||
conns_attrs = vmap(extract_conn_attrs)(conns)
|
||||
|
||||
def cond_fun(carry):
|
||||
values, idx = carry
|
||||
@@ -105,7 +92,7 @@ class DefaultGenome(BaseGenome):
|
||||
def otherwise():
|
||||
conn_indices = u_conns[:, i]
|
||||
hit_attrs = attach_with_inf(conns_attrs, conn_indices)
|
||||
ins = jax.vmap(self.conn_gene.forward, in_axes=(None, 0, 0))(
|
||||
ins = vmap(self.conn_gene.forward, in_axes=(None, 0, 0))(
|
||||
state, hit_attrs, values
|
||||
)
|
||||
|
||||
@@ -130,85 +117,14 @@ class DefaultGenome(BaseGenome):
|
||||
else:
|
||||
return self.output_transform(vals[self.output_idx])
|
||||
|
||||
def update_by_batch(self, state, batch_input, transformed):
|
||||
|
||||
if self.input_transform is not None:
|
||||
batch_input = jax.vmap(self.input_transform)(batch_input)
|
||||
|
||||
cal_seqs, nodes, conns, u_conns = transformed
|
||||
|
||||
batch_size = batch_input.shape[0]
|
||||
batch_ini_vals = jnp.full((batch_size, self.max_nodes), jnp.nan)
|
||||
batch_ini_vals = batch_ini_vals.at[:, self.input_idx].set(batch_input)
|
||||
nodes_attrs = jax.vmap(extract_node_attrs)(nodes)
|
||||
conns_attrs = jax.vmap(extract_conn_attrs)(conns)
|
||||
|
||||
def cond_fun(carry):
|
||||
batch_values, nodes_attrs_, conns_attrs_, idx = carry
|
||||
return (idx < self.max_nodes) & (cal_seqs[idx] != I_INF)
|
||||
|
||||
def body_func(carry):
|
||||
batch_values, nodes_attrs_, conns_attrs_, idx = carry
|
||||
i = cal_seqs[idx]
|
||||
|
||||
def input_node():
|
||||
batch, new_attrs = self.node_gene.update_input_transform(
|
||||
state, nodes_attrs_[i], batch_values[:, i]
|
||||
)
|
||||
return (
|
||||
batch_values.at[:, i].set(batch),
|
||||
nodes_attrs_.at[i].set(new_attrs),
|
||||
conns_attrs_,
|
||||
)
|
||||
|
||||
def otherwise():
|
||||
|
||||
conn_indices = u_conns[:, i]
|
||||
hit_attrs = attach_with_inf(conns_attrs, conn_indices)
|
||||
batch_ins, new_conn_attrs = jax.vmap(
|
||||
self.conn_gene.update_by_batch,
|
||||
in_axes=(None, 0, 1),
|
||||
out_axes=(1, 0),
|
||||
)(state, hit_attrs, batch_values)
|
||||
|
||||
batch_z, new_node_attrs = self.node_gene.update_by_batch(
|
||||
state,
|
||||
nodes_attrs_[i],
|
||||
batch_ins,
|
||||
is_output_node=jnp.isin(i, self.output_idx),
|
||||
)
|
||||
|
||||
return (
|
||||
batch_values.at[:, i].set(batch_z),
|
||||
nodes_attrs_.at[i].set(new_node_attrs),
|
||||
conns_attrs_.at[conn_indices].set(new_conn_attrs),
|
||||
)
|
||||
|
||||
# the val of input nodes is obtained by the task, not by calculation
|
||||
(batch_values, nodes_attrs_, conns_attrs_) = jax.lax.cond(
|
||||
jnp.isin(i, self.input_idx),
|
||||
input_node,
|
||||
otherwise,
|
||||
)
|
||||
|
||||
return batch_values, nodes_attrs_, conns_attrs_, idx + 1
|
||||
|
||||
batch_vals, nodes_attrs, conns_attrs, _ = jax.lax.while_loop(
|
||||
cond_fun, body_func, (batch_ini_vals, nodes_attrs, conns_attrs, 0)
|
||||
def network_dict(self, state, nodes, conns):
|
||||
network = super().network_dict(state, nodes, conns)
|
||||
topo_order, topo_layers = topological_sort_python(
|
||||
set(network["nodes"]), set(network["conns"])
|
||||
)
|
||||
|
||||
nodes = jax.vmap(set_node_attrs)(nodes, nodes_attrs)
|
||||
conns = jax.vmap(set_conn_attrs)(conns, conns_attrs)
|
||||
|
||||
new_transformed = (cal_seqs, nodes, conns, u_conns)
|
||||
|
||||
if self.output_transform is None:
|
||||
return batch_vals[:, self.output_idx], new_transformed
|
||||
else:
|
||||
return (
|
||||
jax.vmap(self.output_transform)(batch_vals[:, self.output_idx]),
|
||||
new_transformed,
|
||||
)
|
||||
network["topo_order"] = topo_order
|
||||
network["topo_layers"] = topo_layers
|
||||
return network
|
||||
|
||||
def sympy_func(
|
||||
self,
|
||||
@@ -241,7 +157,8 @@ class DefaultGenome(BaseGenome):
|
||||
|
||||
input_idx = self.get_input_idx()
|
||||
output_idx = self.get_output_idx()
|
||||
order, _ = topological_sort_python(set(network["nodes"]), set(network["conns"]))
|
||||
order = network["topo_order"]
|
||||
|
||||
hidden_idx = [
|
||||
i for i in network["nodes"] if i not in input_idx and i not in output_idx
|
||||
]
|
||||
@@ -260,8 +177,12 @@ class DefaultGenome(BaseGenome):
|
||||
for i in order:
|
||||
|
||||
if i in input_idx:
|
||||
nodes_exprs[symbols[-i - 1]] = symbols[-i - 1] # origin equal to its symbol
|
||||
nodes_exprs[symbols[i]] = sympy_input_transform[i - min(input_idx)](symbols[-i - 1]) # normed i
|
||||
nodes_exprs[symbols[-i - 1]] = symbols[
|
||||
-i - 1
|
||||
] # origin equal to its symbol
|
||||
nodes_exprs[symbols[i]] = sympy_input_transform[i - min(input_idx)](
|
||||
symbols[-i - 1]
|
||||
) # normed i
|
||||
|
||||
else:
|
||||
in_conns = [c for c in network["conns"] if c[1] == i]
|
||||
@@ -325,3 +246,73 @@ class DefaultGenome(BaseGenome):
|
||||
output_exprs,
|
||||
forward_func,
|
||||
)
|
||||
|
||||
def visualize(
|
||||
self,
|
||||
network,
|
||||
rotate=0,
|
||||
reverse_node_order=False,
|
||||
size=(300, 300, 300),
|
||||
color=("blue", "blue", "blue"),
|
||||
save_path="network.svg",
|
||||
save_dpi=800,
|
||||
**kwargs,
|
||||
):
|
||||
import networkx as nx
|
||||
from matplotlib import pyplot as plt
|
||||
|
||||
nodes_list = list(network["nodes"])
|
||||
conns_list = list(network["conns"])
|
||||
input_idx = self.get_input_idx()
|
||||
output_idx = self.get_output_idx()
|
||||
|
||||
topo_order, topo_layers = network["topo_order"], network["topo_layers"]
|
||||
node2layer = {
|
||||
node: layer for layer, nodes in enumerate(topo_layers) for node in nodes
|
||||
}
|
||||
if reverse_node_order:
|
||||
topo_order = topo_order[::-1]
|
||||
|
||||
G = nx.DiGraph()
|
||||
|
||||
if not isinstance(size, tuple):
|
||||
size = (size, size, size)
|
||||
if not isinstance(color, tuple):
|
||||
color = (color, color, color)
|
||||
|
||||
for node in topo_order:
|
||||
if node in input_idx:
|
||||
G.add_node(node, subset=node2layer[node], size=size[0], color=color[0])
|
||||
elif node in output_idx:
|
||||
G.add_node(node, subset=node2layer[node], size=size[2], color=color[2])
|
||||
else:
|
||||
G.add_node(node, subset=node2layer[node], size=size[1], color=color[1])
|
||||
|
||||
for conn in conns_list:
|
||||
G.add_edge(conn[0], conn[1])
|
||||
pos = nx.multipartite_layout(G)
|
||||
|
||||
def rotate_layout(pos, angle):
|
||||
angle_rad = np.deg2rad(angle)
|
||||
cos_angle, sin_angle = np.cos(angle_rad), np.sin(angle_rad)
|
||||
rotated_pos = {}
|
||||
for node, (x, y) in pos.items():
|
||||
rotated_pos[node] = (
|
||||
cos_angle * x - sin_angle * y,
|
||||
sin_angle * x + cos_angle * y,
|
||||
)
|
||||
return rotated_pos
|
||||
|
||||
rotated_pos = rotate_layout(pos, rotate)
|
||||
|
||||
node_sizes = [n["size"] for n in G.nodes.values()]
|
||||
node_colors = [n["color"] for n in G.nodes.values()]
|
||||
|
||||
nx.draw(
|
||||
G,
|
||||
pos=rotated_pos,
|
||||
node_size=node_sizes,
|
||||
node_color=node_colors,
|
||||
**kwargs,
|
||||
)
|
||||
plt.savefig(save_path, dpi=save_dpi)
|
||||
|
||||
@@ -1,56 +0,0 @@
|
||||
import jax, jax.numpy as jnp
|
||||
from .default import DefaultGenome
|
||||
|
||||
|
||||
class DenseInitialize(DefaultGenome):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
assert self.max_nodes >= self.num_inputs + self.num_outputs
|
||||
assert self.max_conns >= self.num_inputs * self.num_outputs
|
||||
|
||||
def initialize(self, state, randkey):
|
||||
|
||||
k1, k2 = jax.random.split(randkey, num=2)
|
||||
|
||||
input_idx, output_idx = self.input_idx, self.output_idx
|
||||
input_size = len(input_idx)
|
||||
output_size = len(output_idx)
|
||||
|
||||
nodes = jnp.full(
|
||||
(self.max_nodes, self.node_gene.length), jnp.nan, dtype=jnp.float32
|
||||
)
|
||||
|
||||
nodes = nodes.at[input_idx, 0].set(input_idx)
|
||||
nodes = nodes.at[output_idx, 0].set(output_idx)
|
||||
|
||||
total_idx = input_size + output_size
|
||||
rand_keys_n = jax.random.split(k1, num=total_idx)
|
||||
|
||||
node_attr_func = jax.vmap(self.node_gene.new_random_attrs, in_axes=(None, 0))
|
||||
node_attrs = node_attr_func(state, rand_keys_n)
|
||||
nodes = nodes.at[:total_idx, 1:].set(node_attrs)
|
||||
|
||||
conns = jnp.full(
|
||||
(self.max_conns, self.conn_gene.length), jnp.nan, dtype=jnp.float32
|
||||
)
|
||||
|
||||
input_to_output_ids, output_ids = jnp.meshgrid(
|
||||
input_idx, output_idx, indexing="ij"
|
||||
)
|
||||
total_conns = input_size * output_size
|
||||
conns = conns.at[:total_conns, :2].set(
|
||||
jnp.column_stack([input_to_output_ids.flatten(), output_ids.flatten()])
|
||||
)
|
||||
|
||||
rand_keys_c = jax.random.split(k2, num=total_conns)
|
||||
conns_attr_func = jax.vmap(
|
||||
self.conn_gene.new_random_attrs,
|
||||
in_axes=(
|
||||
None,
|
||||
0,
|
||||
),
|
||||
)
|
||||
conns_attrs = conns_attr_func(state, rand_keys_c)
|
||||
conns = conns.at[:total_conns, 2:].set(conns_attrs)
|
||||
|
||||
return nodes, conns
|
||||
@@ -1,70 +0,0 @@
|
||||
import jax, jax.numpy as jnp
|
||||
from .default import DefaultGenome
|
||||
|
||||
|
||||
class HiddenInitialize(DefaultGenome):
|
||||
def __init__(self, hidden_cnt=8, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.hidden_cnt = hidden_cnt
|
||||
|
||||
def initialize(self, state, randkey):
|
||||
|
||||
k1, k2 = jax.random.split(randkey, num=2)
|
||||
|
||||
input_idx, output_idx = self.input_idx, self.output_idx
|
||||
input_size = len(input_idx)
|
||||
output_size = len(output_idx)
|
||||
|
||||
hidden_idx = jnp.arange(
|
||||
input_size + output_size, input_size + output_size + self.hidden_cnt
|
||||
)
|
||||
nodes = jnp.full(
|
||||
(self.max_nodes, self.node_gene.length), jnp.nan, dtype=jnp.float32
|
||||
)
|
||||
|
||||
nodes = nodes.at[input_idx, 0].set(input_idx)
|
||||
nodes = nodes.at[output_idx, 0].set(output_idx)
|
||||
nodes = nodes.at[hidden_idx, 0].set(hidden_idx)
|
||||
|
||||
total_idx = input_size + output_size + self.hidden_cnt
|
||||
rand_keys_n = jax.random.split(k1, num=total_idx)
|
||||
|
||||
node_attr_func = jax.vmap(self.node_gene.new_random_attrs, in_axes=(None, 0))
|
||||
node_attrs = node_attr_func(state, rand_keys_n)
|
||||
nodes = nodes.at[:total_idx, 1:].set(node_attrs)
|
||||
|
||||
conns = jnp.full(
|
||||
(self.max_conns, self.conn_gene.length), jnp.nan, dtype=jnp.float32
|
||||
)
|
||||
|
||||
input_to_hidden_ids, hidden_ids = jnp.meshgrid(
|
||||
input_idx, hidden_idx, indexing="ij"
|
||||
)
|
||||
total_input_to_hidden_conns = input_size * self.hidden_cnt
|
||||
conns = conns.at[:total_input_to_hidden_conns, :2].set(
|
||||
jnp.column_stack([input_to_hidden_ids.flatten(), hidden_ids.flatten()])
|
||||
)
|
||||
|
||||
hidden_to_output_ids, output_ids = jnp.meshgrid(
|
||||
hidden_idx, output_idx, indexing="ij"
|
||||
)
|
||||
total_hidden_to_output_conns = self.hidden_cnt * output_size
|
||||
conns = conns.at[
|
||||
total_input_to_hidden_conns : total_input_to_hidden_conns
|
||||
+ total_hidden_to_output_conns,
|
||||
:2,
|
||||
].set(jnp.column_stack([hidden_to_output_ids.flatten(), output_ids.flatten()]))
|
||||
|
||||
total_conns = total_input_to_hidden_conns + total_hidden_to_output_conns
|
||||
rand_keys_c = jax.random.split(k2, num=total_conns)
|
||||
conns_attr_func = jax.vmap(
|
||||
self.conn_gene.new_random_attrs,
|
||||
in_axes=(
|
||||
None,
|
||||
0,
|
||||
),
|
||||
)
|
||||
conns_attrs = conns_attr_func(state, rand_keys_c)
|
||||
conns = conns.at[:total_conns, 2:].set(conns_attrs)
|
||||
|
||||
return nodes, conns
|
||||
3
tensorneat/algorithm/neat/genome/operations/__init__.py
Normal file
3
tensorneat/algorithm/neat/genome/operations/__init__.py
Normal file
@@ -0,0 +1,3 @@
|
||||
from .crossover import BaseCrossover, DefaultCrossover
|
||||
from .mutation import BaseMutation, DefaultMutation
|
||||
from .distance import BaseDistance, DefaultDistance
|
||||
@@ -0,0 +1,2 @@
|
||||
from .base import BaseCrossover
|
||||
from .default import DefaultCrossover
|
||||
@@ -0,0 +1,12 @@
|
||||
from tensorneat.common import StatefulBaseClass, State
|
||||
|
||||
|
||||
class BaseCrossover(StatefulBaseClass):
|
||||
|
||||
def setup(self, state=State(), genome = None):
|
||||
assert genome is not None, "genome should not be None"
|
||||
self.genome = genome
|
||||
return state
|
||||
|
||||
def __call__(self, state, randkey, nodes1, nodes2, conns1, conns2):
|
||||
raise NotImplementedError
|
||||
@@ -0,0 +1,87 @@
|
||||
import jax
|
||||
from jax import vmap, numpy as jnp
|
||||
|
||||
from .base import BaseCrossover
|
||||
from ...utils import (
|
||||
extract_node_attrs,
|
||||
extract_conn_attrs,
|
||||
set_node_attrs,
|
||||
set_conn_attrs,
|
||||
)
|
||||
|
||||
|
||||
class DefaultCrossover(BaseCrossover):
|
||||
def __call__(self, state, randkey, 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!)
|
||||
"""
|
||||
randkey1, randkey2 = jax.random.split(randkey, 2)
|
||||
randkeys1 = jax.random.split(randkey1, self.genome.max_nodes)
|
||||
randkeys2 = jax.random.split(randkey2, self.genome.max_conns)
|
||||
|
||||
# crossover nodes
|
||||
keys1, keys2 = nodes1[:, 0], nodes2[:, 0]
|
||||
# make homologous genes align in nodes2 align with nodes1
|
||||
nodes2 = self.align_array(keys1, keys2, nodes2, is_conn=False)
|
||||
|
||||
# For not homologous genes, use the value of nodes1(winner)
|
||||
# For homologous genes, use the crossover result between nodes1 and nodes2
|
||||
node_attrs1 = vmap(extract_node_attrs)(nodes1)
|
||||
node_attrs2 = vmap(extract_node_attrs)(nodes2)
|
||||
|
||||
new_node_attrs = jnp.where(
|
||||
jnp.isnan(node_attrs1) | jnp.isnan(node_attrs2), # one of them is nan
|
||||
node_attrs1, # not homologous genes or both nan, use the value of nodes1(winner)
|
||||
vmap(self.genome.node_gene.crossover, in_axes=(None, 0, 0, 0))(
|
||||
state, randkeys1, node_attrs1, node_attrs2
|
||||
), # homologous or both nan
|
||||
)
|
||||
new_nodes = vmap(set_node_attrs)(nodes1, new_node_attrs)
|
||||
|
||||
# crossover connections
|
||||
con_keys1, con_keys2 = conns1[:, :2], conns2[:, :2]
|
||||
conns2 = self.align_array(con_keys1, con_keys2, conns2, is_conn=True)
|
||||
|
||||
conns_attrs1 = vmap(extract_conn_attrs)(conns1)
|
||||
conns_attrs2 = vmap(extract_conn_attrs)(conns2)
|
||||
|
||||
new_conn_attrs = jnp.where(
|
||||
jnp.isnan(conns_attrs1) | jnp.isnan(conns_attrs2),
|
||||
conns_attrs1, # not homologous genes or both nan, use the value of conns1(winner)
|
||||
vmap(self.genome.conn_gene.crossover, in_axes=(None, 0, 0, 0))(
|
||||
state, randkeys2, conns_attrs1, conns_attrs2
|
||||
), # homologous or both nan
|
||||
)
|
||||
new_conns = vmap(set_conn_attrs)(conns1, new_conn_attrs)
|
||||
|
||||
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 consider carefully before 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
|
||||
@@ -0,0 +1,2 @@
|
||||
from .base import BaseDistance
|
||||
from .default import DefaultDistance
|
||||
15
tensorneat/algorithm/neat/genome/operations/distance/base.py
Normal file
15
tensorneat/algorithm/neat/genome/operations/distance/base.py
Normal file
@@ -0,0 +1,15 @@
|
||||
from tensorneat.common import StatefulBaseClass, State
|
||||
|
||||
|
||||
class BaseDistance(StatefulBaseClass):
|
||||
|
||||
def setup(self, state=State(), genome = None):
|
||||
assert genome is not None, "genome should not be None"
|
||||
self.genome = genome
|
||||
return state
|
||||
|
||||
def __call__(self, state, nodes1, nodes2, conns1, conns2):
|
||||
"""
|
||||
The distance between two genomes
|
||||
"""
|
||||
raise NotImplementedError
|
||||
105
tensorneat/algorithm/neat/genome/operations/distance/default.py
Normal file
105
tensorneat/algorithm/neat/genome/operations/distance/default.py
Normal file
@@ -0,0 +1,105 @@
|
||||
from jax import vmap, numpy as jnp
|
||||
|
||||
from .base import BaseDistance
|
||||
from ...utils import extract_node_attrs, extract_conn_attrs
|
||||
|
||||
|
||||
class DefaultDistance(BaseDistance):
|
||||
def __init__(
|
||||
self,
|
||||
compatibility_disjoint: float = 1.0,
|
||||
compatibility_weight: float = 0.4,
|
||||
):
|
||||
self.compatibility_disjoint = compatibility_disjoint
|
||||
self.compatibility_weight = compatibility_weight
|
||||
|
||||
def __call__(self, state, nodes1, nodes2, conns1, conns2):
|
||||
"""
|
||||
The distance between two genomes
|
||||
"""
|
||||
d = self.node_distance(state, nodes1, nodes2) + self.conn_distance(
|
||||
state, conns1, conns2
|
||||
)
|
||||
return d
|
||||
|
||||
def node_distance(self, state, nodes1, nodes2):
|
||||
"""
|
||||
The distance of the nodes part for 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
|
||||
fr_attrs = vmap(extract_node_attrs)(fr)
|
||||
sr_attrs = vmap(extract_node_attrs)(sr)
|
||||
hnd = vmap(self.genome.node_gene.distance, in_axes=(None, 0, 0))(
|
||||
state, fr_attrs, sr_attrs
|
||||
) # homologous node distance
|
||||
hnd = jnp.where(jnp.isnan(hnd), 0, hnd)
|
||||
homologous_distance = jnp.sum(hnd * intersect_mask)
|
||||
|
||||
val = (
|
||||
non_homologous_cnt * self.compatibility_disjoint
|
||||
+ homologous_distance * self.compatibility_weight
|
||||
)
|
||||
|
||||
val = jnp.where(max_cnt == 0, 0, val / max_cnt) # normalize
|
||||
|
||||
return val
|
||||
|
||||
def conn_distance(self, state, conns1, conns2):
|
||||
"""
|
||||
The distance of the conns part for two genomes
|
||||
"""
|
||||
con_cnt1 = jnp.sum(~jnp.isnan(conns1[:, 0]))
|
||||
con_cnt2 = jnp.sum(~jnp.isnan(conns2[:, 0]))
|
||||
max_cnt = jnp.maximum(con_cnt1, con_cnt2)
|
||||
|
||||
cons = jnp.concatenate((conns1, conns2), 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)
|
||||
|
||||
fr_attrs = vmap(extract_conn_attrs)(fr)
|
||||
sr_attrs = vmap(extract_conn_attrs)(sr)
|
||||
hcd = vmap(self.genome.conn_gene.distance, in_axes=(None, 0, 0))(
|
||||
state, fr_attrs, sr_attrs
|
||||
) # homologous connection distance
|
||||
hcd = jnp.where(jnp.isnan(hcd), 0, hcd)
|
||||
homologous_distance = jnp.sum(hcd * intersect_mask)
|
||||
|
||||
val = (
|
||||
non_homologous_cnt * self.compatibility_disjoint
|
||||
+ homologous_distance * self.compatibility_weight
|
||||
)
|
||||
|
||||
val = jnp.where(max_cnt == 0, 0, val / max_cnt) # normalize
|
||||
|
||||
return val
|
||||
@@ -0,0 +1,2 @@
|
||||
from .base import BaseMutation
|
||||
from .default import DefaultMutation
|
||||
12
tensorneat/algorithm/neat/genome/operations/mutation/base.py
Normal file
12
tensorneat/algorithm/neat/genome/operations/mutation/base.py
Normal file
@@ -0,0 +1,12 @@
|
||||
from tensorneat.common import StatefulBaseClass, State
|
||||
|
||||
|
||||
class BaseMutation(StatefulBaseClass):
|
||||
|
||||
def setup(self, state=State(), genome = None):
|
||||
assert genome is not None, "genome should not be None"
|
||||
self.genome = genome
|
||||
return state
|
||||
|
||||
def __call__(self, state, randkey, genome, nodes, conns, new_node_key):
|
||||
raise NotImplementedError
|
||||
292
tensorneat/algorithm/neat/genome/operations/mutation/default.py
Normal file
292
tensorneat/algorithm/neat/genome/operations/mutation/default.py
Normal file
@@ -0,0 +1,292 @@
|
||||
import jax
|
||||
from jax import vmap, numpy as jnp
|
||||
from . import BaseMutation
|
||||
from tensorneat.common import (
|
||||
fetch_first,
|
||||
fetch_random,
|
||||
I_INF,
|
||||
check_cycles,
|
||||
)
|
||||
from ...utils import (
|
||||
unflatten_conns,
|
||||
add_node,
|
||||
add_conn,
|
||||
delete_node_by_pos,
|
||||
delete_conn_by_pos,
|
||||
extract_node_attrs,
|
||||
extract_conn_attrs,
|
||||
set_node_attrs,
|
||||
set_conn_attrs,
|
||||
)
|
||||
|
||||
|
||||
class DefaultMutation(BaseMutation):
|
||||
def __init__(
|
||||
self,
|
||||
conn_add: float = 0.2,
|
||||
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, state, randkey, genome, nodes, conns, new_node_key):
|
||||
k1, k2 = jax.random.split(randkey)
|
||||
|
||||
nodes, conns = self.mutate_structure(
|
||||
state, k1, genome, nodes, conns, new_node_key
|
||||
)
|
||||
nodes, conns = self.mutate_values(state, k2, genome, nodes, conns)
|
||||
|
||||
return nodes, conns
|
||||
|
||||
def mutate_structure(self, state, randkey, genome, nodes, conns, new_node_key):
|
||||
def mutate_add_node(key_, nodes_, conns_):
|
||||
"""
|
||||
add a node while do not influence the output of the network
|
||||
"""
|
||||
|
||||
remain_node_space = jnp.isnan(nodes_[:, 0]).sum()
|
||||
remain_conn_space = jnp.isnan(conns_[:, 0]).sum()
|
||||
i_key, o_key, idx = self.choose_connection_key(
|
||||
key_, conns_
|
||||
) # choose a connection
|
||||
|
||||
def successful_add_node():
|
||||
# remove the original connection and record its attrs
|
||||
original_attrs = extract_conn_attrs(conns_[idx])
|
||||
new_conns = delete_conn_by_pos(conns_, idx)
|
||||
|
||||
# add a new node with identity attrs
|
||||
new_nodes = add_node(
|
||||
nodes_, new_node_key, genome.node_gene.new_identity_attrs(state)
|
||||
)
|
||||
|
||||
# add two new connections
|
||||
# first is with identity attrs
|
||||
new_conns = add_conn(
|
||||
new_conns,
|
||||
i_key,
|
||||
new_node_key,
|
||||
genome.conn_gene.new_identity_attrs(state),
|
||||
)
|
||||
# second is with the origin attrs
|
||||
new_conns = add_conn(
|
||||
new_conns,
|
||||
new_node_key,
|
||||
o_key,
|
||||
original_attrs,
|
||||
)
|
||||
|
||||
return new_nodes, new_conns
|
||||
|
||||
return jax.lax.cond(
|
||||
(idx == I_INF) | (remain_node_space < 1) | (remain_conn_space < 2),
|
||||
lambda: (nodes_, conns_), # do nothing
|
||||
successful_add_node,
|
||||
)
|
||||
|
||||
def mutate_delete_node(key_, nodes_, conns_):
|
||||
"""
|
||||
delete a node
|
||||
"""
|
||||
# randomly choose a node
|
||||
key, idx = self.choose_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 = 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_INF, # no available node to delete
|
||||
lambda: (nodes_, conns_), # do nothing
|
||||
successful_delete_node,
|
||||
)
|
||||
|
||||
def mutate_add_conn(key_, nodes_, conns_):
|
||||
"""
|
||||
add a connection while do not influence the output of the network
|
||||
"""
|
||||
|
||||
remain_conn_space = jnp.isnan(conns_[:, 0]).sum()
|
||||
|
||||
# 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.choose_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.choose_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_INF
|
||||
|
||||
def nothing():
|
||||
return nodes_, conns_
|
||||
|
||||
def successful():
|
||||
# add a connection with zero attrs
|
||||
return nodes_, add_conn(
|
||||
conns_, i_key, o_key, genome.conn_gene.new_zero_attrs(state)
|
||||
)
|
||||
|
||||
if genome.network_type == "feedforward":
|
||||
u_conns = unflatten_conns(nodes_, conns_)
|
||||
conns_exist = u_conns != I_INF
|
||||
is_cycle = check_cycles(nodes_, conns_exist, from_idx, to_idx)
|
||||
|
||||
return jax.lax.cond(
|
||||
is_already_exist | is_cycle | (remain_conn_space < 1),
|
||||
nothing,
|
||||
successful,
|
||||
)
|
||||
|
||||
elif genome.network_type == "recurrent":
|
||||
return jax.lax.cond(
|
||||
is_already_exist | (remain_conn_space < 1),
|
||||
nothing,
|
||||
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.choose_connection_key(key_, conns_)
|
||||
|
||||
return jax.lax.cond(
|
||||
idx == I_INF,
|
||||
lambda: (nodes_, conns_), # nothing
|
||||
lambda: (nodes_, delete_conn_by_pos(conns_, idx)), # success
|
||||
)
|
||||
|
||||
k1, k2, k3, k4 = jax.random.split(randkey, num=4)
|
||||
r1, r2, r3, r4 = jax.random.uniform(k1, shape=(4,))
|
||||
|
||||
def nothing(_, nodes_, conns_):
|
||||
return nodes_, conns_
|
||||
|
||||
if self.node_add > 0:
|
||||
nodes, conns = jax.lax.cond(
|
||||
r1 < self.node_add, mutate_add_node, nothing, k1, nodes, conns
|
||||
)
|
||||
|
||||
if self.node_delete > 0:
|
||||
nodes, conns = jax.lax.cond(
|
||||
r2 < self.node_delete, mutate_delete_node, nothing, k2, nodes, conns
|
||||
)
|
||||
|
||||
if self.conn_add > 0:
|
||||
nodes, conns = jax.lax.cond(
|
||||
r3 < self.conn_add, mutate_add_conn, nothing, k3, nodes, conns
|
||||
)
|
||||
|
||||
if self.conn_delete > 0:
|
||||
nodes, conns = jax.lax.cond(
|
||||
r4 < self.conn_delete, mutate_delete_conn, nothing, k4, nodes, conns
|
||||
)
|
||||
|
||||
return nodes, conns
|
||||
|
||||
def mutate_values(self, state, randkey, genome, nodes, conns):
|
||||
k1, k2 = jax.random.split(randkey)
|
||||
nodes_randkeys = jax.random.split(k1, num=genome.max_nodes)
|
||||
conns_randkeys = jax.random.split(k2, num=genome.max_conns)
|
||||
|
||||
node_attrs = vmap(extract_node_attrs)(nodes)
|
||||
new_node_attrs = vmap(genome.node_gene.mutate, in_axes=(None, 0, 0))(
|
||||
state, nodes_randkeys, node_attrs
|
||||
)
|
||||
new_nodes = vmap(set_node_attrs)(nodes, new_node_attrs)
|
||||
|
||||
conn_attrs = vmap(extract_conn_attrs)(conns)
|
||||
new_conn_attrs = vmap(genome.conn_gene.mutate, in_axes=(None, 0, 0))(
|
||||
state, conns_randkeys, conn_attrs
|
||||
)
|
||||
new_conns = vmap(set_conn_attrs)(conns, new_conn_attrs)
|
||||
|
||||
# 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 choose_node_key(
|
||||
self,
|
||||
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 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(key, mask)
|
||||
key = jnp.where(idx != I_INF, nodes[idx, 0], jnp.nan)
|
||||
return key, idx
|
||||
|
||||
def choose_connection_key(self, key, conns):
|
||||
"""
|
||||
Randomly choose a connection key from the given connections.
|
||||
:return: i_key, o_key, idx
|
||||
"""
|
||||
|
||||
idx = fetch_random(key, ~jnp.isnan(conns[:, 0]))
|
||||
i_key = jnp.where(idx != I_INF, conns[idx, 0], jnp.nan)
|
||||
o_key = jnp.where(idx != I_INF, conns[idx, 1], jnp.nan)
|
||||
|
||||
return i_key, o_key, idx
|
||||
@@ -1,11 +1,11 @@
|
||||
from typing import Callable
|
||||
|
||||
import jax, jax.numpy as jnp
|
||||
from utils import unflatten_conns
|
||||
from .utils import unflatten_conns
|
||||
|
||||
from . import BaseGenome
|
||||
from ..gene import BaseNodeGene, BaseConnGene, DefaultNodeGene, DefaultConnGene
|
||||
from ..ga import BaseMutation, BaseCrossover, DefaultMutation, DefaultCrossover
|
||||
from ..gene import DefaultNodeGene, DefaultConnGene
|
||||
from .operations import DefaultMutation, DefaultCrossover
|
||||
|
||||
|
||||
class RecurrentGenome(BaseGenome):
|
||||
@@ -17,13 +17,13 @@ class RecurrentGenome(BaseGenome):
|
||||
self,
|
||||
num_inputs: int,
|
||||
num_outputs: int,
|
||||
max_nodes: int,
|
||||
max_conns: int,
|
||||
node_gene: BaseNodeGene = DefaultNodeGene(),
|
||||
conn_gene: BaseConnGene = DefaultConnGene(),
|
||||
mutation: BaseMutation = DefaultMutation(),
|
||||
crossover: BaseCrossover = DefaultCrossover(),
|
||||
activate_time: int = 10,
|
||||
max_nodes = 50,
|
||||
max_conns = 100,
|
||||
node_gene=DefaultNodeGene(),
|
||||
conn_gene=DefaultConnGene(),
|
||||
mutation=DefaultMutation(),
|
||||
crossover=DefaultCrossover(),
|
||||
activate_time=10,
|
||||
output_transform: Callable = None,
|
||||
):
|
||||
super().__init__(
|
||||
|
||||
109
tensorneat/algorithm/neat/genome/utils.py
Normal file
109
tensorneat/algorithm/neat/genome/utils.py
Normal file
@@ -0,0 +1,109 @@
|
||||
import jax
|
||||
from jax import vmap, numpy as jnp
|
||||
|
||||
from tensorneat.common import fetch_first, I_INF
|
||||
|
||||
|
||||
def unflatten_conns(nodes, conns):
|
||||
"""
|
||||
transform the (C, CL) connections to (N, N), which contains the idx of the connection in conns
|
||||
connection length, N means the number of nodes, C means the number of connections
|
||||
returns the unflatten connection indices with shape (N, N)
|
||||
"""
|
||||
N = nodes.shape[0] # max_nodes
|
||||
C = conns.shape[0] # max_conns
|
||||
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)
|
||||
|
||||
# Is interesting that jax use clip when attach data in array
|
||||
# however, it will do nothing when setting values in an array
|
||||
# put the index of connections in the unflatten array
|
||||
unflatten = (
|
||||
jnp.full((N, N), I_INF, dtype=jnp.int32)
|
||||
.at[i_idxs, o_idxs]
|
||||
.set(jnp.arange(C, dtype=jnp.int32))
|
||||
)
|
||||
|
||||
return unflatten
|
||||
|
||||
|
||||
def valid_cnt(nodes_or_conns):
|
||||
return jnp.sum(~jnp.isnan(nodes_or_conns[:, 0]))
|
||||
|
||||
|
||||
def extract_node_attrs(node):
|
||||
"""
|
||||
node: Array(NL, )
|
||||
extract the attributes of a node
|
||||
"""
|
||||
return node[1:] # 0 is for idx
|
||||
|
||||
|
||||
def set_node_attrs(node, attrs):
|
||||
"""
|
||||
node: Array(NL, )
|
||||
attrs: Array(NL-1, )
|
||||
set the attributes of a node
|
||||
"""
|
||||
return node.at[1:].set(attrs) # 0 is for idx
|
||||
|
||||
|
||||
def extract_conn_attrs(conn):
|
||||
"""
|
||||
conn: Array(CL, )
|
||||
extract the attributes of a connection
|
||||
"""
|
||||
return conn[2:] # 0, 1 is for in-idx and out-idx
|
||||
|
||||
|
||||
def set_conn_attrs(conn, attrs):
|
||||
"""
|
||||
conn: Array(CL, )
|
||||
attrs: Array(CL-2, )
|
||||
set the attributes of a connection
|
||||
"""
|
||||
return conn.at[2:].set(attrs) # 0, 1 is for in-idx and out-idx
|
||||
|
||||
|
||||
def add_node(nodes, new_key: int, attrs):
|
||||
"""
|
||||
Add a new node to the genome.
|
||||
The new node will place at the first NaN row.
|
||||
"""
|
||||
exist_keys = nodes[:, 0]
|
||||
pos = fetch_first(jnp.isnan(exist_keys))
|
||||
new_nodes = nodes.at[pos, 0].set(new_key)
|
||||
return new_nodes.at[pos, 1:].set(attrs)
|
||||
|
||||
|
||||
def delete_node_by_pos(nodes, pos):
|
||||
"""
|
||||
Delete a node from the genome.
|
||||
Delete the node by its pos in nodes.
|
||||
"""
|
||||
return nodes.at[pos].set(jnp.nan)
|
||||
|
||||
|
||||
def add_conn(conns, i_key, o_key, attrs):
|
||||
"""
|
||||
Add a new connection to the genome.
|
||||
The new connection will place at the first NaN row.
|
||||
"""
|
||||
con_keys = conns[:, 0]
|
||||
pos = fetch_first(jnp.isnan(con_keys))
|
||||
new_conns = conns.at[pos, 0:2].set(jnp.array([i_key, o_key]))
|
||||
return new_conns.at[pos, 2:].set(attrs)
|
||||
|
||||
|
||||
def delete_conn_by_pos(conns, pos):
|
||||
"""
|
||||
Delete a connection from the genome.
|
||||
Delete the connection by its idx.
|
||||
"""
|
||||
return conns.at[pos].set(jnp.nan)
|
||||
Reference in New Issue
Block a user