{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RebindableSyntax #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ViewPatterns #-}
module Julia where
import Data.Array.Accelerate hiding ( fromInteger )
import Data.Array.Accelerate.Data.Complex
import Prelude ( fromInteger )
import qualified Prelude as P
julia
:: forall a. (Num a, RealFloat a, FromIntegral Int a)
=> Int -- ^ image width
-> Int -- ^ image height
-> (Exp Float -> Exp (Complex a) -> Exp (Complex a)) -- ^ iteration function
-> Acc (Scalar Float) -- ^ current time
-> Acc (Scalar a) -- ^ centre x
-> Acc (Scalar a) -- ^ centre y
-> Acc (Scalar a) -- ^ view width
-> Acc (Scalar Int32) -- ^ iteration limit
-> Acc (Scalar a) -- ^ divergence radius
-> Acc (Array DIM2 (Complex a, Int32))
julia screenX screenY next (the -> t) (the -> x0) (the -> y0) (the -> width) (the -> limit) (the -> radius) =
generate
(constant (Z :. screenY :. screenX))
(\ix -> let z0 = complexOfPixel ix
zn = while (\zi -> snd zi < limit
&& dot (fst zi) < radius)
step
(T2 z0 0)
in zn)
where
complexOfPixel :: Exp DIM2 -> Exp (Complex a)
complexOfPixel (unlift -> Z :. y :. x) =
let
height = P.fromIntegral screenY / P.fromIntegral screenX * width
xmin = x0 - width / 2
ymin = y0 - height / 2
--
re = xmin + (fromIntegral x * width) / fromIntegral (constant screenX)
im = ymin + (fromIntegral y * height) / fromIntegral (constant screenY)
in
lift (re :+ im)
-- Divergence condition
--
dot :: Exp (Complex a) -> Exp a
dot (unlift -> x :+ y) = x*x + y*y
-- Take a single step of the recurrence relation
--
step :: Exp (Complex a, Int32) -> Exp (Complex a, Int32)
step (T2 z i) = T2 (next t z) (i + 1)