Permalink
Find file
Fetching contributors…
Cannot retrieve contributors at this time
40 lines (34 sloc) 1.23 KB
{-# 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 #-}