SolverStencil.hs

{-# LANGUAGE BangPatterns, TemplateHaskell, QuasiQuotes, FlexibleContexts #-}
module SolverStencil
        (solveLaplace)
where   
import Data.Array.Repa                          as A
import Data.Array.Repa.Stencil                  as A
import Data.Array.Repa.Stencil.Dim2             as A
import qualified Data.Array.Repa.Shape          as S
import Language.Haskell.TH
import Language.Haskell.TH.Quote


-- | Solver for the Laplace equation.
solveLaplace
        :: Monad m
        => Int                  -- ^ Number of iterations to use.
        -> Array U DIM2 Double  -- ^ Boundary value mask.
        -> Array U DIM2 Double  -- ^ Boundary values.
        -> Array U DIM2 Double  -- ^ Initial state.
        -> m (Array U DIM2 Double)

solveLaplace !steps !arrBoundMask !arrBoundValue !arrInit
 = go steps arrInit
 where  go 0 !arr = return arr
        go n !arr 
         = do   arr'    <- relaxLaplace arr
                go (n - 1) arr'

        relaxLaplace arr
         = computeP
         $ A.szipWith (+) arrBoundValue
         $ A.szipWith (*) arrBoundMask
         $ A.smap (/ 4)
         $ mapStencil2 (BoundConst 0)
            [stencil2|   0 1 0
                         1 0 1 
                         0 1 0 |] arr

{-# NOINLINE solveLaplace #-}