accelerate-examples-0.13.0.0: examples/mandelbrot/Mandel.hs
{-# LANGUAGE ScopedTypeVariables #-}
--
-- A Mandelbrot set generator.
-- Originally submitted by Simon Marlow as part of Issue #49.
--
module Mandel (
-- Types
View, Render, Bitmap,
-- Pretty pictures
mandelbrot, prettyRGBA,
) where
import Prelude as P
import Data.Array.Accelerate as A
import Data.Array.Accelerate.IO as A
-- Types -----------------------------------------------------------------------
-- Current view into the complex plane
type View a = (a, a, a, a)
-- Complex numbers
type Complex a = (a, a)
type ComplexPlane a = Array DIM2 (Complex a)
-- Image data
type Bitmap = Array DIM2 RGBA32
-- Action to render a frame
type Render a = Scalar (View a) -> Bitmap
-- Mandelbrot Set --------------------------------------------------------------
-- Compute the mandelbrot as repeated application of the recurrence relation:
--
-- Z_{n+1} = c + Z_n^2
--
-- This is applied until the iteration limit is reached, regardless of whether
-- or not the pixel has diverged. This implies wasted work, but no thread
-- divergence.
--
mandelbrot
:: forall a. (Elt a, IsFloating a)
=> Int
-> Int
-> Int
-> Acc (Scalar (View a))
-> Acc (Array DIM2 (Complex a,Int))
mandelbrot screenX screenY depth view
= foldr ($) zs0 (P.take depth (repeat go))
where
cs = genPlane screenX screenY view
zs0 = mkinit cs
go :: Acc (Array DIM2 (Complex a, Int)) -> Acc (Array DIM2 (Complex a, Int))
go = A.zipWith iter cs
genPlane :: (Elt a, IsFloating a) => Int -> Int -> Acc (Scalar (View a)) -> Acc (ComplexPlane a)
genPlane screenX screenY view
= generate (constant (Z:.screenY:.screenX))
(\ix -> let pr = unindex2 ix
x = A.fromIntegral (A.snd pr :: Exp Int)
y = A.fromIntegral (A.fst pr :: Exp Int)
in
lift ( xmin + (x * sizex) / viewx
, ymin + (y * sizey) / viewy))
where
(xmin,ymin,xmax,ymax) = unlift (the view)
sizex = xmax - xmin
sizey = ymax - ymin
viewx = constant (P.fromIntegral screenX)
viewy = constant (P.fromIntegral screenY)
next :: (Elt a, IsNum a) => Exp (Complex a) -> Exp (Complex a) -> Exp (Complex a)
next c z = c `plus` (z `times` z)
plus :: (Elt a, IsNum a) => Exp (Complex a) -> Exp (Complex a) -> Exp (Complex a)
plus = lift2 f
where f :: (Elt a, IsNum a) => (Exp a, Exp a) -> (Exp a, Exp a) -> (Exp a, Exp a)
f (x1,y1) (x2,y2) = (x1+x2, y1+y2)
times :: forall a. (Elt a, IsNum a) => Exp (Complex a) -> Exp (Complex a) -> Exp (Complex a)
times = lift2 f
where f :: (Elt a, IsNum a) => (Exp a, Exp a) -> (Exp a, Exp a) -> (Exp a, Exp a)
f (x,y) (x',y') = (x*x'-y*y', x*y'+y*x')
dot :: forall a. (Elt a, IsNum a) => Exp (Complex a) -> Exp a
dot = lift1 f
where f :: (Elt a, IsNum a) => (Exp a, Exp a) -> Exp a
f (x,y) = x*x + y*y
iter :: forall a. (Elt a, IsNum a) => Exp (Complex a) -> Exp (Complex a, Int) -> Exp (Complex a, Int)
iter c zi = f (A.fst zi) (A.snd zi)
where
f :: Exp (Complex a) -> Exp Int -> Exp (Complex a, Int)
f z i =
let z' = next c z
in (dot z' >* 4) ? ( zi , lift (z', i+1) )
mkinit :: Elt a => Acc (ComplexPlane a) -> Acc (Array DIM2 (Complex a, Int))
mkinit cs = A.zip cs (A.fill (A.shape cs) 0)
-- Rendering -------------------------------------------------------------------
prettyRGBA :: forall a. (Elt a, IsFloating a) => Exp Int -> Exp (Complex a, Int) -> Exp RGBA32
prettyRGBA lIMIT s =
let cmax = A.fromIntegral lIMIT
c = A.fromIntegral (A.snd s)
in
c ==* cmax ? ( 0xFF000000, escapeToColour (cmax - c) )
-- Directly convert the iteration count on escape to a colour. The base set
-- (x,y,z) yields a dark background with light highlights.
--
escapeToColour :: Exp Int -> Exp RGBA32
escapeToColour m = constant 0xFFFFFFFF - (packRGBA32 $ lift (x,y,z,w))
where
x = constant 0
w = A.fromIntegral (3 * m)
z = A.fromIntegral (5 * m)
y = A.fromIntegral (7 * m)
{--
-- A simple colour scheme
--
prettyRGBA :: Elt a => Exp Int -> Exp (Complex a, Int) -> Exp RGBA32
prettyRGBA lIMIT s' = r + g + b + a
where
s = A.snd s'
t = A.fromIntegral $ ((lIMIT - s) * 255) `quot` lIMIT
r = (t `rem` 128 + 64) * 0x1000000
g = (t * 2 `rem` 128 + 64) * 0x10000
b = (t * 3 `rem` 256 ) * 0x100
a = 0xFF
--}
{--
prettyRGBA :: forall a. (Elt a, IsFloating a) => Exp Int -> Exp (Complex a, Int) -> Exp RGBA32
prettyRGBA lIMIT s =
let cmax = A.fromIntegral lIMIT :: Exp a
c = A.fromIntegral (A.snd s)
in
c >* 0.98 * cmax ? ( 0xFF000000, rampColourHotToCold 0 cmax c )
-- Standard Hot-to-Cold hypsometric colour ramp. Colour sequence is
-- Red, Yellow, Green, Cyan, Blue
--
rampColourHotToCold
:: (Elt a, IsFloating a)
=> Exp a -- ^ minimum value of the range
-> Exp a -- ^ maximum value of the range
-> Exp a -- ^ data value
-> Exp RGBA32
rampColourHotToCold vmin vmax vNotNorm
= let v = vmin `A.max` vNotNorm `A.min` vmax
dv = vmax - vmin
--
result = v <* vmin + 0.28 * dv
? ( lift ( constant 0.0
, 4 * (v-vmin) / dv
, constant 1.0
, constant 1.0 )
, v <* vmin + 0.5 * dv
? ( lift ( constant 0.0
, constant 1.0
, 1 + 4 * (vmin + 0.25 * dv - v) / dv
, constant 1.0 )
, v <* vmin + 0.75 * dv
? ( lift ( 4 * (v - vmin - 0.5 * dv) / dv
, constant 1.0
, constant 0.0
, constant 1.0 )
, lift ( constant 1.0
, 1 + 4 * (vmin + 0.75 * dv - v) / dv
, constant 0.0
, constant 1.0 )
)))
in
rgba32OfFloat result
--}