tony/tensorneat-mend
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update some files for save

Browse Source
This commit is contained in:
root
2024-07-15 14:25:51 +08:00
parent f032564a43
commit 6edf083d4f
12 changed files with 110 additions and 111 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
src/tensorneat/genome/base.py
View File

@@ -96,9 +96,9 @@ class BaseGenome(StatefulBaseClass):
def setup(self, state=State()):
state = self.node_gene.setup(state)
state = self.conn_gene.setup(state)
state = self.mutation.setup(state, self)
state = self.crossover.setup(state, self)
state = self.distance.setup(state, self)
state = self.mutation.setup(state)
state = self.crossover.setup(state)
state = self.distance.setup(state)
return state
def transform(self, state, nodes, conns):
@@ -114,13 +114,13 @@ class BaseGenome(StatefulBaseClass):
raise NotImplementedError
def execute_mutation(self, state, randkey, nodes, conns, new_node_key):
return self.mutation(state, randkey, nodes, conns, new_node_key)
return self.mutation(state, self, randkey, nodes, conns, new_node_key)
def execute_crossover(self, state, randkey, nodes1, conns1, nodes2, conns2):
return self.crossover(state, randkey, nodes1, conns1, nodes2, conns2)
return self.crossover(state, self, randkey, nodes1, conns1, nodes2, conns2)
def execute_distance(self, state, nodes1, conns1, nodes2, conns2):
return self.distance(state, nodes1, conns1, nodes2, conns2)
return self.distance(state, self, nodes1, conns1, nodes2, conns2)
def initialize(self, state, randkey):
k1, k2 = jax.random.split(randkey) # k1 for nodes, k2 for conns

2
src/tensorneat/genome/gene/node/bias.py
View File

@@ -13,7 +13,7 @@ from tensorneat.common import (
get_func_name
)
from . import BaseNode
from .base import BaseNode
class BiasNode(BaseNode):

7
src/tensorneat/genome/operations/crossover/base.py
View File

@@ -3,10 +3,5 @@ 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):
def __call__(self, state, genome, randkey, nodes1, nodes2, conns1, conns2):
raise NotImplementedError

10
src/tensorneat/genome/operations/crossover/default.py
View File

@@ -11,14 +11,14 @@ from ...utils import (
class DefaultCrossover(BaseCrossover):
def __call__(self, state, randkey, nodes1, conns1, nodes2, conns2):
def __call__(self, state, genome, 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)
randkeys1 = jax.random.split(randkey1, genome.max_nodes)
randkeys2 = jax.random.split(randkey2, genome.max_conns)
# crossover nodes
keys1, keys2 = nodes1[:, 0], nodes2[:, 0]
@@ -33,7 +33,7 @@ class DefaultCrossover(BaseCrossover):
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))(
vmap(genome.node_gene.crossover, in_axes=(None, 0, 0, 0))(
state, randkeys1, node_attrs1, node_attrs2
), # homologous or both nan
)
@@ -49,7 +49,7 @@ class DefaultCrossover(BaseCrossover):
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))(
vmap(genome.conn_gene.crossover, in_axes=(None, 0, 0, 0))(
state, randkeys2, conns_attrs1, conns_attrs2
), # homologous or both nan
)

7
src/tensorneat/genome/operations/distance/base.py
View File

@@ -3,12 +3,7 @@ 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):
def __call__(self, state, genome, nodes1, nodes2, conns1, conns2):
"""
The distance between two genomes
"""

14
src/tensorneat/genome/operations/distance/default.py
View File

@@ -13,16 +13,16 @@ class DefaultDistance(BaseDistance):
self.compatibility_disjoint = compatibility_disjoint
self.compatibility_weight = compatibility_weight
def __call__(self, state, nodes1, conns1, nodes2, conns2):
def __call__(self, state, genome, nodes1, conns1, nodes2, conns2):
"""
The distance between two genomes
"""
d = self.node_distance(state, nodes1, nodes2) + self.conn_distance(
state, conns1, conns2
d = self.node_distance(state, genome, nodes1, nodes2) + self.conn_distance(
state, genome, conns1, conns2
)
return d
def node_distance(self, state, nodes1, nodes2):
def node_distance(self, state, genome, nodes1, nodes2):
"""
The distance of the nodes part for two genomes
"""
@@ -50,7 +50,7 @@ class DefaultDistance(BaseDistance):
# 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))(
hnd = vmap(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)
@@ -65,7 +65,7 @@ class DefaultDistance(BaseDistance):
return val
def conn_distance(self, state, conns1, conns2):
def conn_distance(self, state, genome, conns1, conns2):
"""
The distance of the conns part for two genomes
"""
@@ -89,7 +89,7 @@ class DefaultDistance(BaseDistance):
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))(
hcd = vmap(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)

7
src/tensorneat/genome/operations/mutation/base.py
View File

@@ -3,10 +3,5 @@ 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, nodes, conns, new_node_key):
def __call__(self, state, genome, randkey, nodes, conns, new_node_key):
raise NotImplementedError

42
src/tensorneat/genome/operations/mutation/default.py
View File

@@ -33,17 +33,17 @@ class DefaultMutation(BaseMutation):
self.node_add = node_add
self.node_delete = node_delete
def __call__(self, state, randkey, nodes, conns, new_node_key):
def __call__(self, state, genome, randkey, nodes, conns, new_node_key):
k1, k2 = jax.random.split(randkey)
nodes, conns = self.mutate_structure(
state, k1, nodes, conns, new_node_key
state, genome, k1, nodes, conns, new_node_key
)
nodes, conns = self.mutate_values(state, k2, nodes, conns)
nodes, conns = self.mutate_values(state, genome, k2, nodes, conns)
return nodes, conns
def mutate_structure(self, state, randkey, nodes, conns, new_node_key):
def mutate_structure(self, state, genome, randkey, nodes, conns, new_node_key):
def mutate_add_node(key_, nodes_, conns_):
"""
add a node while do not influence the output of the network
@@ -62,7 +62,7 @@ class DefaultMutation(BaseMutation):
# add a new node with identity attrs
new_nodes = add_node(
nodes_, new_node_key, self.genome.node_gene.new_identity_attrs(state)
nodes_, new_node_key, genome.node_gene.new_identity_attrs(state)
)
# add two new connections
@@ -71,7 +71,7 @@ class DefaultMutation(BaseMutation):
new_conns,
i_key,
new_node_key,
self.genome.conn_gene.new_identity_attrs(state),
genome.conn_gene.new_identity_attrs(state),
)
# second is with the origin attrs
new_conns = add_conn(
@@ -97,8 +97,8 @@ class DefaultMutation(BaseMutation):
key, idx = self.choose_node_key(
key_,
nodes_,
self.genome.input_idx,
self.genome.output_idx,
genome.input_idx,
genome.output_idx,
allow_input_keys=False,
allow_output_keys=False,
)
@@ -136,8 +136,8 @@ class DefaultMutation(BaseMutation):
i_key, from_idx = self.choose_node_key(
k1_,
nodes_,
self.genome.input_idx,
self.genome.output_idx,
genome.input_idx,
genome.output_idx,
allow_input_keys=True,
allow_output_keys=True,
)
@@ -146,8 +146,8 @@ class DefaultMutation(BaseMutation):
o_key, to_idx = self.choose_node_key(
k2_,
nodes_,
self.genome.input_idx,
self.genome.output_idx,
genome.input_idx,
genome.output_idx,
allow_input_keys=False,
allow_output_keys=True,
)
@@ -161,10 +161,10 @@ class DefaultMutation(BaseMutation):
def successful():
# add a connection with zero attrs
return nodes_, add_conn(
conns_, i_key, o_key, self.genome.conn_gene.new_zero_attrs(state)
conns_, i_key, o_key, genome.conn_gene.new_zero_attrs(state)
)
if self.genome.network_type == "feedforward":
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)
@@ -175,7 +175,7 @@ class DefaultMutation(BaseMutation):
successful,
)
elif self.genome.network_type == "recurrent":
elif genome.network_type == "recurrent":
return jax.lax.cond(
is_already_exist | (remain_conn_space < 1),
nothing,
@@ -183,7 +183,7 @@ class DefaultMutation(BaseMutation):
)
else:
raise ValueError(f"Invalid network type: {self.genome.network_type}")
raise ValueError(f"Invalid network type: {genome.network_type}")
def mutate_delete_conn(key_, nodes_, conns_):
# randomly choose a connection
@@ -223,19 +223,19 @@ class DefaultMutation(BaseMutation):
return nodes, conns
def mutate_values(self, state, randkey, nodes, conns):
def mutate_values(self, state, genome, randkey, nodes, conns):
k1, k2 = jax.random.split(randkey)
nodes_randkeys = jax.random.split(k1, num=self.genome.max_nodes)
conns_randkeys = jax.random.split(k2, num=self.genome.max_conns)
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(self.genome.node_gene.mutate, in_axes=(None, 0, 0))(
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(self.genome.conn_gene.mutate, in_axes=(None, 0, 0))(
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)