Permalink
Find file
Fetching contributors…
Cannot retrieve contributors at this time
128 lines (103 sloc) 4.17 KB
{-# LANGUAGE BangPatterns, RankNTypes, FlexibleContexts #-}
-- | Solver for the Laplace equation
-- You supply a BMP files specifying the boundary conditions.
-- The output is written back to another BMP file.
--
import Data.Array.Repa.Algorithms.Pixel
import Data.Array.Repa.Algorithms.ColorRamp
import Data.Array.Repa.IO.BMP
import Data.Array.Repa.IO.Timing
import System.Environment
import Data.Word
import Control.Monad
import SolverGet as SG
import SolverStencil as SS
import Data.Array.Repa as R
import Prelude as P
type Solver m
= 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)
solvers
= [ ("get", SG.solveLaplace)
, ("stencil", SS.solveLaplace) ]
main :: IO ()
main
= do args <- getArgs
case args of
[strSolver, steps, fileInput, fileOutput]
-> let Just solver = lookup strSolver solvers
in laplace solver (read steps) fileInput fileOutput
_ -> do
putStr usage
return ()
-- | Command line usage information.
usage :: String
usage = unlines
[ "Usage: laplace <solver> <iterations> <input.bmp> <output.bmp>"
, ""
, " iterations :: Int Number of iterations to use in the solver."
, " input.bmp :: FileName Uncompressed RGB24 or RGBA32 BMP file for initial and boundary values."
, " output.bmp :: FileName BMP file to write output to."
, ""
, " solver = one of " P.++ show (P.map fst solvers)
, ""
, " Format of input file:"
, " Boundary values are indicated in greyscale,"
, " ie from the list [(x, x, x) | x <- [0 .. 255]]"
, " Non-boundary values are indicated in blue,"
, " ie (0, 0, 255)"
, " Any other pixel value is an error."
, ""
, " NOTE: For GHC 7.0.3, this runs better when you turn off the parallel"
, " garbage collector. Run with +RTS -qg"
, "" ]
-- | Solve it.
laplace :: Solver IO
-> Int -- ^ Number of iterations to use.
-> FilePath -- ^ Input file.
-> FilePath -- ^ Output file
-> IO ()
laplace solve steps fileInput fileOutput
= do
-- Load up the file containing boundary conditions.
arrImage <- liftM (either (error . show) id)
$ readImageFromBMP fileInput
arrBoundValue <- computeP $ R.map slurpBoundValue arrImage
arrBoundMask <- computeP $ R.map slurpBoundMask arrImage
let arrInitial = arrBoundValue
-- Run the Laplace solver and print how long it took.
(arrFinal, t) <- time $ solve steps arrBoundMask arrBoundValue arrInitial
putStr (prettyTime t)
-- Write out the result to a file.
arrImageOut <- computeP
$ R.map rgb8OfDouble
$ R.map (rampColorHotToCold 0.0 1.0) arrFinal
writeImageToBMP fileOutput arrImageOut
-- | Extract the boundary value from a RGB triple.
slurpBoundValue :: (Word8, Word8, Word8) -> Double
{-# INLINE slurpBoundValue #-}
slurpBoundValue (!r, !g, !b)
-- A non-boundary value.
| r == 0 && g == 0 && b == 255
= 0
-- A boundary value.
| (r == g) && (r == b)
= fromIntegral (fromIntegral r :: Int) / 255
| otherwise
= error $ "Unhandled pixel value in input " P.++ show (r, g, b)
-- | Extract boundary mask from a RGB triple.
slurpBoundMask :: (Word8, Word8, Word8) -> Double
{-# INLINE slurpBoundMask #-}
slurpBoundMask (!r, !g, !b)
-- A non-boundary value.
| r == 0 && g == 0 && b == 255
= 1
-- A boundary value.
| (r == g) && (r == b)
= 0
| otherwise
= error $ "Unhandled pixel value in input " P.++ show (r, g, b)