change a lot a lot a lot!!!!!!!
This commit is contained in:
wls2002
@@ -1,5 +0,0 @@
|
||||
from .basic import initialize_genomes
|
||||
from .mutate import create_mutate
|
||||
from .distance import create_distance
|
||||
from .crossover import crossover
|
||||
from .graph import topological_sort
|
||||
@@ -1,111 +0,0 @@
|
||||
from typing import Type, Tuple
|
||||
|
||||
import numpy as np
|
||||
import jax
|
||||
from jax import Array, numpy as jnp
|
||||
|
||||
from algorithm import State
|
||||
from ..gene import BaseGene
|
||||
from algorithm.utils import fetch_first
|
||||
|
||||
|
||||
def initialize_genomes(state: State, gene_type: Type[BaseGene]):
|
||||
o_nodes = np.full((state.N, state.NL), np.nan, dtype=np.float32) # original nodes
|
||||
o_conns = np.full((state.C, state.CL), np.nan, dtype=np.float32) # original connections
|
||||
|
||||
input_idx = state.input_idx
|
||||
output_idx = state.output_idx
|
||||
new_node_key = max([*input_idx, *output_idx]) + 1
|
||||
|
||||
o_nodes[input_idx, 0] = input_idx
|
||||
o_nodes[output_idx, 0] = output_idx
|
||||
o_nodes[new_node_key, 0] = new_node_key
|
||||
o_nodes[np.concatenate([input_idx, output_idx]), 1:] = jax.device_get(gene_type.new_node_attrs(state))
|
||||
o_nodes[new_node_key, 1:] = jax.device_get(gene_type.new_node_attrs(state))
|
||||
|
||||
input_conns = np.c_[input_idx, np.full_like(input_idx, new_node_key)]
|
||||
o_conns[input_idx, 0:2] = input_conns # in key, out key
|
||||
o_conns[input_idx, 2] = True # enabled
|
||||
o_conns[input_idx, 3:] = jax.device_get(gene_type.new_conn_attrs(state))
|
||||
|
||||
output_conns = np.c_[np.full_like(output_idx, new_node_key), output_idx]
|
||||
o_conns[output_idx, 0:2] = output_conns # in key, out key
|
||||
o_conns[output_idx, 2] = True # enabled
|
||||
o_conns[output_idx, 3:] = jax.device_get(gene_type.new_conn_attrs(state))
|
||||
|
||||
# repeat origin genome for P times to create population
|
||||
pop_nodes = np.tile(o_nodes, (state.P, 1, 1))
|
||||
pop_conns = np.tile(o_conns, (state.P, 1, 1))
|
||||
|
||||
return jax.device_put([pop_nodes, pop_conns])
|
||||
|
||||
|
||||
def count(nodes: Array, cons: Array):
|
||||
"""
|
||||
Count how many nodes and connections are in the genome.
|
||||
"""
|
||||
node_cnt = jnp.sum(~jnp.isnan(nodes[:, 0]))
|
||||
cons_cnt = jnp.sum(~jnp.isnan(cons[:, 0]))
|
||||
return node_cnt, cons_cnt
|
||||
|
||||
|
||||
def add_node(nodes: Array, cons: Array, new_key: int, attrs: Array) -> Tuple[Array, Array]:
|
||||
"""
|
||||
Add a new node to the genome.
|
||||
The new node will place at the first NaN row.
|
||||
"""
|
||||
exist_keys = nodes[:, 0]
|
||||
idx = fetch_first(jnp.isnan(exist_keys))
|
||||
nodes = nodes.at[idx, 0].set(new_key)
|
||||
nodes = nodes.at[idx, 1:].set(attrs)
|
||||
return nodes, cons
|
||||
|
||||
|
||||
def delete_node(nodes: Array, cons: Array, node_key: Array) -> Tuple[Array, Array]:
|
||||
"""
|
||||
Delete a node from the genome. Only delete the node, regardless of connections.
|
||||
Delete the node by its key.
|
||||
"""
|
||||
node_keys = nodes[:, 0]
|
||||
idx = fetch_first(node_keys == node_key)
|
||||
return delete_node_by_idx(nodes, cons, idx)
|
||||
|
||||
|
||||
def delete_node_by_idx(nodes: Array, cons: Array, idx: Array) -> Tuple[Array, Array]:
|
||||
"""
|
||||
Delete a node from the genome. Only delete the node, regardless of connections.
|
||||
Delete the node by its idx.
|
||||
"""
|
||||
nodes = nodes.at[idx].set(np.nan)
|
||||
return nodes, cons
|
||||
|
||||
|
||||
def add_connection(nodes: Array, cons: Array, i_key: Array, o_key: Array, enable: bool, attrs: Array) -> Tuple[
|
||||
Array, Array]:
|
||||
"""
|
||||
Add a new connection to the genome.
|
||||
The new connection will place at the first NaN row.
|
||||
"""
|
||||
con_keys = cons[:, 0]
|
||||
idx = fetch_first(jnp.isnan(con_keys))
|
||||
cons = cons.at[idx, 0:3].set(jnp.array([i_key, o_key, enable]))
|
||||
cons = cons.at[idx, 3:].set(attrs)
|
||||
return nodes, cons
|
||||
|
||||
|
||||
def delete_connection(nodes: Array, cons: Array, i_key: Array, o_key: Array) -> Tuple[Array, Array]:
|
||||
"""
|
||||
Delete a connection from the genome.
|
||||
Delete the connection by its input and output node keys.
|
||||
"""
|
||||
idx = fetch_first((cons[:, 0] == i_key) & (cons[:, 1] == o_key))
|
||||
return delete_connection_by_idx(nodes, cons, idx)
|
||||
|
||||
|
||||
def delete_connection_by_idx(nodes: Array, cons: Array, idx: Array) -> Tuple[Array, Array]:
|
||||
"""
|
||||
Delete a connection from the genome.
|
||||
Delete the connection by its idx.
|
||||
"""
|
||||
cons = cons.at[idx].set(np.nan)
|
||||
return nodes, cons
|
||||
@@ -1,66 +0,0 @@
|
||||
import jax
|
||||
from jax import jit, Array, numpy as jnp
|
||||
|
||||
|
||||
def crossover(randkey, nodes1: Array, conns1: Array, nodes2: Array, conns2: Array):
|
||||
"""
|
||||
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 = 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, crossover_gene(randkey_1, nodes1, nodes2))
|
||||
|
||||
# crossover connections
|
||||
con_keys1, con_keys2 = conns1[:, :2], conns2[:, :2]
|
||||
cons2 = align_array(con_keys1, con_keys2, conns2, True)
|
||||
new_cons = jnp.where(jnp.isnan(conns1) | jnp.isnan(cons2), conns1, crossover_gene(randkey_2, conns1, cons2))
|
||||
|
||||
return new_nodes, new_cons
|
||||
|
||||
|
||||
def align_array(seq1: Array, seq2: Array, ar2: Array, is_conn: bool) -> Array:
|
||||
"""
|
||||
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(rand_key: Array, g1: Array, g2: Array) -> Array:
|
||||
"""
|
||||
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)
|
||||
@@ -1,76 +0,0 @@
|
||||
from typing import Dict, Type
|
||||
|
||||
from jax import Array, numpy as jnp, vmap
|
||||
|
||||
from ..gene import BaseGene
|
||||
|
||||
|
||||
def create_distance(config: Dict, gene_type: Type[BaseGene]):
|
||||
def node_distance(state, nodes1: Array, nodes2: Array):
|
||||
"""
|
||||
Calculate the distance between nodes of two genomes.
|
||||
"""
|
||||
# statistics nodes count of 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
|
||||
hnd = vmap(gene_type.distance_node, in_axes=(None, 0, 0))(state, fr, sr)
|
||||
hnd = jnp.where(jnp.isnan(hnd), 0, hnd)
|
||||
homologous_distance = jnp.sum(hnd * intersect_mask)
|
||||
|
||||
val = non_homologous_cnt * config['compatibility_disjoint'] + homologous_distance * config[
|
||||
'compatibility_weight']
|
||||
|
||||
return jnp.where(max_cnt == 0, 0, val / max_cnt) # avoid zero division
|
||||
|
||||
def connection_distance(state, cons1: Array, cons2: Array):
|
||||
"""
|
||||
Calculate the distance between connections of two genomes.
|
||||
Similar process as node_distance.
|
||||
"""
|
||||
con_cnt1 = jnp.sum(~jnp.isnan(cons1[:, 0]))
|
||||
con_cnt2 = jnp.sum(~jnp.isnan(cons2[:, 0]))
|
||||
max_cnt = jnp.maximum(con_cnt1, con_cnt2)
|
||||
|
||||
cons = jnp.concatenate((cons1, cons2), 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)
|
||||
hcd = vmap(gene_type.distance_conn, in_axes=(None, 0, 0))(state, fr, sr)
|
||||
hcd = jnp.where(jnp.isnan(hcd), 0, hcd)
|
||||
homologous_distance = jnp.sum(hcd * intersect_mask)
|
||||
|
||||
val = non_homologous_cnt * config['compatibility_disjoint'] + homologous_distance * config[
|
||||
'compatibility_weight']
|
||||
|
||||
return jnp.where(max_cnt == 0, 0, val / max_cnt)
|
||||
|
||||
def distance(state, nodes1, conns1, nodes2, conns2):
|
||||
return node_distance(state, nodes1, nodes2) + connection_distance(state, conns1, conns2)
|
||||
|
||||
return distance
|
||||
|
||||
@@ -1,67 +0,0 @@
|
||||
"""
|
||||
Some graph algorithm implemented in jax.
|
||||
Only used in feed-forward networks.
|
||||
"""
|
||||
|
||||
import jax
|
||||
from jax import jit, Array, numpy as jnp
|
||||
|
||||
from algorithm.utils import fetch_first, I_INT
|
||||
|
||||
|
||||
@jit
|
||||
def topological_sort(nodes: Array, conns: Array) -> Array:
|
||||
"""
|
||||
a jit-able version of topological_sort!
|
||||
"""
|
||||
|
||||
in_degree = jnp.where(jnp.isnan(nodes[:, 0]), jnp.nan, jnp.sum(conns, axis=0))
|
||||
res = jnp.full(in_degree.shape, I_INT)
|
||||
|
||||
def cond_fun(carry):
|
||||
res_, idx_, in_degree_ = carry
|
||||
i = fetch_first(in_degree_ == 0.)
|
||||
return i != I_INT
|
||||
|
||||
def body_func(carry):
|
||||
res_, idx_, in_degree_ = carry
|
||||
i = fetch_first(in_degree_ == 0.)
|
||||
|
||||
# add to res and flag it is already in it
|
||||
res_ = res_.at[idx_].set(i)
|
||||
in_degree_ = in_degree_.at[i].set(-1)
|
||||
|
||||
# decrease in_degree of all its children
|
||||
children = conns[i, :]
|
||||
in_degree_ = jnp.where(children, in_degree_ - 1, in_degree_)
|
||||
return res_, idx_ + 1, in_degree_
|
||||
|
||||
res, _, _ = jax.lax.while_loop(cond_fun, body_func, (res, 0, in_degree))
|
||||
return res
|
||||
|
||||
|
||||
@jit
|
||||
def check_cycles(nodes: Array, conns: Array, from_idx, to_idx) -> Array:
|
||||
"""
|
||||
Check whether a new connection (from_idx -> to_idx) will cause a cycle.
|
||||
"""
|
||||
|
||||
conns = conns.at[from_idx, to_idx].set(True)
|
||||
|
||||
visited = jnp.full(nodes.shape[0], False)
|
||||
new_visited = visited.at[to_idx].set(True)
|
||||
|
||||
def cond_func(carry):
|
||||
visited_, new_visited_ = carry
|
||||
end_cond1 = jnp.all(visited_ == new_visited_) # no new nodes been visited
|
||||
end_cond2 = new_visited_[from_idx] # the starting node has been visited
|
||||
return jnp.logical_not(end_cond1 | end_cond2)
|
||||
|
||||
def body_func(carry):
|
||||
_, visited_ = carry
|
||||
new_visited_ = jnp.dot(visited_, conns)
|
||||
new_visited_ = jnp.logical_or(visited_, new_visited_)
|
||||
return visited_, new_visited_
|
||||
|
||||
_, visited = jax.lax.while_loop(cond_func, body_func, (visited, new_visited))
|
||||
return visited[from_idx]
|
||||
@@ -1,205 +0,0 @@
|
||||
from typing import Dict, Tuple, Type
|
||||
|
||||
import jax
|
||||
from jax import Array, numpy as jnp, vmap
|
||||
|
||||
from algorithm import State
|
||||
from .basic import add_node, add_connection, delete_node_by_idx, delete_connection_by_idx
|
||||
from .graph import check_cycles
|
||||
from algorithm.utils import fetch_random, fetch_first, I_INT, unflatten_connections
|
||||
from ..gene import BaseGene
|
||||
|
||||
|
||||
def create_mutate(config: Dict, gene_type: Type[BaseGene]):
|
||||
"""
|
||||
Create function to mutate a single genome
|
||||
"""
|
||||
|
||||
def mutate_structure(state: State, randkey, nodes, conns, new_node_key):
|
||||
|
||||
def mutate_add_node(key_, nodes_, conns_):
|
||||
i_key, o_key, idx = choice_connection_key(key_, nodes_, conns_)
|
||||
|
||||
def nothing():
|
||||
return nodes_, conns_
|
||||
|
||||
def successful_add_node():
|
||||
# disable the connection
|
||||
aux_nodes, aux_conns = nodes_, conns_
|
||||
|
||||
# set enable to false
|
||||
aux_conns = aux_conns.at[idx, 2].set(False)
|
||||
|
||||
# add a new node
|
||||
aux_nodes, aux_conns = add_node(aux_nodes, aux_conns, new_node_key, gene_type.new_node_attrs(state))
|
||||
|
||||
# add two new connections
|
||||
aux_nodes, aux_conns = add_connection(aux_nodes, aux_conns, i_key, new_node_key, True,
|
||||
gene_type.new_conn_attrs(state))
|
||||
aux_nodes, aux_conns = add_connection(aux_nodes, aux_conns, new_node_key, o_key, True,
|
||||
gene_type.new_conn_attrs(state))
|
||||
|
||||
return aux_nodes, aux_conns
|
||||
|
||||
# if from_idx == I_INT, that means no connection exist, do nothing
|
||||
new_nodes, new_conns = jax.lax.cond(idx == I_INT, nothing, successful_add_node)
|
||||
|
||||
return new_nodes, new_conns
|
||||
|
||||
def mutate_delete_node(key_, nodes_, conns_):
|
||||
# TODO: Do we really need to delete a node?
|
||||
# randomly choose a node
|
||||
key, idx = choice_node_key(key_, nodes_, config['input_idx'], config['output_idx'],
|
||||
allow_input_keys=False, allow_output_keys=False)
|
||||
def nothing():
|
||||
return nodes_, conns_
|
||||
|
||||
def successful_delete_node():
|
||||
# delete the node
|
||||
aux_nodes, aux_cons = delete_node_by_idx(nodes_, conns_, idx)
|
||||
|
||||
# delete all connections
|
||||
aux_cons = jnp.where(((aux_cons[:, 0] == key) | (aux_cons[:, 1] == key))[:, None],
|
||||
jnp.nan, aux_cons)
|
||||
|
||||
return aux_nodes, aux_cons
|
||||
|
||||
return jax.lax.cond(idx == I_INT, nothing, successful_delete_node)
|
||||
|
||||
def mutate_add_conn(key_, nodes_, conns_):
|
||||
# randomly choose two nodes
|
||||
k1_, k2_ = jax.random.split(key_, num=2)
|
||||
i_key, from_idx = choice_node_key(k1_, nodes_, config['input_idx'], config['output_idx'],
|
||||
allow_input_keys=True, allow_output_keys=True)
|
||||
o_key, to_idx = choice_node_key(k2_, nodes_, config['input_idx'], config['output_idx'],
|
||||
allow_input_keys=False, allow_output_keys=True)
|
||||
|
||||
con_idx = fetch_first((conns_[:, 0] == i_key) & (conns_[:, 1] == o_key))
|
||||
|
||||
def nothing():
|
||||
return nodes_, conns_
|
||||
|
||||
def successful():
|
||||
new_nodes, new_cons = add_connection(nodes_, conns_, i_key, o_key, True, gene_type.new_conn_attrs(state))
|
||||
return new_nodes, new_cons
|
||||
|
||||
def already_exist():
|
||||
new_cons = conns_.at[con_idx, 2].set(True)
|
||||
return nodes_, new_cons
|
||||
|
||||
is_already_exist = con_idx != I_INT
|
||||
|
||||
if config['network_type'] == 'feedforward':
|
||||
u_cons = unflatten_connections(nodes_, conns_)
|
||||
cons_exist = jnp.where(~jnp.isnan(u_cons[0, :, :]), True, False)
|
||||
is_cycle = check_cycles(nodes_, cons_exist, from_idx, to_idx)
|
||||
|
||||
choice = jnp.where(is_already_exist, 0, jnp.where(is_cycle, 1, 2))
|
||||
return jax.lax.switch(choice, [already_exist, nothing, successful])
|
||||
|
||||
elif config['network_type'] == 'recurrent':
|
||||
return jax.lax.cond(is_already_exist, already_exist, successful)
|
||||
|
||||
else:
|
||||
raise ValueError(f"Invalid network type: {config['network_type']}")
|
||||
|
||||
def mutate_delete_conn(key_, nodes_, conns_):
|
||||
# randomly choose a connection
|
||||
i_key, o_key, idx = choice_connection_key(key_, nodes_, conns_)
|
||||
|
||||
def nothing():
|
||||
return nodes_, conns_
|
||||
|
||||
def successfully_delete_connection():
|
||||
return delete_connection_by_idx(nodes_, conns_, idx)
|
||||
|
||||
return jax.lax.cond(idx == I_INT, 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(k, n, c):
|
||||
return n, c
|
||||
|
||||
nodes, conns = jax.lax.cond(r1 < config['node_add_prob'], mutate_add_node, no, k1, nodes, conns)
|
||||
|
||||
nodes, conns = jax.lax.cond(r2 < config['node_delete_prob'], mutate_delete_node, no, k2, nodes, conns)
|
||||
|
||||
nodes, conns = jax.lax.cond(r3 < config['conn_add_prob'], mutate_add_conn, no, k3, nodes, conns)
|
||||
|
||||
nodes, conns = jax.lax.cond(r4 < config['conn_delete_prob'], mutate_delete_conn, no, k4, nodes, conns)
|
||||
|
||||
return nodes, conns
|
||||
|
||||
def mutate_values(state: State, randkey, 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])
|
||||
|
||||
nodes_attrs, conns_attrs = nodes[:, 1:], conns[:, 3:]
|
||||
|
||||
new_nodes_attrs = vmap(gene_type.mutate_node, in_axes=(None, 0, 0))(state, nodes_attrs, nodes_keys)
|
||||
new_conns_attrs = vmap(gene_type.mutate_conn, in_axes=(None, 0, 0))(state, conns_attrs, conns_keys)
|
||||
|
||||
# nan nodes not changed
|
||||
new_nodes_attrs = jnp.where(jnp.isnan(nodes_attrs), jnp.nan, new_nodes_attrs)
|
||||
new_conns_attrs = jnp.where(jnp.isnan(conns_attrs), jnp.nan, new_conns_attrs)
|
||||
|
||||
new_nodes = nodes.at[:, 1:].set(new_nodes_attrs)
|
||||
new_conns = conns.at[:, 3:].set(new_conns_attrs)
|
||||
|
||||
return new_nodes, new_conns
|
||||
|
||||
def mutate(state, randkey, nodes, conns, new_node_key):
|
||||
k1, k2 = jax.random.split(randkey)
|
||||
|
||||
nodes, conns = mutate_structure(state, k1, nodes, conns, new_node_key)
|
||||
nodes, conns = mutate_values(state, k2, nodes, conns)
|
||||
|
||||
return nodes, conns
|
||||
|
||||
return mutate
|
||||
|
||||
|
||||
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]:
|
||||
"""
|
||||
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_keys:
|
||||
:param output_keys:
|
||||
: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_keys))
|
||||
|
||||
if not allow_output_keys:
|
||||
mask = jnp.logical_and(mask, ~jnp.isin(node_keys, output_keys))
|
||||
|
||||
idx = fetch_random(rand_key, mask)
|
||||
key = jnp.where(idx != I_INT, nodes[idx, 0], jnp.nan)
|
||||
return key, idx
|
||||
|
||||
|
||||
def choice_connection_key(rand_key: Array, nodes: Array, cons: Array) -> Tuple[Array, Array, Array]:
|
||||
"""
|
||||
Randomly choose a connection key from the given connections.
|
||||
:param rand_key:
|
||||
:param nodes:
|
||||
:param cons:
|
||||
:return: i_key, o_key, idx
|
||||
"""
|
||||
|
||||
idx = fetch_random(rand_key, ~jnp.isnan(cons[:, 0]))
|
||||
i_key = jnp.where(idx != I_INT, cons[idx, 0], jnp.nan)
|
||||
o_key = jnp.where(idx != I_INT, cons[idx, 1], jnp.nan)
|
||||
|
||||
return i_key, o_key, idx
|
||||
Reference in New Issue
Block a user