Ticket #3960: TestMat.hs

{-# LANGUAGE PArr, ParallelListComp #-}
{-# OPTIONS -fvectorise #-}

module TestMat where

import qualified Prelude as P
import Data.Array.Parallel
import Data.Array.Parallel.Prelude
import Data.Array.Parallel.Prelude.Double as D
import Data.Array.Parallel.Prelude.Int as I

snd (_, b) = b

fst (a, _) = a

type Vector = [:Double:]
type SparseMatrix = [:[:(Int, Double):]:]
type DenseMatrix = [:[:Double:]:]

dotp :: Vector -> Vector -> Double
dotp v1 v2 = D.sumP (zipWithP (D.*) v1 v2)

mvMult :: SparseMatrix -> Vector -> Vector 
mvMult mat vec = [: D.sumP [: v D.* (vec !: n)  | (n, v) <- v1 :] | v1 <- mat:]

vp :: Vector -> Vector -> Vector 
vp a b = zipWithP (D.+) a b

vt :: Vector -> Vector -> Vector 
vt a b = zipWithP (D.*) a b

middle ib u1 ic u2 = vp (mvMult ib u1) (mvMult ic u2)
firstfn mv eta1 ib u1 ic u2 = vt (vp mv eta1) (middle ib u1 ic u2)
lastfn nv eta1 eta2 = vt (vp nv eta1) eta2

tmpfn mv eta1 ib u1 ic u2 nv eta2 = vp (firstfn mv eta1 ib u1 ic u2) (lastfn nv eta1 eta2)