Created
January 31, 2010 13:49
-
-
Save jrockway/291074 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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