67 lines
2.0 KiB
Python
67 lines
2.0 KiB
Python
import jax.numpy as jnp
|
|
from ..gene import BaseNodeGene, BaseConnGene, DefaultNodeGene, DefaultConnGene
|
|
from utils import fetch_first
|
|
|
|
|
|
class BaseGenome:
|
|
|
|
network_type = None
|
|
|
|
def __init__(
|
|
self,
|
|
num_inputs: int,
|
|
num_outputs: int,
|
|
max_nodes: int,
|
|
max_conns: int,
|
|
node_gene: BaseNodeGene = DefaultNodeGene(),
|
|
conn_gene: BaseConnGene = DefaultConnGene(),
|
|
):
|
|
self.num_inputs = num_inputs
|
|
self.num_outputs = num_outputs
|
|
self.input_idx = jnp.arange(num_inputs)
|
|
self.output_idx = jnp.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
|
|
|
|
def transform(self, nodes, conns):
|
|
raise NotImplementedError
|
|
|
|
def forward(self, inputs, transformed):
|
|
raise NotImplementedError
|
|
|
|
def add_node(self, 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(self, 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(self, conns, i_key, o_key, enable: bool, 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:3].set(jnp.array([i_key, o_key, enable]))
|
|
return new_conns.at[pos, 3:].set(attrs)
|
|
|
|
def delete_conn_by_pos(self, conns, pos):
|
|
"""
|
|
Delete a connection from the genome.
|
|
Delete the connection by its idx.
|
|
"""
|
|
return conns.at[pos].set(jnp.nan)
|