mandelbrot set

Mandelbrot.hs

{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE BangPatterns #-}

module Main where

import System.Environment (getArgs)
import Text.Read (readMaybe)
import Control.Monad (when,unless)
import Data.Maybe (isNothing,fromJust,fromMaybe)
import Data.Complex
import Codec.Picture
import Data.Vector (Vector,(!))
import qualified Data.Vector as V

main :: IO ()
main = do
    args <- getArgs
    when (length args >= 2) $ do
        let
            iterations = readMaybe $ head args
            bound      = readMaybe $ args !! 1
        unless (isNothing iterations || isNothing bound) $ do
            let
                f'    = f (fromJust iterations) (fromJust bound)
                image = generateImage f' width height
            writePng "mandelbrot.png" image

width, height :: Int
width =  1000
height = 1000

f :: Int -> Double -> Int -> Int -> PixelRGB8
f i b x' y' = g i b $ (x-c1)/c3:+(y-c2)/c3
    where
        x       = fromIntegral x'
        y       = fromIntegral y'
        width'  = fromIntegral width
        c1      = width'*(1/2+1/8)
        c2      = width'/2
        c3      = width'/(5/2)

g :: Int -> Double -> Complex Double -> PixelRGB8
g iterations bound c =
    let
        g' z          = z^(2 :: Int)+c
        xs            = V.iterateN (iterations+1) g' 0
        outOfBounds x = isNaN val || val > bound
            where val = magnitude x
        colorize i    = pallete ! round (len*i'/iter :: Double)
            where
                len  = fromIntegral $ length pallete-1
                i'   = fromIntegral i
                iter = fromIntegral iterations
    in colorize . fromMaybe 0 $ V.findIndex outOfBounds xs

pallete :: Vector PixelRGB8
pallete = (\(r',g',b') -> PixelRGB8 r' g' b') <$> V.concatMap (V.generate 255)
    [(\i -> let i' = fI i in (0, 0,     i'))
    ,(\i -> let i' = fI i in (0, i',    255-i'))
    ,(\i -> let i' = fI i in (i',255-i',0))]
    where fI = fromIntegral