Haskell Mandelbrot

mandelbrot.hs

import Data.Complex
import Data.Bits
import Data.Binary.Put
import qualified Data.ByteString.Lazy as BL
import Data.Word
import System.IO
import Control.Exception
import Char
import Ix

dibHeader :: Int -> Int -> Put
dibHeader w h = do
	putWord32le 40
	putWord32le (fromIntegral w)
	putWord32le (fromIntegral h)
	putWord16le 1
	putWord16le 32
	putWord32le 0
	putWord32le (fromIntegral (w * h))
	putWord32le 72
	putWord32le 72
	putWord32le 0
	putWord32le 0

bmpHeader :: Int -> Int -> Put
bmpHeader w h = do
	putWord8 (fromIntegral (ord 'B'))
	putWord8 (fromIntegral (ord 'M'))
	putWord32le (fromIntegral (w * h + hsize))
	putWord16le 0
	putWord16le 0
	putWord32le (fromIntegral hsize)
	dibHeader w h
		where hsize = 54

orbit :: (Complex Float) -> (Complex Float) -> (Complex Float) -> Int -> (Int, (Complex Float))
orbit c z zn i
	| i <= 0 = (i, z)
	| magnitude z >= 4 = (i, z)
	| otherwise = orbit c (z^2 + c) z (i - 1)

color :: Int -> (Int, (Complex Float)) -> (Float, Float, Float)
color mi (i, o)
	| i == 0 = (0, 0, 0)
	| otherwise = (c, c, c)
		where c = fromIntegral(mi - i) / fromIntegral(mi)

mandel :: Int -> Int -> Int -> [[(Float, Float, Float)]]
mandel w h i = [[color i (orbit ((fromIntegral(4*y) / fromIntegral(h) - 2) :+
	                         (fromIntegral(4*x) / fromIntegral(w) - 2))
	                        0 0 i) | y <- range(0, h - 1)] | x <- range(0, w - 1)]

genMandel :: Int -> Int -> Put
genMandel w h = sequence_ (map (mapM_ (\(r, g, b) -> do
	putWord8 (truncate (r*255))
	putWord8 (truncate (g*255))
	putWord8 (truncate (b*255))
	putWord8 255))
	(mandel w h 100))

run :: Int -> Int -> IO ()
run w h = BL.writeFile "mandelbrot.bmp" (runPut (bmpHeader w h >> genMandel w h))

main = run 512 512