Just draw something on the screen!

Nostalgia for times past sometimes leads to frustration and anger. Yes, times are harder for would-be pixel artists, but nothing a little work can't solve. Let's do fractals in OCaml!

Suppose you've already installed the wonderful findlib (vielen Dank, Herr Stolpmann!). Suppose, furthermore, you've succeded in compiling OcamlSDL. This might look like a heavy prerequisite to fulfill before showing pretty graphics, but it's a job you have to do only once. Then, something like plotting Newton fractals requires less than ninety lines of code, and the result is this:

I find natural to abstract iteration patterns into higher-order functions. In this case, centering and scaling the pixel plane over the complex plane (the so-called screen-space transformation) is tedious but repetitive code that could be reused. Furthermore, the details of striding, color conversion and such are the same from program to program, where the only thing that varies is the actual calculation performed on each point:

let iter_surface ?(mag=1.0) f g surface =
  let info   = Sdlvideo.surface_info surface
  and pixels = Sdlvideo.pixel_data_32 surface in
  let size   = min info.Sdlvideo.w info.Sdlvideo.h in
  let scale  = mag /. float size
  and cx, cy = info.Sdlvideo.w / 2, info.Sdlvideo.h / 2 in
  for i = 0 to info.Sdlvideo.h - 1 do
    let b = i * info.Sdlvideo.pitch / 4
    and y = float (cy - i) *. scale in
    for j = 0 to info.Sdlvideo.w - 1 do
      let x = float (j - cx) *. scale in
      let c = f x y in
      pixels.{b + j} <- Sdlvideo.map_RGB surface c
    done;
    g i
  done

The magnification mag specifies how many real units the shortest dimension spans. The outer loop iterates with i for every scanline, whose base address into the pixels array is b. The inner loop iterates with j the index to successive pixels, and from both i and j the coordinates x, y are computed. With those the function f is called, and the result c (an RGB triple of type Sdlvideo.color) is mapped and assigned to the pixel. At the end of every scanline, the function g is called so that the main program can display progress.

The simplest SDL event loop waits for the user to quit the application:

let rec event_loop () =
  match Sdlevent.wait_event () with
  | Sdlevent.QUIT -> ()
  | _ -> event_loop ()

Anything more complicated is easy to accomodate. Now, there are a myriad schemes for color rendering of mathematical data, but using the color wheel is very common, especially since the complex argument maps naturally to a hue value. Converting HSV colors to RGB is somewhat involved, but the following code is entirely reusable:

let quomod d x =
  let f, e = modf (float d *. (x -. floor x)) in truncate e, f

let byte_of_frac x =
  let mask x = lnot (x asr 30) land 0xff in
  let n = truncate (ldexp x 8) in
  (n land (mask n))  lor  mask (n - 256)

let hsv_to_rgb (h, s, v) =
  let epsilon = 1e-3 in
  let w = byte_of_frac v in 
  if abs_float s < epsilon then (w, w, w) else
  let (e, f) = quomod 6 h in
  let x = byte_of_frac (v *. (1. -. s))
  and y = byte_of_frac (v *. (1. -. s *. f))
  and z = byte_of_frac (v *. (1. -. s *. (1. -. f))) in
  match e with
  | 0 -> (w, z, x)
  | 1 -> (y, w, x)
  | 2 -> (x, w, z)
  | 3 -> (x, y, w)
  | 4 -> (z, x, w)
  | 5 -> (w, x, y)
  | _ -> assert false

quomod scales the fractional part of x by the integer d, and splits the result into an entier part e and a fractional part f. This clasifies the hue into a sextant on which the fractional part linearly interpolates. To map a real number in the interval [0… 1] into [0… 256), byte_of_frac scales the fraction by 2⁸ and clips the result to a byte by constructing the appropriate bit-mask.

The code so far doesn't incorporate any specific functionality, and could be shipped off to a standard module to use again and again. The particular application that I had in mind, Newton fractals, are the plot on the complex plane of the basins of attraction to the complex roots of a polynomial, as found by applying the Newton method.

Polynomials and their derivatives are especially easy to evaluate at a point in one pass using Horner's method, when represented as the list of their coefficients with the independent term at the head, a so-called dense representation of polynomials:

open Complex

let eval p x =
  List.fold_left (fun (y, dy) p ->
    (add p (mul y x), add y (mul dy x)))
    (zero, zero) p

With that, the Newton iteration is very simple:

let nr_root ?(max_iter=32) ?(epsilon=1e-6) p x =
  let rec go (x, i as res) =
    if i = max_iter then res else
    let y, dy = eval p x in
    if norm y < epsilon then res else
    let dx = div y dy in
    go (sub x dx, succ i)
  in go (x, 0)

The root finder returns the number of iterations required, up to a maximum and/or specified precision, together with the closest estimate to the root. This allows me to "color" the fractal according to the time needed to converge to a root. With the chosen representation, the polynomial zⁿ - 1 (which represents, by definition, the n-th roots of unity) is straightforward to construct:

let pow n x =
  let rec go i l =
    if i = 0 then l else go (pred i) (zero :: l)
  in go n [x]

let nth_roots n = (neg one) :: pow (n - 1) one

These twenty lines are all the mathematical machinery required. It only remains to put all the pieces together:

let () =
  let max_iter = 128 in
  Sdl.init [`EVERYTHING];
  let screen = Sdlvideo.set_video_mode ~w:640 ~h:480 ~bpp:32 [`SWSURFACE] in
  Sdlwm.set_caption ~title:"nrfractal" ~icon:"";
  let p = nth_roots 3 in
  iter_surface ~mag:2.0 (fun x y ->
    let (z, n) = nr_root ~max_iter p { re = x; im = y } in
    let hue = Complex.arg z *. 0.159154943091895346 +. 0.5
    and vlu = float (max_iter - n) /. float max_iter in
    hsv_to_rgb (hue, 1., vlu) )
    (fun i -> if i land 7 = 0 then Sdlvideo.flip screen)
    screen;
  Sdlvideo.flip screen;
  event_loop ();
  Sdl.quit ()

The iteration limit max_iter controls the fractal detail. I first initialize the SDL library and create a big enough window, with a descriptive title. Then, I build a polynomial p representing the third roots of unity, and invoke iter_surface with a function that finds p's root starting from the complex value corresponding to the coordinates of the point to be plotted. I convert the complex angle of the returned approximation z to a hue value (by dividing by 2π and normalizing to the range [0… 1]); analogously, the number of iterations determines the vividness, or value of the color. This is converted to an RGB triple and returned. Every eight scanlines the screen is refreshed with the contents of the backing buffer. Once every point is plotted, the event loop is entered until the user decides to quit.

So, not counting the generic, reusable code, less than 40 lines of code suffice to just draw something on the screen. Have fun!