public
Last active

  • Download Gist
mandelbrot.hs
Haskell
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128
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 ()

Please sign in to comment on this gist.

Something went wrong with that request. Please try again.