programming_rust/mandelbrot/src/main.rs

207 lines
6.8 KiB
Rust

extern crate num;
extern crate image;
extern crate crossbeam;
use num::Complex;
use std::str::FromStr;
use image::ColorType;
use image::pnm::PNMEncoder;
use image::pnm::{PNMSubtype, SampleEncoding};
use std::io::Write;
use std::fs::File;
/// Try to determine if a complex number is a member of the Mandelbrot
/// set, using at most 'limit' iterations to decide.
///
/// If 'c' is not a member, return 'Some(i)' where 'i' is the number of
/// iterations it took for 'c' to leave the circle of radius 2
/// centered on the origin. If 'c' seems to be a member (if we
/// reached the iteration limit without being able to prove 'c' is not
/// a member), return 'None'
///
fn escape_time(c: Complex<f64>, limit: u32) -> Option<u32> {
let mut z = Complex{ re: 0.0, im: 0.0 };
for i in 0..limit {
z = z * z + c;
if z.norm_sqr() > 4.0 {
return Some(i);
}
}
None
}
fn parse_pair<T: FromStr>(s: &str, separator: char) -> Option<(T, T)> {
match s.find(separator) {
None => None,
Some(index) => {
match (T::from_str(&s[..index]), T::from_str(&s[index+1..])) {
(Ok(l), Ok(r)) => Some((l, r)),
_ => None
}
}
}
}
fn parse_complex(s: &str) -> Option<Complex<f64>> {
match parse_pair(s, ',') {
Some((re, im)) => Some(Complex{ re, im }),
None => None
}
}
/// Contains the definitions of two planes: an integral cartesian plane,
/// and a complex cartesian plane. Used to map points on one to the
/// other.
struct Plane {
bounds: (usize, usize),
upper_left: Complex<f64>,
lower_right: Complex<f64>,
width: f64,
height: f64
}
impl Plane where {
/// Define a mapping between the coordinates of an integral
/// cartesian plane and a complex cartesian plane.
pub fn new(width: usize, height: usize, ulp: Complex<f64>, lrp: Complex<f64>) -> Plane {
let bound_width = width as f64;
let bound_height = height as f64;
Plane {
bounds: (width, height),
upper_left: ulp,
lower_right: lrp,
width: (lrp.re - ulp.re) / bound_width,
height: (lrp.im - ulp.im) / bound_height,
}
}
/// Given the row and column of a pixel on the integral cartesian plane,
/// return an complex number that corresponds to the equivalent location
/// mapped to the complex cartesian plane.
pub fn pixel_to_point(&self, pixel: (usize, usize)) -> Complex<f64> {
Complex{
re: self.upper_left.re + pixel.0 as f64 * self.width,
im: self.upper_left.im + pixel.1 as f64 * self.height
}
}
pub fn render(&self, pixels: &mut[u8]) {
assert!(pixels.len() == self.bounds.0 * self.bounds.1);
for row in 0..self.bounds.1 {
for column in 0..self.bounds.0 {
let point = self.pixel_to_point((column, row));
pixels[row * self.bounds.0 + column] =
match escape_time(point, 255) {
None => 0,
Some(count) => 255 - count as u8
};
}
}
}
pub fn subplane(&self, top: usize, height: usize) -> Plane {
let ulc = self.pixel_to_point((0, top));
let lrc = self.pixel_to_point((self.bounds.0, top + height));
Plane::new(self.bounds.0, height, ulc, lrc)
}
pub fn render_concurrent(&self, pixels: &mut[u8], threads: usize) {
let rows_per_band = self.bounds.1 / threads + 1;
{
let bands: Vec<&mut [u8]> = pixels.chunks_mut(rows_per_band * self.bounds.0).collect();
crossbeam::scope(|spawner| {
for (i, band) in bands.into_iter().enumerate() {
let top = rows_per_band * i;
let height = band.len() / self.bounds.0;
let subplane = self.subplane(top, height);
spawner.spawn(move || {
subplane.render(band)
});
}
});
}
}
}
fn write_image(filename: &str, pixels: &[u8], bounds: (usize, usize)) -> Result<(), std::io::Error> {
let output = File::create(filename)?;
let mut encoder = PNMEncoder::new(output).with_subtype(PNMSubtype::Graymap(SampleEncoding::Binary));
encoder.encode(pixels, bounds.0 as u32 ,bounds.1 as u32, ColorType::Gray(8))?;
Ok(())
}
pub fn main() {
let args: Vec<String> = std::env::args().collect();
if args.len() != 5 {
writeln!(std::io::stderr(), "Usage: mandelbrot FILE PIXELS UPPERLEFT LOWERRIGHT").unwrap();
writeln!(std::io::stderr(), "Exaple: {} mandel.png 1000x750 -1.20,0.35 -1,02.0", args[0]).unwrap();
std::process::exit(1);
}
let bounds = parse_pair(&args[2], 'x').expect("Error parsing image dimensions");
let upper_left = parse_complex(&args[3]).expect("Error parsing upper left hand point");
let lower_right = parse_complex(&args[4]).expect("Error parsing lower right hand point");
let mut pixels = vec![0; bounds.0 * bounds.1];
let plane = Plane::new(bounds.0, bounds.1, upper_left, lower_right);
plane.render_concurrent(&mut pixels, 8);
write_image(&args[1], &pixels, bounds).expect("Error writing PNM file");
}
/*
fn pixel_to_point(
pixel: (usize, usize),
bounds: (usize, usize),
upper_left: Complex<f64>,
lower_right: Complex<f64>) -> Complex<f64>
{
let (width, height) = (lower_right.re - upper_left.re,
lower_right.im - upper_left.im);
Complex {
re: upper_left.re + pixel.0 as f64 * width / bounds.0 as f64,
im: upper_left.im + pixel.1 as f64 * height / bounds.1 as f64
}
}
*/
#[test]
fn test_parse_pair() {
assert_eq!(parse_pair::<i32>("", ','), None);
assert_eq!(parse_pair::<i32>("10", ','), None);
assert_eq!(parse_pair::<i32>(",10", ','), None);
assert_eq!(parse_pair::<i32>("10,20xy", ','), None);
assert_eq!(parse_pair::<i32>("0.5x", ','), None);
assert_eq!(parse_pair::<i32>("10,20", ','), Some((10, 20)));
assert_eq!(parse_pair::<f64>("0.5x", ','), None);
assert_eq!(parse_pair::<f64>("0.5x1.5", 'x'), Some((0.5, 1.5)));
}
#[test]
fn test_parse_complex() {
assert_eq!(parse_complex("1.25,-0.0625"), Some(Complex{ re: 1.25, im: -0.0625 }));
assert_eq!(parse_complex(",-0.0625"), None);
}
/*
#[test]
fn test_pixel_to_point() {
assert_eq!(pixel_to_point((25, 75), (100, 100),
Complex { re: -1.0, im: 1.0 },
Complex { re: 1.0, im: -1.0 }), Complex { re: -0.5, im: -0.5 });
}
*/
#[test]
fn test_plane() {
let plane = Plane::new(100, 100, Complex{ re: -1.0, im: 1.0 }, Complex{ re: 1.0, im: -1.0 });
assert_eq!(plane.pixel_to_point((25, 75)), Complex{ re: -0.5, im: -0.5 });
}