add gene type RNN

test/unit/test_graphs.py

@@ -0,0 +1,32 @@
+import jax.numpy as jnp
+
+from algorithm.neat.genome.graph import topological_sort, check_cycles
+from algorithm.neat.utils import I_INT
+
+nodes = jnp.array([
+    [0],
+    [1],
+    [2],
+    [3],
+    [jnp.nan]
+])
+
+# {(0, 2), (1, 2), (1, 3), (2, 3)}
+conns = jnp.array([
+    [0, 0, 1, 0, 0],
+    [0, 0, 1, 1, 0],
+    [0, 0, 0, 1, 0],
+    [0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0]
+])
+
+
+def test_topological_sort():
+    assert jnp.all(topological_sort(nodes, conns) == jnp.array([0, 1, 2, 3, I_INT]))
+
+
+def test_check_cycles():
+    assert check_cycles(nodes, conns, 3, 2)
+    assert ~check_cycles(nodes, conns, 2, 3)
+    assert ~check_cycles(nodes, conns, 0, 3)
+    assert ~check_cycles(nodes, conns, 1, 0)

test/unit/test_utils.py

@@ -1,7 +1,5 @@
-import pytest
-import jax
-
-from algorithm.neat.utils import *
+import jax.numpy as jnp
+from algorithm.neat.utils import unflatten_connections
 
 
 def test_unflatten():
@@ -13,7 +11,6 @@ def test_unflatten():
         [jnp.nan, jnp.nan, jnp.nan, jnp.nan]
     ])
 
-
     conns = jnp.array([
         [0, 1, True, 0.1, 0.11],
         [0, 2, False, 0.2, 0.22],
@@ -33,4 +30,4 @@ def test_unflatten():
     mask = mask.at[:, [0, 0, 1, 1], [1, 2, 2, 3]].set(False)
 
     # Ensure all other places are jnp.nan
-    assert jnp.all(jnp.isnan(res[mask]))
+    assert jnp.all(jnp.isnan(res[mask]))