diff --git a/src/agent.rs b/src/agent.rs
new file mode 100644
index 0000000..87792c0
--- /dev/null
+++ b/src/agent.rs
@@ -0,0 +1,107 @@
+use crate::{
+    grid::{Grid, PopulationConfig},
+    util::wrap,
+};
+
+use rand::{seq::SliceRandom, Rng};
+use std::f32::consts::TAU;
+
+// A single Physarum agent. The x and y positions are continuous, hence we use floating point numbers instead of integers.
+#[derive(Debug)]
+pub struct Agent {
+    pub x: f32,
+    pub y: f32,
+    pub angle: f32,
+    pub population_id: usize,
+    pub i: usize,
+}
+
+impl Agent {
+    // Construct a new agent with random parameters.
+    pub fn new<R: Rng + ?Sized>(width: usize, height: usize, id: usize, rng: &mut R, i: usize) -> Self {
+        let (x, y, angle) = rng.gen::<(f32, f32, f32)>();
+        Agent {
+            x: x * width as f32,
+            y: y * height as f32,
+            angle: angle * TAU,
+            population_id: id,
+            i,
+        }
+    }
+
+    #[inline]
+    pub fn tick(&mut self, grid: &Grid) {
+        let (width, height) = (grid.width, grid.height);
+        let PopulationConfig {
+            sensor_distance,
+            sensor_angle,
+            rotation_angle,
+            step_distance,
+            ..
+        } = grid.config;
+
+        let xc = self.x + fastapprox::faster::cos(self.angle) * sensor_distance;
+        let yc = self.y + fastapprox::faster::sin(self.angle) * sensor_distance;
+
+        let agent_add_sens = self.angle + sensor_angle;
+        let agent_sub_sens = self.angle - sensor_angle;
+
+        let xl = self.x + fastapprox::faster::cos(agent_sub_sens) * sensor_distance;
+        let yl = self.y + fastapprox::faster::sin(agent_sub_sens) * sensor_distance;
+        let xr = self.x + fastapprox::faster::cos(agent_add_sens) * sensor_distance;
+        let yr = self.y + fastapprox::faster::sin(agent_add_sens) * sensor_distance;
+
+        // We sense from the buffer because this is where we previously combined data from all the grid.
+        let center = grid.get_buf(xc, yc);
+        let left = grid.get_buf(xl, yl);
+        let right = grid.get_buf(xr, yr);
+
+        // Rotate and move logic
+        let mut rng = rand::thread_rng();
+        let mut direction: f32 = 0.0;
+
+        if (center > left) && (center > right) {
+            direction = 0.0;
+        } else if (center < left) && (center < right) {
+            direction = *[-1.0, 1.0].choose(&mut rng).unwrap();
+        } else if left < right {
+            direction = 1.0;
+        } else if right < left {
+            direction = -1.0;
+        }
+
+        let delta_angle = rotation_angle * direction;
+
+        self.angle = wrap(self.angle + delta_angle, TAU);
+        self.x = wrap(
+            self.x + step_distance * fastapprox::faster::cos(self.angle),
+            width as f32,
+        );
+        self.y = wrap(
+            self.y + step_distance * fastapprox::faster::sin(self.angle),
+            height as f32,
+        );
+    }
+}
+
+impl Clone for Agent {
+    fn clone(&self) -> Agent {
+        Agent {
+            x: self.x,
+            y: self.y,
+            angle: self.angle,
+            population_id: self.population_id,
+            i: self.i,
+        }
+    }
+}
+
+impl PartialEq for Agent {
+    fn eq(&self, other: &Self) -> bool {
+        self.x == other.x
+            && self.y == other.y
+            && self.angle == other.angle
+            && self.population_id == other.population_id
+            && self.i == other.i
+    }
+}
\ No newline at end of file
diff --git a/src/grid.rs b/src/grid.rs
index 7e718f9..1710ef7 100644
--- a/src/grid.rs
+++ b/src/grid.rs
@@ -1,6 +1,6 @@
 use crate::{
     blur::Blur,
-    model::Agent,
+    agent::Agent,
 };
 
 use rand::{distributions::Uniform, Rng};
@@ -86,11 +86,11 @@ pub struct Grid {
     pub width: usize,
     pub height: usize,
 
-    data: Vec<f32>,
+    pub data: Vec<f32>,
 
     // Scratch space for the blur operation.
-    buf: Vec<f32>,
-    blur: Blur,
+    pub buf: Vec<f32>,
+    pub blur: Blur,
     pub agents: Vec<Agent>
 }
 
diff --git a/src/imgdata.rs b/src/imgdata.rs
index a91c221..05126ab 100644
--- a/src/imgdata.rs
+++ b/src/imgdata.rs
@@ -1,8 +1,67 @@
 use crate::{grid::Grid, palette::Palette};
 
+use itertools::multizip;
+
+// Stores data that is located in grids that is used for image generation
+pub struct ThinGridData {
+    pub width: usize,
+    pub height: usize,
+    pub data: Vec<f32>,
+}
+
+
+impl Clone for ThinGridData {
+    fn clone(&self) -> ThinGridData {
+        ThinGridData {
+            width: self.width.clone(),
+            height: self.height.clone(),
+            data: self.data.clone(),
+        }
+    }
+}
+
+impl ThinGridData {
+    // Convert Grid to ThinGridData
+    pub fn from_grid(in_grid: &Grid) -> Self {
+        return ThinGridData {
+            width: in_grid.width.clone(),
+            height: in_grid.height.clone(),
+            data: in_grid.data.clone(),
+        }
+    }
+
+    #[allow(dead_code)]
+    pub fn from_grid_vec(in_grids: Vec<Grid>) -> Vec<Self> {
+        return in_grids.iter().map(|grid|{
+            return Self::from_grid(grid);
+        }).collect();
+    }
+
+    // from grid.rs (needed in image gen)
+    #[allow(dead_code)]
+    pub fn data(&self) -> &[f32] {
+        &self.data
+    }
+
+    // from grid.rs (needed in image gen)
+    #[allow(dead_code)]
+    pub fn quantile(&self, fraction: f32) -> f32 {
+        let index = if (fraction - 1.0_f32).abs() < f32::EPSILON {
+            self.data.len() - 1
+        } else {
+            (self.data.len() as f32 * fraction) as usize
+        };
+        let mut sorted = self.data.clone();
+        sorted
+            .as_mut_slice()
+            .select_nth_unstable_by(index, |a, b| a.partial_cmp(b).unwrap());
+        sorted[index]
+    }
+}
+
 // Class for storing data that will be used to create images
 pub struct ImgData {
-    pub grids: Vec<Grid>,
+    pub grids: Vec<ThinGridData>,
     pub palette: Palette,
     pub iteration: i32,
 }
@@ -18,11 +77,46 @@ impl Clone for ImgData {
 }
 
 impl ImgData {
-    pub fn new(in_grids: Vec<Grid>, in_palette: Palette, in_iteration: i32) -> Self {
+    pub fn new(in_grids: Vec<ThinGridData>, in_palette: Palette, in_iteration: i32) -> Self {
         ImgData {
             grids: in_grids,
             palette: in_palette,
             iteration: in_iteration,
         }
     }
+
+    #[inline]
+    pub fn save_to_image(&self) {
+        let (width, height) = (self.grids[0].width, self.grids[0].height);
+        let mut img = image::RgbImage::new(width as u32, height as u32);
+    
+        let max_values: Vec<_> = self
+            .grids
+            .iter()
+            .map(|grid| grid.quantile(0.999) * 1.5)
+            .collect();
+    
+        for y in 0..height {
+            for x in 0..width {
+                let i = y * width + x;
+                let (mut r, mut g, mut b) = (0.0_f32, 0.0_f32, 0.0_f32);
+                for (grid, max_value, color) in
+                    multizip((&self.grids, &max_values, &self.palette.colors))
+                {
+                    let mut t = (grid.data()[i] / max_value).clamp(0.0, 1.0);
+                    t = t.powf(1.0 / 2.2); // gamma correction
+                    r += color.0[0] as f32 * t;
+                    g += color.0[1] as f32 * t;
+                    b += color.0[2] as f32 * t;
+                }
+                r = r.clamp(0.0, 255.0);
+                g = g.clamp(0.0, 255.0);
+                b = b.clamp(0.0, 255.0);
+                img.put_pixel(x as u32, y as u32, image::Rgb([r as u8, g as u8, b as u8]));
+            }
+        }
+    
+        img.save(format!("./tmp/out_{}.png", self.iteration).as_str())
+            .unwrap();
+    }
 }
diff --git a/src/lib.rs b/src/lib.rs
index 9b05422..fee30ca 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -4,4 +4,5 @@ mod imgdata; // for storing image data
 mod math;
 pub mod model;
 mod palette;
-mod util; // for math things
\ No newline at end of file
+mod util; // for math things
+mod agent;
\ No newline at end of file
diff --git a/src/main.rs b/src/main.rs
index 369c4d5..1c03d80 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,15 +2,17 @@ use physarum::model;
 
 fn main() {
     // # of iterations to go through
-    let n_iterations = 254;
+    // let n_iterations = 1024;
+    let n_iterations = 128;
 
     // Size of grid and pictures
-    let (width, height) = (256, 256);
+    // let (width, height) = (256, 256);
+    let (width, height) = (512, 512);
     // let (width, height) = (1024, 1024);
 
     // # of agents
-    // let n_particles = 1 << 24;
-    let n_particles = 1 << 16;
+    let n_particles = 1 << 10;
+    // let n_particles = 1 << 16;
     println!("n_particles: {}", n_particles);
 
     let diffusivity = 1;
diff --git a/src/model.rs b/src/model.rs
index b0a3f3b..9b0c1b2 100644
--- a/src/model.rs
+++ b/src/model.rs
@@ -1,122 +1,18 @@
 use crate::{
-    grid::{combine, Grid, PopulationConfig},
-    imgdata::ImgData,
+    grid::{combine, Grid},
+    imgdata::{ImgData, ThinGridData},
     palette::{random_palette, Palette},
-    util::wrap,
+    agent::Agent,
 };
 
 use indicatif::{ParallelProgressIterator, ProgressBar, ProgressStyle};
-use itertools::multizip;
-use rand::{seq::SliceRandom, Rng};
+// use rand::Rng;
 use rand_distr::{Distribution, Normal};
 use rayon::{iter::ParallelIterator, prelude::*};
-use std::{f32::consts::TAU, path::Path, time::Instant};
-
-// A single Physarum agent. The x and y positions are continuous, hence we use floating point numbers instead of integers.
-#[derive(Debug)]
-pub struct Agent {
-    x: f32,
-    y: f32,
-    angle: f32,
-    population_id: usize,
-    i: usize,
-}
-
-impl Agent {
-    // Construct a new agent with random parameters.
-    fn new<R: Rng + ?Sized>(width: usize, height: usize, id: usize, rng: &mut R, i: usize) -> Self {
-        let (x, y, angle) = rng.gen::<(f32, f32, f32)>();
-        Agent {
-            x: x * width as f32,
-            y: y * height as f32,
-            angle: angle * TAU,
-            population_id: id,
-            i,
-        }
-    }
-
-    #[inline]
-    pub fn tick(&mut self, grid: &Grid) {
-        let (width, height) = (grid.width, grid.height);
-        let PopulationConfig {
-            sensor_distance,
-            sensor_angle,
-            rotation_angle,
-            step_distance,
-            ..
-        } = grid.config;
-
-        let xc = self.x + fastapprox::faster::cos(self.angle) * sensor_distance;
-        let yc = self.y + fastapprox::faster::sin(self.angle) * sensor_distance;
-
-        let agent_add_sens = self.angle + sensor_angle;
-        let agent_sub_sens = self.angle - sensor_angle;
-
-        let xl = self.x + fastapprox::faster::cos(agent_sub_sens) * sensor_distance;
-        let yl = self.y + fastapprox::faster::sin(agent_sub_sens) * sensor_distance;
-        let xr = self.x + fastapprox::faster::cos(agent_add_sens) * sensor_distance;
-        let yr = self.y + fastapprox::faster::sin(agent_add_sens) * sensor_distance;
-
-        // We sense from the buffer because this is where we previously combined data from all the grid.
-        let center = grid.get_buf(xc, yc);
-        let left = grid.get_buf(xl, yl);
-        let right = grid.get_buf(xr, yr);
-
-        // Rotate and move logic
-        let mut rng = rand::thread_rng();
-        let mut direction: f32 = 0.0;
-
-        if (center > left) && (center > right) {
-            direction = 0.0;
-        } else if (center < left) && (center < right) {
-            direction = *[-1.0, 1.0].choose(&mut rng).unwrap();
-        } else if left < right {
-            direction = 1.0;
-        } else if right < left {
-            direction = -1.0;
-        }
-
-        let delta_angle = rotation_angle * direction;
-
-        self.angle = wrap(self.angle + delta_angle, TAU);
-        self.x = wrap(
-            self.x + step_distance * fastapprox::faster::cos(self.angle),
-            width as f32,
-        );
-        self.y = wrap(
-            self.y + step_distance * fastapprox::faster::sin(self.angle),
-            height as f32,
-        );
-    }
-}
-
-impl Clone for Agent {
-    fn clone(&self) -> Agent {
-        Agent {
-            x: self.x,
-            y: self.y,
-            angle: self.angle,
-            population_id: self.population_id,
-            i: self.i,
-        }
-    }
-}
-
-impl PartialEq for Agent {
-    fn eq(&self, other: &Self) -> bool {
-        self.x == other.x
-            && self.y == other.y
-            && self.angle == other.angle
-            && self.population_id == other.population_id
-            && self.i == other.i
-    }
-}
+use std::{path::Path, time::Instant};
 
 // Top-level simulation class.
 pub struct Model {
-    // Physarum agents.
-    // agents: Vec<Agent>,
-
     // The grid they move on.
     grids: Vec<Grid>,
 
@@ -276,8 +172,14 @@ impl Model {
         );
     }
 
+    fn strip_grid_data(grids: Vec<Grid>) -> Vec<ThinGridData> {
+        return grids.iter().map(|grid| {
+            return ThinGridData::from_grid(grid);
+        }).collect();
+    }
+
     fn save_image_data(&mut self) {
-        let grids = self.grids.clone();
+        let grids = Self::strip_grid_data(self.grids.clone());
         let img_data = ImgData::new(grids, self.palette, self.iteration);
         self.img_data_vec.push(img_data);
         if self.grids[0].width > 1024 && self.grids[0].height > 1024 && self.img_data_vec.len() > 100 {
@@ -312,10 +214,12 @@ impl Model {
             .par_iter()
             .progress_with(pb)
             .for_each(|img| {
-                Self::save_to_image(img.to_owned());
+                // Self::save_to_image(img.to_owned());
+                img.save_to_image();
             });
     }
 
+    /*
     pub fn save_to_image(imgdata: ImgData) {
         let (width, height) = (imgdata.grids[0].width, imgdata.grids[0].height);
         let mut img = image::RgbImage::new(width as u32, height as u32);
@@ -349,4 +253,5 @@ impl Model {
         img.save(format!("./tmp/out_{}.png", imgdata.iteration).as_str())
             .unwrap();
     }
+    */
 }