feat: improve colours in mandelbrot set

src/main.rs

@@ -18,6 +18,7 @@ struct SpriteScale {
     horiz: f64,
     verti: f64,
     step: f64,
+    iterations: u32,
 }
 
 impl Default for SpriteScale {
@@ -27,6 +28,7 @@ impl Default for SpriteScale {
             horiz: 0.0,
             verti: 0.0,
             step: 0.01,
+            iterations: 64,
         }
     }
 }
@@ -40,12 +42,30 @@ fn transform_mandelbrot_system(
     let horizontal = input.pressed(KeyCode::D) as i8 - input.pressed(KeyCode::A) as i8;
     let vertical = input.pressed(KeyCode::W) as i8 - input.pressed(KeyCode::S) as i8;
     let step_mod = input.pressed(KeyCode::R) as i8 - input.pressed(KeyCode::F) as i8;
+    let iter_mod = input.pressed(KeyCode::T) as i8 - input.pressed(KeyCode::G) as i8;
 
     scale.step += (scale.step / 10.0) * step_mod as f64;
     scale.verti += scale.step * vertical as f64;
     scale.horiz += scale.step * horizontal as f64;
     scale.zoom += scale.step * zoom as f64;
 
+    if iter_mod == 1 {
+        scale.iterations += (scale.step * 100.0) as u32;
+
+        if scale.iterations > 512 {
+            scale.iterations = 512;
+        }
+    } else if iter_mod == -1 {
+        let step = (scale.step * 100.0) as u32;
+        let diff = scale.iterations as i64 - step as i64;
+
+        if diff < 32 {
+            scale.iterations = 32;
+        } else {
+            scale.iterations -= step;
+        }
+    }
+
     if scale.zoom < 0.0 {
         // We can't go below 0
         scale.zoom -= scale.step * zoom as f64;
@@ -68,18 +88,15 @@ fn mandelbrot_render_system(
                     let z = scale.zoom;
                     let h = scale.horiz;
                     let v = scale.verti;
+                    let iters = scale.iterations;
 
                     let start = Instant::now();
-                    texture.data = fractal.generate_scaled_image(
+                    texture.data.copy_from_slice(fractal.generate_scaled_image(
                         ((-2.5 * z) + h, (1.0 * z) + h),
                         ((-1.0 * z) - v, (1.0 * z) - v),
-                    );
+                        iters,
-                    let diff = Instant::now() - start;
+                    ));
-                    println!(
+                    let _diff = Instant::now() - start;
-                        "Frametime: {:?} Framerate: {:.2}",
-                        diff,
-                        1.0 / diff.as_secs_f32()
-                    );
                 }
             }
         }

src/mandelbrot.rs

@@ -5,19 +5,7 @@ const MAX_ITERATIONS: u16 = 512;
 
 #[derive(Debug, Clone)]
 pub struct Mandelbrot {
-    iterations: Vec<u16>,
+    texture_buffer: Vec<u8>,
-}
-
-enum Colour {
-    Purple,
-    Violet,
-    Blue,
-    Green,
-    Yellow,
-    Orange,
-    Red,
-    White,
-    Black,
 }
 
 impl Mandelbrot {
@@ -26,7 +14,7 @@ impl Mandelbrot {
 
     pub fn new() -> Self {
         Mandelbrot {
-            iterations: vec![0; Self::IMG_WIDTH * Self::IMG_HEIGHT],
+            texture_buffer: vec![0; (Self::IMG_WIDTH * Self::IMG_HEIGHT) * 4],
         }
     }
 
@@ -39,138 +27,109 @@ impl Mandelbrot {
     }
 
     pub fn generate_image(&mut self) -> Vec<u8> {
-        self.escape_time((-2.5, 1.0), (-1.0, 1.0));
+        self.texture_buffer
-
+            .par_chunks_mut(4)
-        let mut texture_buffer = vec![0; (Self::IMG_WIDTH * Self::IMG_HEIGHT) * 4];
-        for i in 0..texture_buffer.len() {
-            if i % 4 == 0 {
-                let iter_i = i / 4;
-                self.assign_colour(i, iter_i, &mut texture_buffer);
-            }
-        }
-
-        texture_buffer
-    }
-
-    pub fn generate_scaled_image(&mut self, x_bounds: (f64, f64), y_bounds: (f64, f64)) -> Vec<u8> {
-        let mut texture_buffer = vec![0; (Self::IMG_WIDTH * Self::IMG_HEIGHT) * 4];
-
-        self.escape_time(x_bounds, y_bounds);
-
-        for i in 0..texture_buffer.len() {
-            if i % 4 == 0 {
-                let iter_i = i / 4;
-                self.assign_colour(i, iter_i, &mut texture_buffer);
-            }
-        }
-
-        texture_buffer
-    }
-
-    fn escape_time(&mut self, x_bounds: (f64, f64), y_bounds: (f64, f64)) {
-        self.iterations
-            .par_iter_mut()
             .enumerate()
-            .for_each(|(i, value)| {
+            .for_each(|(i, buf)| {
+                let iters = Self::new_escape_time(i, (-2.5, 1.0), (-1.0, 1.0), 64);
+                let normalized_iters = iters / MAX_ITERATIONS as f64;
+                let h = 0.5 + (10.0 * normalized_iters);
+                buf.copy_from_slice(&Self::hsv_to_rgb(h, 0.6, 1.0));
+            });
+
+        self.texture_buffer.clone()
+    }
+
+    pub fn generate_scaled_image(
+        &mut self,
+        x_bounds: (f64, f64),
+        y_bounds: (f64, f64),
+        max_iterations: u32,
+    ) -> &[u8] {
+        self.texture_buffer
+            .par_chunks_mut(4)
+            .enumerate()
+            .for_each(|(i, buf)| {
+                let max_iters = max_iterations as f64;
+
+                let iters = Self::new_escape_time(i, x_bounds, y_bounds, max_iterations);
+                let normalized_iters = iters / max_iters;
+
+                let h = normalized_iters * 350.0;
+                let v = if iters == max_iters { 0.0 } else { 1.0 };
+                buf.copy_from_slice(&Self::hsv_to_rgb(h, 1.0, v))
+            });
+
+        &self.texture_buffer
+    }
+
+    fn new_escape_time(
+        i: usize,
+        x_bounds: (f64, f64),
+        y_bounds: (f64, f64),
+        max_iterations: u32,
+    ) -> f64 {
         let px = i % Self::IMG_WIDTH;
         let py = i / Self::IMG_WIDTH;
-
         let c = Self::coords_to_complex(px, py, x_bounds, y_bounds);
 
-                *value = Self::calc_num_iterations(c);
+        Self::new_calc_num_iters(c, max_iterations)
-            });
     }
 
-    fn find_colour(iteration: u16) -> Colour {
+    #[inline]
-        if iteration == 1 {
+    fn hsv_to_rgb(h: f64, s: f64, v: f64) -> [u8; 4] {
-            return Colour::Black;
+        let c = v * s;
+        let x = c * (1.0 - (((h / 60.0) % 2.0) - 1.0).abs());
+        let m = v - c;
+
+        let (r_prime, g_prime, b_prime) = {
+            if h < 60.0 {
+                (c, x, 0.0)
+            } else if h < 120.0 {
+                (x, c, 0.0)
+            } else if h < 180.0 {
+                (0.0, c, x)
+            } else if h < 240.0 {
+                (0.0, x, c)
+            } else if h < 300.0 {
+                (x, 0.0, c)
+            } else if h < 360.0 {
+                (c, 0.0, x)
+            } else {
+                unreachable!()
+            }
+        };
+
+        let r = ((r_prime + m) * 255.0) as u8;
+        let g = ((g_prime + m) * 255.0) as u8;
+        let b = ((b_prime + m) * 255.0) as u8;
+        let a = 0xFF;
+
+        [r, g, b, a]
     }
 
-        match iteration % 7 {
+    fn new_calc_num_iters(c: Complex<f64>, max_iterations: u32) -> f64 {
-            0 => Colour::Purple,
-            1 => Colour::Violet,
-            2 => Colour::Blue,
-            3 => Colour::Green,
-            4 => Colour::Yellow,
-            5 => Colour::Orange,
-            6 => Colour::Red,
-            _ => unreachable!(),
-        }
-    }
-
-    fn assign_colour(&self, i: usize, iter_i: usize, texture_buffer: &mut [u8]) {
-        match Self::find_colour(self.iterations[iter_i]) {
-            Colour::White => {
-                texture_buffer[i] = 0xFF;
-                texture_buffer[i + 1] = 0xFF;
-                texture_buffer[i + 2] = 0xFF;
-                texture_buffer[i + 3] = 0xFF;
-            }
-            Colour::Black => {
-                texture_buffer[i] = 0x00;
-                texture_buffer[i + 1] = 0x00;
-                texture_buffer[i + 2] = 0x00;
-                texture_buffer[i + 3] = 0xFF;
-            }
-            Colour::Purple => {
-                texture_buffer[i] = 0x94;
-                texture_buffer[i + 1] = 0x00;
-                texture_buffer[i + 2] = 0xD3;
-                texture_buffer[i + 3] = 0xFF;
-            }
-            Colour::Violet => {
-                texture_buffer[i] = 0x4B;
-                texture_buffer[i + 1] = 0x00;
-                texture_buffer[i + 2] = 0x82;
-                texture_buffer[i + 3] = 0xFF;
-            }
-            Colour::Blue => {
-                texture_buffer[i] = 0x00;
-                texture_buffer[i + 1] = 0x00;
-                texture_buffer[i + 2] = 0xFF;
-                texture_buffer[i + 3] = 0xFF;
-            }
-            Colour::Green => {
-                texture_buffer[i] = 0x00;
-                texture_buffer[i + 1] = 0xFF;
-                texture_buffer[i + 2] = 0x00;
-                texture_buffer[i + 3] = 0xFF;
-            }
-            Colour::Yellow => {
-                texture_buffer[i] = 0xFF;
-                texture_buffer[i + 1] = 0xFF;
-                texture_buffer[i + 2] = 0x00;
-                texture_buffer[i + 3] = 0xFF;
-            }
-            Colour::Orange => {
-                texture_buffer[i] = 0xFF;
-                texture_buffer[i + 1] = 0x7F;
-                texture_buffer[i + 2] = 0x00;
-                texture_buffer[i + 3] = 0xFF;
-            }
-            Colour::Red => {
-                texture_buffer[i] = 0xFF;
-                texture_buffer[i + 1] = 0x00;
-                texture_buffer[i + 2] = 0x00;
-                texture_buffer[i + 3] = 0xFF;
-            }
-        }
-    }
-
-    fn calc_num_iterations(c: Complex<f64>) -> u16 {
         let mut z: Complex<f64> = Complex::new(0.0, 0.0);
-        let mut num_iterations: u16 = 0;
+        let mut num_iters: u32 = 0;
 
-        for i in 0..MAX_ITERATIONS {
+        loop {
             if z.norm_sqr() > 4.0 {
-                num_iterations = i;
+                break;
+            }
+
+            if num_iters >= max_iterations {
                 break;
             }
 
             z = (z * z) + c;
+            num_iters += 1;
         }
 
-        num_iterations + 1 // account for off by one error
+        if num_iters < max_iterations {
+            (num_iters as f64) - z.norm() / 2f64.ln()
+        } else {
+            num_iters as f64
+        }
     }
 
     fn coords_to_complex(