Refactor genetic algorithm

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 rand::Rng;
+use factorio_pathfinding::{Connection, PathInput, Pathfinder, examples::HashMapMap};
+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,
                     });
                 }
             }