mandelbrot/mandelbrot.ss

(library (mandelbrot)
  (export make-ppm mandelbrot!
          steps
          make-rgb rgb-r rgb-g rgb-b
          number->rgb rgb-bin rgb-map rgb-lerp)
  (import (chezscheme))

(define-record-type rgb (fields r g b))
(define (number->rgb n)
  (assert (and (integer? n) (<= 0 n)))
  (make-rgb (bitwise-and #xff (bitwise-arithmetic-shift-right n 16))
            (bitwise-and #xff (bitwise-arithmetic-shift-right n 8))
            (bitwise-and #xff (bitwise-arithmetic-shift-right n 0))))
(define (rgb-bin +)
  (case-lambda
    [(a) a]
    [(a b) (make-rgb (+ (rgb-r a) (rgb-r b))
                     (+ (rgb-g a) (rgb-g b))
                     (+ (rgb-b a) (rgb-b b)))]
    [(a . bs) (fold-left (rgb-bin +) a bs)]))

(define (rgb-map c f)
  (make-rgb (f (rgb-r c))
            (f (rgb-g c))
            (f (rgb-b c))))

(define (rgb-mean . rgbs)
  (if (null? rgbs)
      (make-rgb 0 0 0)
      (let* ([sum (apply (rgb-bin +) rgbs)]
             [len (length rgbs)]
             [scale (lambda (c) (inexact->exact (floor (/ c len))))])
        (rgb-map sum scale))))

(define (rgb-lerp i min max)
  ((rgb-bin +) min (rgb-map ((rgb-bin -) max min) (lambda (x) (* i x)))))

;; proc : (x y) -> rgb
(define (make-ppm path width height proc)
  (call-with-port (open-file-output-port path (file-options no-fail))
    (lambda (p)
      (define header (format "P6\n~a ~a\n255\n" width height))
      (put-bytevector p (string->utf8 header))
      (do ([x 0 (mod (+ 1 x) width)]
           [y 0 (if (= (+ 1 x) width) (+ 1 y) y)])
          [(= y height)]
        (let ([color (proc x y)])
          (put-u8 p (rgb-r color))
          (put-u8 p (rgb-g color))
          (put-u8 p (rgb-b color))))))
  'done)


(define (lerp i min max)
  (+ min (* i (- max min))))

(define steps 100)
(define power 2)
(define samples-per-pixel 4)

(define tau (* 4 (acos 0)))

(define (rotate-about origin p angle)
  (+ (* (- p origin)
      (make-rectangular (cos angle) (sin angle)))
     origin))

;; let's rotate the sample?
(define (mandelbrot! path width height color xmin xmax ymin ymax)
  (define pixel-width  (exact->inexact (/ (- xmax xmin) width)))
  (define pixel-height (exact->inexact (/ (- ymax ymin) height)))
  (define sample-rotation (/ tau samples-per-pixel))
  (define (sample-vals center)
    (let ([top (+ center (make-rectangular 0 (- (/ pixel-height 2))))])
      (map (lambda (i) (rotate-about center top (* i sample-rotation)))
           (iota samples-per-pixel))))
  (make-ppm path width height
    (lambda (x y)
      (define ix (lerp (/ x width)  xmin xmax))
      (define iy (lerp (/ y height) ymin ymax))
      (define center (exact->inexact (make-rectangular ix iy)))
      (define (sample p)
        (define (step z) (+ (expt z power) p))
        (let loop ([i 0]
                   [z (step 0)])
          (cond
           [(= i steps) (color -1)]
           [(> (magnitude z) 2) (color i)]
           [else (loop (+ 1 i) (step z))])))
      (define samples (map sample (sample-vals center)))
      (apply rgb-mean samples)))))

(top-level-program
 (import (chezscheme)
         (mandelbrot))

(define (zoom path-fmt x y width steps)
  (for-each
   (lambda (i)
     (let ([idx (+ 1 i)]
           [i (+ 1 i)])
       (mandelbrot!
        (format path-fmt idx)
        width width
        (/ (- x 0.5) i) (/ (+ x 0.5) i)
        (/ (- y 0.5) i) (/ (+ y 0.5) i))))
   (iota steps)))

(define (cli-fail)
  (format (standard-error-port 'line (current-transcoder))
          "usage: ./mandelbrot.ss PATH WIDTH HEIGHT [COLOR:hex+|--lerp INTERIOR:hex START:hex END:hex]\n")
  (exit 1))

(define args (command-line-arguments))

(unless (<= 3 (length args))
  (cli-fail))

(define path   (list-ref args 0))
(define width  (list-ref args 1))
(define height (list-ref args 2))
(set! width (string->number width))
(set! height (string->number height))

(unless (and (integer? width) (integer? height) (positive? width) (positive? height))
  (cli-fail))

(define (color i)
  (define c
    (if (negative? i)
        0
        (floor (* 256 (/ i steps)))))
  (make-rgb c c c))

(when (< 3 (length args))
  (let* ([colors (cdddr args)]
         [lerp? (equal? "--lerp" (car colors))]
         [colors (if lerp? (cdr colors) colors)]
         [nums (map (lambda (n) (string->number n 16)) colors)]
         [_ (unless (for-all (lambda (x) (and (integer? x) (<= 0 x))) nums)
              (cli-fail))]
         [stripes (list->vector (map number->rgb nums))]
         [stripe-count (vector-length stripes)])
    (when (and lerp? (< stripe-count 3))
      (cli-fail))

    (set! color
      (if lerp?
          (lambda (i)
            (if (negative? i)
                (vector-ref stripes 0)
                (rgb-map (rgb-lerp (expt (/ i steps) 1.3)
                                   (vector-ref stripes 1)
                                   (vector-ref stripes 2))
                         floor)))
          (lambda (i) (vector-ref stripes (mod i stripe-count)))))))

(mandelbrot! path width height color -2.2 1.1 -1 1))