WIPHEADmaster

1 files changed, 60 insertions, 0 deletions

diff --git a/test/Examples/Mandel.hs b/test/Examples/Mandel.hs
new file mode 100644
index 0000000..adb116b
--- /dev/null
+++ b/test/Examples/Mandel.hs
@@ -0,0 +1,60 @@
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE ViewPatterns #-}
+module Examples.Mandel (afun) where
+
+import Prelude ()
+import qualified Prelude as P
+import Data.Array.Accelerate
+
+
+type Dims = (Int, Int)
+type Pos = (Double, Double)
+type Viewport = (Dims     -- image size
+                ,Pos      -- midpoint
+                ,Double)  -- complex width of the viewport
+
+type RGB = (Word8, Word8, Word8)
+
+-- Arguments: viewport and maxiter
+-- Result: image in row-major order
+afun :: Acc (Scalar (Viewport, Int)) -> Acc (Matrix RGB)
+afun (the -> T2 viewport maxiter) = mandel viewport maxiter (clrBasic maxiter)
+
+clrBasic :: Exp Int -> Exp Int -> Exp RGB
+clrBasic maxiterI nI =
+    let maxiter = log (toFloating maxiterI)
+        n = log (toFloating nI)
+    in cond (n == maxiter)
+            (T3 0 0 0)
+            (let r = slope (Just (maxiter / 4)) Nothing n maxiter
+                 g = slope Nothing Nothing n maxiter
+                 b = slope Nothing (Just (maxiter * 3 / 4)) n maxiter
+             in T3 r g b)
+  where
+    slope :: Maybe (Exp Double) -> Maybe (Exp Double) -> Exp Double -> Exp Double -> Exp Word8
+    slope mlo mhi x m =
+        (P.maybe ($ 0) (\lo' -> max 0 . ($ lo')) mlo) $ \lo ->
+        (P.maybe ($ m) (\hi' -> min 255 . ($ hi')) mhi) $ \hi ->
+            fromIntegral @Int @Word8 $ round @Double @Int $
+                max 0 $ min 255 $ (x - lo) * 255 / (hi - lo)
+
+mandel :: Exp Viewport -> Exp Int -> (Exp Int -> Exp RGB) -> Acc (Matrix RGB)
+mandel (T3 (T2 w h) (T2 cx cy) cw) maxiter clrscheme =
+    generate (I2 h w) $ \(I2 yi xi) ->
+        let minx = cx - cw / 2
+            ch = toFloating h / toFloating w * cw
+            maxy = cy + ch / 2
+            x = minx + toFloating xi / (toFloating w - 1) * cw
+            y = maxy - toFloating yi / (toFloating h - 1) * ch
+        in clrscheme (mandeliter x y maxiter)
+
+mandeliter :: Exp Double -> Exp Double -> Exp Int -> Exp Int
+mandeliter x y maxiter =
+    let T5 _ _ _ _ n =
+            while (\(T5 _ _ a2 b2 i) -> a2 + b2 < 4 && i < maxiter)
+                  (\(T5 a b a2 b2 i) ->
+                      let a' = a2 - b2 + x
+                          b' = 2 * a * b + y
+                      in T5 a' b' (a'*a') (b'*b') (i + 1))
+                  (T5 x y (x*x) (y*y) 0)
+    in n