Skip to content

Instantly share code, notes, and snippets.

@wyager
Created July 2, 2016 22:48
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
Star You must be signed in to star a gist
Save wyager/1c5879660a61c5cd6afaceb3e928d889 to your computer and use it in GitHub Desktop.
{-# 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