gistfile1.hs

advect :: Int -> Grid -> Grid -> Grid -> Grid -> Double -> IO ()
advect b d@(Grid n _) d0 u v dt = forEachPixel d
                                   (\(i,j) ->
                                    do
                                      uVal <- readVal u (i,j)
                                      vVal <- readVal v (i,j)
                                      let n5 = fromIntegral n + 0.5
                                          x = min n5 (max 0.5 (fromIntegral i - dt0 * uVal))
                                          y = min n5 (max 0.5 (fromIntegral j - dt0 * vVal))
                                          i0 = truncate x
                                          i1 = i0 + 1
                                          j0 = truncate y
                                          j1 = j0 + 1
                                          s1 = x - fromIntegral i0
                                          s0 = 1 - s1
                                          t1 = y - fromIntegral j0
                                          t0 = 1 - t1
                                      xd0 <- readVal d0 (i0,j0)
                                      xd1 <- readVal d0 (i0,j1)
                                      xd2 <- readVal d0 (i1,j0)
                                      xd3 <- readVal d0 (i1,j1)
                                      writeVal d (i,j) (s0*(t0*xd0 + t1*xd1) + s1*(t0*xd2+ t0*xd3)))
                                  >> setBnd b d
                                      where
                                        dt0 = dt * fromIntegral n

project :: Grid -> Grid -> Grid -> Grid -> IO ()
project u@(Grid n _) v p d = forEachPixel u
                                (\(i,j) ->
                                     do
                                       u0 <- readVal u (i+1,j)
                                       u1 <- readVal u (i-1,j)
                                       v0 <- readVal v (i,j+1)
                                       v1 <- readVal v (i,j-1)
                                       writeVal d (i,j) (-0.5 * ((u0-u1+v0-v1) / fromIntegral n))
                                       writeVal p (i,j) 0)
                               >> setBnd 0 d 
                               >> setBnd 0 p 
                               >> linSolve 0 p d 1 4 
                               >> forEachPixel p
                                   (\(i,j) ->
                                    do
                                      (up,down,left,right) <- neighbours p (i,j)
                                      u0 <- readVal u (i,j)
                                      v0 <- readVal v (i,j)
                                      writeVal u (i,j) (u0 - 0.5*fromIntegral n*(down - up))
                                      writeVal v (i,j) (v0 - 0.5*fromIntegral n*(right - left)))
                               >> setBnd 1 u
                               >> setBnd 2 v