http://nineties.github.io/math-seminar/7.html 練習問題(行列計算/ピボット選択まで)をHaskellで書いてみた

mathpro7.hs

{-
 - http://nineties.github.io/math-seminar/7.html
 - 練習問題(行列計算/ピボット選択まで)をHaskellで書いてみた
 -}
import Data.Array

{- 目的：行列aとbの和を求める -}
mat_add :: Num a => Array (Int, Int) a -> Array (Int, Int) a
                 -> Array (Int, Int) a
mat_add a b = array bs [((i,j),a!(i,j)+b!(i,j))|i<-range(l1,l2),j<-range(m1,m2)]
  where ((l1,m1),(l2,m2)) = bounds a
        bs = if ((l1,m1),(l2,m2)) == bounds b
             then ((l1,m1),(l2,m2)) else error "行列の型の不一致"

test_mat_add = (mat_add a b) == r
  where a = array ((0,0),(1,2)) [((0,0),1),((0,1),2),((0,2),3),
                                 ((1,0),4),((1,1),5),((1,2),6)]
        b = array ((0,0),(1,2)) [((0,0),2),((0,1),1),((0,2),4),
                                 ((1,0),5),((1,1),2),((1,2),1)]
        r = array ((0,0),(1,2)) [((0,0),3),((0,1),3),((0,2),7),
                                 ((1,0),9),((1,1),7),((1,2),7)]

{- 目的：行列aのk倍を求める -}
mat_scale :: Num a => a -> Array (Int, Int) a -> Array (Int, Int) a
mat_scale k a = array bs [((i,j),k*a!(i,j))|i<-range(l1,l2),j<-range(m1,m2)]
  where ((l1,m1),(l2,m2)) = bounds a
        bs = ((l1,m1),(l2,m2))

test_mat_scale = (mat_scale 2 a) == r
  where a = array ((0,0),(1,2)) [((0,0),1),((0,1),2),((0,2),3),
                                 ((1,0),4),((1,1),5),((1,2),6)]
        r = array ((0,0),(1,2)) [((0,0),2),((0,1),4),((0,2),6),
                                 ((1,0),8),((1,1),10),((1,2),12)]

{- 目的：行列aとbの積を求める -}
mat_mul :: Num a => Array (Int, Int) a -> Array (Int, Int) a
                 -> Array (Int, Int) a
mat_mul a b = accumArray (+) 0 bs [((i,j),a!(i,k)*b!(k,j))|i<-range(l1,l2),
                                                           j<-range(n1,n2),
                                                           k<-range(m1,m2)]
  where ((l1,m1),(l2,m2)) = bounds a
        ((m1',n1),(m2',n2)) = bounds b
        bs = if (m1,m2) == (m1',m2')
               then ((l1,n1),(l2,n2)) else error "行列の型の不整合"

test_mat_mul = (mat_mul a b) == r
  where a = array ((0,0),(1,2)) [((0,0),1),((0,1),2),((0,2),3),
                                 ((1,0),4),((1,1),5),((1,2),6)]
        b = array ((0,0),(2,2)) [((0,0),2),((0,1),1),((0,2),4),
                                 ((1,0),5),((1,1),2),((1,2),1),
                                 ((2,0),0),((2,1),1),((2,2),3)]
        r = array ((0,0),(1,2)) [((0,0),12),((0,1), 8),((0,2),15),
                                 ((1,0),33),((1,1),20),((1,2),39)]

{- 目的：フィボナッチ数列の第n番目を行列を用いて求める -}
fib :: Int -> Int
fib n = (if n < 1 then a else foldr1 mat_mul $ take n $ repeat a)!(1,0)
  where a = array ((0,0),(1,1)) [((0,0),1),((0,1),1),
                                 ((1,0),1),((1,1),0)]

test_fib = (fib 50) == 12586269025

{- 目的：ガウス消去法を用いて連立一次方程式を解く -}
solve :: Array (Int, Int) Double -> Array (Int, Int) Double
solve = backward_substitution . forward_elimination

test_solve = r!(0,3) > 2.5999 && r!(0,3) <= 2.6 && r!(1,3) == 3.4 && r!(2,3) == -2.8
  where a = array ((0,0),(2,3)) [((0,0),1),((0,1),2),((0,2),3),((0,3),1),
                                 ((1,0),-1),((1,1),3),((1,2),2),((1,3),2),
                                 ((2,0),2),((2,1),1),((2,2),2),((2,3),3)]
        r = solve a

{- 目的：行列aを前進消去する -}
forward_elimination a = snd $ (iterate f (0,a))!!l2
  where ((l1,m1),(l2,m2)) = bounds a
        f (n,b) = let (l,m) = (l1+n,m1+n)
                      g i j = if i<=l||j<m then b!(i,j)
                              else              b!(i,j)-b!(l,j)*b!(i,m)/b!(l,m)
                  in (n+1,array ((l1,m1),(l2,m2))
                                [((i,j),g i j)|i<-range(l1,l2),j<-range(m1,m2)])

{- 目的：行列aを後退代入する -}
backward_substitution a = snd $ (iterate f (0,a))!!(l2+1)
  where ((l1,m1),(l2,m2)) = bounds a
        f (n,b) = let (l,m) = (l2-n,m2-n-1)
                      g i j = if      i==l&&j==m2 then b!(i,j)/b!(i,m)
                              else if i==l&&j==m  then 1
                              else if j==m        then 0
                              else if j==m2       then b!(i,j)-b!(i,m)*b!(l,m2)/b!(l,m)
                              else                     b!(i,j)
                  in (n+1,array ((l1,m1),(l2,m2))
                                [((i,j),g i j)|i<-range(l1,l2),j<-range(m1,m2)])

{- 目的：ガウス消去法(部分ピボット選択)を用いて連立一次方程式を解く -}
solve' :: Array (Int, Int) Double -> Array (Int, Int) Double
solve' = backward_substitution . forward_elimination'

test_solve' = r!(0,3) == 1 && r!(1,3) > 3.332 && r!(1,3) < 3.334 && r!(2,3) == 7
  where a = array ((0,0),(2,3)) [((0,0),2),((0,1),3),((0,2),-1),((0,3),5),
                                 ((1,0),4),((1,1),6),((1,2),-3),((1,3),3),
                                 ((2,0),2),((2,1),-3),((2,2),1),((2,3),-1)]
        r = solve' a

{- 目的：行列aを前進消去(部分ピボット選択)する -}
forward_elimination' a = snd $ (iterate (f.h) (0,a))!!l2
  where ((l1,m1),(l2,m2)) = bounds a
        f (n,b) = let (l,m) = (l1+n,m1+n)
                      g i j = if i<=l||j<m then b!(i,j)
                              else              b!(i,j)-b!(l,j)*b!(i,m)/b!(l,m)
                  in (n+1, array ((l1,m1),(l2,m2))
                                 [((i,j),g i j)|i<-range(l1,l2),j<-range(m1,m2)])
        h (n,b) = let (max_i,_) = foldr1 (\p@(_,x) q@(_,y) -> if x>y&&x/=0 then p else q)
                                         [(i,b!(i,m1+n))|i<-range(l1+n,l2)]
                      g i j = if      i==l1+n  then b!(max_i,j)
                              else if i==max_i then b!(l1+n,j)
                              else                  b!(i,j)
                  in (n,array ((l1,m1),(l2,m2))
                              [((i,j),g i j)|i<-range(l1,l2),j<-range(m1,m2)])