Refactor genetic algorithm

factorio-blueprint-generator/src/assembly.rs

@@ -1,9 +1,8 @@
 use factorio_blueprint::abstraction::{Blueprint, ElectricPoleType, Entity, InserterType};
 use factorio_core::{beltoptions::Beltspeed, prelude::*};
 
-pub fn assembly_line(assembly_machines: usize) -> Blueprint {
+pub fn assembly_line(assembly_machines: usize, recipe: impl AsRef<str>) -> Blueprint {
     let mut blueprint = Blueprint::new();
-
     let mut last = None;
     for i in 0..assembly_machines.div_ceil(3) {
         let top = blueprint.add_entity(Entity::new_electric_pole(
@@ -28,7 +27,7 @@ pub fn assembly_line(assembly_machines: usize) -> Blueprint {
     for i in 0..assembly_machines {
         blueprint.add_entity(Entity::new_production(
             "assembling-machine-3",
-            "iron-gear-wheel",
+            recipe.as_ref(),
             Position::new(3 + 6 * i as PositionType, 9),
             Direction::Up,
             Position::new(6, 6),

factorio-blueprint-generator/src/bin/assembly.rs

@@ -7,11 +7,14 @@ struct Args {
     #[arg(short, long)]
     json: bool,
     assembly_machines: usize,
+    recipe: String,
 }
 
 fn main() {
     let args = Args::parse();
-    let b = BlueprintString::Blueprint(assembly_line(args.assembly_machines).to_blueprint());
+    let b = BlueprintString::Blueprint(
+        assembly_line(args.assembly_machines, args.recipe).to_blueprint(),
+    );
 
     if args.json {
         println!("{}", serde_json::to_string_pretty(&b).unwrap());

factorio-blueprint-generator/src/bin/generate_factory.rs

@@ -2,7 +2,7 @@ use clap::Parser;
 use factorio_blueprint::{BlueprintString, encode};
 use factorio_blueprint_generator::factory::generate_factory;
 use factorio_core::prelude::*;
-use factorio_layout::valid_layout::ValidLayout;
+use factorio_layout::{genetic_algorithm_v1::GeneticAlgorithm, valid_layout::ValidLayout};
 use factorio_pathfinding::belt_finding::ConflictAvoidance;
 use rand::{SeedableRng, rngs::SmallRng};
 
@@ -19,10 +19,19 @@ fn main() {
 
     let l = ValidLayout {
         max_tries: 10,
-        retries: 10,
+        retries: 20,
-        start_size: Position::new(100, 100),
+        start_size: Position::new(20, 20),
         growth: Position::new(5, 5),
     };
+
+    let l = GeneticAlgorithm {
+        mutation_retries: 20,
+        population_size: 10,
+        population_keep: 2,
+        population_new: 2,
+        generations: 20,
+        valid_layout: l,
+    };
     let p = ConflictAvoidance {
         timeout: Some(std::time::Duration::from_millis(20)),
     };

factorio-blueprint-generator/src/factory.rs

@@ -95,10 +95,10 @@ pub fn generate_factory<L: Layouter, P: Pathfinder>(
         .unwrap();
 
     let blueprints = vec![
-        assembly_line(2),
+        assembly_line(2, "iron-gear-wheel"),
-        assembly_line(3),
+        assembly_line(3, "copper-cable"),
-        assembly_line(2),
+        assembly_line(2, "electronic-circuit"),
-        assembly_line(2),
+        assembly_line(2, "inserter"),
     ];
 
     let mut b = Blueprint::new();

factorio-layout/src/genetic_algorithm_v1.rs

@@ -0,0 +1,81 @@
+use crate::{
+    LayoutResult, Layouter,
+    misc::{mutate, path_input_from_blocks_positions, score},
+    valid_layout::ValidLayout,
+};
+
+pub struct GeneticAlgorithm {
+    pub mutation_retries: usize,
+    pub population_size: usize,
+    pub population_keep: usize,
+    pub population_new: usize,
+    pub generations: usize,
+    pub valid_layout: ValidLayout,
+}
+
+impl Layouter for GeneticAlgorithm {
+    fn layout<R: rand::Rng, P: factorio_pathfinding::Pathfinder>(
+        &self,
+        input: &crate::LayoutInput,
+        pathfinder: &P,
+        rng: &mut R,
+    ) -> Option<crate::LayoutResult> {
+        let mut population = Vec::new();
+
+        for _ in 0..self.population_size {
+            loop {
+                if let Some(l) = self.valid_layout.layout(input, pathfinder, rng) {
+                    let score = score(input, &l);
+                    population.push((l, score));
+                    break;
+                }
+            }
+        }
+
+        for g in 0..self.generations {
+            population.sort_by_key(|(_, s)| *s);
+
+            for i in 0..(self.population_size - self.population_keep - self.population_new) {
+                loop {
+                    let parent = &population[i % self.population_keep].0;
+                    if let Some(blocks) = mutate(input, parent, rng) {
+                        let (connections, map) =
+                            path_input_from_blocks_positions(input, parent.size, &blocks);
+
+                        if let Some(paths) =
+                            pathfinder.find_paths(factorio_pathfinding::PathInput {
+                                connections: &connections,
+                                map,
+                            })
+                        {
+                            let r = LayoutResult {
+                                positions: blocks,
+                                path_result: paths,
+                                size: parent.size,
+                            };
+                            let score = score(input, &r);
+                            population.push((r, score));
+                            break;
+                        }
+                    }
+                }
+            }
+
+            for i in 0..self.population_new {
+                loop {
+                    if let Some(l) = self.valid_layout.layout(input, pathfinder, rng) {
+                        let score = score(input, &l);
+                        population[self.population_size - self.population_new + i] = (l, score);
+                        break;
+                    }
+                }
+            }
+
+            println!("completed generation {g}\nscore: {}", population[0].1);
+        }
+
+        population.sort_by_key(|(_, s)| *s);
+
+        population.into_iter().next().map(|(l, _)| l)
+    }
+}

factorio-layout/src/lib.rs

@@ -40,6 +40,7 @@ pub struct LayoutInput {
 
 #[derive(Debug)]
 pub struct LayoutResult {
+    pub size: Position,
     pub positions: Vec<Block>,
     pub path_result: Vec<Vec<PathField>>,
 }

factorio-layout/src/misc.rs

@@ -1,7 +1,7 @@
-use crate::LayoutInput;
+use crate::{LayoutInput, LayoutResult};
 use factorio_core::prelude::*;
-use factorio_pathfinding::{Connection, PathInput, examples::HashMapMap};
+use factorio_pathfinding::{Connection, PathInput, Pathfinder, examples::HashMapMap};
-use rand::Rng;
+use rand::{Rng, seq::SliceRandom};
 
 pub fn initally_set_blocks(
     input: &LayoutInput,
@@ -117,3 +117,154 @@ pub fn path_input_from_blocks_positions(
 
     (connections, map)
 }
+
+pub fn score(input: &LayoutInput, output: &LayoutResult) -> usize {
+    let _ = input;
+    output
+        .path_result
+        .iter()
+        .flatten()
+        .map(|p| p.cost())
+        .sum::<usize>()
+}
+
+pub fn mutate<R: Rng>(
+    input: &LayoutInput,
+    output: &LayoutResult,
+    rng: &mut R,
+) -> Option<Vec<Block>> {
+    let mut blocks = output.positions.clone();
+
+    #[allow(clippy::type_complexity)]
+    let r: &[(
+        &dyn Fn(&LayoutInput, &LayoutResult, &mut Vec<Block>, &mut R) -> bool,
+        _,
+    )] = &[
+        (&mutate_replace::<R>, 30),
+        (&mutate_flip::<R>, 50),
+        (&mutate_jiggle::<R>, 160),
+    ];
+
+    loop {
+        let p = r.choose_weighted(rng, |i| i.1).unwrap().0;
+
+        if p(input, output, &mut blocks, rng) && rng.gen_bool(0.5) {
+            break;
+        }
+    }
+    Some(blocks)
+}
+
+fn mutate_replace<R: Rng>(
+    input: &LayoutInput,
+    output: &LayoutResult,
+    blocks: &mut Vec<Block>,
+    rng: &mut R,
+) -> bool {
+    let _ = input;
+    let i = rng.gen_range(0..blocks.len());
+    let block = blocks[i];
+    let dir = rng.r#gen::<Direction>();
+
+    let size = output.size;
+
+    let pos = match dir {
+        Direction::Up => Position::new(
+            rng.gen_range(0..=(size.x - block.size().x)),
+            rng.gen_range(0..=(size.y - block.size().y)),
+        ),
+        Direction::Right => Position::new(
+            rng.gen_range((block.size().y - 1)..size.x),
+            rng.gen_range(0..=(size.y - block.size().x)),
+        ),
+        Direction::Down => Position::new(
+            rng.gen_range((block.size().x - 1)..size.x),
+            rng.gen_range((block.size().y - 1)..size.y),
+        ),
+        Direction::Left => Position::new(
+            rng.gen_range(0..=(size.x - block.size().y)),
+            rng.gen_range((block.size().x - 1)..size.y),
+        ),
+    };
+
+    let new_block = Block::new(pos, dir, block.size());
+    let new_aabb = new_block.get_aabb();
+
+    if blocks
+        .iter()
+        .enumerate()
+        .all(|(j, b)| i == j || !AABB::collision(b.get_aabb(), new_aabb))
+    {
+        blocks[i] = new_block;
+        true
+    } else {
+        false
+    }
+}
+
+fn mutate_flip<R: Rng>(
+    input: &LayoutInput,
+    output: &LayoutResult,
+    blocks: &mut Vec<Block>,
+    rng: &mut R,
+) -> bool {
+    let _ = output;
+    let _ = input;
+    let i = rng.gen_range(0..blocks.len());
+    let block = blocks[i];
+
+    blocks[i] = Block::new(
+        match block.dir() {
+            Direction::Up => block.pos() + block.size() - Position::new(1, 1),
+            Direction::Right => block.pos() + Position::new(1 - block.size().y, block.size().x - 1),
+            Direction::Down => block.pos() - block.size() + Position::new(1, 1),
+            Direction::Left => block.pos() + Position::new(block.size().y - 1, 1 - block.size().x),
+        },
+        block.dir().reverse(),
+        block.size(),
+    );
+    true
+}
+
+fn mutate_jiggle<R: Rng>(
+    input: &LayoutInput,
+    output: &LayoutResult,
+    blocks: &mut Vec<Block>,
+    rng: &mut R,
+) -> bool {
+    let _ = input;
+    let i = rng.gen_range(0..blocks.len());
+    let block = blocks[i];
+
+    let dir = rng.r#gen::<Direction>();
+    let step = [(1, 10), (2, 5), (3, 5)]
+        .choose_weighted(rng, |i| i.1)
+        .unwrap()
+        .0;
+
+    let new_block = Block::new(
+        block.pos().in_direction(&dir, step),
+        block.dir(),
+        block.size(),
+    );
+    let new_aabb = new_block.get_aabb();
+
+    if new_aabb.min().x < 0
+        || new_aabb.min().y < 0
+        || new_aabb.max().x >= output.size.x
+        || new_aabb.max().y >= output.size.y
+    {
+        return false;
+    }
+
+    if blocks
+        .iter()
+        .enumerate()
+        .all(|(j, b)| j == i || !AABB::collision(b.get_aabb(), new_aabb))
+    {
+        blocks[i] = new_block;
+        true
+    } else {
+        false
+    }
+}

factorio-layout/src/valid_layout.rs

@@ -34,6 +34,7 @@ impl Layouter for ValidLayout {
                     return Some(LayoutResult {
                         positions: blocks,
                         path_result: paths,
+                        size,
                     });
                 }
             }

factorio-pathfinding/src/belt_finding/mod.rs

@@ -47,7 +47,7 @@ impl SinglePathfinder for ConflictAvoidance {
             // c.print_visualization();
 
             if c.remove_all_conflicts(self.timeout) {
-                c.print_visualization();
+                // c.print_visualization();
                 Some(c.get_paths().to_vec())
             } else {
                 None