Banded.hs

{-# LANGUAGE AllowAmbiguousTypes #-}

module Test.Banded (bandedTests) where

import Test.Tasty
import Test.Tasty.HUnit
import qualified Data.Massiv.Array as M

import Numeric.LinearAlgebra.Massiv.Internal
import Numeric.LinearAlgebra.Massiv.BLAS.Level2 (matvec)
import Numeric.LinearAlgebra.Massiv.BLAS.Level3 (matMul, transpose)
import Numeric.LinearAlgebra.Massiv.Solve.Banded
import Test.Types (matApproxEq, vecApproxEq)

bandedTests :: TestTree
bandedTests = testGroup "Banded"
  [ testCase "bandLU + band solve (tridiagonal p=1, q=1)" test_bandLUTridiag
  , testCase "bandLU + band solve (pentadiagonal p=2, q=2)" test_bandLUPentadiag
  , testCase "bandCholesky reconstruction (p=1)" test_bandCholTridiag
  , testCase "bandCholesky reconstruction (p=2)" test_bandCholPentadiag
  , testCase "tridiagSolve SPD system" test_tridiagSolveSPD
  ]

test_bandLUTridiag :: Assertion
test_bandLUTridiag = do
  -- Diagonally dominant tridiagonal 5x5 matrix
  let a = makeMatrix @5 @5 @M.P $ \i j ->
            case abs (i - j) of
              0 -> 4.0 :: Double
              1 -> -1.0
              _ -> 0.0
      b = makeVector @5 @M.P $ \i -> fromIntegral (i + 1) :: Double
      lu = bandLU 1 1 a
      y = bandForwardSub 1 lu b
      x = bandBackSub 1 lu y
      ax = matvec a x
  assertBool "bandLU solve Ax ~= b" (vecApproxEq ax b)

test_bandLUPentadiag :: Assertion
test_bandLUPentadiag = do
  -- Diagonally dominant pentadiagonal 6x6 matrix (p=2, q=2)
  let a = makeMatrix @6 @6 @M.P $ \i j ->
            case abs (i - j) of
              0 -> 10.0 :: Double
              1 -> -2.0
              2 -> 0.5
              _ -> 0.0
      b = makeVector @6 @M.P $ \i -> fromIntegral (i * 2 + 1) :: Double
      lu = bandLU 2 2 a
      y = bandForwardSub 2 lu b
      x = bandBackSub 2 lu y
      ax = matvec a x
  assertBool "bandLU pentadiagonal Ax ~= b" (vecApproxEq ax b)

test_bandCholTridiag :: Assertion
test_bandCholTridiag = do
  -- Symmetric positive definite tridiagonal 5x5 matrix (p=1)
  let a = makeMatrix @5 @5 @M.P $ \i j ->
            case abs (i - j) of
              0 -> 5.0 :: Double
              1 -> 1.0
              _ -> 0.0
      g = bandCholesky 1 a
      gt = transpose g
      ggt = matMul g gt
  assertBool "bandCholesky G*G^T ~= A (p=1)" (matApproxEq a ggt)

test_bandCholPentadiag :: Assertion
test_bandCholPentadiag = do
  -- Symmetric positive definite pentadiagonal 5x5 matrix (p=2)
  let a = makeMatrix @5 @5 @M.P $ \i j ->
            case abs (i - j) of
              0 -> 10.0 :: Double
              1 -> 2.0
              2 -> 0.5
              _ -> 0.0
      g = bandCholesky 2 a
      gt = transpose g
      ggt = matMul g gt
  assertBool "bandCholesky G*G^T ~= A (p=2)" (matApproxEq a ggt)

test_tridiagSolveSPD :: Assertion
test_tridiagSolveSPD = do
  -- 5x5 SPD tridiagonal system
  let d = makeVector @5 @M.P $ \_ -> 4.0 :: Double
      sd = makeVector @5 @M.P $ \i -> if i < 4 then 1.0 else 0.0 :: Double
      b = makeVector @5 @M.P $ \i -> fromIntegral (i + 1) :: Double
      x = tridiagSolve d sd b
      -- Reconstructed explicit matrix
      a = makeMatrix @5 @5 @M.P $ \i j ->
            case abs (i - j) of
              0 -> 4.0 :: Double
              1 -> 1.0
              _ -> 0.0
      ax = matvec a x
  assertBool "tridiagSolve Ax ~= b" (vecApproxEq ax b)