complete normal neat algorithm

algorithm/neat/pipeline.py

@@ -0,0 +1,79 @@
+import time
+from typing import Union, Callable
+
+import jax
+from jax import vmap, jit
+import numpy as np
+
+class Pipeline:
+    """
+    Neat algorithm pipeline.
+    """
+
+    def __init__(self, config, algorithm):
+        self.config = config
+        self.algorithm = algorithm
+        randkey = jax.random.PRNGKey(config['random_seed'])
+        self.state = algorithm.setup(randkey)
+
+        self.best_genome = None
+        self.best_fitness = float('-inf')
+        self.generation_timestamp = time.time()
+
+        self.evaluate_time = 0
+
+        self.forward_func = algorithm.gene_type.create_forward(config)
+        self.batch_forward_func = jit(vmap(self.forward_func, in_axes=(0, None)))
+        self.pop_batch_forward_func = jit(vmap(self.batch_forward_func, in_axes=(None, 0)))
+
+        self.pop_transform_func = jit(vmap(algorithm.gene_type.forward_transform))
+
+    def ask(self):
+        pop_transforms = self.pop_transform_func(self.state.pop_nodes, self.state.pop_conns)
+        return lambda inputs: self.pop_batch_forward_func(inputs, pop_transforms)
+
+    def tell(self, fitness):
+        self.state = self.algorithm.step(self.state, fitness)
+        from algorithm.neat.genome.basic import count
+        # print([count(self.state.pop_nodes[i], self.state.pop_conns[i]) for i in range(self.state.P)])
+
+
+    def auto_run(self, fitness_func, analysis: Union[Callable, str] = "default"):
+        for _ in range(self.config['generation_limit']):
+            forward_func = self.ask()
+
+            fitnesses = fitness_func(forward_func)
+
+            if analysis is not None:
+                if analysis == "default":
+                    self.default_analysis(fitnesses)
+                else:
+                    assert callable(analysis), f"What the fuck you passed in? A {analysis}?"
+                    analysis(fitnesses)
+
+            if max(fitnesses) >= self.config['fitness_threshold']:
+                print("Fitness limit reached!")
+                return self.best_genome
+
+            self.tell(fitnesses)
+        print("Generation limit reached!")
+        return self.best_genome
+
+    def default_analysis(self, fitnesses):
+        max_f, min_f, mean_f, std_f = max(fitnesses), min(fitnesses), np.mean(fitnesses), np.std(fitnesses)
+
+        new_timestamp = time.time()
+        cost_time = new_timestamp - self.generation_timestamp
+        self.generation_timestamp = new_timestamp
+
+        max_idx = np.argmax(fitnesses)
+        if fitnesses[max_idx] > self.best_fitness:
+            self.best_fitness = fitnesses[max_idx]
+            self.best_genome = (self.state.pop_nodes[max_idx], self.state.pop_conns[max_idx])
+
+        member_count = jax.device_get(self.state.species_info[:, 3])
+        species_sizes = [int(i) for i in member_count if i > 0]
+
+        print(f"Generation: {self.state.generation}",
+              f"species: {len(species_sizes)}, {species_sizes}",
+              f"fitness: {max_f}, {min_f}, {mean_f}, {std_f}, Cost time: {cost_time}")