Laplace equation

laplace.hs

{-ラプラス方程式を解く-}

import Data.Array.IO

{-パラメータ設定-}
n = 50 {-行列のサイズ-}
k = (n+2) {-境界を含めた行列のサイズ-}
loop u 0 = return ()
loop u n = do
    writeArray u (1,n) 100
    loop u (n-1)

{-次のuを計算する-}
nu u error i j
    |i == k = return (u,error)
    |j == k = nu u error (i+1) 2
    |otherwise = do
        n <- readArray u (i,j)
        n1 <- readArray u (i+1,j)
        n2 <- readArray u (i-1,j)
        n3 <- readArray u (i,j+1)
        n4 <- readArray u (i,j-1)
        n5 <- readArray u (i,j)
        let r = (0.25*(n1 + n2 + n3 + n4) - n5)
        writeArray u (i,j) (n+r)
        if r > error then
            nu u r i (j+1)
        else
            nu u error i (j+1)

{-誤差が小さくなるまで繰り返す-}
select l = do
    (next,error) <- nu l 0 2 2
    if error < 0.001 then
        return next
    else
        select next

main = do
    u <- newArray ((1,1), (k,k)) 0.0 :: IO (IOUArray (Int, Int) Double)
    loop u k
    uu <- select u
    ans <- getElems uu
    let s = unwords [if (mod i k)/=(n+1) then show (ans!!i) else (show $ ans!!i)++"\n"| i <- [0..(length ans)-1]]
    print "ok"
    writeFile "output.txt" s

出力するファイルはgnuplotで表示させるためのもの。