mandelbrot.ml

(* The Mandelbrot set.

   Compile with:

   ocamlbuild mandelbrot.native

   Example usage:
   
   ./mandelbrot.native --xmin 0.27085 --xmax 0.27100 --ymin 0.004640 --ymax 0.004810 --xres 1000 --maxiter 1024 --file pic.ppm

  The interior of Mandelbrot set is black, the levels are gray.
  If you have very many levels, the picture is likely going to be quite
  dark. You can postprocess it to fix the palette. For instance,
  with ImageMagick you can do (assuming the picture was saved to pic.ppm):

  convert -normalize pic.ppm pic.png

  The resulting pic.png is still gray, but the levels will be nicer. You
  can also add colors, for instance:

  convert -negate -normalize -fill blue -tint 100 pic.ppm pic.png

  See http://www.imagemagick.org/Usage/color_mods/ for what ImageMagick
  can do. It can do a lot.
*)

let main = 
  let xmin = ref (-2.5) in
  let xmax = ref 0.5 in
  let ymin = ref (-1.5) in
  let ymax = ref (1.5) in
  let xres = ref 512 in
  let yres = ref 512 in
  let maxiter = ref 1000 in
  let filename = ref None in

  let options =
    Arg.align [
      ("--xmin", Arg.Set_float xmin, "<float> the left edge") ;
      ("--xmax", Arg.Set_float xmax, "<float> the right edge") ;
      ("--ymin", Arg.Set_float ymin, "<float> the bottom edge") ;
      ("--ymax", Arg.Set_float ymax, "<float> the top edge") ;
      ("--xres", Arg.Int (fun k ->
        xres := k ;
        yres := int_of_float (float k *. (!ymax -. !ymin) /. (!xmax -. !xmin))
       ), "<int> resolution in x direction") ;
      ("--maxiter", Arg.Int (fun k -> maxiter := max 256 k), "<int> number of iterations") ;
      ("--file", Arg.String (fun t -> filename := Some t), "<file.ppm> output file")
    ] in
  let anon = (fun _ -> raise (Arg.Bad "unexpected command-line argument")) in
  let usage = "Usage: mandelbrot --xmin <float> --xmax <float> --ymin <float> --ymax <float> --xres <int> --maxiter <int> --file <file.ppm>" in

    (* Parse command line *)
    Arg.parse options anon usage ;
    
    let fh =
      (match !filename with
        | None -> 
          print_endline "Please specify an output file." ;
          Arg.usage options usage ;
          exit 1
        | Some f -> open_out_bin f)
    in

      (* Print the PPM header *)
      Printf.fprintf fh "P6\n# Mandelbrot, xmin=%f, xmax=%f, ymin=%f, ymax=%f, maxiter=%d\n%d\n%d\n%d\n"
        !xmin !xmax !ymin !ymax !maxiter !xres !yres !maxiter ;
      
      (* Precompute the pixel width and height *)
      let dx = (!xmax -. !xmin) /. float !xres in
      let dy = (!ymax -. !ymin) /. float !yres in
      for j = 0 to !yres - 1 do
        let y = !ymax -. float j *. dy in
        for i = 0 to !xres - 1 do
          let x = !xmin +. float i *. dx in
          let u = ref 0.0 in
          let v = ref 0.0 in
          let u2 = ref (!u *. !u) in
          let v2 = ref (!v *. !v) in
          let k = ref 0 in
            while !k < !maxiter && !u2 +. !v2 < 4.0 do
              v := 2.0 *. !u *. !v +. y ;
              u := !u2 -. !v2 +. x ;
              u2 := !u *. !u ;
              v2 := !v *. !v ;
              incr k
            done ;
            (* Compute the pixel color and write it to file *)
            if !k >= !maxiter then
              (* interior *)
              for i = 0 to 5 do output_byte fh 0 done
            else begin
              (* exterior *)
              let a = char_of_int (!k lsr 8) in
              let b = char_of_int (!k land 255) in
              let s = "012345" in
                s.[0] <- a ; s.[1] <- b ;
                s.[2] <- a ; s.[3] <- b ;
                s.[4] <- a ; s.[5] <- b ;
                output_string fh s
            end
        done
      done ;
      close_out fh