mandelbrot

  1use std::io::Write;
  2
  3use num::Complex;
  4use rayon::prelude::*;
  5
  6const CHECK_ITERS: i64 = 128;
  7
  8const X_DIMEN: i64 = 16000;
  9const ASPECT_RATIO: f64 = 16.0 / 9.0;
 10
 11type N = Complex<f64>;
 12
 13#[derive(Debug, Default)]
 14struct Color {
 15    r: f64,
 16    g: f64,
 17    b: f64,
 18}
 19
 20impl Color {
 21    pub fn bytes(&self) -> [u8; 6] {
 22        let r = (self.r * 65535.).round() as u16;
 23        let g = (self.g * 65535.).round() as u16;
 24        let b = (self.b * 65535.).round() as u16;
 25
 26        [
 27            (r >> 8) as u8,
 28            (r & 0xFF) as u8,
 29            (g >> 8) as u8,
 30            (g & 0xFF) as u8,
 31            (b >> 8) as u8,
 32            (b & 0xFF) as u8,
 33        ]
 34    }
 35}
 36
 37impl std::ops::Add<f64> for Color {
 38    type Output = Color;
 39
 40    fn add(self, rhs: f64) -> Self::Output {
 41        let Color { r, g, b } = self;
 42        Color {
 43            r: r + rhs,
 44            g: g + rhs,
 45            b: b + rhs,
 46        }
 47    }
 48}
 49
 50impl std::ops::Div<isize> for Color {
 51    type Output = Self;
 52
 53    fn div(self, rhs: isize) -> Self::Output {
 54        let Color { r, g, b } = self;
 55
 56        Color {
 57            r: r / rhs as f64,
 58            g: g / rhs as f64,
 59            b: b / rhs as f64,
 60        }
 61    }
 62}
 63
 64impl std::ops::AddAssign for Color {
 65    fn add_assign(&mut self, rhs: Self) {
 66        self.r += rhs.r;
 67        self.g += rhs.g;
 68        self.b += rhs.b;
 69    }
 70}
 71
 72fn in_set(c: N) -> Option<i64> {
 73    let mut z = N::new(0.0, 0.0);
 74    for i in 0..CHECK_ITERS {
 75        z = z * z + c;
 76        if z.norm_sqr() > 4.0 {
 77            return Some(i);
 78        }
 79    }
 80
 81    None
 82}
 83
 84fn color_for_iters(iters: i64) -> Color {
 85    let s = 0.5f64;
 86    let l = 0.5f64;
 87    // FIXME: ensure h is between 0 and 360;
 88    let h = iters as f64 / (CHECK_ITERS as f64) * 360.0;
 89
 90    assert!(h >= 0.0);
 91    assert!(h < 360.0);
 92
 93    let c = (1.0 - (2.0 * l - 1.0).abs()) * s;
 94    let x = c * (1.0 - ((h / 60.0) % 2.0 - 1.0).abs());
 95
 96    let m = l - c / 2.0;
 97
 98    let intermediate = if h >= 300.0 {
 99        Color { r: c, g: 0.0, b: x }
100    } else if h >= 240.0 {
101        Color { r: x, g: 0.0, b: c }
102    } else if h >= 180.0 {
103        Color { r: 0.0, g: x, b: c }
104    } else if h >= 120.0 {
105        Color { r: 0.0, g: c, b: x }
106    } else if h >= 60.0 {
107        Color { r: x, g: c, b: 0.0 }
108    } else if h >= 0.0 {
109        Color { r: c, g: x, b: 0.0 }
110    } else {
111        panic!("h out of range");
112    };
113
114    intermediate + m
115}
116
117fn main() {
118    let x_start = -3.0;
119    let x_end = 3.0;
120    let y_dimen: i64 = (X_DIMEN as f64 / ASPECT_RATIO).round() as i64;
121    let y_start = (x_end - x_start) / ASPECT_RATIO / 2.0;
122    let substeps = 16;
123
124    let top_left = N::new(x_start, y_start);
125
126    let step = (x_end - x_start) / X_DIMEN as f64;
127    let substep = step / substeps as f64;
128
129    let f = std::fs::File::create("image.ppm").unwrap();
130    let mut w = std::io::BufWriter::new(f);
131    writeln!(w, "P6\n{X_DIMEN} {y_dimen}\n65535").unwrap();
132
133    let rows: Vec<_> = (0..y_dimen)
134        .into_par_iter()
135        .map(|y| {
136            let mut row = Vec::new();
137            for x in 0..X_DIMEN {
138                let mut color = Color::default();
139
140                let inner_top_left = top_left + N::new(x as f64 * step, -y as f64 * step);
141                for i in 0..substeps {
142                    for j in 0..substeps {
143                        let point = inner_top_left + N::new(substep * i as f64, substep * j as f64);
144                        color += in_set(point).map_or_else(Color::default, color_for_iters)
145                            / (substeps * substeps);
146                    }
147                }
148
149                row.push(color);
150            }
151
152            row
153        })
154        .collect();
155
156    for row in rows {
157        for color in row {
158            w.write_all(&color.bytes()).unwrap();
159        }
160    }
161}