more customizable colors

mandelbrot.ss

@@ -1,8 +1,18 @@
 (library (mandelbrot)
-  (export make-ppm mandelbrot!)
+  (export make-ppm mandelbrot!
+          steps
+          make-rgb rgb-r rgb-g rgb-b
+          number->rgb)
   (import (chezscheme))
 
-;; proc : (x y) -> (u8 u8 u8)
+  (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))))
+
+;; 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)
@@ -11,19 +21,19 @@
       (do ([x 0 (mod (+ 1 x) width)]
            [y 0 (if (= (+ 1 x) width) (+ 1 y) y)])
           [(= y height)]
-        (let-values ([(r g b) (proc x y)])
-          (put-u8 p r)
-          (put-u8 p g)
-          (put-u8 p b)))))
+        (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 20)
-(define power 4)
+(define steps 100)
+(define power 2)
 
-(define (mandelbrot! path width height xmin xmax ymin ymax)
+(define (mandelbrot! path width height color xmin xmax ymin ymax)
   (make-ppm path width height
     (lambda (x y)
       (define ix (lerp (/ x width)  xmin xmax))
@@ -33,12 +43,8 @@
       (let loop ([i 0]
                  [z (step 0)])
         (cond
-         [(= i steps) (values 255 255 255)]
-         [(> (magnitude z) 2)
-          (let ([i (/ (+ 1 i) steps)])
-            (values (floor (lerp i 0 256))
-                    (floor (lerp i 256 0))
-                    (floor (lerp i 128 256))))]
+         [(= i steps) (color -1)]
+         [(> (magnitude z) 2) (color i)]
          [else (loop (+ 1 i) (step z))]))))))
 
 (top-level-program
@@ -59,19 +65,37 @@
 
 (define (cli-fail)
   (format (standard-error-port 'line (current-transcoder))
-          "usage: ./mandelbrot.ss path width height\n"))
+          "usage: ./mandelbrot.ss path width height [color:hex*]\n")
+  (exit 1))
 
 (define args (command-line-arguments))
 
-(unless (= 3 (length args))
+(unless (<= 3 (length args))
   (cli-fail))
 
-(define-values (path width height) (apply values args))
+(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))
 
-(mandelbrot! path width height -1.35 1.05 -1.2 1.2))
+(define (color i)
+  (define c (floor (* 256 (/ (+ 1 i) steps))))
+  (make-rgb c c c))
+
+(when (< 3 (length args))
+  (let* ([stripes
+          (let* ([colors (cdddr args)]
+                 [nums (map (lambda (n) (string->number n 16)) colors)]
+                 [_ (unless (for-all (lambda (x) (and (integer? x) (<= 0 x))) nums)
+                      (cli-fail))]
+                 [colors (map number->rgb nums)])
+            (list->vector colors))]
+         [stripe-count (vector-length stripes)])
+    (set! color (lambda (i) (vector-ref stripes (mod i stripe-count))))))
+
+(mandelbrot! path width height color -2.2 1.1 -1 1))