Create a gist now

Instantly share code, notes, and snippets.

import Control.Applicative
import Control.Concurrent
import Control.Monad
import Data.Complex
import Data.List (find)
import Graphics.EasyRaster.GTK
import System.Environment
import System.IO
-- | Type representing points in an image
type Point = (Coord, Coord)
-- | Type representing colors of pixels
type Color = (Channel, Channel, Channel)
-- | Communicate results between threads.
data Result a = Result { getResult :: a } -- | A result
| Done -- | Process is done
deriving (Eq)
infinity = 1.0/0.0
-- | Some nice-looking colors. (Actually they are hideous, but I am
-- lazy...)
colors :: [Color]
colors = [(0,0,255),(0,255,0),(0,255,255),(255,0,0),
(255,0,255),(255,255,0),(0,0,128),(0,128,0),(0,128,128),(128,0,0),
(128,0,128),(128,128,0),(128,128,255),(128,255,128),
(128,255,255),(255,128,128),
(255,128,255),(255,255,128),(128,128,128)]
-- | Calculate the sequence for a given c.
iterateMandelbrot :: Complex Double -> [Complex Double]
iterateMandelbrot c = iterate (\z -> z ** 2 + c) c
-- | Calculate whether a given complex number "c" is in the Mandelbrot
-- set. Return 0 if the point is in the set, number of iterations
-- until divergence otherwise.
inMandelbrotSet :: Complex Double -> Int
inMandelbrotSet c = case find (\(i,z) -> magnitude z == infinity) (take 50 zs) of
Just (i,_) -> i
Nothing -> 0
where zs = zip [1..] (iterateMandelbrot c)
-- | Convert a pixel position in the image buffer to a coordinate on
-- the complex plane. (One axis.)
toPlaneCoord :: Double -> Double -> Int -> Int -> Double
toPlaneCoord min max x xmax =
(s * ((fromIntegral x) / (fromIntegral xmax))) + min
where s = max - min
-- | Given the maximum x and y bounds of the image, return a function
-- that calculates the color for any point on the image.
mandelbrotSet :: Int -> Int -> Point -> Color
mandelbrotSet xmax ymax (x,y) = case inMandelbrotSet c of
0 -> (0,0,0)
x -> colors !! (x `mod` 19)
where c = toPlaneCoord (-2) 0.5 x xmax :+ toPlaneCoord (-1) 1 y ymax
-- | Called when a thread is finished; kills the thread and notifies
-- the coordinator of its completion.
imDone mvar = do
putMVar mvar True
myThreadId >>= killThread
-- | Given a pair of coordinates representing the size of the image,
-- and a function from coordinates to colors, run the function for
-- each point in the image and return the generated image.
doIb :: Point -> (Point -> Color) -> IO ImageBuffer
doIb pmax f = do
ib <- ib_new pmax
todo <- newChan
results <- newChan
done <- newEmptyMVar
-- consume colors to write, write them to image
forkIO $ forever $ do
result <- readChan results
when (result == Done) $ imDone done
let (p,c) = getResult result
let (rowNum,colNum) = p
( ib_write_pixel ib p ) $! c
when (colNum `mod` 20 == 0 && rowNum == 0) $ print colNum
-- consume list of pixels, produce colors to write
let numThreads = 6
replicateM_ numThreads $ forkIO $ forever $ do
ps <- readChan todo
when (null ps) $ imDone done
forM_ ps (\p -> do
let fp = f $! p
writeChan results $! (fp `seq` Result (p,fp)))
-- produce list of coordinates
forM (ib_coords2 ib) $ \ps -> do
writeChan todo ps
-- ensure workers kill themselves when they get to the end of the chan
replicateM_ numThreads $ (writeChan todo [])
putStrLn "workers started"
-- wait for workers to finish
replicateM_ numThreads (takeMVar done)
putStrLn "workers done"
-- wait for image writer to finish
writeChan results Done
takeMVar done
return ib
-- | Main function; generate the image of size xmax, ymax and save it
-- to my homedir as a PNG.
mainMandelbrot xmax ymax = do
ib <- doIb (xmax,ymax) (mandelbrotSet xmax ymax)
putStrLn "Writing image..."
ib_save ib IFT_PNG "/home/jon/out.png"
putStrLn "done"
return ()
-- | Entry point. Accepts two args on the command-line, xmax and
-- ymax.
main = do
init_system
xmax:ymax:[] <- getArgs
mainMandelbrot (read xmax) (read ymax)
return ()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment