Save and load neural networks

.gitignore

@@ -0,0 +1,4 @@
+.DS_Store
+/.vscode
+/target
+output.brain

Cargo.toml

@@ -10,3 +10,5 @@ egui-macroquad = "0.15.0"
 macroquad = "0.3.25"
 once_cell = "1.17.1"
 rand = "0.8.5"
+serde = { version = "1.0", features = ["derive"] }
+bincode = "1.3.3"

src/agent.rs

@@ -30,6 +30,12 @@ impl Bird {
         }
     }
 
+    pub fn with_net(net: &Net) -> Self {
+        let mut new_bird = Bird::new();
+        new_bird.brain = net.clone();
+        new_bird
+    }
+
     pub fn with_brain(other: &Bird) -> Self {
         let mut new_bird = Bird::new();
         new_bird.brain = other.brain.clone();
@@ -107,6 +113,18 @@ impl Bird {
         self.brain.mutate();
     }
 
+    pub fn save(&self) {
+        self.brain
+            .save_to_disk("output.brain")
+            .expect("Failed to save brain file");
+    }
+
+    pub fn load(&self) -> Net {
+        self.brain
+            .load_from_disk("output.brain")
+            .expect("Failed to load brain")
+    }
+
     fn mark_dead(&mut self) {
         self.is_dead = true;
     }

src/editor.rs

@@ -10,6 +10,8 @@ pub struct Settings {
     pub is_frame_skip: bool,
     pub is_show_egui: bool,
     pub show_one_bird: bool,
+    pub save: bool,
+    pub load: bool,
 }
 
 pub struct Editor {
@@ -25,6 +27,8 @@ impl Settings {
             is_frame_skip: false,
             is_show_egui: true,
             show_one_bird: false,
+            save: false,
+            load: false,
         }
     }
 }
@@ -83,6 +87,12 @@ impl Editor {
                             if ui.add(egui::Button::new("Restart")).clicked() {
                                 self.settings.is_restart = true;
                             }
+                            if ui.add(egui::Button::new("Save")).clicked() {
+                                self.settings.save = true;
+                            }
+                            if ui.add(egui::Button::new("Load")).clicked() {
+                                self.settings.load = true;
+                            }
                         });
                 });
         });

src/main.rs

@@ -36,6 +36,14 @@ async fn main() {
                 stats = simulation.update(&editor.settings).unwrap_or(stats);
             }
         }
+        if editor.settings.save {
+            editor.settings.save = false;
+            simulation.save();
+        }
+        if editor.settings.load {
+            editor.settings.load = false;
+            simulation.load();
+        }
 
         if is_key_pressed(KeyCode::Q) {
             break;

src/nn.rs

@@ -1,14 +1,17 @@
 use macroquad::rand::gen_range;
+use serde::{Deserialize, Serialize};
+use std::fs::File;
+use std::io::{self, BufReader, BufWriter};
 
 use crate::*;
 
-#[derive(Clone)]
+#[derive(Clone, Serialize, Deserialize)]
 pub struct Net {
     n_inputs: usize,
     layers: Vec<Layer>,
 }
 
-#[derive(Clone)]
+#[derive(Clone, Serialize, Deserialize)]
 struct Layer {
     nodes: Vec<Vec<f64>>,
 }
@@ -57,6 +60,18 @@ impl Net {
     pub fn mutate(&mut self) {
         self.layers.iter_mut().for_each(|l| l.mutate());
     }
+
+    pub fn save_to_disk(&self, filename: &str) -> io::Result<()> {
+        let file = File::create(filename)?;
+        let writer = BufWriter::new(file);
+        bincode::serialize_into(writer, self).map_err(|e| io::Error::new(io::ErrorKind::Other, e))
+    }
+
+    pub fn load_from_disk(&self, filename: &str) -> io::Result<Self> {
+        let file = File::open(filename)?;
+        let reader = BufReader::new(file);
+        bincode::deserialize_from(reader).map_err(|e| io::Error::new(io::ErrorKind::Other, e))
+    }
 }
 
 impl Layer {

src/simulation.rs

@@ -68,6 +68,24 @@ impl Simulation {
         self.birds.iter().for_each(|b| b.draw());
     }
 
+    pub fn save(&self) {
+        let bird = self.birds.iter().filter(|b| !b.is_dead).next().unwrap();
+        bird.save();
+    }
+
+    pub fn load(&mut self) -> () {
+        let bird = self.birds.first().unwrap();
+        let net = bird.load();
+        let mut new_birds = Vec::new();
+
+        for _ in 0..NUM_BIRDS {
+            let mut new_bird = Bird::with_net(&net);
+            new_bird.mutate();
+            new_birds.push(new_bird);
+        }
+        self.birds = new_birds;
+    }
+
     fn selection(&self) -> Vec<Bird> {
         let mut rng = thread_rng();
         let gene_pool = self.calc_fitness();