{-# LANGUAGE TypeOperators #-}
--
-- A Mandelbrot set generator. Submitted by Simon Marlow as part of Issue #49.
--
import Config
import Control.Monad
import Control.Exception
import Foreign.Ptr
import Foreign.ForeignPtr
import System.IO.Unsafe
import System.Environment
import Data.Array.Accelerate.Array.Data ( ptrsOfArrayData )
import Data.Array.Accelerate.Array.Sugar ( Array(..) )
import Prelude as P
import Data.Array.Accelerate as A hiding ( size )
import Data.Array.Accelerate.IO
import qualified Data.Array.Repa as R
import qualified Data.Array.Repa.Repr.ForeignPtr as R
import Data.Array.Repa.IO.DevIL
-- Mandelbrot Set --------------------------------------------------------------
-- <<types
type F = Float
type Complex = (F,F)
type ComplexPlane = Array DIM2 Complex
-- >>
-- <<mandelbrot
mandelbrot :: F -> F -> F -> F -> Int -> Int -> Int
-> Acc (Array DIM2 (Complex,Int))
mandelbrot x y x' y' screenX screenY max_depth
= iterate go zs0 !! max_depth -- <4>
where
cs = genPlane x y x' y' screenX screenY -- <1>
zs0 = mkinit cs -- <2>
go :: Acc (Array DIM2 (Complex,Int))
-> Acc (Array DIM2 (Complex,Int))
go = A.zipWith iter cs -- <3>
-- >>
-- <<genPlane
genPlane :: F -> F
-> F -> F
-> Int
-> Int
-> Acc ComplexPlane
-- >>
genPlane lowx lowy highx highy viewx viewy
= generate (constant (Z:.viewy:.viewx))
(\ix -> let pr = unindex2 ix
x = A.fromIntegral (A.fst pr)
y = A.fromIntegral (A.snd pr)
in
lift ( elowx + (x * exsize) / eviewx
, elowy + (y * eysize) / eviewy))
where
elowx, elowy, exsize, eysize, eviewx, eviewy :: Exp F
elowx = constant lowx
elowy = constant lowy
exsize = constant (highx - lowx)
eysize = constant (highy - lowy)
eviewx = constant (P.fromIntegral viewx)
eviewy = constant (P.fromIntegral viewy)
-- <<next
next :: Exp Complex -> Exp Complex -> Exp Complex
next c z = c `plus` (z `times` z)
-- >>
-- <<plus
plus :: Exp Complex -> Exp Complex -> Exp Complex
plus = lift2 f
where f :: (Exp F, Exp F) -> (Exp F, Exp F) -> (Exp F, Exp F)
f (x1,y1) (x2,y2) = (x1+x2,y1+y2)
-- >>
-- <<times
times :: Exp Complex -> Exp Complex -> Exp Complex
times = lift2 f
where f :: (Exp F, Exp F) -> (Exp F, Exp F) -> (Exp F, Exp F)
f (ax,ay) (bx,by) = (ax*bx-ay*by, ax*by+ay*bx)
-- >>
-- <<dot
dot :: Exp Complex -> Exp F
dot = lift1 f
where f :: (Exp F, Exp F) -> Exp F
f (x,y) = x*x + y*y
-- >>
-- <<iter
iter :: Exp Complex -> Exp (Complex,Int) -> Exp (Complex,Int)
iter c p =
let
(z,i) = unlift p :: (Exp Complex, Exp Int) -- <1>
z' = next c z -- <2>
in
(dot z' >* 4.0) ? -- <3>
( p -- <4>
, lift (z', i+1) -- <5>
)
-- >>
-- <<mkinit
mkinit :: Acc ComplexPlane -> Acc (Array DIM2 (Complex,Int))
mkinit cs = A.map (lift1 f) cs
where f :: (Exp F, Exp F) -> ((Exp F, Exp F), Exp Int)
f (x,y) = ((x,y),0)
-- >>
-- Rendering -------------------------------------------------------------------
type RGBA = Word32
prettyRGBA :: Exp Int -> Exp (Complex, Int) -> Exp RGBA
prettyRGBA lIMIT s' = r + g + b + a
where
(_, s) = unlift s' :: (Exp (F, F), Exp Int)
t = A.fromIntegral $ ((lIMIT - s) * 255) `quot` lIMIT
r = (t `mod` 128 + 64)
g = (t * 2 `mod` 128 + 64) * 0x100
b = (t * 3 `mod` 256 ) * 0x10000
a = 0xFF000000
makePicture :: Options -> Int -> Acc (Array DIM2 (Complex, Int))
-> R.Array R.F R.DIM3 Word8
makePicture opt limit zs = R.fromForeignPtr (R.Z R.:. h R.:. w R.:. 4) rawData
where
arrPixels = run opt $ A.map (prettyRGBA (constant limit)) zs
(Z:.h:.w) = arrayShape arrPixels
{-# NOINLINE rawData #-}
rawData = let (Array _ adata) = arrPixels
((), ptr) = ptrsOfArrayData adata
in
unsafePerformIO $ newForeignPtr_ (castPtr ptr)
-- Main ------------------------------------------------------------------------
main :: IO ()
main
= do (config, nops) <- processArgs =<< getArgs
let size = optSize config
limit = optLimit config
--
x = -0.25 -- should get this from command line as well
y = -1.0
x' = 0.0
y' = -0.75
--
image = makePicture config limit
$ mandelbrot x y x' y' size size limit
runIL $ writeImage "out.png" (RGBA image)