Skip to content

GitLab

Commit 94384bea authored by Christoph Ruegge's avatar Christoph Ruegge
Browse files

.hs for lecture 25 (PDEs)

parent 0dc04170
Hide whitespace changes
Inline Side-by-side
lecture/25 poisson equation (PDE)/01-hmatrix.hs
View file @ 94384bea
{-# OPTIONS_GHC -Odph -rtsopts -threaded -fno-liberate-case -funfolding-use-threshold1000 -funfolding-keeness-factor1000 -fllvm -optlo-O3 #-}
{-# OPTIONS_GHC -fdefer-type-errors #-}
{-# LANGUAGE BangPatterns #-}
import Data.Array.Repa hiding (map)
import Data.Array.Repa.Eval (fromList)
import Data.Array.Repa
import Data.Maybe (fromJust)
import Graphics.Rendering.Plot
import qualified Graphics.Rendering.Plot as P
import Numeric.LinearAlgebra.Helpers
import Numeric.LinearAlgebra.Repa
import Numeric.LinearAlgebra.Repa.Conversion
import Numeric.LinearAlgebra.Repa (linearSolve)
import Numeric.LinearAlgebra.Repa.Conversion (repa2hm)
type Interval = (Double, Double, Int)
poisson :: Int -> (Double -> Double -> Double) -> Mat Double
poisson n f = computeS . reshape gridshape . fromJust . linearSolve laplace $ asColumn rhs
where
-- coordinates
h = 1 / fromIntegral n :: Double
xs :: Array U DIM1 Double
xs = computeS $ fromFunction (Z:.n) $ \(Z:.(!j)) -> (fromIntegral j + 1/2) * h
-- right hand side
gridshape = Z:.n:.n :: DIM2
rhs = vec [ f (unsafeLinearIndex xs jx) (unsafeLinearIndex xs jy)
| i <- [0..n*n-1] , let Z:.(!jx):.(!jy) = fromIndex gridshape i]
poisson :: (Double -> Double -> Double) -> Interval -> Interval -> ([(Double, Double)], Mat Double)
poisson f (x0, x1, nx) (y0, y1, ny) = (mesh, solution)
where
-- coordinates
x = computeS $ linspace nx (x0, x1) :: Vec Double
hx = x ! (Z:.1)
y = computeS $ linspace ny (y0, y1) :: Vec Double
hy = y ! (Z:.1)
-- mesh
gridshape = Z:.nx:.ny :: DIM2
mesh = [ (unsafeLinearIndex x jx, unsafeLinearIndex y jy)
| i <- [0..nx*ny-1]
, let Z:.(!jx):.(!jy) = fromIndex gridshape i
]
-- laplace matrix
matshape = Z:.(nx*ny):.(nx*ny) :: DIM2
laplace = computeS $ fromFunction matshape idxf :: Mat Double
idxf :: DIM2 -> Double
idxf (Z:.(!j):.(!k))
| j == k = 4*c
| xj == xk && abs (yj - yk) == 1 = -c
| yj == yk && abs (xj - xk) == 1 = -c
| otherwise = 0
where
Z:.(!xj):.(!yj) = fromIndex gridshape j :: DIM2
Z:.(!xk):.(!yk) = fromIndex gridshape k :: DIM2
c = 1/(hx*hy) :: Double
-- right hand side
rhs = vec $ map (uncurry f) mesh
-- solve linear system
solution = computeS . reshape gridshape . fromJust . linearSolve laplace . asColumn $ rhs :: Mat Double
-- Both the right hand side and the solution *should be* 2-d arrays of
-- size (n x n), but linear algebra routines require them to be vectors.
-- So we need a way to convert from a 1-d vector index j (0 <= j <= n^2-1)
-- to a 2-d coordinate pair (xj, yj) (0 <= xj, yj <= n-1). Instead of
-- reinventing it, we use repa's own functions to convert between linear
-- and multi-dimensional indices:
--
-- fromIndex gridshape :: Int -> DIM2
--
-- converts 1-d to 2-d. The resulting vector is converted back to a 2-d
-- array using
--
-- reshape gridshape :: Array r DIM1 Double -> Array D DIM2 Double
matshape = Z:.(n^2):.(n^2) :: DIM2
laplace :: Mat Double
laplace = computeS $ fromFunction matshape computeElem
where computeElem :: DIM2 -> Double
computeElem (Z:.(!j):.(!k))
| j == k = 4*c
| xj == xk && abs (yj - yk) == 1 = -c
| yj == yk && abs (xj - xk) == 1 = -c
| otherwise = 0
where
Z:.(!xj):.(!yj) = fromIndex gridshape j :: DIM2
Z:.(!xk):.(!yk) = fromIndex gridshape k :: DIM2
c = 1/h**2
main :: IO ()
main = do
-- let (mesh, solution) = poisson (\x y -> x*y^(4::Int)) (0, 1, 50) (0, 1, 50)
-- let (xs, ys) = unzip mesh
-- saveMatrix "xmat.txt" "%f" $ fromList (Z:.50:.50) xs
-- saveMatrix "ymat.txt" "%f" $ fromList (Z:.50:.50) ys
-- saveMatrix "zmat.txt" "%f" solution
let (_, solution) = poisson (\x y -> x*y^(4::Int)) (0, 1, 50) (0, 1, 50)
writeFigure PNG "01.png" (500, 500) $ do
setPlots 1 1
withPlot (1, 1) $ setDataset $ repa2hm solution
let solution = poisson 50 (\x y -> x*y**4)
P.writeFigure P.PNG "01.png" (500, 500) $ do
P.setPlots 1 1
P.withPlot (1, 1) $ P.setDataset $ repa2hm solution
lecture/25 poisson equation (PDE)/02-repa.hs
View file @ 94384bea
{-# OPTIONS_GHC -Odph -rtsopts -threaded -fno-liberate-case -funfolding-use-threshold1000 -funfolding-keeness-factor1000 -fllvm -optlo-O3 #-}
{-# OPTIONS_GHC -fdefer-type-errors #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}
import Data.Array.Repa as R hiding ((++))
import Graphics.Rendering.Plot
import Numeric.LinearAlgebra.Repa.Conversion
import CG
import CG (runCG)
import Data.Array.Repa
import qualified Graphics.Rendering.Plot as P
import Numeric.LinearAlgebra.Repa.Conversion (repa2hm)
poissonOperator :: Int -> Array U DIM2 Double -> Array D DIM2 Double
poissonOperator n arr = fromFunction ext getElem
where ext = extent arr
get idx = if inShape ext idx then index arr idx else 0
getElem (Z:.i:.j) = ( - get (Z:.(i-1):.j) - get (Z:.(i+1):.j)
- get (Z:.i:.(j-1)) - get (Z:.i:.(j+1))
+ 4*get (Z:.i:.j) )
* fromIntegral (n^(2::Int))
poisson :: Int -> (Double -> Double -> Double) -> IO (Array U DIM2 Double)
poisson n f = runCG 1e-6 laplace =<< computeP rhs
where
-- coordinates
h = 1 / fromIntegral n :: Double
xs :: Array U DIM1 Double
xs = computeS $ fromFunction (Z:.n) $ \(Z:.(!j)) -> (fromIntegral j + 1/2) * h
-- right hand side
gridshape = Z:.n:.n :: DIM2
rhs = fromFunction gridshape $ \(Z:.(!jx):.(!jy)) ->
f (unsafeLinearIndex xs jx) (unsafeLinearIndex xs jy)
-- laplace operator
laplace :: Array U DIM2 Double -> Array D DIM2 Double
laplace arr = fromFunction ext computeElem
where
!ext = extent arr
get !idx | inShape ext idx = unsafeIndex arr idx
| otherwise = 0
computeElem (Z:.(!jx):.(!jy)) = ( 4*get (Z:.jx:.jy)
- get (Z:.(jx-1):.jy) - get (Z:.(jx+1):.jy)
- get (Z:.jx:.(jy-1)) - get (Z:.jx:.(jy+1)) )
/ h**2
rhs :: Int -> Array U DIM2 Double
rhs n = computeS $ fromFunction (Z:.n:.n) $ \(Z:.i:.j) -> f (toCoord i) (toCoord j)
where h = 1 / fromIntegral n
toCoord k = h*fromIntegral k + 0.5
f x y = x*(y^(4::Int))
runCG :: Shape sh => Double -> (Array U sh Double -> Array D sh Double) -> Array U sh Double -> IO (Array U sh Double)
runCG tol op rhs = do
let initial = computeS $ fromFunction (extent rhs) (const 0)
let steps' = cg op rhs initial
let r20 = cgr2 $ head steps'
let steps = takeUntil (\x -> sqrt (cgr2 x / r20) < tol) steps'
result <- process (zip [(1::Int)..] steps) $ \(n, cgs) ->
putStrLn $ show n ++ " " ++ show (sqrt $ cgr2 cgs / r20)
return $ cgx . snd $ result
main :: IO ()
main = do
solution <- let n = 50 in runCG 1e-6 (poissonOperator n) (rhs n)
writeFigure PNG "02.png" (500, 500) $ do
setPlots 1 1
withPlot (1, 1) $ setDataset $ repa2hm $ copyS solution
solution <- poisson 50 (\x y -> x*y**4)
P.writeFigure P.PNG "02.png" (500, 500) $ do
P.setPlots 1 1
P.withPlot (1, 1) $ P.setDataset $ repa2hm $ copyS solution
lecture/25 poisson equation (PDE)/03-stencils.hs
View file @ 94384bea
import Data.Array.Repa.Stencil.Dim2
import Data.Array.Repa.Stencil
{-# OPTIONS_GHC -Odph -rtsopts -threaded -fno-liberate-case -funfolding-use-threshold1000 -funfolding-keeness-factor1000 -fllvm -optlo-O3 #-}
poissonStencil :: Integer -> Stencil DIM2 Double
poissonStencil n = makeStencil2 3 3 getElem
where val = fromIntegral n^2
getElem (Z:. -1 :. 0) = Just $ -val
getElem (Z:. 1 :. 0) = Just $ -val
getElem (Z:. 0 :. -1) = Just $ -val
getElem (Z:. 0 :. 1) = Just $ -val
getElem (Z:. 0 :. 0) = Just $ 4*val
getElem _ = Nothing
poissonOperator :: Integer -> Array U DIM2 Double -> Array D DIM2 Double
poissonOperator n = delay . mapStencil2 (BoundConst 0) (poissonStencil n)
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE QuasiQuotes #-}
-- solution <- let n = 50 in runCG 1e-3 (poissonOperator n) (rhs n)
-- Operations like `laplace` in 02-repa.hs are a frequent pattern: each element
-- of the resulting array is a linear combination of the surrounding elements of
-- the input array.
--
-- Repa has special means for constructing such operations in 2-d using
-- stencils. Stencils are constructed with
--
-- makeStencil2 :: Num a => Int -> Int -> (DIM2 -> Maybe a) -> Stencil DIM2 a
--
-- and applied to 2-d arrays with
--
-- mapStencil2 :: Source r a => Boundary a -> Stencil DIM2 a -> Array r DIM2 a -> Array PC5 DIM2 a
--
-- `Boundary a` is a type that handles the boundary conditions. It can be
-- - `BoundFixed x` returns a fixed value `x :: a` at the boundary,
-- - `BoundConst x` assumes pixels outside the array have the value `x :: a`,
-- - `BoundClamp` extends the boundary values to outside the boundary.
--
-- `PC5` is a special representation for stencil results; it is defined as (no
-- need to understand this)
--
-- type PC5 = P C (P (S D) (P (S D) (P (S D) (P (S D) X))))
--
-- `PC5` arrays can be converted to `D` arrays using the polymorphic function
--
-- delay :: (Shape sh, Source r e) => Array r sh e -> Array D sh e
{-# LANGUAGE QuasiQuotes #-}
import CG (runCG)
import Data.Array.Repa as R
import Data.Array.Repa.Stencil
import Data.Array.Repa.Stencil.Dim2
import qualified Graphics.Rendering.Plot as P
import Numeric.LinearAlgebra.Repa.Conversion (repa2hm)
poissonStencil :: Stencil DIM2 Double
poissonStencil = [stencil2| 0 -1 0
-1 4 -1
0 -1 0 |]
poisson :: Int -> (Double -> Double -> Double) -> IO (Array U DIM2 Double)
poisson n f = runCG 1e-6 laplace =<< computeP rhs
where
-- coordinates
h = 1 / fromIntegral n :: Double
xs :: Array U DIM1 Double
xs = computeS $ fromFunction (Z:.n) $ \(Z:.(!j)) -> (fromIntegral j + 1/2) * h
-- right hand side
gridshape = Z:.n:.n :: DIM2
rhs = fromFunction gridshape $ \(Z:.(!jx):.(!jy)) ->
f (unsafeLinearIndex xs jx) (unsafeLinearIndex xs jy)
-- laplace stencil
laplaceStencil = makeStencil2 3 3 computeElem :: Stencil DIM2 Double
where computeElem (Z:. -1 :. 0) = Just $ -c
computeElem (Z:. 1 :. 0) = Just $ -c
computeElem (Z:. 0 :. -1) = Just $ -c
computeElem (Z:. 0 :. 1) = Just $ -c
computeElem (Z:. 0 :. 0) = Just $ 4*c
computeElem _ = Nothing
c = 1/h**2
-- laplace operator
laplace :: Array U DIM2 Double -> Array D DIM2 Double
laplace = delay . mapStencil2 (BoundConst 0) laplaceStencil
-- {-# LANGUAGE QuasiQuotes #-}
-- There is special syntax for stencils. The laplaceStencil above can be
-- produced more easily by
--
-- laplaceStencil = [stencil2| 0 -1 0
-- -1 4 -1
-- 0 -1 0 |]
-- laplace = R.map (/h**2) . mapStencil2 (BoundConst 0) laplaceStencil
--
-- This syntax requires the QuasiQuotes language extension. The stencil is
-- constructed at compile time and can therefore not depend on runtime values.
-- That's why we have to use R.map to do the division by h**2 here.
-- poissonStencil :: Stencil DIM2 Double
-- poissonStencil n = [stencil2| 0 (-val) 0
-- (-val) (4*val) (-val)
-- 0 (-val) 0 |]
-- where val = fromIntegral n^2
main :: IO ()
main = do
solution <- poisson 50 (\x y -> x*y**4)
P.writeFigure P.PNG "03.png" (500, 500) $ do
P.setPlots 1 1
P.withPlot (1, 1) $ P.setDataset $ repa2hm $ copyS solution
lecture/25 poisson equation (PDE)/CG.hs → lecture/25 poisson equation (PDE)/lib/CG.hs
View file @ 94384bea
{-# OPTIONS_GHC -Odph -rtsopts -threaded -fno-liberate-case -funfolding-use-threshold1000 -funfolding-keeness-factor1000 -fllvm -optlo-O3 #-}
{-# OPTIONS_GHC -fdefer-type-errors #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}
-- CG is a simple, efficient iterative solver for linear systems
--
-- A x = y
--
-- with a positive definite matrix A, i.e. (x^T A x >= 0) for all x.
module CG where
import Control.Monad (foldM)
import Control.Monad.Identity (runIdentity)
import Data.Array.Repa as R hiding ((++))
-- CGState is the state of a CG iteration, i.e. everything that needs to be
-- passed to the next iteration.
data CGState sh = CGState { cgx :: !(Array U sh Double)
, cgp :: !(Array U sh Double)
, cgr :: !(Array U sh Double)
, cgr2 :: !Double
}
-- CG takes a function implementing a linear operator, a right hand side and an
-- initial guess and returns the (lazy, infinte) list of the CGStates of all
-- iterations. This way, the iterations can be decoupled from the stopping
-- criterion and the output of diagnostic information.
--
-- Computations use Repa's parallel computation in a straight-forward way and
-- needs to be performed in a monad. We use the Identity monad which does
-- essentially nothing except for providing the sequencing. All operations are
-- performed strictly using BangPatterns.
cg :: Shape sh => (Array U sh Double -> Array D sh Double) -> Array U sh Double -> Array U sh Double-> [CGState sh]
cg op rhs initial =
runIdentity $ do
!rInit <- computeP $ rhs -^ op initial
!r2Init <- normSquaredP rInit
return $ iterate cgStep (CGState initial rInit rInit r2Init)
where normSquaredP = sumAllP . R.map (^(2::Int))
where normSquaredP = sumAllP . R.map (**2)
scale a = R.map (* a)
-- the main iteration; maps a CGState to another CGState
cgStep (CGState x p r r2) =
cgStep (CGState !x !p !r !r2) =
-- We need a monad to avoid nested parallelism; the Identity
-- works fine.
runIdentity $ do
-- need inline type annotation here because computeP is polymorphic in its
-- output type (syntax like this requires ScopedTypeVariables)
-- We need inline type annotation here because computeP is
-- polymorphic in its output type, and the type of q is not
-- fixed by any other expression below. Syntax like this
-- requires ScopedTypeVariables.
!(q :: Array U sh Double) <- computeP $ op p
!qp <- sumAllP $ q *^ p
let alpha = r2 / qp
......@@ -39,11 +60,41 @@ cg op rhs initial =
!p' <- computeP $ r' +^ scale beta p
return $ CGState x' p' r' r2'
-- takeUntil is like takeWhile, but also returns the final iterate (no need to
-- waste it).
takeUntil :: (a -> Bool) -> [a] -> [a]
takeUntil _ [] = []
takeUntil predicate (x:xs)
| predicate x = [x]
| otherwise = x : takeUntil predicate xs
-- We want to be able to output some information about the iterates, i.e. run an
-- IO action over the list of iterates, before returning the final one. One
-- option would be something like
--
-- fmap last . forM iterates $ \it -> do [...]
--
-- However, this would perform the `fmap last` only *after* the IO actions are
-- done, therfore retaining the entire list of iterates in memory.
--
-- We want all but the final iterate to be garbage collected directly after
-- printing information. This can be done using a monadic fold, where we do not
-- actually accumulate a result but only return the last value.
process :: Monad m => [a] -> (a -> m ()) -> m a
process xs f = foldM (\_ x -> f x >> return x) undefined xs
-- runCG calls cg with initial guess 0, a stopping rule based on residuals
-- relative to the first initial guess, and outputs the residual while
-- iterating -- therefore, it returns a result in the IO monad.
runCG :: Shape sh => Double -> (Array U sh Double -> Array D sh Double) -> Array U sh Double -> IO (Array U sh Double)
runCG tol op rhs = do
let initial = computeS $ fromFunction (extent rhs) (const 0)
let steps' = cg op rhs initial
let r20 = cgr2 $ head steps'
let steps = takeUntil (\x -> sqrt (cgr2 x / r20) < tol) steps'
result <- process (zip [(1::Int)..] steps) $ \(n, cgs) ->
putStrLn $ show n ++ " " ++ show (sqrt $ cgr2 cgs / r20)
return $ cgx . snd $ result
lecture/25 poisson equation (PDE)/poisson-equation.cabal
View file @ 94384bea
......@@ -3,31 +3,40 @@ version: 0.1.0.0
author: Christoph Ruegge
maintainer: c.ruegge@math.uni-goettingen.de
build-type: Simple
-- extra-source-files:
cabal-version: >=1.10
library
default-language: Haskell2010
hs-source-dirs: lib
exposed-modules: CG
build-depends: base >=4.8 && <4.9
, mtl
, repa
executable 01-hmatrix
default-language: Haskell2010
main-is: 01-hmatrix.hs
build-depends: base >=4.8 && <4.9
, lahelpers
, plot
, repa
, repa-linear-algebra
default-language: Haskell2010
, poisson-equation
executable 02-repa
default-language: Haskell2010
main-is: 02-repa.hs
build-depends: base >=4.8 && <4.9
, lahelpers
, mtl
, plot
, repa
, repa-linear-algebra
default-language: Haskell2010
, poisson-equation
library
exposed-modules: CG
executable 03-stencils
default-language: Haskell2010
main-is: 03-stencils.hs
build-depends: base >=4.8 && <4.9
, mtl
, plot
, repa
default-language: Haskell2010
, repa-linear-algebra
, poisson-equation
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment