Created
July 2, 2016 22:48
Star
You must be signed in to star a gist
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
{-# LANGUAGE QuasiQuotes, ScopedTypeVariables, TypeOperators, | |
FlexibleContexts, TypeFamilies, GeneralizedNewtypeDeriving, | |
StandaloneDeriving, DataKinds, FlexibleInstances #-} | |
module Main (main, kernel) where | |
import Prelude hiding ( traverse,) | |
import Data.Array.Repa as R hiding ((++)) | |
import Data.Array.Repa.Repr.Vector (toVector, computeVectorS) | |
import Data.Array.Repa.Stencil (Stencil, Boundary(BoundClamp,BoundConst)) | |
import Data.Array.Repa.Stencil.Dim2 (stencil2, makeStencil2, mapStencil2) | |
import Data.Array.Repa.Eval (Elt) | |
import GHC.Float (double2Float) | |
import Codec.Picture (Image(Image), PixelF, DynamicImage(ImageYF), savePngImage) | |
import Data.Vector.Generic (convert) | |
import Control.Monad.Identity (runIdentity) | |
{-# INLINE kernel #-} | |
kernel :: Stencil DIM2 Double | |
kernel = [stencil2| 1 2 1 | |
2 -12 2 | |
1 2 1 |] | |
{-# INLINE kern #-} | |
kern :: Monad m => Vals -> m Vals | |
kern vals = computeP $ mapStencil2 (BoundConst 0) kernel vals | |
type Height = Double | |
type Vals = Array U DIM2 Double | |
type Velocity = Double | |
type Vels = Array U DIM2 Double | |
type Stiffness = Double | |
type T = Double | |
move :: Monad m => T -> Stiffness -> (Vals, Vels) -> m (Vals,Vels) | |
move timestep stiffness (vals, vels) = do | |
diff <- kern vals | |
let dv = R.map (* stiffness) diff | |
vels' <- computeP $ dv +^ vels | |
let dx = R.map (* timestep) vels' | |
vals' <- computeP $ dx +^ vals | |
return (vals', vels') | |
initialVals :: Vals | |
initialVals = computeS $ fromFunction (Z :. 100 :. 100) init | |
where init (Z :. x :. y) = (exp $ negate $ fromIntegral $ (x-50)^2 + (y-50)^2) | |
initialVels :: Vels | |
initialVels = computeS $ fromFunction (Z :. 100 :. 100) (const 0) | |
toImage :: Vals -> Image PixelF | |
toImage vals = Image w h (convert floatVec) | |
where | |
maximum :: Height | |
maximum = index (foldS max 0 $ foldS max 0 $ vals) Z | |
minimum :: Height | |
minimum = index (foldS min 0 $ foldS min 0 $ vals) Z | |
toDimensionless :: Height -> Float | |
toDimensionless x = double2Float $ (x - minimum) / (maximum - minimum) | |
floatMat :: Array D DIM2 Float | |
floatMat = R.map toDimensionless vals | |
floatVec = toVector $ computeVectorS floatMat | |
(Z :. w :. h) = extent floatMat | |
applyN 0 f x = return x | |
applyN n f x = f x >>= applyN (n-1) f | |
main = do | |
--n <- readLn | |
let step = move 0.001 0.1 | |
let steps = applyN 100 step | |
let save n vals = savePngImage ("out" ++ show n ++ ".png") $ ImageYF $ toImage vals | |
let results = scanl (>>=) (return (initialVals, initialVels)) $ replicate 5 steps | |
let outputs = runIdentity $ sequence results | |
mapM_ (\(n,(vals,vels)) -> save n vals) $ zip [1..] outputs |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment