Basically, I'm comparably new to Haskell, which means lacking both knowledge of advanced principles, available libraries and, maybe, some syntax.
For educational purposes I hacked together a code implementing Stam-style fluid dynamics solver. For now the code is running in constant space and linear time, which is exactly what's needed. But feeling remains that code I wrote is rather ugly.
I have serious concerns about whether this code is possible to read at all, but hope that somebody could point out most obvious idiosyncrasies.
import Data.Array.Unboxed
import Data.List
import Text.Printf
n = 20
timeStep = 0.1
numSteps = 100
indicies = [[0..n + 1], [1..n], [2..n-1]]
type Matrix = UArray (Int, Int) Float
data Bound = XHard | YHard | Soft
data State = State !Matrix !Matrix !Matrix
applyNTimes :: Int -> (a -> a) -> a -> a
applyNTimes times f val = foldl' (\m i -> f m) val [1..times]
showMatrix :: Matrix -> String
showMatrix m = foldl (++) "" [showRow m i ++ "\n" | i <- head indicies] where showRow m i = foldl (++) "" [printf "%.3f " (m ! (i, j)) | j <- head indicies]
showState :: State -> String
showState (State u v d) = showMatrix d
createMatrix :: [Int] -> ((Int, Int) -> Float) -> Matrix
createMatrix ind f = array ((0, 0), (n + 1, n + 1)) [((i, j), f (i, j)) | j <- ind, i <- ind] // [(ix, 0) | ix <- other]
where other = [(i, j) | i <- head indicies, j <- head indicies] \\ [(i, j) | i <- ind, j <- ind]
zeroMatrix :: Matrix
zeroMatrix = createMatrix (head indicies) (\ (i, j) -> 0.0)
add :: Matrix -> Matrix -> Matrix
add a b = createMatrix (head indicies) (add_cell a b) where add_cell a b ix = a ! ix + b ! ix
addPecularSource :: State -> State
addPecularSource (State u v d) = State u v (d `add` createMatrix (head indicies) spike) where spike ix = if ix == (n `div` 2, n `div` 2) then 0.1 else 0.0
advance :: State -> State
advance s = solveDensity . addPecularSource . solveVelocity $ s
main :: IO()
main = putStrLn . showState . applyNTimes numSteps advance $ State z z z where z = zeroMatrix
solveDensity :: State -> State
solveDensity (State u v d) = State u v d' where d' = advect Soft (diffuse d diffusion) u v where diffusion = 0.0 :: Float
setCorners :: Matrix -> Matrix
setCorners a = a // [((n + 1, n + 1), (a ! (n , n + 1) + a ! (n + 1, n)) / 2),
((n + 1, 0), (a ! (n , 0) + a ! (n + 1, 1)) / 2),
((0, n + 1), (a ! (0, n) + a ! (1, n + 1)) / 2),
((0, 0), (a ! (1 , 0) + a ! (0, 1)) / 2)]
diffuse :: Matrix -> Float -> Matrix
diffuse c0 diff = linearSolver c0 a (1 + 4 * a) where a = timeStep * diff * fromIntegral n * fromIntegral n
setBoundary :: Bound -> Matrix -> Matrix
setBoundary b m = setCorners (setSides b m)
setSides :: Bound -> Matrix -> Matrix
setSides XHard m = setSides Soft m // ([((0, i), - m ! (0, i)) | i <- indicies !! 1] ++ [((n + 1, i), - m ! (n + 1, i)) | i <- indicies !! 1])
setSides YHard m = setSides Soft m // ([((i, 0), - m ! (i, 0)) | i <- indicies !! 1] ++ [((i, n + 1), - m ! (i, n + 1)) | i <- indicies !! 1])
setSides Soft m = m // ([((0, i), m ! (1, i)) | i <- indicies !! 1] ++ [((n + 1, i), m ! (n, i)) | i <- indicies !! 1] ++
[((i, 0), m ! (i, 1)) | i <- indicies !! 1] ++ [((i, n + 1), m ! (i, n)) | i <- indicies !! 1])
linearSolver :: Matrix -> Float -> Float -> Matrix
linearSolver x0 a c = applyNTimes 20 (iterFun x0 a c) zeroMatrix
where iterFun x0 a c m = setBoundary Soft (createMatrix (indicies !! 1) (update m x0 a c) )
where update m x0 a c (i, j) = (a * (m ! (i - 1, j) + m ! (i + 1, j) + m ! (i, j - 1) + m ! (i, j + 1)) + x0 ! (i, j)) / c
curl :: Matrix -> Matrix -> Matrix
curl u v = createMatrix (indicies !! 1) gen where gen (i, j) = (u ! (i, j + 1) - u ! (i, j - 1) - v ! (i + 1, j) + v ! (i - 1, j)) / 2
archCell :: Matrix -> Float -> Float -> (Int, Int) -> Float
archCell d gravity push (i, j) = gravity * d ! (i, j) + push * (d ! (i, j) - ambient d)
where ambient d = sum [d ! (i, j) | i <- indicies !! 1, j <- indicies !! 1]
archimedes :: Matrix -> Matrix
archimedes d = createMatrix (head indicies) (archCell d gravity push)
where (gravity, push) = (0.001 :: Float, 0.005 :: Float)
solveVelocity :: State -> State
solveVelocity (State u v d) = State u' v' d
where (u', v') = project (advect YHard uProjected uProjected vProjected) (advect XHard vProjected uProjected vProjected)
where (uProjected, vProjected) = project uDiffused vDiffused
where (uDiffused, vDiffused) = (diffuse uVorticityConfined viscosity, diffuse (vVorticityConfined `add` archimedes d) viscosity)
where ((uVorticityConfined, vVorticityConfined), viscosity) = (vorticityConfimentForce u v, 0.0)
advectCell :: Matrix -> Matrix -> Matrix -> (Int, Int) -> Float
advectCell d0 du dv (i, j) = s0 * t0 * d0 ! (floor x, floor y) +
s0 * t1 * d0 ! (floor x, floor y + 1) +
s1 * t0 * d0 ! (floor x + 1, floor y) +
s1 * t1 * d0 ! (floor x + 1, floor y + 1)
where (s0, t0, s1, t1, x, y) = (1.0 - x' + xx, 1.0 - y' + yy, x' - xx, y' - yy, x', y')
where (x', y', xx, yy) = (x, y, fromIntegral (floor x), fromIntegral (floor y))
where (x, y) = (sort [0.5 + fromIntegral n, 0.5, fromIntegral i - dt0 * du ! (i, j)] !! 1,
sort [0.5 + fromIntegral n, 0.5, fromIntegral j - dt0 * dv ! (i, j)] !! 1)
where dt0 = timeStep * fromIntegral n
advect :: Bound -> Matrix -> Matrix -> Matrix -> Matrix
advect b d0 du dv = setBoundary b (createMatrix (indicies !! 1) (advectCell d0 du dv))
projectCellX :: Matrix -> Matrix -> (Int, Int) -> Float
projectCellX x p (i, j) = x ! (i, j) - 0.5 * fromIntegral n * (p!(i + 1, j) - p!(i - 1, j))
projectCellY :: Matrix -> Matrix -> (Int, Int) -> Float
projectCellY y p (i, j) = y ! (i, j) - 0.5 * fromIntegral n * (p!(i, j + 1) - p!(i, j - 1))
project :: Matrix -> Matrix -> (Matrix, Matrix)
project x y = (setBoundary YHard x', setBoundary XHard y')
where (x', y') = (createMatrix (indicies !! 1) (projectCellX x p), createMatrix (indicies !! 1) (projectCellY y p))
where p = linearSolver (setBoundary Soft diverg) 1 4
where diverg = createMatrix (indicies !! 1) (computeDiv x y)
where computeDiv x y (i, j) = -0.5 * (x ! (i + 1, j) - x ! (i - 1, j) + y ! (i, j + 1) - y ! (i, j - 1)) / fromIntegral n
dwdx :: Matrix -> (Int, Int) -> Float
dwdx c (i, j) = (abs (c ! (i + 1, j)) - abs (c ! (i - 1, j))) / 2
dwdy :: Matrix -> (Int, Int) -> Float
dwdy c (i, j) = (abs (c ! (i, j + 1)) - abs (c ! (i, j - 1))) / 2
len :: Matrix -> (Int, Int) -> Float
len c (i, j) = sqrt(dwdx c (i, j) * dwdx c (i, j) + dwdy c (i, j) * dwdy c (i, j)) + eps where eps = 1e-6
funX :: Matrix -> (Int, Int) -> Float
funX c (i, j) = - c ! (i, j) * dwdy c (i, j) / len c (i, j)
funY :: Matrix -> (Int, Int) -> Float
funY c (i, j) = c ! (i, j) * dwdx c (i, j) / len c (i, j)
vorticityConfimentForce :: Matrix -> Matrix -> (Matrix, Matrix)
vorticityConfimentForce u v = (createMatrix (indicies !! 2) (funX (curl u v)), createMatrix (indicies !! 2) (funY (curl u v)))
GHC 6.12.3 was used. Bang-patterns were added via -XBangPatterns.
My question is: how to change this code to be more comprehensible and idiomatic?
