modifying

algorithms/neat/genome/genome.py

@@ -0,0 +1,181 @@
+"""
+Vectorization of genome representation.
+
+Utilizes Tuple[nodes: Array(N, 5), connections: Array(C, 4)] to encode the genome, where:
+nodes: [key, bias, response, act, agg]
+connections: [in_key, out_key, weight, enable]
+N: Maximum number of nodes in the network.
+C: Maximum number of connections in the network.
+"""
+
+from typing import Tuple, Dict
+
+import numpy as np
+from numpy.typing import NDArray
+from jax import jit, numpy as jnp
+
+from .utils import fetch_first
+
+
+def initialize_genomes(N: int, C: int, config: Dict) -> Tuple[NDArray, NDArray]:
+    """
+    Initialize genomes with default values.
+
+    Args:
+        N (int): Maximum number of nodes in the network.
+        C (int): Maximum number of connections in the network.
+        config (Dict): Configuration dictionary.
+
+    Returns:
+        Tuple[NDArray, NDArray, NDArray, NDArray]: pop_nodes, pop_connections, input_idx, and output_idx arrays.
+    """
+    # Reserve one row for potential mutation adding an extra node
+    assert config['num_inputs'] + config['num_outputs'] + 1 <= N, \
+        f"Too small N: {N} for input_size: {config['num_inputs']} and output_size: {config['num_inputs']}!"
+
+    assert config['num_inputs'] * config['num_outputs'] + 1 <= C, \
+        f"Too small C: {C} for input_size: {config['num_inputs']} and output_size: {config['num_outputs']}!"
+
+    pop_nodes = np.full((config['pop_size'], N, 5), np.nan)
+    pop_cons = np.full((config['pop_size'], C, 4), np.nan)
+    input_idx = config['input_idx']
+    output_idx = config['output_idx']
+
+    pop_nodes[:, input_idx, 0] = input_idx
+    pop_nodes[:, output_idx, 0] = output_idx
+
+    # pop_nodes[:, output_idx, 1] = config['bias_init_mean']
+    pop_nodes[:, output_idx, 1] = np.random.normal(loc=config['bias_init_mean'], scale=config['bias_init_std'],
+                                                   size=(config['pop_size'], 1))
+    pop_nodes[:, output_idx, 2] = np.random.normal(loc=config['response_init_mean'], scale=config['response_init_std'],
+                                                   size=(config['pop_size'], 1))
+    pop_nodes[:, output_idx, 3] = np.random.choice(config['activation_options'], size=(config['pop_size'], 1))
+    pop_nodes[:, output_idx, 4] = np.random.choice(config['aggregation_options'], size=(config['pop_size'], 1))
+
+    grid_a, grid_b = np.meshgrid(input_idx, output_idx)
+    grid_a, grid_b = grid_a.flatten(), grid_b.flatten()
+
+    p = config['num_inputs'] * config['num_outputs']
+    pop_cons[:, :p, 0] = grid_a
+    pop_cons[:, :p, 1] = grid_b
+    pop_cons[:, :p, 2] = np.random.normal(loc=config['weight_init_mean'], scale=config['weight_init_std'],
+                                                   size=(config['pop_size'], p))
+    pop_cons[:, :p, 3] = 1
+
+    return pop_nodes, pop_cons
+
+
+def expand_single(nodes: NDArray, cons: NDArray, new_N: int, new_C: int) -> Tuple[NDArray, NDArray]:
+    """
+    Expand a single genome to accommodate more nodes or connections.
+    :param nodes: (N, 5)
+    :param cons:  (C, 4)
+    :param new_N:
+    :param new_C:
+    :return: (new_N, 5), (new_C, 4)
+    """
+    old_N, old_C = nodes.shape[0], cons.shape[0]
+    new_nodes = np.full((new_N, 5), np.nan)
+    new_nodes[:old_N, :] = nodes
+
+    new_cons = np.full((new_C, 4), np.nan)
+    new_cons[:old_C, :] = cons
+
+    return new_nodes, new_cons
+
+
+def expand(pop_nodes: NDArray, pop_cons: NDArray, new_N: int, new_C: int) -> Tuple[NDArray, NDArray]:
+    """
+    Expand the population to accommodate more nodes or connections.
+    :param pop_nodes: (pop_size, N, 5)
+    :param pop_cons:  (pop_size, C, 4)
+    :param new_N:
+    :param new_C:
+    :return: (pop_size, new_N, 5), (pop_size, new_C, 4)
+    """
+    pop_size, old_N, old_C = pop_nodes.shape[0], pop_nodes.shape[1], pop_cons.shape[1]
+
+    new_pop_nodes = np.full((pop_size, new_N, 5), np.nan)
+    new_pop_nodes[:, :old_N, :] = pop_nodes
+
+    new_pop_cons = np.full((pop_size, new_C, 4), np.nan)
+    new_pop_cons[:, :old_C, :] = pop_cons
+
+    return new_pop_nodes, new_pop_cons
+
+
+@jit
+def count(nodes: NDArray, cons: NDArray) -> Tuple[NDArray, NDArray]:
+    """
+    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
+
+
+@jit
+def add_node(nodes: NDArray, cons: NDArray, new_key: int,
+             bias: float = 0.0, response: float = 1.0, act: int = 0, agg: int = 0) -> Tuple[NDArray, NDArray]:
+    """
+    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].set(jnp.array([new_key, bias, response, act, agg]))
+    return nodes, cons
+
+
+@jit
+def delete_node(nodes: NDArray, cons: NDArray, node_key: int) -> Tuple[NDArray, NDArray]:
+    """
+    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)
+
+
+@jit
+def delete_node_by_idx(nodes: NDArray, cons: NDArray, idx: int) -> Tuple[NDArray, NDArray]:
+    """
+    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
+
+
+@jit
+def add_connection(nodes: NDArray, cons: NDArray, i_key: int, o_key: int,
+                   weight: float = 1.0, enabled: bool = True) -> Tuple[NDArray, NDArray]:
+    """
+    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].set(jnp.array([i_key, o_key, weight, enabled]))
+    return nodes, cons
+
+
+@jit
+def delete_connection(nodes: NDArray, cons: NDArray, i_key: int, o_key: int) -> Tuple[NDArray, NDArray]:
+    """
+    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)
+
+
+@jit
+def delete_connection_by_idx(nodes: NDArray, cons: NDArray, idx: int) -> Tuple[NDArray, NDArray]:
+    """
+    Delete a connection from the genome.
+    Delete the connection by its idx.
+    """
+    cons = cons.at[idx].set(np.nan)
+    return nodes, cons