Multigrid Methods with Repa

Multigrid Methods with Repa

Introduction

We implement the multigrid and full multigrid schemes using the Haskell parallel array library Repa. Whilst this file is a literate Haskell program it omits some preliminaries. The full code can be found on GitHub at multiGrid.

Repa (short for ‘regular parallel arrays’) is a Haskell library for efficiently calculating with arrays functionally. It allows parallel computation to be expressed with a single primitive (computeP). This is possible because we only use immutable arrays so calculations are deterministic without any need to worry about locks/mutual exclusion/race conditions/etc.

The multigrid and full multigrid schemes are used for approximating solutions of partial differential equations but they are strategies rather than solvers. They can be used to drastically improve more basic solvers provided convergence and error analysis are considered for the basic solver.

We only consider linear partial differential equations and the Poisson equation in particular where linear partial differential equations can be summarised in the form

A u = -f

Here, A is a linear operator (such as the laplace operator \nabla^2 or higher order or lower order partial derivatives or any linear combination of these), f is given and u is the function we want to solve for, and both are defined on a region \Omega. The multigrid scheme starts with a fine grid \Omega^h (where h is the grid spacing) on which we want to obtain an approximate solution by finite difference methods. The scheme involves the use of a coarser grid (e.g. \Omega^{2h}) and, recursively, a stack of coarser and coarser grids to apply error correction on approximations. The full multigrid scheme also uses coarser grids to improve initial approximations used by multigrid.

Outline of Coarse Grid Error Correction

Assume we have a basic method to approximate solutions which we call ‘approximately solve’. Then the scheme for coarse grid correction of errors is

  1. Take an initial guess v_0 on the fine grid (\Omega^h) and use it to approximately solve A u = -f to get a new approximation v_1.

    We want to estimate errors

    e = u - v_1

    We do not have u to calculate them, but the errors satisfy

    A e = A u - A v_1 = - f - A v_1

    The negative of these values A e are called the residuals (r) and these we can calculate.

  2. Compute residuals r = A v_1 + f
  3. Move to the coarse grid \Omega^{2h} and (recursively) solve A e = -r

    (This is a problem of the same form as we started with, but it requires solving with restrictions of A and r for the coarse grid.)

    We use zeros as initial guess for errors and this recursion results in an approximation for errors e^{2h},

  4. Interpolate coarse errors (e^{2h}) into the fine grid to get e^h and get a corrected guess

    v^h = v_1 + e^h

  5. Approximately solve A u = -f again (on the fine grid), but now starting with v^h to get a final approximation.

So a basic requirement of the multigrid scheme is to be able to move values to a coarser grid (restriction) and from a coarser grid to a finer grid (interpolation). We will write functions for these below. First we give a quick summary of Repa and operations we will be using, which experienced Repa users may prefer to skip.

Background Repa Summary

Repa array types have an extent (shape) and an indication of representation as well as a content type. For example Array U DIM2 Double is the type of a two dimensional array of Doubles. The U indicates a ‘manifest’ representation which means that the contents are fully calculated rather than delayed. A delayed representation is expressed with D rather than U. The only other representation type we use explicitly is (TR PC5) which is the type of a partitioned, array resulting from a stencil mapping operation. Non-manifest arrays are useful for enabling the compiler to combine (fuse) operations to optimise after inlining code. Some operations require the argument array to be manifest, so to make a delayed array manifest we can use

  • computeP which employes parallel evaluation on the array elements. The type of computeP requires that it is used within a monad. This is to prevent us accidentally writing nested calls of computeP.

The extent DIM2 abbreviates Z :. Int :. Int (where Z represents a base shape for the empty zero dimensional array). The integers give the sizes in each dimension. We can have higher dimensions (adding more with :.) but will only be using DIM2 here. Values of an extent type are used to express the sizes of an array in each dimensions and also used as indexes into an array where indexing starts at 0.

We only use a few other Repa operations

  • fromListUnboxed takes an extent and a list of values to create a new manifest array.

  • traverse takes 3 arguments to produce a new delayed array. The first is the (old) array. The second is a function which produces the extent of the new array when given the extent of the old array. The third is a function to calculate any item in the new array [when supplied with an operation (get) to retrieve values from the old array and an index in the new array for the item to calculate].

  • szipWith is analgous to zipWith for lists, taking a binary operation to combine two arrays. (There is also a zipWith for arrays, but the szipWith version is tuned for partitioned arrays)

  • smap is analgous to map for lists, mapping an operation over an array. (There is also a map for arrays, but the smap version is tuned for partitioned arrays)

  • mapStencil2 takes a boundary handling option (we only use BoundClamp here), a stencil, and a two dimensional array. It maps (convolves) the stencil over the array to produce a new (partititioned) array.

There is a convenient way of writing down stencils which are easy to visualize. We will see examples later, but this format requires pragmas

> {-# LANGUAGE TemplateHaskell     #-}
> {-# LANGUAGE QuasiQuotes         #-}

Interpolation (Prolongation) and Restriction

To interpolate from coarse to fine \Omega^{2 h} \rightarrow \Omega^h we want a (linear) mapping that is full rank. E.g identity plus averaging neighbours.

We can restrict by injection \Omega^h \rightarrow \Omega^{2 h} or take some weighting of neighbours as well. A common requirement is that this is also linear and a constant times the transpose of the interpolation when these are expressed as matrix operations.

Repa examples of restriction and interpolation in 2D.

We will be using DIM2 grids and work with odd extents so that border points remain on the border when moving between grids. A common method for restriction is to take a weighted average of fine neighbours with each coarse grid point and we can easily express this as a stencil mapping. We calculate one sixteenth of the results from mapping this stencil

> restStencil :: Stencil DIM2 Double
> restStencil =   [stencil2|  1 2 1         
>                             2 4 2 
>                             1 2 1 |]

I.e we map this stencil then divide by 16 and take the coarse array of items where both the row and column are even. Taking the coarse grid items can be done with an array traversal

> {-# INLINE coarsen #-}  
> coarsen :: Array U DIM2 Double -> Array D DIM2 Double
> coarsen !arr 
>  = traverse arr   -- i+1 and j+1 to deal with odd extents for arr correctly
>             (\ (e :. i :. j) -> (e :. (i+1) `div` 2 :. (j+1) `div` 2))
>             (\ get (e :. i :. j) -> get (e :. 2*i :. 2*j))

Here the second argument for traverse – the function to calculate the new extent from the old – adds one before the division to ensures that odd extents are dealt with appropriately.
The inline pragma and bang pattern (!) on the argument are needed for good optimisation.

Coarsening after mapping the above stencil works well with odd extents but is not so good for even extents. For completeness we allow for even extents as well and in this case it is more appropriate to calculate one quarter of the results from mapping this stencil

> sum4Stencil :: Stencil DIM2 Double    
> sum4Stencil = [stencil2|  0 0 0         
>                           0 1 1 
>                           0 1 1 |] 

We treat mixed even and odd extents this way as well so we can express restriction for all cases as

> {-# INLINE restrict #-}       
> restrict :: Array U DIM2 Double -> Array D DIM2 Double
> restrict !arr 
>   | odd n && odd m
>     = coarsen
>       $ smap (/16)  
>       $ mapStencil2 BoundClamp restStencil arr
>   | otherwise
>     = coarsen
>       $ smap (/4)  
>       $ mapStencil2 BoundClamp sum4Stencil arr
>   where _ :. m :. n = extent arr

For interpolation in the case of odd extents we want to distribute coarse to fine values according to the “give-to” stencil

1/4 1/2 1/4
1/2  1  1/2
1/4 1/2 1/4

This means that the fine value at the centre becomes the coarse value at the corresponding position (if you picture the coarse grid spaced out to overlay the fine grid). The fine values around this central value inherit proportions of the coarse value as indicated by the give-to stencil (along with proportions from other neighbouring coarse values).

Odd Extent Interpolation

Odd Extent Interpolation

For the even extent case we simply make four copies of each coarse value using a traverse

> {-# INLINE inject4 #-} 
> inject4 :: Source r a => Array r DIM2 a -> Array D DIM2 a
> inject4 !arr 
>  = traverse arr -- mod 2s deal with odd extents
>             (\ (e :. i :. j) -> (e :. 2*i - (i `mod` 2) :. 2*j - (j `mod` 2)))
>             (\get (e :. i :. j) -> get(e :. i `div` 2 :. j `div` 2))
Even Extent Interpolation

Even Extent Interpolation

Again, we just use the latter version to cover cases with mixed even and odd extents. There is no primitive for “give-to” stencils but it is easy enough to define what we want with a traverse.

> {-# INLINE interpolate #-} 
> interpolate :: Array U DIM2 Double -> Array D DIM2 Double
> interpolate !arr 
>   | odd m && odd n
>        = traverse arr 
>                   (\ (e :. i :. j) -> (e :. 2*i - (i `mod` 2) :. 2*j - (j `mod` 2)))
>                   (\get (e :. i :. j) -> case () of
>                    _ | even i && even j     -> get(e :. i `div` 2 :. j `div` 2)
>                      | even i               -> (0.5)*(get(e :. i `div` 2 :. j `div` 2)
>                                                       + get(e :. i `div` 2 :. (j `div` 2)+1)) 
>                       -- odd i
>                      | even j               -> (0.5)*(get(e :. i `div` 2 :. j `div` 2)
>                                                       + get(e :. (i `div` 2)+1 :. j `div` 2)) 
>                       -- odd i and j
>                      | otherwise            -> (0.25)*(get(e :. i `div` 2 :. j `div` 2)
>                                                       + get(e :. i `div` 2 :. (j `div` 2)+1)
>                                                       + get(e :. (i `div` 2)+1 :. j `div` 2)
>                                                       + get(e :. (i `div` 2)+1 :. (j `div` 2)+1))
>                   )
>   | otherwise = inject4 arr
>  where _ :. n :. m = extent arr

The uses of mod 2 in the new size calculation are there to ensure that odd extents remain odd.

An alternative interpolation calculation

As an aside, we found that the above interpolation for odd extents can be implemented with a stencil convolution using sum4Stencil (which we defined above) after applying inject4 and then dividing by 4.
So we could define (for odd extents)

interpolate' :: Monad m
                => Array U DIM2 Double -> m (Array (TR PC5) DIM2 Double) 

interpolate' arr = 
   do fineArr <- computeP $ inject4 arr
      return $ smap (/4) 
             $ mapStencil2 BoundClamp sum4Stencil fineArr

We have to make the intermediate fineArr manifest (using computeP) before mapping the stencil over it which is why a monad is used. The final array is not manifest but a structured array resulting from the stencil map. By not making this manifest, we allow the computation to be combined with other computations improving inlining optimisation opportunities.

To illustrate the effect of interpolate', suppose the coarse array content is just

a b c
d e f
g h i

Then the injected array content will be

a a b b c c
a a b b c c
d d e e f f
d d e e f f
g g h h i i
g g h h i i

After mapping the stencil and dividing by 4 we have (assuming top left is (0,0))

at (1,1) (a+b+d+e)/4               (odd row,  odd column)
at (2,1) (d+e+d+e)/4 = (d+e)/2     (even row, odd column)
at (1,2) (b+b+e+e)/4 = (b+e)/2     (odd row,  even column)
at (2,2) (e+e+e+e)/4 = e           (even row, even column)

This even handles the bottom and right boundaries as in the original interpolation.

Slightly more generally, after inject4 and for any w x y z then convolution with the stencil

0 0 0
0 w x
0 y z

will produce the same as interpolating with the “give-to” stencil

  z     (z+y)    y   
(z+x) (z+y+x+w) (y+w)           
  x     (x+w)    w      

Conversely after inject4, the give-to stencil has to have this form for a stencil convolution to produce the same results.

Since the stencil version does require making an intermediate array manifest it is not clear at face value which is more efficient, so we will stick with interpolate.

We note that for even extents, the combination of an interpolation followed by a restriction is an identity (inject4 followed by coarsen). For odd extents, the combination preserves internal values but not values on the boarder. This will not be a problem if boundary conditions of the problem are used to update boarders of the array.

MultiGrid with Coarse Grid Correction Scheme

The problem to solve for u is

A u = -f

where A is a linear operator. This linear operator could be represented as a matrix or as a matrix multiplication, but this is not the most efficient way of implementing a solver. In the multigrid scheme described above we need knowledge of A to implement an approximate solver, and also to calculate residuals. There are ways to implement these latter two operations with some mathematical analysis of the linear operator using finite differences. We will therefore be using an approximator and a residual calculator directly rather than a direct representation of A as a parameter.

Let us assume that the approximate solver is a single step improvement which we want to iterate a few times. We can write down the iterator (iterateSolver)

> {-# INLINE iterateSolver #-}
> iterateSolver !opA !steps !arrInit
>  = go steps arrInit
>   where
>     go 0 !arr = return arr
>     go n !arr 
>       = do   arr' <- computeP $ opA arr
>              go (n - 1) arr'

In the definition of multiGrid we make use of a function to create an array of zeros of a given shape (extent)

> {-# INLINE zeroArray #-}
> zeroArray  :: Shape sh => sh -> Array U sh Double
> zeroArray !sh = fromListUnboxed sh $ replicate (size sh) 0.0

and a function to determine when we have the coarsest grid (the bases case for odd extents will be 3 by 3)

> {-# INLINE coarsest #-}
> coarsest :: Array U DIM2 Double -> Bool
> coarsest !arr = m<4 where (Z :. _ :. m) = extent arr

We also use some global parameters to indicate the number of iteration steps (to simplify expressions by passing fewer arguments)

> steps1,steps2 :: Int
> steps1 = 5       -- number of iterations for first step in multiGrid
> steps2 = 5       -- number of iterations for last step in multiGrid

To write down a first version of multiGrid, we temporarily ignore boundary conditions. We will also assume that the grid spacing is the same in each dimension and use just one parameter (h) as the grid spacing.

{- version ignoring boundary conditions -}
multiGrid approxOp residualOp h f uInit 
 = if coarsest uInit
   then 
     do computeP $ approxOp h f uInit
   else
     do v  <- iterateSolver (approxOp h f) steps1 uInit 
        r  <- computeP $ residualOp h f v                 -- calculate fine residuals
        r' <- computeP $ restrict r                       -- move to coarser grid                          
        err <- multiGrid approxOp residualOp (2*h) r' $ zeroArray (extent r')
        vC <- computeP $ szipWith (+) v  
                       $ interpolate err                  -- correct with errors on fine grid
        iterateSolver (approxOp h f) steps2 vC            -- solve again with improved approximation

The parameters are

  1. approxOp – an approximate solver to be iterated.
  2. residualOp – a residual calculator
  3. h – the grid spacing
  4. f – a representation for the (negative of) the function on the right hand side of the equation
  5. uInit – an array representing a suitable first guess to start from.

Note that both approxOp and residualOp need to be passed h and f as well as an array when they are used. Also, the recursive call of multiGrid requres the two function parameters to work at each coarseness of grid. The grid spacing doubles and the array of residuals has to be converted to the coarser grid for the recursive call.

Next we consider how we deal with the boundary conditions. We will only deal with the Dirichlet type of boundary conditions here.

MultiGrid with Dirichlet boundary conditions

For u as above, Dirichlet boundary conditions have the form u = g for some function g on the boundary \partial \Omega. We will take the standard approach of representing the boundary conditions using a mask array (with 0’s indicating boundary positions and 1’s at non-boundary positions) along with a boundary values array which has 0’s at non-boundary positions. This will enable us to adapt these two arrays for different coarsenesses of grid.

We are assuming we are working in just two dimensions so u=u(x,y). We can thus assume two dimensional arrays represent f and the boundary mask and boundary values.

We can now define multiGrid with boundary conditions as:

 multiGrid approxOp residualOp h f boundMask boundValues uInit 
  = if coarsest uInit
    then 
      do computeP $ approxOp h f boundMask boundValues uInit
    else
      do v  <- iterateSolver (approxOp h f boundMask boundValues) steps1 uInit 
         r  <- computeP $ residualOp h f boundMask v      -- calculate fine residuals
         boundMask' <- computeP $ coarsen boundMask       -- move to coarser grid
         r' <- computeP $ szipWith (*) boundMask' $ restrict r                            
         let zeros = zeroArray (extent r')
         err <- multiGrid approxOp residualOp (2*h) r' boundMask' zeros zeros
         vC <- computeP $ szipWith (+) v  
                        $ szipWith (*) boundMask 
                        $ interpolate err                 -- correct with errors on fine grid
         iterateSolver (approxOp h f boundMask boundValues) steps2 vC 

In the base case (when the grid size is 3*3) there is only one value at the centre point which needs to be evaluated (assuming the surrounding 8 points are given by the boundary conditions). This will be solved exactly with a single step of the approximate solver. The recursive call uses zeros for both the boundValues (redundantly) as well the initial guess. We note that the residual calculation makes use of the mask but not the boundary values as these should be zero for residuals. The restriction of residuals for the coarse grid also requires applying the mask adapted for the coarse grid. The interpolation of errors uses the (fine version of the) mask but the boundary values will already be set in v which is added to the errors.

Full MultiGrid

Before looking at specific examples, we can now write down the algorithm for the full multigrid scheme which extends multigrid as follows.

Before calling multiGrid we precalculate the initial guess using a coarser grid and interpolate the result. The precalculation on the coarser grid involves the same full multigrid process recursively on coarser and coarser grids.

fullMG approxOp residualOp h f boundMask boundValues
 = if coarsest boundValues
     then do computeP $ approxOp h f boundMask boundValues boundValues
         -- an approximation with  1 interior point will be exact using boundValues as initial
     else do  
         -- recursively precalculate for the coarser level
            f' <- computeP $ restrict f
            boundMask' <- computeP $ coarsen boundMask
            boundValues' <- computeP $ coarsen boundValues
            v' <- fullMG approxOp residualOp (2*h) f' boundMask' boundValues'
         -- move to finer level     
            v <- computeP $ szipWith (+) boundValues  
                           $ szipWith (*) boundMask 
                           $ interpolate v'  
         -- solve for finer level
            multiGrid approxOp residualOp h f boundMask boundValues v

Note that in the base case we need an initial guess for the coarsest initial array. We simply use the boundValues array for this. The interpolation of v requires resetting the boundary values.

Improved inlining using modules

The previous (higher order) versions of multiGrid and fullMG which take approxOp and residualOp as arguments can by substantially improved by importing the operations in a module and specialising the definitions for the imported functions. Thus we define a module which imports a separate module (defining approxOp and residualOp) and which exports

> multiGrid :: Monad m =>
>              Double
>              -> Array U DIM2 Double
>              -> Array U DIM2 Double
>              -> Array U DIM2 Double
>              -> Array U DIM2 Double
>              -> m (Array U DIM2 Double)

> multiGrid !h !f !boundMask !boundValues !uInit 
>  = if coarsest uInit
>    then 
>     do computeP $ approxOp h f boundMask boundValues uInit
>    else
>     do v  <- iterateSolver (approxOp h f boundMask boundValues) steps1 uInit
>        r  <- computeP $ residualOp h f boundMask v
>        boundMask' <- computeP $ coarsen boundMask
>        r'  <- computeP $ szipWith (*) boundMask' $ restrict r 
>        let zeros = zeroArray (extent r')
>        err  <- multiGrid (2*h) r' boundMask' zeros zeros
>        vC   <- computeP $ szipWith (+) v  
>                            $ szipWith (*) boundMask 
>                            $ interpolate err
>        iterateSolver (approxOp h f boundMask boundValues) steps2 vC 

and the main full multigrid operation

> fullMG :: Monad m =>
>           Double
>           -> Array U DIM2 Double
>           -> Array U DIM2 Double
>           -> Array U DIM2 Double
>           -> m (Array U DIM2 Double)

> fullMG !h !f !boundMask !boundValues
>  = if coarsest boundValues
>      then do computeP $ approxOp h f boundMask boundValues boundValues
>      else do 
>             f'  <- computeP $ restrict f
>             boundMask' <- computeP $ coarsen boundMask
>             boundValues' <- computeP $ coarsen boundValues
>             v'  <- fullMG (2*h) f' boundMask' boundValues' 
>             v   <- computeP $ szipWith (*) boundValues  
>                                $ szipWith (*) boundMask 
>                                $ interpolate v'
>             multiGrid h f boundMask boundValues v

These versions perform significantly better when optimised.

Poisson’s Equation

For the two dimensional case, where u = u(x,y), Poisson’s equation has the form

\nabla^2 u = -f(x,y)

for (x,y) \in \Omega subject to the Dirichlet boundary condition u(x,y) = g(x,y) for (x,y) \in \partial \Omega

So the linear operator here is the Laplace operator

\Delta = \nabla^2 = {\partial^2 \over\partial x^2} + {\partial^2 \over\partial y^2}

We need to make use of finite difference analysis to implement the approximation step and the residual calculation for Poisson’s equation.

Finite difference analysis with the central difference operator leads to the following (five point difference) formula approximating Poisson’s equation

u_{i-1,j} + u_{i+1,j} + u_{i,j-1} + u_{i,j+1} - 4u_{i,j} = -h^2 f_{ij}

Rewriting the above as

u_{ij} = \frac{1}{4}  \{ u_{i-1,j} + u_{i+1,j} + u_{i,j-1} + u_{i,j+1} + h^2 f_{ij} \}

gives us an algorithm (approxOp) for the approximating iteration step. This is known as the the Jacobi method. We map a suitable stencil to add the four neighbours of each item, then we add corresponding elements of the array f multiplied by h^2, then we divide by 4 and reset the boundary with the mask and values.

> {-# INLINE approxOp #-} 
> approxOp :: Double
>             -> Array U DIM2 Double
>             -> Array U DIM2 Double
>             -> Array U DIM2 Double
>             -> Array U DIM2 Double
>             -> Array (TR PC5) DIM2 Double               

> approxOp !h !f !boundMask !boundValues !arr
>        =   szipWith (+) boundValues
>          $ szipWith (*) boundMask
>          $ smap (/4)
>          $ szipWith (+) (R.map (*hSq) f )
>          $ mapStencil2 (BoundConst 0) 
>             [stencil2|   0 1 0
>                          1 0 1 
>                          0 1 0 |] arr
>          where hSq = h*h

We have chosen to leave the resulting array in a non-manifest form so that inlining can combine this operation with any subsequent operations to optimise the combination.

To calculate residuals from an estimation v, we can use the same formula as above to approximate r = \nabla^2 v + f. Rearranging we have
r_{ij} = \frac{1}{h^2} \{ v_{i-1,j} + v_{i+1,j} + v_{i,j-1} + v_{i,j+1} - 4 v_{ij} \} + f_{ij}
This leads us to implement residualOp with a five point stencil that adds four neighbours and subtracts four times the middle item. After this we divide all items by h^2 and then add f and reset the boundary.

> {-# INLINE residualOp #-}
> residualOp :: Double
>               -> Array U DIM2 Double
>               -> Array U DIM2 Double
>               -> Array U DIM2 Double
>               -> Array (TR PC5) DIM2 Double

> residualOp !h !f !boundMask !v
>        =   szipWith (*) boundMask
>          $ szipWith (+) f
>          $ smap (*hFactor)
>          $ mapStencil2 (BoundConst 0) 
>             [stencil2|   0 1  0
>                          1 -4 1 
>                          0 1  0 |] v
>          where hFactor = 1/(h*h)

As noted earlier, the boundary is reset simply with the mask. The boundary values should all be zero for residuals, so we can drop any resetting of boundary values as the mask will have done that already.

The above two functions (approxOP and residualOp) are put in a module PoissonOps to be imported by the module MultiGrid_Poisson.

Setting up to run an example

This example is taken from (Iserles 2009 p.162). On the unit square

f(x,y) = -\{x^2+y^2\}

The Dirichlet boundary conditions are

u(x,0) = 0 \ \ \ \text{and}\ \ \     u(x,1) = \frac{1}{2} x^2 \ \ \    \text{for}\ \ \      0 \le x \le 1

and

u(0,y) = \sin {\pi x} \ \ \ \text{and}\ \ \     u(1,y) = e^{\pi x} \sin {\pi y} + \frac{y^2}{2} \ \ \    \text{for}\ \ \      0 \le y \le 1

We define the following to calculate values from the number of grids to be used.

> fineGridSpaces:: Int -> Int
> fineGridSpaces gridStackSize
>   = 2^gridStackSize

> fineGridShape:: Int -> DIM2
> fineGridShape gridStackSize 
>   = (Z :. n :. n) where n = 1+fineGridSpaces gridStackSize

We will set a global value for the sides of the area (the unit square) and also set a stack size of 9 grids so the finest will have 2^9+1 by 2^9+1 (i.e. 513 by 513) grid points

> distance :: Double
> distance = 1.0        -- length of sides for square area
> gridStackSize :: Int
> gridStackSize = 9     -- number of grids to be used including finest and coarsest

> intervals = fineGridSpaces gridStackSize
> shapeInit = fineGridShape gridStackSize

> hInit :: Double
> hInit = distance / fromIntegral intervals -- initial finest grid spacing

The initial arrays are defined using

> boundMask :: Array U DIM2 Double
> boundMask = 
>  fromListUnboxed shapeInit $ concat 
>   [edgeRow,
>    take ((intervals-1)*(intervals+1)) (cycle mainRow),
>    edgeRow
>   ]
>   where edgeRow = replicate (intervals+1) 0.0
>         mainRow = 0.0: replicate (intervals-1) 1.0  Prelude.++ [0.0]
  

> coordList :: [Double]
> coordList = Prelude.map ((hInit*) . fromIntegral) [0..intervals] 

> boundValues :: Array U DIM2 Double
> boundValues =
>   fromListUnboxed shapeInit $ concat
>      [Prelude.map (\j -> sin (pi*j)) coordList, 
>       concat $ Prelude.map mainRow $ tail $ init coordList,
>       Prelude.map (\j -> exp pi * sin (pi*j) + 0.5*j^2) coordList
>      ] 
>      where mainRow i = replicate intervals 0.0 Prelude.++ [0.5*i^2]

> fInit :: Array U DIM2 Double
> fInit =                        -- negative of RHS of Poisson Equation 
>  fromListUnboxed shapeInit $ concat $ Prelude.map row coordList    
>  where row i = Prelude.map (item i) coordList
>        item i j = -(i^2 + j^2)

> uInit :: Array U DIM2 Double
> uInit = boundValues

We are now in a position to calculate

> test1 = multiGrid hInit fInit boundMask boundValues uInit

and

> test2 = fullMG hInit fInit boundMask boundValues

as well as comparing with iterations of the basic approximation operation on its own

> solverTest :: Monad m => Int -> m(Array U DIM2 Double)
> solverTest n  = iterateSolver (approxOp hInit fInit boundMask boundValues) n uInit

We note that this particular example of Poisson’s equation does have a known exact solution

u(x,y) = e^{\pi x} \sin {\pi y} + \frac{1}{2} (xy)^2

So we can calculate an array with values for the exact solution and look at the errors. We just look at the maximum error here.

> exact :: Array U DIM2 Double
>  exact =
>  fromListUnboxed shapeInit $
>  concat $ Prelude.map row coordList    
>  where row i = Prelude.map (item i) coordList
>        item i j = exp (pi*i) * sin (pi*j) + 0.5*i^2*j^2

> maxError :: Monad m => m(Array U DIM2 Double) -> m (Double)
> maxError test = 
>   do ans <- test
>      err :: Array U DIM2 Double 
>          <- computeP
>                $ R.map abs
>                $ R.zipWith (-) ans exact
>      foldAllS max 0.0 err

A note on errors

The Jacobi method used for the approximation step is known to converge very slowly, but fullMG significantly improves convergence rates. The literature analysing multigrid and full multigrid shows that the errors are reduced rapidly if the approximation method smooths the high frequency errors quickly (error components with a wavelength less than half the grid size). Unfortunately the Jacobi method does not do this. Gauss-Seidel and SOR methods do smooth high frequency errors but their usual implementation relies on in-place updating of arrays which is problematic for parallel evaluation. However, although multiGrid with the Jacobi approxOp does not reduce errors rapidly, fullMG does (with only gridStackSize=7 and step1=step2=3). This shows that the improved initial guess using fullMG outweighs the slow reduction of errors in multiGrid.

Some results

Clearly this code is designed to be run on multi-cores to use the parallelism. However, even using a single processor, we can see the impressive performance of Repa with optimisations on some simple tests of the Poisson example (running on a 2.13 GHz Intel Core 2 Duo).

Results for fullMG (steps1 = 5, steps2 = 5)
Grid Stack Fine Grid cpu Time(MS) Max Error
6 65*65 22 2.1187627917047536e-3
7 129*129 83 5.321669730857792e-4
8 257*257 276 1.3324309721163274e-4
9 513*513 1417 3.332619984242058e-5
10 1025*1025 6017 8.332744698691386e-6
11 2049*2049 26181 2.0832828990791086e-6

So increasing the grid stack size by 1 roughly quadruples the size of the fine grid, increases runtime by a similar factor, and improves the error by a factor of 4 as well.

As a comparison, multiGrid alone (with the Jacobi approxOp) has high errors and increasing grid stack size does not help.

Results for multiGrid (steps1 = 5, steps2 = 5)
Grid Stack cpu Time(MS) Max Error
6 14 1.116289597967647
7 49 1.1682431240072333
8 213 1.1912130219418806
9 1118 1.2018233367410431
10 4566 1.206879281001907
11 20696 1.209340584664126

To see the problems with the basic Jacobi method on its own we note that after 400 iteration steps (with grid stack size of 6) the max error is 5.587133822631921. Increasing the grid stack size just makes this worse.

Compiler options

As we had observed in previous work Repa Laplace and SOR the large amount of inline optimisation seems to trigger a ghc compiler bug warning. This is resolved by adding the compiler option -fsimpl-tick-factor=1000 (along with -O2).

Some variations

We also tried two variations on opproxOp.

The first was to add in an extra relaxation step which weights the old value with the new value at each step using a weighting \omega between 0 and 1. This slightly increases runtime, but does not improve errors for fullMG. It’s purpose is to improve smoothing of errors, but as we have seen although this may be significant for multiGrid it is not for fullMG. The improved starting points are much more significant for the latter. Note that successive over relaxing (SOR) with \omega between 1 and 2 is not suitable for use with the Jacobi method in general as this is likely to diverge.

omega :: Double
omega = 2/3

{-# INLINE approxOp #-} 
approxOp :: Double
            -> Array U DIM2 Double
            -> Array U DIM2 Double
            -> Array U DIM2 Double
            -> Array U DIM2 Double
            -> Array (TR PC5) DIM2 Double               

approxOp !h !f !boundMask !boundValues !arr
       = szipWith weightedSum arr         -- extra relaxation step
         $ szipWith (+) boundValues
         $ szipWith (*) boundMask
         $ smap (/4)
         $ szipWith (+) (R.map (*hSq) f )
         $ mapStencil2 (BoundConst 0) 
            [stencil2|   0 1 0
                         1 0 1 
                         0 1 0 |] arr
         where hSq = h*h
               weightedSum old new = (1-omega) * old + omega * new

The second variation is based on the ‘modified 9 point scheme’ method discussed in (Iserles 2009 p.162). This has slower runtimes than the 5 point Jacobi scheme as it involves two stencil mappings (one across f as well as u). It had no noticable improvement on error reduction for fullMG.

-- modified 9 point scheme
approxOp ::  Double
                -> Array U DIM2 Double
                -> Array U DIM2 Double
                -> Array U DIM2 Double
                -> Array U DIM2 Double
                -> Array (TR PC5) DIM2 Double
            
approxOp !h !f !boundMask !boundValues !arr
       =   R.szipWith (+) boundValues
         $ R.szipWith (*) boundMask
         $ R.smap (/20)
         $ R.szipWith (+) 
            (R.map (*hFactor) 
             $ mapStencil2 (BoundConst 0) 
                  [stencil2|   0 1 0
                               1 8 1 
                               0 1 0 |] f )
         $ mapStencil2 (BoundConst 0) 
            [stencil2|   1 4 1
                         4 0 4 
                         1 4 1 |] arr
         where hFactor = h*h/2

This algorithm was obtained from rearranging the formula

\frac{-10}{3}u_{ij} + \frac{1}{6} \{ {u_{i-1,j-1} + 4 u_{i-1,j} + u_{i-1,j+1} + 4 u_{i,j-1} + 4 u_{i,j+1} + u_{i+1,j-1}+ 4 u_{i+1,j} + u_{i+1,j+1} \} } = \frac{h^2}{12} { \{ 8 f_{i,j} + f_{i-1,j} + f_{i,j-1} + f_{i,j+1} + f_{i+1,j+1} \} }

Parallel Processing

The important point to note is that this code is ready to run using parallel processing. Results using multi-core processors will be posted in a separate blog.

Acknowledgements and References

Iserles (2009) gives some analysis of the multigrid method and this book also provided the example used (and the modified 9 point scheme). There are numerous other finite methods books covering multigrid and some devoted to it. Online there are some tutorial slides by Briggs part1 and part2 and some by Stüben amongst many others. A paper by Lippmeier, Keller, and Peyton Jones discussing Repa design and optimisation for stencil convolution in Haskell can be found here.

Iserles, Arieh. 2009. A First Course in Numerical Analysis of Differential Equations. Second edition. CUP.

Advertisement

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s