diff --git a/gloss-examples.cabal b/gloss-examples.cabal --- a/gloss-examples.cabal +++ b/gloss-examples.cabal @@ -1,5 +1,5 @@ Name: gloss-examples -Version: 1.7.4.2 +Version: 1.7.4.4 License: MIT License-file: LICENSE Author: Ben Lippmeier @@ -253,9 +253,9 @@ base == 4.*, gloss == 1.7.*, vector == 0.9.*, - repa == 3.1.*, - repa-io == 3.1.*, - repa-algorithms == 3.1.*, + repa == 3.2.*, + repa-io == 3.2.*, + repa-algorithms == 3.2.*, ghc-prim Main-is: Main.hs other-modules: @@ -263,7 +263,7 @@ Solve.Density Solve.Velocity Stage.Advection Stage.Boundary Stage.Diffusion Stage.Linear Stage.Project Stage.Sources - hs-source-dirs: raster/Fluid + hs-source-dirs: raster/Fluid/src-repa ghc-options: -Wall -threaded -rtsopts -eventlog -Odph -fno-liberate-case @@ -278,7 +278,7 @@ Build-depends base == 4.*, gloss == 1.7.*, - repa == 3.1.*, + repa == 3.2.*, Main-is: Main.hs hs-source-dirs: raster/Snow ghc-options: @@ -293,7 +293,7 @@ Build-depends base == 4.*, gloss == 1.7.*, - repa >= 3.1.4.2 && < 3.2 + repa >= 3.2.1 && < 3.3 Main-is: Main.hs other-modules: Solver hs-source-dirs: raster/Mandel diff --git a/raster/Fluid/Args.hs b/raster/Fluid/Args.hs deleted file mode 100644 --- a/raster/Fluid/Args.hs +++ /dev/null @@ -1,180 +0,0 @@ - -module Args where -import Config -import Data.Array.Repa as R -import Data.Array.Repa.Algorithms.Pixel as R -import Data.Array.Repa.IO.BMP as R -import System.Console.GetOpt -import Data.IORef -import Prelude as P -import Control.Monad - --- | Command line options. -loadConfig :: [String] -> IO Config -loadConfig args - = do - batchModeArg <- newIORef False - maxStepsArg <- newIORef 0 - widthArg <- newIORef 100 - scaleArg <- newIORef 5 - rateArg <- newIORef 25 - deltaArg <- newIORef 0.1 - diffArg <- newIORef 0 - viscArg <- newIORef 0 - densArg <- newIORef 100 - velArg <- newIORef (20, 20) - - densityBmpArg <- newIORef Nothing - velocityBmpArg <- newIORef Nothing - - let setWidthArg arg = writeIORef widthArg (read arg) - let setScaleArg arg = writeIORef scaleArg (read arg) - let setDeltaArg arg = writeIORef deltaArg (read arg) - let setDiffArg arg = writeIORef diffArg (read arg) - let setViscArg arg = writeIORef viscArg (read arg) - let setDensArg arg = writeIORef densArg (read arg) - let setVelArg arg = let a = read arg in writeIORef velArg (a, a) - let setRate arg = writeIORef rateArg (read arg) - let setMaxSteps arg = writeIORef maxStepsArg (read arg) - let setBatchArg = writeIORef batchModeArg True - let setDensityBMP arg = writeIORef densityBmpArg (Just arg) - let setVelocityBMP arg = writeIORef velocityBmpArg (Just arg) - - let options :: [OptDescr (IO ())] - options - = [Option [] ["batch"] (NoArg setBatchArg ) - "Run a fixed number of steps instead of displaying in a window.", - - Option [] ["max"] (ReqArg setMaxSteps "INT") - "Quit after this number of steps.", - - Option [] ["width"] (ReqArg setWidthArg "INT") - "Size of simulation (100)", - - Option [] ["scale"] (ReqArg setScaleArg "INT") - "Width of cell in window (4)", - - Option [] ["rate"] (ReqArg setRate "INT") - "Frame rate of simulator (25)", - - Option [] ["delta"] (ReqArg setDeltaArg "FLOAT") - "Length of time steps (0.1)", - - Option [] ["diffusion"] (ReqArg setDiffArg "FLOAT") - "Diffusion rate for the density (0)", - - Option [] ["viscosity"] (ReqArg setViscArg "FLOAT") - "Viscosity rate for the velocity (0)", - - Option [] ["user-density"] (ReqArg setDensArg "FLOAT") - "Magnitude of a user inserted density (100)", - - Option [] ["user-velocity"] (ReqArg setVelArg "FLOAT") - "Magnitude of a user inserted velocity (20)", - - Option [] ["bmp-density"] (ReqArg setDensityBMP "FILE.bmp") - "File for initial fluid density", - - Option [] ["bmp-velocity"] (ReqArg setVelocityBMP "FILE.bmp") - "File for initial fluid velocity" - ] - - case getOpt RequireOrder options args of - (actions, [], []) - -> foldl (>>) (return ()) actions - - ( _, nonOpts, []) - -> error $ "unrecognized arguments: " P.++ unwords nonOpts - - ( _, _, msgs) - -> error $ concat msgs P.++ usageInfo "Usage: fluid [OPTION...]" options - - - batchMode <- readIORef batchModeArg - maxSteps <- readIORef maxStepsArg - width <- readIORef widthArg - let height = width - scale <- readIORef scaleArg - rate <- readIORef rateArg - delta <- readIORef deltaArg - diff <- readIORef diffArg - visc <- readIORef viscArg - dens <- readIORef densArg - vel <- readIORef velArg - densityBmp <- readIORef densityBmpArg - velocityBmp <- readIORef velocityBmpArg - - - -- Load the initial desity bmp if we were given one. - initialDensity - <- case densityBmp of - -- No density file given, so just set the field to zero. - Nothing - -> return - $ R.fromListUnboxed (Z :. height :. width) - $ replicate (height * width) 0 - - -- Load density from a .bmp, using the luminance as - -- the scalar density value. - Just filePath - -> do result <- readImageFromBMP filePath - let arr = case result of - Right arr' -> arr' - Left err -> error $ show err - - let Z :. height' :. width' - = extent arr - - when (height /= height' || width /= width') - $ error "Fluid: bmp size does not match --width argument" - - density <- computeUnboxedP - $ R.map floatLuminanceOfRGB8 arr - return density - - - -- Load the initial velocity bmp if we were given one - initialVelocity - <- case velocityBmp of - -- No velocity file given, so just set the field to zero. - Nothing - -> return - $ R.fromListUnboxed (Z :. height :. width) - $ replicate (height * width) (0, 0) - - -- Load - Just filePath - -> do result <- readImageFromBMP filePath - let arr = case result of - Right arr' -> arr' - Left err -> error $ show err - - let Z :. height' :. width' - = extent arr - - when (height /= height' || width /= width') - $ error "Fluid: bmp size does not match --width argument" - - let {-# INLINE conv #-} - conv (r, g, _b) - = let r' = fromIntegral (-128 + fromIntegral r :: Int) - g' = fromIntegral (-128 + fromIntegral g :: Int) - in (r' * 0.0001, g' * 0.0001) - - velocity <- computeUnboxedP $ R.map conv arr - return velocity - - return Config - { configRate = rate - , configWindowSize = (scale * width, scale * width) - , configMaxSteps = maxSteps - , configBatchMode = batchMode - , configModelSize = (width, width) - , configDelta = delta - , configDiffusion = diff - , configViscosity = visc - , configDensity = dens - , configVelocity = vel - , configInitialDensity = initialDensity - , configInitialVelocity = initialVelocity } - diff --git a/raster/Fluid/Config.hs b/raster/Fluid/Config.hs deleted file mode 100644 --- a/raster/Fluid/Config.hs +++ /dev/null @@ -1,54 +0,0 @@ - -module Config - ( Config (..) - , configScale) -where -import Model -import Data.Array.Repa - -data Config - = Config - { -- | Size of window in pixels - configWindowSize :: (Int, Int) - - -- | Simulation rate (frames per second) - , configRate :: Int - - -- | Maximum number of steps in simulation - , configMaxSteps :: Int - - -- | Whether to run in batch-mode, non-interactively. - , configBatchMode :: Bool - - -- | Number of cells in model. - , configModelSize :: (Int, Int) - - -- | Time delta per step. - , configDelta :: !Delta - - -- | Diffusion rate. - , configDiffusion :: !Float - - -- | Fluid viscosity. - , configViscosity :: !Float - - -- | Magnitude of density to add with user interface. - , configDensity :: !Float - - -- | Magnitude of velocity to add with user interface. - , configVelocity :: !(Float, Float) - - -- | Initial density array - , configInitialDensity :: !(Array U DIM2 Float) - - -- | BMP file to use as the initial velocity - , configInitialVelocity :: !(Array U DIM2 (Float, Float)) - } - - -configScale :: Config -> (Float, Float) -configScale config - = let (windowWidth, windowHeight) = configWindowSize config - (modelWidth, modelHeight) = configModelSize config - in ( fromIntegral windowWidth / fromIntegral modelWidth - , fromIntegral windowHeight / fromIntegral modelHeight) diff --git a/raster/Fluid/FieldElt.hs b/raster/Fluid/FieldElt.hs deleted file mode 100644 --- a/raster/Fluid/FieldElt.hs +++ /dev/null @@ -1,95 +0,0 @@ -{-# LANGUAGE FlexibleInstances, BangPatterns #-} - --- | Field operations -module FieldElt where - -class FieldElt a where - zero :: a - - -- | Add all elements. - (~+~) :: a -> a -> a - - -- | Subtract elements from second from first. - (~-~) :: a -> a -> a - - -- | Multiply all elements. - (~*~) :: a -> Float -> a - - -- | Divide first elemends by the second. - (~/~) :: a -> Float -> a - - -- | Negate all mambers. - negate :: a -> a - - -- | Masked addition. - addIf :: Bool -> a -> a -> a - - -- | Masked filter, members with their corresponding flags set - -- to False get zero. - useIf :: Bool -> a -> a - - -instance FieldElt Float where - zero = 0 - {-# INLINE zero #-} - - (~+~) a b = a + b - {-# INLINE (~+~) #-} - - (~-~) a b = a - b - {-# INLINE (~-~) #-} - - (~*~) a b = a * b - {-# INLINE (~*~) #-} - - (~/~) a b = a / b - {-# INLINE (~/~) #-} - - negate a = 0 ~-~ a - {-# INLINE negate #-} - - addIf True a b = a + b - addIf False _ b = b - {-# INLINE addIf #-} - - useIf True a = a - useIf False _ = 0 - {-# INLINE useIf #-} - - -instance FieldElt (Float, Float) where - zero = (0, 0) - {-# INLINE zero #-} - - (~+~) (a1, a2) (b1, b2) = (c1, c2) - where !c1 = a1 + b1 - !c2 = a2 + b2 - {-# INLINE (~+~) #-} - - (~-~) (a1, a2) (b1, b2) = (c1, c2) - where !c1 = a1 - b1 - !c2 = a2 - b2 - {-# INLINE (~-~) #-} - - (~*~) (a1, a2) b = (c1, c2) - where !c1 = a1 * b - !c2 = a2 * b - {-# INLINE (~*~) #-} - - (~/~) (a1, a2) b = (c1, c2) - where !c1 = a1 / b - !c2 = a2 / b - {-# INLINE (~/~) #-} - - addIf True a b = a ~+~ b - addIf False _ b = b - {-# INLINE addIf #-} - - useIf True a = a - useIf False _ = (0, 0) - {-# INLINE useIf #-} - - negate (a1, a2) - = (~-~) (0, 0) (a1, a2) - {-# INLINE negate #-} - diff --git a/raster/Fluid/Main.hs b/raster/Fluid/Main.hs deleted file mode 100644 --- a/raster/Fluid/Main.hs +++ /dev/null @@ -1,89 +0,0 @@ - --- | Stable fluid flow-solver. --- Based on "Real-time Fluid Dynamics for Games", Jos Stam, Game developer conference, 2003. --- Implementation by Ben Lambert-Smith. --- Converted to Repa 3 by Ben Lippmeier. --- -{-# LANGUAGE ScopedTypeVariables #-} -module Main (main) -where -import Solve.Density -import Solve.Velocity -import Model -import UserEvent -import Args -import Config - -import Graphics.Gloss -import Graphics.Gloss.Interface.IO.Game -import System.Mem -import System.Environment (getArgs) -import Data.Array.Repa as R -import Data.Array.Repa.IO.BMP as R - - -main :: IO () -main - = do -- Parse the command-line arguments. - args <- getArgs - config <- loadConfig args - - -- Setup the initial fluid model. - let model = initModel - (configInitialDensity config) - (configInitialVelocity config) - - case configBatchMode config of - False -> runInteractive config model - True -> runBatchMode config model - - --- | Run the simulation interactively. -runInteractive :: Config -> Model -> IO () -runInteractive config model0 - = playIO (InWindow "Stam's stable fluid. Use left-click right-drag to add density / velocity." - (configWindowSize config) - (20, 20)) - black - (configRate config) - model0 - (pictureOfModel (configScale config)) - (\event model -> return $ userEvent config event model) - (\_ -> stepFluid config) - - --- | Run in batch mode and dump a .bmp of the final state. -runBatchMode :: Config -> Model -> IO () -runBatchMode config model - | stepsPassed model > configMaxSteps config - , configMaxSteps config > 0 - = do outputBMP $ densityField model - return () - - | otherwise - = do model' <- stepFluid config model - runBatchMode config model' - - --- Function to step simulator one step forward in time -stepFluid :: Config -> Model -> IO Model -stepFluid config m@(Model df ds vf vs cl sp cb) - | sp > configMaxSteps config - , configMaxSteps config > 0 - = case configBatchMode config of - True -> return m - False -> error "Finished simulation" - - | otherwise - = do performGC - vf' <- velocitySteps config vf vs - df' <- densitySteps config df ds vf' - return $ Model df' Nothing vf' Nothing cl (sp + 1) cb - - --- Writes bitmap data to test batch-mode ran correctly -outputBMP :: DensityField -> IO () -outputBMP df - = do arr <- computeUnboxedP $ R.map pixel8OfDensity df - R.writeImageToBMP "./output.bmp" arr - diff --git a/raster/Fluid/Model.hs b/raster/Fluid/Model.hs deleted file mode 100644 --- a/raster/Fluid/Model.hs +++ /dev/null @@ -1,139 +0,0 @@ -{-# LANGUAGE ScopedTypeVariables, BangPatterns #-} -module Model - ( Delta - , Rate - , Field - , DensityField - , VelocityField - , Source (..) - , CurrentButton (..) - - , Model (..) - , initModel - - , pictureOfModel - , pixel32OfDensity - , pixel8OfDensity) -where -import Graphics.Gloss -import Data.Array.Repa as R -import Data.Array.Repa.Repr.ForeignPtr as R -import Data.Bits -import Data.Word -import Unsafe.Coerce - - --- | Time delta (seconds) -type Delta = Float - --- | Time rate (or period) (1/seconds) -type Rate = Float - - --- | A 2d field. -type Field a - = Array U DIM2 a - --- | Scalar density field. -type DensityField - = Field Float - --- | Vector velocity field. -type VelocityField - = Field (Float, Float) - - --- | Button being pressed by the user. -data CurrentButton - = LeftButton - | RightButton - | None - --- | A density source added by the user. -data Source a - = Source DIM2 a - - --- Model ---------------------------------------------------------------------- --- | The world model. -data Model - = Model - { densityField :: DensityField - , densitySource :: Maybe (Source Float) - , velocityField :: VelocityField - , velocitySource :: Maybe (Source (Float, Float)) - , clickLoc :: Maybe (Int, Int) - , stepsPassed :: Int - , currButton :: CurrentButton - } - - --- | Creates an initial blank model -initModel - :: Array U DIM2 Float - -> Array U DIM2 (Float, Float) - -> Model - -initModel density velocity - | extent density /= extent velocity - = error "Fluid: initModel density and velocity extents do not match" - - | otherwise - = Model - { densityField = density - , densitySource = Nothing - , velocityField = velocity - , velocitySource = Nothing - , clickLoc = Nothing - , stepsPassed = 0 - , currButton = None } -{-# INLINE initModel #-} - - --- Picture -------------------------------------------------------------------- --- | Function to convert the Model into a Bitmap for displaying in Gloss -pictureOfModel :: Monad m => (Float, Float) -> Model -> m Picture -pictureOfModel (scaleX, scaleY) m - = let (Z :. width' :. height') = R.extent $ densityField m - width = fromIntegral width' - height = fromIntegral height' - - in do - (arrDensity :: Array F DIM2 Word32) - <- computeP $ R.map pixel32OfDensity $ densityField m - - return $ Scale scaleX scaleY - $ bitmapOfForeignPtr width height - (R.toForeignPtr $ unsafeCoerce arrDensity) - False -{-# NOINLINE pictureOfModel #-} - - --- | Converts Float value to Word32 for pixel data -pixel32OfDensity :: Float -> Word32 -pixel32OfDensity f - = let !fsat - | f < 0 = 0 - | f >= 1 = 1 - | otherwise = f - - !x = truncate $ fsat * 255 - !a = 255 - - in unsafeShiftL x 24 .|. unsafeShiftL x 16 - .|. unsafeShiftL x 8 .|. a -{-# INLINE pixel32OfDensity #-} - - --- | Converts Float value to a tuple of pixel components. -pixel8OfDensity :: Float -> (Word8, Word8, Word8) -pixel8OfDensity f - = let !fsat - | f < 0 = 0 - | f >= 1 = 1 - | otherwise = f - - !x = truncate $ fsat * 255 - in (x, x, x) -{-# INLINE pixel8OfDensity #-} - diff --git a/raster/Fluid/Solve/Density.hs b/raster/Fluid/Solve/Density.hs deleted file mode 100644 --- a/raster/Fluid/Solve/Density.hs +++ /dev/null @@ -1,24 +0,0 @@ - -module Solve.Density - (densitySteps) -where -import Stage.Diffusion -import Stage.Advection -import Stage.Sources -import Config -import Model - --- | Run the stages for processing the density field in one time step -densitySteps - :: Config - -> DensityField - -> Maybe (Source Float) - -> VelocityField - -> IO DensityField - -densitySteps config df ds vf - = {-# SCC "Solve.densitySteps" #-} - do df1 <- addSources (configDelta config) (configDensity config) ds df - df2 <- diffusion (configDelta config) (configDiffusion config) df1 - df' <- advection (configDelta config) vf df2 - return df' diff --git a/raster/Fluid/Solve/Velocity.hs b/raster/Fluid/Solve/Velocity.hs deleted file mode 100644 --- a/raster/Fluid/Solve/Velocity.hs +++ /dev/null @@ -1,33 +0,0 @@ - -module Solve.Velocity - (velocitySteps) -where -import Stage.Diffusion -import Stage.Advection -import Stage.Sources -import Stage.Project -import Model -import Config - --- The pass that sets boundary conditions is buggy and --- currently disabled. --- import Stage.Boundary - -velocitySteps - :: Config - -> VelocityField - -> Maybe (Source (Float, Float)) - -> IO VelocityField - -velocitySteps config vf vs - = {-# SCC "Solve.velocitySteps" #-} - do vf1 <- addSources (configDelta config) (configVelocity config) vs vf - vf2 <- diffusion (configDelta config) (configViscosity config) vf1 --- vf3 <- setBoundary vf2 - vf4 <- project vf2 --- vf5 <- setBoundary vf4 - vf6 <- advection (configDelta config) vf4 vf4 --- vf7 <- setBoundary vf6 - vf8 <- project vf6 --- vf' <- setBoundary vf8 - return vf8 diff --git a/raster/Fluid/Stage/Advection.hs b/raster/Fluid/Stage/Advection.hs deleted file mode 100644 --- a/raster/Fluid/Stage/Advection.hs +++ /dev/null @@ -1,103 +0,0 @@ -{-# LANGUAGE BangPatterns #-} -module Stage.Advection - (advection) -where -import Model -import FieldElt -import Data.Array.Repa as R -import Data.Array.Repa.Unsafe as R -import Data.Vector.Unboxed (Unbox) -import Debug.Trace - - --- | Apply a velocity field to another field. --- Both fields must have the same extent. -advection - :: (FieldElt a, Unbox a) - => Delta - -> VelocityField - -> Field a - -> IO (Field a) - -advection !delta velField field - = {-# SCC "advection" #-} - velField `deepSeqArray` field `deepSeqArray` - do traceEventIO "Fluid: advection" - computeP $ unsafeTraverse field id (advectElem delta velField) - -{-# SPECIALIZE advection - :: Delta - -> VelocityField -> Field Float - -> IO (Field Float) #-} - -{-# SPECIALIZE advection - :: Delta - -> VelocityField -> Field (Float, Float) - -> IO (Field (Float, Float)) #-} - - --- | Compute the new field value at the given location. -advectElem - :: (FieldElt a, Unbox a) - => Delta -- ^ Time delta (in seconds) - -> VelocityField -- ^ Velocity field that moves the source field. - -> (DIM2 -> a) -- ^ Get an element from the source field. - -> DIM2 -- ^ Compute the new value at this index. - -> a - -advectElem !delta !velField !get !pos@(Z:. j :. i) - = velField `deepSeqArray` - (((d00 ~*~ t0) ~+~ (d01 ~*~ t1)) ~*~ s0) - ~+~ (((d10 ~*~ t0) ~+~ (d11 ~*~ t1)) ~*~ s1) - where - _ :. _ :. width' = R.extent velField - !width = fromIntegral width' - - -- helper values - !dt0 = delta * width - !(u, v) = velField `unsafeIndex` pos - - -- backtrack densities to point based on velocity field - -- and make sure they are in field - !x = checkLocation width $ fromIntegral i - dt0 * u - !y = checkLocation width $ fromIntegral j - dt0 * v - - -- calculate discrete locations surrounding point - !i0 = truncate x - !i1 = i0 + 1 - - !j0 = truncate y - !j1 = j0 + 1 - - -- calculate ratio point is between the discrete locations - !s1 = x - fromIntegral i0 - !s0 = 1 - s1 - - !t1 = y - fromIntegral j0 - !t0 = 1 - t1 - - -- grab values from grid surrounding advected point - !d00 = get (Z:. j0 :. i0) - !d01 = get (Z:. j1 :. i0) - !d10 = get (Z:. j0 :. i1) - !d11 = get (Z:. j1 :. i1) - - -{-# SPECIALIZE advectElem - :: Delta - -> VelocityField -> (DIM2 -> Float) - -> DIM2 -> Float #-} - -{-# SPECIALIZE advectElem - :: Delta - -> VelocityField -> (DIM2 -> (Float,Float)) - -> DIM2 -> (Float,Float) #-} - - --- | Wrap an index back into the simulation area if it is outside. -checkLocation :: Float -> Float -> Float -checkLocation !width !x - | x < 0.5 = 0.5 - | x > width - 1.5 = width - 1.5 - | otherwise = x -{-# INLINE checkLocation #-} diff --git a/raster/Fluid/Stage/Boundary.hs b/raster/Fluid/Stage/Boundary.hs deleted file mode 100644 --- a/raster/Fluid/Stage/Boundary.hs +++ /dev/null @@ -1,116 +0,0 @@ - -{-# LANGUAGE BangPatterns #-} -module Stage.Boundary - (setBoundary) -where -import Model -import Data.Array.Repa as R -import Control.Monad -import Debug.Trace -import Config -import FieldElt - --- | Apply boundary conditions to a velocity field. -setBoundary :: Config -> VelocityField -> IO VelocityField -setBoundary config f - = let (width, _) = configModelSize config - in (rebuild width f <=< setBoundary' width <=< grabBorders width) f - - --- | Takes the original VelocityField and the array of edges and replaces --- edge values with new values -rebuild :: Int -> VelocityField -> VelocityField -> IO VelocityField -rebuild width field edges - = field `deepSeqArray` edges `deepSeqArray` - do traceEventIO "Fluid: rebuild" - computeUnboxedP $ backpermuteDft field (rebuildPosMap width) edges -{-# INLINE rebuild #-} - - -rebuildPosMap :: Int -> DIM2 -> Maybe DIM2 -rebuildPosMap !width (Z:.j:.i) - | j == 0 - = Just (Z:.0:.i) - - | j == width - 1 - = Just (Z:.1:.i) - - | i == 0 - = if j == 0 then Just (Z:. 0 :. 0) - else if j == end then Just (Z:. 1 :. 0) - else Just (Z:. 2 :. j) - - | i == width - 1 - = if j == 0 then Just (Z:. 0 :. (width-1)) - else if j == end then Just (Z:. 1 :. (width-1)) - else Just (Z:. 3 :. j) - - | otherwise = Nothing - where end = width - 1 -{-# INLINE rebuildPosMap #-} - - - --- | Grabs the border elements of the VelocityField and outputs them as --- one array, for ease of adding back into the original VelocityField later -grabBorders :: Int -> VelocityField -> IO VelocityField -grabBorders width f - = f `deepSeqArray` - do traceEventIO "Fluid: grabBorders" - computeUnboxedP $ backpermute (Z:. 4 :. width) (edgeCases width) f -{-# INLINE grabBorders #-} - - --- | Map a position in the edges array to what they were in the original --- array. -edgeCases :: Int -> DIM2 -> DIM2 -edgeCases width (Z:.j:.i) - | j == 0 = (Z:.0 :.i) - | j == 1 = (Z:.(width-1) :.i) - | j == 2 = (Z:.i :.0) - | j == 3 = (Z:.i :.(width-1)) - | otherwise = error "Incorrect coordinate given in setBoundary" -{-# INLINE edgeCases #-} - - - -setBoundary' :: Int -> VelocityField -> IO VelocityField -setBoundary' width e - = e `deepSeqArray` - do traceEventIO "Fluid: setBoundary'" - computeUnboxedP $ traverse e id (revBoundary width) -{-# INLINE setBoundary' #-} - - --- | Based on position in edges array set the velocity accordingly -revBoundary :: Int -> (DIM2 -> (Float,Float)) -> DIM2 -> (Float,Float) -revBoundary width loc pos@(Z:.j:.i) - | j == 0 - = if i == 0 then grabCornerCase loc (Z:.2:.1) (Z:.0:.1) - else if i == end then grabCornerCase loc (Z:.0:.(width-2)) (Z:.3:.1) - else (-p1,p2) - | j == 1 - = if i == 0 then grabCornerCase loc (Z:.2:.(width-2)) (Z:.1:.1) - else if i == end then grabCornerCase loc (Z:.1:.(width-2)) (Z:.3:.(width-2)) - else (-p1,p2) - - | j == 2 = (p1,-p2) - | j == 3 = (p1,-p2) - - | otherwise = error "Fluid: revBoundary" - where (p1,p2) = loc pos - end = width - 1 - - -{-# INLINE revBoundary #-} - --- | Corner cases are special and are calculated with this function -grabCornerCase :: (DIM2 -> (Float, Float)) -> DIM2 -> DIM2 -> (Float, Float) -grabCornerCase loc pos1 pos2 - = (p1 * q1, p2 * q2) ~*~ 0.5 - where (p1,p2) = loc pos1 - (q1,q2) = loc pos2 -{-# INLINE grabCornerCase #-} - - - diff --git a/raster/Fluid/Stage/Diffusion.hs b/raster/Fluid/Stage/Diffusion.hs deleted file mode 100644 --- a/raster/Fluid/Stage/Diffusion.hs +++ /dev/null @@ -1,40 +0,0 @@ -{-# LANGUAGE BangPatterns #-} -module Stage.Diffusion - (diffusion) -where -import Model -import FieldElt -import Stage.Linear -import Data.Array.Repa as R -import Data.Array.Repa.Eval as R -import Data.Vector.Unboxed -import Debug.Trace - - --- | Diffuse a field at a certain rate. -diffusion - :: (FieldElt a, Num a, Elt a, Unbox a) - => Delta -> Rate - -> Field a - -> IO (Field a) -diffusion !delta !rate field - = {-# SCC diffusion #-} - field `deepSeqArray` - let _ :. _ :. width' = R.extent field - !width = fromIntegral width' - !a = delta * rate * width * width - !c = 1 + 4 * a - !repeats = 20 - in do - traceEventIO "Fluid: diffusion" - linearSolver field field a c repeats - -{-# SPECIALIZE diffusion - :: Delta -> Rate - -> Field Float - -> IO (Field Float) #-} - -{-# SPECIALIZE diffusion - :: Delta -> Rate - -> Field (Float, Float) - -> IO (Field (Float, Float)) #-} diff --git a/raster/Fluid/Stage/Linear.hs b/raster/Fluid/Stage/Linear.hs deleted file mode 100644 --- a/raster/Fluid/Stage/Linear.hs +++ /dev/null @@ -1,88 +0,0 @@ -{-# LANGUAGE FlexibleContexts, BangPatterns #-} -module Stage.Linear - (linearSolver) -where -import Model -import FieldElt as E -import Data.Array.Repa as R -import Data.Array.Repa.Stencil as R -import Data.Array.Repa.Stencil.Dim2 as R -import Data.Array.Repa.Eval as R -import Data.Vector.Unboxed -import Debug.Trace - - -linearSolver - :: (FieldElt a, Repr U a, Unbox a, Elt a, Num a) - => Field a -- ^ Original field. - -> Field a -- ^ Current field. - -> Float - -> Float - -> Int -- ^ Number of iterations to apply. - -> IO (Field a) - -linearSolver origField curField !a !c !i - -- If nothing would change by running the solver, then skip it. - | 0 <- a = return origField - - -- The solver has finished its loop - | 0 <- i = return curField - - -- Do one iteration - | otherwise - = origField `deepSeqArray` curField `deepSeqArray` - do let !c' = 1/c - let {-# INLINE zipFunc #-} - zipFunc !orig !new - = new ~+~ (orig ~*~ c') - - traceEventIO "Fluid: linear solver mapStencil" - newField <- {-# SCC "linearSolver.mapStencil" #-} - computeUnboxedP - $ R.czipWith zipFunc origField - $ mapStencil2 (BoundConst E.zero) (linearSolverStencil a c) curField - - linearSolver origField newField a c (i - 1) - -{-# SPECIALIZE linearSolver - :: Field Float - -> Field Float - -> Float -> Float -> Int - -> IO (Field Float) #-} - -{-# SPECIALIZE linearSolver - :: Field (Float, Float) - -> Field (Float, Float) - -> Float -> Float -> Int - -> IO (Field (Float, Float)) #-} - - - --- | Stencil function for the linear solver. -linearSolverStencil - :: FieldElt a - => Float -> Float -> Stencil DIM2 a - -linearSolverStencil a c - = StencilStatic (Z:.3:.3) E.zero - (\ix val acc -> - case linearSolverCoeffs a c ix of - Nothing -> acc - Just coeff -> acc ~+~ (val ~*~ coeff)) -{-# INLINE linearSolverStencil #-} - - --- | Linear solver stencil kernel. -linearSolverCoeffs - :: Float -> Float -> DIM2 -> Maybe Float - -linearSolverCoeffs a c (Z:.j:.i) - | i == 1, j == 0 = Just (a/c) - | i == -1, j == 0 = Just (a/c) - | i == 0, j == 1 = Just (a/c) - | i == 0, j == -1 = Just (a/c) - | otherwise = Nothing -{-# INLINE linearSolverCoeffs #-} - - - diff --git a/raster/Fluid/Stage/Project.hs b/raster/Fluid/Stage/Project.hs deleted file mode 100644 --- a/raster/Fluid/Stage/Project.hs +++ /dev/null @@ -1,68 +0,0 @@ -{-# LANGUAGE BangPatterns #-} -module Stage.Project - (project) -where -import Model -import FieldElt -import Stage.Linear -import Data.Array.Repa as R -import Data.Array.Repa.Unsafe as R -import Debug.Trace - - -project :: Field (Float, Float) -> IO (Field (Float, Float)) -project field - = {-# SCC project #-} - field `deepSeqArray` - do let _ :. _ :. width = extent field - let !repeats = 20 - - traceEventIO "Fluid: project gradient" - divergence <- {-# SCC "project.genDiv" #-} - computeUnboxedP - $ fromFunction (Z:. width :. width) (genDivergence width field) - - traceEventIO "Fluid: project linear solver" - p <- {-# SCC "project.linearSolver" #-} - linearSolver divergence divergence 1 4 repeats - - traceEventIO "Fluid: project apply" - f' <- {-# SCC "project.apply" #-} - computeUnboxedP - $ unsafeTraverse field id (projectElem width p) - - return f' -{-# NOINLINE project #-} - - --- | Subtract a gradient field from the regular field to --- create a mass-conserving field. -projectElem - :: Int -- ^ Width of model. - -> Field Float - -> (DIM2 -> (Float, Float)) -- ^ Get data from the regular field. - -> DIM2 -- ^ Compute the value at this point. - -> (Float, Float) - -projectElem !width !p !get !pos@(Z:.j:.i) - = get pos ~-~ (0.5 * width' * (p0 - p1), - 0.5 * width' * (p2 - p3)) - where - !width' = fromIntegral width - !p0 = useIf (i < width - 1) (p `unsafeIndex` (Z :. j :. i+1)) - !p1 = useIf (i > 0) (p `unsafeIndex` (Z :. j :. i-1)) - !p2 = useIf (j < width - 1) (p `unsafeIndex` (Z :. j+1 :. i )) - !p3 = useIf (j > 0) (p `unsafeIndex` (Z :. j-1 :. i )) -{-# INLINE projectElem #-} - - --- | Get an approximation of the gradient at this point. -genDivergence :: Int -> VelocityField -> DIM2 -> Float -genDivergence !width !f (Z :. j :. i) - = (-0.5 * ((u0 - u1) + (v0 - v1))) / fromIntegral width - where - (u0, _) = useIf (i < width - 1) (f `unsafeIndex` (Z:. j :. i+1)) - (u1, _) = useIf (i > 0) (f `unsafeIndex` (Z:. j :. i-1)) - ( _, v0) = useIf (j < width - 1) (f `unsafeIndex` (Z:. j+1 :. i )) - ( _, v1) = useIf (j > 0) (f `unsafeIndex` (Z:. j-1 :. i )) -{-# INLINE genDivergence #-} diff --git a/raster/Fluid/Stage/Sources.hs b/raster/Fluid/Stage/Sources.hs deleted file mode 100644 --- a/raster/Fluid/Stage/Sources.hs +++ /dev/null @@ -1,79 +0,0 @@ -{-# LANGUAGE BangPatterns, FlexibleInstances #-} -module Stage.Sources - (addSources) -where -import Model -import FieldElt -import Data.Array.Repa as R -import Data.Array.Repa.Unsafe as R -import Data.Vector.Unboxed (Unbox) -import Debug.Trace - - --- | Addition of forces stage for simulation -addSources - :: (FieldElt a, FieldSource a, Unbox a) - => Delta -- ^ Time delta. - -> a -- ^ Value to insert. - -> Maybe (Source a) - -> Field a - -> IO (Field a) - -addSources !delta !value (Just (Source aim mul)) field - = {-# SCC addSources #-} - field `deepSeqArray` - do traceEventIO "Fluid: addSources" - computeP $ unsafeTraverse field id (insertSource delta value aim mul) - -addSources _ _ Nothing field - = return field - - -insertSource - :: (FieldElt a, FieldSource a) - => Delta - -> a -- ^ Value to insert - -> DIM2 -> a - -> (DIM2 -> a) - -> DIM2 - -> a - -insertSource !delta !value !aim !mul locate !pos - | aim == pos = addSource delta value (locate pos) mul - | otherwise = locate pos -{-# INLINE insertSource #-} - - -{-# SPECIALIZE addSources - :: Delta - -> Float - -> Maybe (Source Float) - -> Field Float - -> IO (Field Float) #-} - -{-# SPECIALIZE addSources - :: Delta - -> (Float, Float) - -> Maybe (Source (Float, Float)) - -> Field (Float, Float) - -> IO (Field (Float, Float)) #-} - - --- FieldSource ---------------------------------------------------------------- -class FieldSource a where - addSource :: Delta -> a -> a -> a -> a - -instance FieldSource Float where - addSource !delta !value !a !mul - = a ~+~ (value * delta * mul) - {-# INLINE addSource #-} - -instance FieldSource (Float, Float) where - addSource !delta (newA, newB) (a,b) (mulA, mulB) - = ( a + (newA * delta * (-mulA)) - , b + (newB * delta * (-mulB))) - {-# INLINE addSource #-} - - - - diff --git a/raster/Fluid/UserEvent.hs b/raster/Fluid/UserEvent.hs deleted file mode 100644 --- a/raster/Fluid/UserEvent.hs +++ /dev/null @@ -1,136 +0,0 @@ -module UserEvent - (userEvent) -where -import Config -import Model as M -import Data.Array.Repa as R -import Graphics.Gloss.Interface.Pure.Game as G - - --- | Handle user events for the Gloss `playIO` wrapper. -userEvent :: Config -> Event -> Model -> Model -userEvent config - (EventKey key keyState mods (x, y)) - (Model df ds vf vs cl sp _cb) - - -- Add velocity --------------------------------------------- - | MouseButton G.RightButton <- key - , Down <- keyState - , (x',y') <- windowToModel config (x,y) - = Model { densityField = df - , densitySource = ds - , velocityField = vf - , velocitySource = vs - , clickLoc = Just (x',y') - , stepsPassed = sp - , currButton = M.RightButton - } - - -- Accept shift-leftbutton for people with trackpads - | MouseButton G.LeftButton <- key - , Down <- keyState - , Down <- shift mods - , (x',y') <- windowToModel config (x,y) - = Model { densityField = df - , densitySource = ds - , velocityField = vf - , velocitySource = vs - , clickLoc = Just (x',y') - , stepsPassed = sp - , currButton = M.RightButton - } - - | MouseButton G.RightButton <- key - , Up <- keyState - , Just (locX, locY) <- cl - , (x',y') <- windowToModel config (x,y) - = Model { densityField = df - , densitySource = ds - , velocityField = vf - , velocitySource = Just (Source (Z:.locY:.locX) - (fromIntegral (locX-x'),fromIntegral (locY-y'))) - , clickLoc = Nothing - , stepsPassed = sp - , currButton = M.None - } - - - -- Add density ---------------------------------------------- - | MouseButton G.LeftButton <- key - , Down <- keyState - , (x',y') <- windowToModel config (x, y) - = Model { densityField = df - , densitySource = Just (Source (Z:.y':.x') 1) - , velocityField = vf - , velocitySource = vs - , clickLoc = cl - , stepsPassed = sp - , currButton = M.LeftButton - } - - | MouseButton G.LeftButton <- key - , Up <- keyState - = Model { densityField = df - , densitySource = ds - , velocityField = vf - , velocitySource = vs - , clickLoc = cl - , stepsPassed = sp - , currButton = M.None - } - - - -- Reset model ---------------------------------------------- - | Char 'r' <- key - , Down <- keyState - = initModel (configInitialDensity config) - (configInitialVelocity config) - - -- Quit program --------------------------------------------- - | Char 'q' <- key - , Down <- keyState - = error "Quitting" - - -userEvent config (EventMotion (x, y)) - (Model df _ds vf vs cl sp M.LeftButton) - | (x',y') <- windowToModel config (x, y) - = Model { densityField = df - , densitySource = Just (Source (Z:.y':.x') 1) - , velocityField = vf - , velocitySource = vs - , clickLoc = cl - , stepsPassed = sp - , currButton = M.LeftButton - } - -userEvent config (EventMotion (x, y)) - (Model df ds vf _vs (Just (clx, cly)) sp M.RightButton) - | (x', y') <- windowToModel config (x,y) - = Model { densityField = df - , densitySource = ds - , velocityField = vf - , velocitySource = Just (Source (Z:.y':.x') - (fromIntegral (clx-x'), fromIntegral (cly-y'))) - , clickLoc = Just (x',y') - , stepsPassed = sp - , currButton = M.RightButton - } - -userEvent _ _ m = m - --- Converts a window location to the corresponding location in the --- simulation. -windowToModel :: Config -> (Float, Float) -> (Int, Int) -windowToModel config (x, y) = (x', y') - where (scaleX, scaleY) = configScale config - (windowWidth, windowHeight) = configWindowSize config - - x' = round ((x + (fromIntegral windowWidth / 2)) / scaleX) - y' = round ((y + (fromIntegral windowHeight / 2)) / scaleY) - - - - - - diff --git a/raster/Fluid/src-repa/Args.hs b/raster/Fluid/src-repa/Args.hs new file mode 100644 --- /dev/null +++ b/raster/Fluid/src-repa/Args.hs @@ -0,0 +1,230 @@ + +module Args where +import Config +import Data.Array.Repa as R +import Data.Array.Repa.Algorithms.Pixel as R +import Data.Array.Repa.IO.BMP as R +import System.Console.GetOpt +import Data.IORef +import Prelude as P +import Control.Monad + + +-- | Command line options. +loadConfig :: [String] -> IO Config +loadConfig args + = do + batchModeArg <- newIORef False + benchModeArg <- newIORef False + maxStepsArg <- newIORef 0 + widthArg <- newIORef 100 + itersArg <- newIORef 40 + scaleArg <- newIORef 5 + rateArg <- newIORef 25 + deltaArg <- newIORef 0.1 + diffArg <- newIORef 0.00001 + viscArg <- newIORef 0 + densArg <- newIORef 100 + velArg <- newIORef (20, 20) + + densityBmpArg <- newIORef Nothing + velocityBmpArg <- newIORef Nothing + + let setWidthArg arg = writeIORef widthArg (read arg) + let setScaleArg arg = writeIORef scaleArg (read arg) + let setItersArg arg = writeIORef itersArg (read arg) + let setDeltaArg arg = writeIORef deltaArg (read arg) + let setDiffArg arg = writeIORef diffArg (read arg) + let setViscArg arg = writeIORef viscArg (read arg) + let setDensArg arg = writeIORef densArg (read arg) + let setVelArg arg = let a = read arg in writeIORef velArg (a, a) + let setRate arg = writeIORef rateArg (read arg) + let setMaxSteps arg = writeIORef maxStepsArg (read arg) + let setBatchArg = writeIORef batchModeArg True + let setBenchArg = writeIORef benchModeArg True + let setDensityBMP arg = writeIORef densityBmpArg (Just arg) + let setVelocityBMP arg = writeIORef velocityBmpArg (Just arg) + + let options :: [OptDescr (IO ())] + options + = [Option [] ["batch"] (NoArg setBatchArg ) + "Run a fixed number of steps instead of displaying in a window.", + + Option [] ["bench"] (NoArg setBenchArg ) + "Set standard initial conditions for benchmarking.", + + Option [] ["max"] (ReqArg setMaxSteps "INT") + "Quit after this number of steps.", + + Option [] ["width"] (ReqArg setWidthArg "INT") + "Size of simulation (100)", + + Option [] ["iters"] (ReqArg setItersArg "INT") + "Iterations for the linear solver (20)", + + Option [] ["scale"] (ReqArg setScaleArg "INT") + "Width of cell in window (5)", + + Option [] ["rate"] (ReqArg setRate "INT") + "Frame rate of simulator (25)", + + Option [] ["delta"] (ReqArg setDeltaArg "FLOAT") + "Length of time steps (0.1)", + + Option [] ["diffusion"] (ReqArg setDiffArg "FLOAT") + "Diffusion rate for the density (0)", + + Option [] ["viscosity"] (ReqArg setViscArg "FLOAT") + "Viscosity rate for the velocity (0)", + + Option [] ["user-density"] (ReqArg setDensArg "FLOAT") + "Magnitude of a user inserted density (100)", + + Option [] ["user-velocity"] (ReqArg setVelArg "FLOAT") + "Magnitude of a user inserted velocity (20)", + + Option [] ["bmp-density"] (ReqArg setDensityBMP "FILE.bmp") + "File for initial fluid density", + + Option [] ["bmp-velocity"] (ReqArg setVelocityBMP "FILE.bmp") + "File for initial fluid velocity" + ] + + case getOpt RequireOrder options args of + (actions, [], []) + -> foldl (>>) (return ()) actions + + ( _, nonOpts, []) + -> error $ "unrecognized arguments: " P.++ unwords nonOpts + + ( _, _, msgs) + -> error $ concat msgs P.++ usageInfo "Usage: fluid [OPTION...]" options + + + batchMode <- readIORef batchModeArg + benchMode <- readIORef benchModeArg + maxSteps <- readIORef maxStepsArg + width <- readIORef widthArg + let height = width + iters <- readIORef itersArg + scale <- readIORef scaleArg + rate <- readIORef rateArg + delta <- readIORef deltaArg + diff <- readIORef diffArg + visc <- readIORef viscArg + dens <- readIORef densArg + vel <- readIORef velArg + densityBmp <- readIORef densityBmpArg + velocityBmp <- readIORef velocityBmpArg + + + -- Load the initial desity bmp if we were given one. + let mkInitialDensity + -- Load density from a .bmp, using the luminance as + -- the scalar density value. + | Just filePath <- densityBmp + = do result <- readImageFromBMP filePath + let arr = case result of + Right arr' -> arr' + Left err -> error $ show err + + let Z :. height' :. width' + = extent arr + + when (height /= height' || width /= width') + $ error "Fluid: bmp size does not match --width argument" + + density <- computeUnboxedP + $ R.map floatLuminanceOfRGB8 arr + return density + + | benchMode + = let width' = fromIntegral width + yc = fromIntegral (width `div` 2) + xc = fromIntegral (width `div` 2) + + in return + $ R.fromListUnboxed (Z :. height :. width) + $ [ let x' = fromIntegral (x - 1) + y' = fromIntegral (y - 1) + xk1 = cos (10 * (x' - xc) / width') + yk1 = cos (10 * (y' - yc) / width') + d1 = xk1 * yk1 + in if (d1 < 0) then 0 else d1 + | y <- [1..width] + , x <- [1..width] ] + + -- No density file given, so just set the field to zero. + | otherwise + = return + $ R.fromListUnboxed (Z :. height :. width) + $ replicate (height * width) 0 + + + initialDensity <- mkInitialDensity + + + -- Load the initial velocity bmp if we were given one + let mkInitialVelocity + -- Load + | Just filePath <- velocityBmp + = do result <- readImageFromBMP filePath + let arr = case result of + Right arr' -> arr' + Left err -> error $ show err + + let Z :. height' :. width' + = extent arr + + when (height /= height' || width /= width') + $ error "Fluid: bmp size does not match --width argument" + + let {-# INLINE conv #-} + conv (r, g, _b) + = let r' = fromIntegral (-128 + fromIntegral r :: Int) + g' = fromIntegral (-128 + fromIntegral g :: Int) + in (r' * 0.0001, g' * 0.0001) + + velocity <- computeUnboxedP $ R.map conv arr + return velocity + + | benchMode + = let width' = fromIntegral width + yc = fromIntegral (width `div` 2) + xc = fromIntegral (width `div` 2) + + in return + $ R.fromListUnboxed (Z :. height :. width) + $ [ let x' = fromIntegral x + y' = fromIntegral y + xk2 = cos (15 * (x' - xc) / width') + yk2 = cos (15 * (y' - yc) / width') + d2 = xk2 * yk2 / 5 + in (0, d2) + | y <- [0..width-1] + , x <- [0..width-1] ] + + -- No velocity file given, so just set the field to zero. + | otherwise + = return + $ R.fromListUnboxed (Z :. height :. width) + $ replicate (height * width) (0, 0) + + initialVelocity <- mkInitialVelocity + + + return Config + { configRate = rate + , configWindowSize = (scale * width, scale * width) + , configMaxSteps = maxSteps + , configBatchMode = batchMode + , configModelSize = (width, width) + , configIters = iters + , configDelta = delta + , configDiffusion = diff + , configViscosity = visc + , configDensity = dens + , configVelocity = vel + , configInitialDensity = initialDensity + , configInitialVelocity = initialVelocity } + diff --git a/raster/Fluid/src-repa/Config.hs b/raster/Fluid/src-repa/Config.hs new file mode 100644 --- /dev/null +++ b/raster/Fluid/src-repa/Config.hs @@ -0,0 +1,57 @@ + +module Config + ( Config (..) + , configScale) +where +import Model +import Data.Array.Repa + +data Config + = Config + { -- | Size of window in pixels + configWindowSize :: (Int, Int) + + -- | Simulation rate (frames per second) + , configRate :: Int + + -- | Maximum number of steps in simulation + , configMaxSteps :: Int + + -- | Whether to run in batch-mode, non-interactively. + , configBatchMode :: Bool + + -- | Number of cells in model. + , configModelSize :: (Int, Int) + + -- | Number of iterations to use in the linear solver + , configIters :: !Int + + -- | Time delta per step. + , configDelta :: !Delta + + -- | Diffusion rate. + , configDiffusion :: !Float + + -- | Fluid viscosity. + , configViscosity :: !Float + + -- | Magnitude of density to add with user interface. + , configDensity :: !Float + + -- | Magnitude of velocity to add with user interface. + , configVelocity :: !(Float, Float) + + -- | Initial density array + , configInitialDensity :: !(Array U DIM2 Float) + + -- | BMP file to use as the initial velocity + , configInitialVelocity :: !(Array U DIM2 (Float, Float)) + } + + +configScale :: Config -> (Float, Float) +configScale config + = let (windowWidth, windowHeight) = configWindowSize config + (modelWidth, modelHeight) = configModelSize config + in ( fromIntegral windowWidth / fromIntegral modelWidth + , fromIntegral windowHeight / fromIntegral modelHeight) diff --git a/raster/Fluid/src-repa/FieldElt.hs b/raster/Fluid/src-repa/FieldElt.hs new file mode 100644 --- /dev/null +++ b/raster/Fluid/src-repa/FieldElt.hs @@ -0,0 +1,95 @@ +{-# LANGUAGE FlexibleInstances, BangPatterns #-} + +-- | Field operations +module FieldElt where + +class Show a => FieldElt a where + zero :: a + + -- | Add all elements. + (~+~) :: a -> a -> a + + -- | Subtract elements from second from first. + (~-~) :: a -> a -> a + + -- | Multiply all elements. + (~*~) :: a -> Float -> a + + -- | Divide first elemends by the second. + (~/~) :: a -> Float -> a + + -- | Negate all mambers. + negate :: a -> a + + -- | Masked addition. + addIf :: Bool -> a -> a -> a + + -- | Masked filter, members with their corresponding flags set + -- to False get zero. + useIf :: Bool -> a -> a + + +instance FieldElt Float where + zero = 0 + {-# INLINE zero #-} + + (~+~) a b = a + b + {-# INLINE (~+~) #-} + + (~-~) a b = a - b + {-# INLINE (~-~) #-} + + (~*~) a b = a * b + {-# INLINE (~*~) #-} + + (~/~) a b = a / b + {-# INLINE (~/~) #-} + + negate a = 0 ~-~ a + {-# INLINE negate #-} + + addIf True a b = a + b + addIf False _ b = b + {-# INLINE addIf #-} + + useIf True a = a + useIf False _ = 0 + {-# INLINE useIf #-} + + +instance FieldElt (Float, Float) where + zero = (0, 0) + {-# INLINE zero #-} + + (~+~) (a1, a2) (b1, b2) = (c1, c2) + where !c1 = a1 + b1 + !c2 = a2 + b2 + {-# INLINE (~+~) #-} + + (~-~) (a1, a2) (b1, b2) = (c1, c2) + where !c1 = a1 - b1 + !c2 = a2 - b2 + {-# INLINE (~-~) #-} + + (~*~) (a1, a2) b = (c1, c2) + where !c1 = a1 * b + !c2 = a2 * b + {-# INLINE (~*~) #-} + + (~/~) (a1, a2) b = (c1, c2) + where !c1 = a1 / b + !c2 = a2 / b + {-# INLINE (~/~) #-} + + addIf True a b = a ~+~ b + addIf False _ b = b + {-# INLINE addIf #-} + + useIf True a = a + useIf False _ = (0, 0) + {-# INLINE useIf #-} + + negate (a1, a2) + = (~-~) (0, 0) (a1, a2) + {-# INLINE negate #-} + diff --git a/raster/Fluid/src-repa/Main.hs b/raster/Fluid/src-repa/Main.hs new file mode 100644 --- /dev/null +++ b/raster/Fluid/src-repa/Main.hs @@ -0,0 +1,93 @@ + +-- | Stable fluid flow-solver. +-- Based on "Real-time Fluid Dynamics for Games", Jos Stam, Game developer conference, 2003. +-- Implementation by Ben Lambert-Smith. +-- Converted to Repa 3 by Ben Lippmeier. +-- +{-# LANGUAGE ScopedTypeVariables #-} +module Main (main) +where +import Solve.Density +import Solve.Velocity +import Model +import UserEvent +import Args +import Config + +import Graphics.Gloss +import Graphics.Gloss.Interface.IO.Game +import System.Mem +import System.Environment (getArgs) +import Data.Array.Repa.IO.Timing as R +import Prelude as P + + +main :: IO () +main + = do -- Parse the command-line arguments. + args <- getArgs + config <- loadConfig args + + -- Setup the initial fluid model. + let model = initModel + (configInitialDensity config) + (configInitialVelocity config) + + performGC + case configBatchMode config of + False -> runInteractive config model + True + -> do (_, elapsed) <- time $ do result <- runBatchMode config model + result `seq` return () + putStrLn $ show $ wallTime milliseconds elapsed + + +-- | Run the simulation interactively. +runInteractive :: Config -> Model -> IO () +runInteractive config model0 + = playIO (InWindow "Stam's stable fluid. Use left-click right-drag to add density / velocity." + (configWindowSize config) + (20, 20)) + black + (configRate config) + model0 + (pictureOfModel (configScale config)) + (\event model -> return $ userEvent config event model) + (\_ -> stepFluid config) + + +-- | Run in batch mode and dump a .bmp of the final state. +runBatchMode :: Config -> Model -> IO () +runBatchMode config model + | stepsPassed model >= configMaxSteps config + = do -- outputBMP $ densityField model + -- outputPPM (stepsPassed model) "density" 1 (densityField model) + -- outputPPM (stepsPassed model) "velctyU" 10 (R.computeS $ R.map fst $ velocityField model) + -- outputPPM (stepsPassed model) "velctyV" 10 (R.computeS $ R.map snd $ velocityField model) + return () + + | otherwise + = do model' <- stepFluid config model + runBatchMode config model' + + +-- Function to step simulator one step forward in time +stepFluid :: Config -> Model -> IO Model +stepFluid config m@(Model df ds vf vs cl step cb) + | step >= configMaxSteps config + , configMaxSteps config > 0 + = case configBatchMode config of + True -> return m + False -> error "Finished simulation" + + | otherwise + = do performGC + vf' <- velocitySteps config step vf vs + df' <- densitySteps config step df ds vf' + return $ Model df' Nothing vf' Nothing cl (step + 1) cb + + + + + + diff --git a/raster/Fluid/src-repa/Model.hs b/raster/Fluid/src-repa/Model.hs new file mode 100644 --- /dev/null +++ b/raster/Fluid/src-repa/Model.hs @@ -0,0 +1,183 @@ +{-# LANGUAGE ScopedTypeVariables, BangPatterns #-} +module Model + ( Delta + , Rate + , Field + , DensityField + , VelocityField + , SourceDensity (..) + , CurrentButton (..) + + , Model (..) + , initModel + + , pictureOfModel + , pixel32OfDensity + , pixel8OfDensity + + , outputBMP + , outputPPM) +where +import Data.Array.Repa as R +import Data.Array.Repa.IO.BMP as R +import Data.Array.Repa.Repr.ForeignPtr as R +import Data.Bits +import Data.Word +import Unsafe.Coerce +import Prelude as P +import qualified Graphics.Gloss as G + + +-- | Time delta (seconds) +type Delta = Float + +-- | Time rate (or period) (1/seconds) +type Rate = Float + + +-- | A 2d field. +type Field a + = Array U DIM2 a + +-- | Scalar density field. +type DensityField + = Field Float + +-- | Vector velocity field. +type VelocityField + = Field (Float, Float) + + +-- | Button being pressed by the user. +data CurrentButton + = LeftButton + | RightButton + | None + +-- | A density source added by the user. +data SourceDensity a + = SourceDensity DIM2 a + + +-- Model ---------------------------------------------------------------------- +-- | The world model. +data Model + = Model + { densityField :: DensityField + , densitySource :: Maybe (SourceDensity Float) + , velocityField :: VelocityField + , velocitySource :: Maybe (SourceDensity (Float, Float)) + , clickLoc :: Maybe (Int, Int) + , stepsPassed :: Int + , currButton :: CurrentButton + } + + +-- | Creates an initial blank model +initModel + :: Array U DIM2 Float + -> Array U DIM2 (Float, Float) + -> Model + +initModel density velocity + | extent density /= extent velocity + = error "Fluid: initModel density and velocity extents do not match" + + | otherwise + = Model + { densityField = density + , densitySource = Nothing + , velocityField = velocity + , velocitySource = Nothing + , clickLoc = Nothing + , stepsPassed = 0 + , currButton = None } +{-# INLINE initModel #-} + + +-- Picture -------------------------------------------------------------------- +-- | Function to convert the Model into a Bitmap for displaying in Gloss +pictureOfModel :: Monad m => (Float, Float) -> Model -> m G.Picture +pictureOfModel (scaleX, scaleY) m + = let (Z :. width' :. height') = R.extent $ densityField m + width = fromIntegral width' + height = fromIntegral height' + + in do + (arrDensity :: Array F DIM2 Word32) + <- computeP $ R.map pixel32OfDensity $ densityField m + + return $ G.Scale scaleX scaleY + $ G.bitmapOfForeignPtr width height + (R.toForeignPtr $ unsafeCoerce arrDensity) + False +{-# NOINLINE pictureOfModel #-} + + +-- | Converts Float value to Word32 for pixel data +pixel32OfDensity :: Float -> Word32 +pixel32OfDensity f + = let !fsat + | f < 0 = 0 + | f >= 1 = 1 + | otherwise = f + + !x = truncate $ fsat * 255 + !a = 255 + + in unsafeShiftL x 24 .|. unsafeShiftL x 16 + .|. unsafeShiftL x 8 .|. a +{-# INLINE pixel32OfDensity #-} + + +-- | Converts Float value to a tuple of pixel components. +pixel8OfDensity :: Float -> (Word8, Word8, Word8) +pixel8OfDensity f + = let !fsat + | f < 0 = 0 + | f >= 1 = 1 + | otherwise = f + + !x = truncate $ fsat * 255 + in (x, x, x) +{-# INLINE pixel8OfDensity #-} + + +-- Dump ----------------------------------------------------------------------- +-- Writes bitmap data to test batch-mode ran correctly +outputBMP :: DensityField -> IO () +outputBMP df + = do arr <- computeUnboxedP $ R.map pixel8OfDensity df + R.writeImageToBMP "./out.bmp" arr + + +outputPPM :: Int -> String -> Float -> Array U DIM2 Float -> IO () +outputPPM step name scale df + = do let (Z :. h :. w) = extent df + let step' + = replicate (4 - length (show step)) '0' P.++ show step + + let getVal x y + = let v = df R.! (Z :. y :. x) + in truncate (v * 255 * scale) :: Int + + let showVal x y + = let v = getVal x y + in replicate (4 - length (show v)) ' ' P.++ show v + + let mx + = maximum [getVal x y | x <- [0..w-1], y <- [0..h-1]] + + let out = unlines $ + [ "P2" + , show h P.++ " " P.++ show w + , show mx] + P.++ [ concat [ showVal x y P.++ " " + | x <- [0..w - 1]] + | y <- [h - 1, h - 2 .. 0]] + + writeFile ("out/" P.++ step' P.++ "-" P.++ name P.++ ".ppm") out + + + + diff --git a/raster/Fluid/src-repa/Solve/Density.hs b/raster/Fluid/src-repa/Solve/Density.hs new file mode 100644 --- /dev/null +++ b/raster/Fluid/src-repa/Solve/Density.hs @@ -0,0 +1,30 @@ + +module Solve.Density + (densitySteps) +where +import Stage.Diffusion +import Stage.Advection +import Stage.Sources +import Config +import Model + +-- | Run the stages for processing the density field in one time step +densitySteps + :: Config + -> Int + -> DensityField + -> Maybe (SourceDensity Float) + -> VelocityField + -> IO DensityField + +densitySteps config step df ds vf + = {-# SCC "Solve.densitySteps" #-} + do df1 <- addSources (configDelta config) (configDensity config) + ds df + + df2 <- diffusion (configIters config) (configDelta config) (configDiffusion config) + df1 + + df' <- advection (configDelta config) vf df2 + + return df' diff --git a/raster/Fluid/src-repa/Solve/Velocity.hs b/raster/Fluid/src-repa/Solve/Velocity.hs new file mode 100644 --- /dev/null +++ b/raster/Fluid/src-repa/Solve/Velocity.hs @@ -0,0 +1,41 @@ + +module Solve.Velocity + (velocitySteps) +where +import Stage.Diffusion +import Stage.Advection +import Stage.Sources +import Stage.Project +import Model +import Config + +-- The pass that sets boundary conditions is buggy and +-- currently disabled. +-- import Stage.Boundary +velocitySteps + :: Config + -> Int + -> VelocityField + -> Maybe (SourceDensity (Float, Float)) + -> IO VelocityField + +velocitySteps config _step vf vs + = {-# SCC "Solve.velocitySteps" #-} + do + vf1 <- addSources (configDelta config) (configVelocity config) + vs vf + + vf2 <- diffusion (configIters config) (configDelta config) (configViscosity config) + vf1 +-- vf3 <- setBoundary vf2 + + vf4 <- project (configIters config) vf2 +-- vf5 <- setBoundary vf4 + + vf6 <- advection (configDelta config) vf4 vf4 +-- vf7 <- setBoundary vf6 + + vf8 <- project (configIters config) vf6 +-- vf' <- setBoundary vf8 + + return vf8 diff --git a/raster/Fluid/src-repa/Stage/Advection.hs b/raster/Fluid/src-repa/Stage/Advection.hs new file mode 100644 --- /dev/null +++ b/raster/Fluid/src-repa/Stage/Advection.hs @@ -0,0 +1,117 @@ +{-# LANGUAGE BangPatterns, ScopedTypeVariables #-} +module Stage.Advection + (advection) +where +import Model +import FieldElt +import Data.Array.Repa as R +import Data.Array.Repa.Unsafe as R +import Data.Vector.Unboxed (Unbox) +import Text.Printf + +-- | Apply a velocity field to another field. +-- Both fields must have the same extent. +advection + :: (FieldElt a, Unbox a, Show a) + => Delta + -> VelocityField + -> Field a + -> IO (Field a) + +advection !delta velField field + = {-# SCC "advection" #-} + velField `deepSeqArray` field `deepSeqArray` + do computeP $ unsafeTraverse field id (advectElem delta velField) + +{-# SPECIALIZE advection + :: Delta + -> VelocityField -> Field Float + -> IO (Field Float) #-} + +{-# SPECIALIZE advection + :: Delta + -> VelocityField -> Field (Float, Float) + -> IO (Field (Float, Float)) #-} + + +-- | Compute the new field value at the given location. +advectElem + :: (FieldElt a, Unbox a) + => Delta -- ^ Time delta (in seconds) + -> VelocityField -- ^ Velocity field that moves the source field. + -> (DIM2 -> a) -- ^ Get an element from the source field. + -> DIM2 -- ^ Compute the new value at this index. + -> a + +advectElem !delta !velField !get !pos@(Z:. j :. i) + = velField `deepSeqArray` + (((d00 ~*~ t0) ~+~ (d01 ~*~ t1)) ~*~ s0) + ~+~ (((d10 ~*~ t0) ~+~ (d11 ~*~ t1)) ~*~ s1) + where + _ :. _ :. width' = R.extent velField + !width = fromIntegral width' + + -- helper values + !dt0 = delta * width + !(u, v) = velField `unsafeIndex` pos + + -- backtrack densities to point based on velocity field + -- and make sure they are in field + !x = fromIntegral i - dt0 * u + !y = fromIntegral j - dt0 * v + + !x' | x < -0.5 = -0.5 + | x > width + 0.5 = width + 0.5 + | otherwise = x + + !y' | y < -0.5 = -0.5 + | y > width + 0.5 = width + 0.5 + | otherwise = y + + -- calculate discrete locations surrounding point + !i0 = truncate (x' + 1) - 1 + !i1 = i0 + 1 + + !j0 = truncate (y' + 1) - 1 + !j1 = j0 + 1 + + -- calculate ratio point is between the discrete locations + !s1 = x' - fromIntegral i0 + !s0 = 1 - s1 + + !t1 = y' - fromIntegral j0 + !t0 = 1 - t1 + + get' ix@(Z :. jj :. ii) + | ii < 0 = zero + | ii >= width' = zero + | jj < 0 = zero + | jj >= width' = zero + | otherwise = get ix + + -- grab values from grid surrounding advected point + !d00 = get' (Z:. j0 :. i0) + !d01 = get' (Z:. j1 :. i0) + !d10 = get' (Z:. j0 :. i1) + !d11 = get' (Z:. j1 :. i1) + + +{-# SPECIALIZE advectElem + :: Delta + -> VelocityField -> (DIM2 -> Float) + -> DIM2 -> Float #-} + +{-# SPECIALIZE advectElem + :: Delta + -> VelocityField -> (DIM2 -> (Float,Float)) + -> DIM2 -> (Float,Float) #-} + +{- +-- | Wrap an index back into the simulation area if it is outside. +checkLocation :: Float -> Float -> Float +checkLocation !width !x + | x < 0.5 = 0.5 + | x > width - 1.5 = width - 1.5 + | otherwise = x +{-# INLINE checkLocation #-} +-} diff --git a/raster/Fluid/src-repa/Stage/Boundary.hs b/raster/Fluid/src-repa/Stage/Boundary.hs new file mode 100644 --- /dev/null +++ b/raster/Fluid/src-repa/Stage/Boundary.hs @@ -0,0 +1,115 @@ + +{-# LANGUAGE BangPatterns #-} +module Stage.Boundary + (setBoundary) +where +import Model +import Data.Array.Repa as R +import Control.Monad +import Debug.Trace +import Config +import FieldElt + +-- | Apply boundary conditions to a velocity field. +setBoundary :: Config -> VelocityField -> IO VelocityField +setBoundary config f + = let (width, _) = configModelSize config + in (rebuild width f <=< setBoundary' width <=< grabBorders width) f + + +-- | Takes the original VelocityField and the array of edges and replaces +-- edge values with new values +rebuild :: Int -> VelocityField -> VelocityField -> IO VelocityField +rebuild width field edges + = field `deepSeqArray` edges `deepSeqArray` + do computeUnboxedP $ backpermuteDft field (rebuildPosMap width) edges +{-# INLINE rebuild #-} + + +rebuildPosMap :: Int -> DIM2 -> Maybe DIM2 +rebuildPosMap !width (Z:.j:.i) + | j == 0 + = Just (Z:.0:.i) + + | j == width - 1 + = Just (Z:.1:.i) + + | i == 0 + = if j == 0 then Just (Z:. 0 :. 0) + else if j == end then Just (Z:. 1 :. 0) + else Just (Z:. 2 :. j) + + | i == width - 1 + = if j == 0 then Just (Z:. 0 :. (width-1)) + else if j == end then Just (Z:. 1 :. (width-1)) + else Just (Z:. 3 :. j) + + | otherwise = Nothing + where end = width - 1 +{-# INLINE rebuildPosMap #-} + + + +-- | Grabs the border elements of the VelocityField and outputs them as +-- one array, for ease of adding back into the original VelocityField later +grabBorders :: Int -> VelocityField -> IO VelocityField +grabBorders width f + = f `deepSeqArray` + do traceEventIO "Fluid: grabBorders" + computeUnboxedP $ backpermute (Z:. 4 :. width) (edgeCases width) f +{-# INLINE grabBorders #-} + + +-- | Map a position in the edges array to what they were in the original +-- array. +edgeCases :: Int -> DIM2 -> DIM2 +edgeCases width (Z:.j:.i) + | j == 0 = (Z:.0 :.i) + | j == 1 = (Z:.(width-1) :.i) + | j == 2 = (Z:.i :.0) + | j == 3 = (Z:.i :.(width-1)) + | otherwise = error "Incorrect coordinate given in setBoundary" +{-# INLINE edgeCases #-} + + + +setBoundary' :: Int -> VelocityField -> IO VelocityField +setBoundary' width e + = e `deepSeqArray` + do traceEventIO "Fluid: setBoundary'" + computeUnboxedP $ traverse e id (revBoundary width) +{-# INLINE setBoundary' #-} + + +-- | Based on position in edges array set the velocity accordingly +revBoundary :: Int -> (DIM2 -> (Float,Float)) -> DIM2 -> (Float,Float) +revBoundary width loc pos@(Z:.j:.i) + | j == 0 + = if i == 0 then grabCornerCase loc (Z:.2:.1) (Z:.0:.1) + else if i == end then grabCornerCase loc (Z:.0:.(width-2)) (Z:.3:.1) + else (-p1,p2) + | j == 1 + = if i == 0 then grabCornerCase loc (Z:.2:.(width-2)) (Z:.1:.1) + else if i == end then grabCornerCase loc (Z:.1:.(width-2)) (Z:.3:.(width-2)) + else (-p1,p2) + + | j == 2 = (p1,-p2) + | j == 3 = (p1,-p2) + + | otherwise = error "Fluid: revBoundary" + where (p1,p2) = loc pos + end = width - 1 + + +{-# INLINE revBoundary #-} + +-- | Corner cases are special and are calculated with this function +grabCornerCase :: (DIM2 -> (Float, Float)) -> DIM2 -> DIM2 -> (Float, Float) +grabCornerCase loc pos1 pos2 + = (p1 * q1, p2 * q2) ~*~ 0.5 + where (p1,p2) = loc pos1 + (q1,q2) = loc pos2 +{-# INLINE grabCornerCase #-} + + + diff --git a/raster/Fluid/src-repa/Stage/Diffusion.hs b/raster/Fluid/src-repa/Stage/Diffusion.hs new file mode 100644 --- /dev/null +++ b/raster/Fluid/src-repa/Stage/Diffusion.hs @@ -0,0 +1,39 @@ +{-# LANGUAGE BangPatterns #-} +module Stage.Diffusion + (diffusion) +where +import Model +import FieldElt +import Stage.Linear +import Data.Array.Repa as R +import Data.Array.Repa.Eval as R +import Data.Vector.Unboxed + + +-- | Diffuse a field at a certain rate. +diffusion + :: (FieldElt a, Num a, Elt a, Unbox a) + => Int + -> Delta + -> Rate + -> Field a + -> IO (Field a) +diffusion !iters !delta !rate field + = {-# SCC diffusion #-} + field `deepSeqArray` + let _ :. _ :. width' = R.extent field + !width = fromIntegral width' + !a = delta * rate * width * width + !c = 1 + 4 * a + in do + linearSolver field field a c iters + +{-# SPECIALIZE diffusion + :: Int -> Delta -> Rate + -> Field Float + -> IO (Field Float) #-} + +{-# SPECIALIZE diffusion + :: Int -> Delta -> Rate + -> Field (Float, Float) + -> IO (Field (Float, Float)) #-} diff --git a/raster/Fluid/src-repa/Stage/Linear.hs b/raster/Fluid/src-repa/Stage/Linear.hs new file mode 100644 --- /dev/null +++ b/raster/Fluid/src-repa/Stage/Linear.hs @@ -0,0 +1,86 @@ +{-# LANGUAGE FlexibleContexts, BangPatterns #-} +module Stage.Linear + (linearSolver) +where +import Model +import FieldElt as E +import Data.Array.Repa as R +import Data.Array.Repa.Stencil as R +import Data.Array.Repa.Stencil.Dim2 as R +import Data.Array.Repa.Eval as R +import Data.Vector.Unboxed +import Prelude as P + + +linearSolver + :: (FieldElt a, Source U a, Unbox a, Elt a, Num a, Show a) + => Field a -- ^ Original field. + -> Field a -- ^ Current field. + -> Float + -> Float + -> Int -- ^ Number of iterations to apply. + -> IO (Field a) + +linearSolver origField curField !a !c !iters + -- If nothing would change by running the solver, then skip it. + | 0 <- a = return origField + + -- The solver has finished its loop + | 0 <- iters = return curField + + -- Do one iteration + | otherwise + = origField `deepSeqArray` curField `deepSeqArray` + do let !c' = 1/c + let {-# INLINE zipFunc #-} + zipFunc !orig !new + = (orig ~+~ (new ~*~ a)) ~*~ c' + + newField <- {-# SCC "linearSolver.mapStencil" #-} + computeUnboxedP + $ R.szipWith zipFunc origField + $ mapStencil2 (BoundConst 0) linearSolverStencil curField + + -- TODO: this boundConst thing is costing a fair bit. + -- Do something about the branches introduced into the core code. + + linearSolver origField newField a c (iters - 1) + +{-# SPECIALIZE linearSolver + :: Field Float + -> Field Float + -> Float -> Float -> Int + -> IO (Field Float) #-} + +{-# SPECIALIZE linearSolver + :: Field (Float, Float) + -> Field (Float, Float) + -> Float -> Float -> Int + -> IO (Field (Float, Float)) #-} + + + +-- | Stencil function for the linear solver. +linearSolverStencil + :: FieldElt a + => Stencil DIM2 a + +linearSolverStencil + = StencilStatic (Z:.3:.3) E.zero + (\ix val acc -> + case linearSolverCoeffs ix of + Nothing -> acc + Just coeff -> acc ~+~ (val ~*~ coeff)) +{-# INLINE linearSolverStencil #-} + + +-- | Linear solver stencil kernel. +linearSolverCoeffs :: DIM2 -> Maybe Float +linearSolverCoeffs (Z:.j:.i) + | i == 1, j == 0 = Just 1 + | i == -1, j == 0 = Just 1 + | i == 0, j == 1 = Just 1 + | i == 0, j == -1 = Just 1 + | otherwise = Nothing +{-# INLINE linearSolverCoeffs #-} + diff --git a/raster/Fluid/src-repa/Stage/Project.hs b/raster/Fluid/src-repa/Stage/Project.hs new file mode 100644 --- /dev/null +++ b/raster/Fluid/src-repa/Stage/Project.hs @@ -0,0 +1,66 @@ +{-# LANGUAGE BangPatterns #-} +module Stage.Project + (project) +where +import Model +import FieldElt +import Stage.Linear +import Data.Array.Repa as R +import Data.Array.Repa.Unsafe as R +import Prelude as P + + +project :: Int -> Field (Float, Float) -> IO (Field (Float, Float)) +project iters field + = {-# SCC project #-} + field `deepSeqArray` + do let _ :. _ :. width = extent field + + divergence <- {-# SCC "project.genDiv" #-} + computeUnboxedP + $ fromFunction (Z:. width :. width) (genDivergence width field) + + p <- {-# SCC "project.linearSolver" #-} + linearSolver divergence divergence 1 4 iters + + f' <- {-# SCC "project.apply" #-} + computeUnboxedP + $ unsafeTraverse field id (projectElem width p) + + return f' +{-# NOINLINE project #-} + + +-- | Subtract a gradient field from the regular field to +-- create a mass-conserving field. +projectElem + :: Int -- ^ Width of model. + -> Field Float + -> (DIM2 -> (Float, Float)) -- ^ Get data from the regular field. + -> DIM2 -- ^ Compute the value at this point. + -> (Float, Float) + +projectElem !width !p !get !pos@(Z:.j:.i) + = get pos ~-~ (0.5 * width' * (p0 - p1), + 0.5 * width' * (p2 - p3)) + where + !width' = fromIntegral width + !p0 = useIf (i < width - 1) (p `unsafeIndex` (Z :. j :. i+1)) + !p1 = useIf (i > 0) (p `unsafeIndex` (Z :. j :. i-1)) + !p2 = useIf (j < width - 1) (p `unsafeIndex` (Z :. j+1 :. i )) + !p3 = useIf (j > 0) (p `unsafeIndex` (Z :. j-1 :. i )) +{-# INLINE projectElem #-} + + +-- | Get an approximation of the gradient at this point. +genDivergence :: Int -> VelocityField -> DIM2 -> Float +genDivergence !width !f (Z :. j :. i) + = (-0.5 * ((u0 - u1) + (v0 - v1))) / fromIntegral width + where + (u0, _) = useIf (i < width - 1) (f `unsafeIndex` (Z:. j :. i+1)) + (u1, _) = useIf (i > 0) (f `unsafeIndex` (Z:. j :. i-1)) + ( _, v0) = useIf (j < width - 1) (f `unsafeIndex` (Z:. j+1 :. i )) + ( _, v1) = useIf (j > 0) (f `unsafeIndex` (Z:. j-1 :. i )) +{-# INLINE genDivergence #-} + + diff --git a/raster/Fluid/src-repa/Stage/Sources.hs b/raster/Fluid/src-repa/Stage/Sources.hs new file mode 100644 --- /dev/null +++ b/raster/Fluid/src-repa/Stage/Sources.hs @@ -0,0 +1,77 @@ +{-# LANGUAGE BangPatterns, FlexibleInstances #-} +module Stage.Sources + (addSources) +where +import Model +import FieldElt +import Data.Array.Repa as R +import Data.Array.Repa.Unsafe as R +import Data.Vector.Unboxed (Unbox) + + +-- | Addition of forces stage for simulation +addSources + :: (FieldElt a, FieldSource a, Unbox a) + => Delta -- ^ Time delta. + -> a -- ^ Value to insert. + -> Maybe (SourceDensity a) + -> Field a + -> IO (Field a) + +addSources !delta !value (Just (SourceDensity aim mul)) field + = {-# SCC addSources #-} + field `deepSeqArray` + do computeP $ unsafeTraverse field id (insertSource delta value aim mul) + +addSources _ _ Nothing field + = return field + + +insertSource + :: (FieldElt a, FieldSource a) + => Delta + -> a -- ^ Value to insert + -> DIM2 -> a + -> (DIM2 -> a) + -> DIM2 + -> a + +insertSource !delta !value !aim !mul locate !pos + | aim == pos = addSource delta value (locate pos) mul + | otherwise = locate pos +{-# INLINE insertSource #-} + + +{-# SPECIALIZE addSources + :: Delta + -> Float + -> Maybe (SourceDensity Float) + -> Field Float + -> IO (Field Float) #-} + +{-# SPECIALIZE addSources + :: Delta + -> (Float, Float) + -> Maybe (SourceDensity (Float, Float)) + -> Field (Float, Float) + -> IO (Field (Float, Float)) #-} + + +-- FieldSource ---------------------------------------------------------------- +class FieldSource a where + addSource :: Delta -> a -> a -> a -> a + +instance FieldSource Float where + addSource !delta !value !a !mul + = a ~+~ (value * delta * mul) + {-# INLINE addSource #-} + +instance FieldSource (Float, Float) where + addSource !delta (newA, newB) (a,b) (mulA, mulB) + = ( a + (newA * delta * (-mulA)) + , b + (newB * delta * (-mulB))) + {-# INLINE addSource #-} + + + + diff --git a/raster/Fluid/src-repa/UserEvent.hs b/raster/Fluid/src-repa/UserEvent.hs new file mode 100644 --- /dev/null +++ b/raster/Fluid/src-repa/UserEvent.hs @@ -0,0 +1,136 @@ +module UserEvent + (userEvent) +where +import Config +import Model as M +import Data.Array.Repa as R +import Graphics.Gloss.Interface.Pure.Game as G + + +-- | Handle user events for the Gloss `playIO` wrapper. +userEvent :: Config -> Event -> Model -> Model +userEvent config + (EventKey key keyState mods (x, y)) + (Model df ds vf vs cl sp _cb) + + -- Add velocity --------------------------------------------- + | MouseButton G.RightButton <- key + , Down <- keyState + , (x',y') <- windowToModel config (x,y) + = Model { densityField = df + , densitySource = ds + , velocityField = vf + , velocitySource = vs + , clickLoc = Just (x',y') + , stepsPassed = sp + , currButton = M.RightButton + } + + -- Accept shift-leftbutton for people with trackpads + | MouseButton G.LeftButton <- key + , Down <- keyState + , Down <- shift mods + , (x',y') <- windowToModel config (x,y) + = Model { densityField = df + , densitySource = ds + , velocityField = vf + , velocitySource = vs + , clickLoc = Just (x',y') + , stepsPassed = sp + , currButton = M.RightButton + } + + | MouseButton G.RightButton <- key + , Up <- keyState + , Just (locX, locY) <- cl + , (x',y') <- windowToModel config (x,y) + = Model { densityField = df + , densitySource = ds + , velocityField = vf + , velocitySource = Just (SourceDensity (Z:.locY:.locX) + (fromIntegral (locX-x'),fromIntegral (locY-y'))) + , clickLoc = Nothing + , stepsPassed = sp + , currButton = M.None + } + + + -- Add density ---------------------------------------------- + | MouseButton G.LeftButton <- key + , Down <- keyState + , (x',y') <- windowToModel config (x, y) + = Model { densityField = df + , densitySource = Just (SourceDensity (Z:.y':.x') 1) + , velocityField = vf + , velocitySource = vs + , clickLoc = cl + , stepsPassed = sp + , currButton = M.LeftButton + } + + | MouseButton G.LeftButton <- key + , Up <- keyState + = Model { densityField = df + , densitySource = ds + , velocityField = vf + , velocitySource = vs + , clickLoc = cl + , stepsPassed = sp + , currButton = M.None + } + + + -- Reset model ---------------------------------------------- + | Char 'r' <- key + , Down <- keyState + = initModel (configInitialDensity config) + (configInitialVelocity config) + + -- Quit program --------------------------------------------- + | Char 'q' <- key + , Down <- keyState + = error "Quitting" + + +userEvent config (EventMotion (x, y)) + (Model df _ds vf vs cl sp M.LeftButton) + | (x',y') <- windowToModel config (x, y) + = Model { densityField = df + , densitySource = Just (SourceDensity (Z:.y':.x') 1) + , velocityField = vf + , velocitySource = vs + , clickLoc = cl + , stepsPassed = sp + , currButton = M.LeftButton + } + +userEvent config (EventMotion (x, y)) + (Model df ds vf _vs (Just (clx, cly)) sp M.RightButton) + | (x', y') <- windowToModel config (x,y) + = Model { densityField = df + , densitySource = ds + , velocityField = vf + , velocitySource = Just (SourceDensity (Z:.y':.x') + (fromIntegral (clx-x'), fromIntegral (cly-y'))) + , clickLoc = Just (x',y') + , stepsPassed = sp + , currButton = M.RightButton + } + +userEvent _ _ m = m + +-- Converts a window location to the corresponding location in the +-- simulation. +windowToModel :: Config -> (Float, Float) -> (Int, Int) +windowToModel config (x, y) = (x', y') + where (scaleX, scaleY) = configScale config + (windowWidth, windowHeight) = configWindowSize config + + x' = round ((x + (fromIntegral windowWidth / 2)) / scaleX) + y' = round ((y + (fromIntegral windowHeight / 2)) / scaleY) + + + + + +