module Bench.Vector.Algo.Tridiag ( tridiag ) where

import Data.Vector.Unboxed as V

tridiag :: (Vector Double, Vector Double, Vector Double, Vector Double)
            -> Vector Double
{-# NOINLINE tridiag #-}
tridiag :: (Vector Double, Vector Double, Vector Double, Vector Double)
-> Vector Double
tridiag (Vector Double
as,Vector Double
bs,Vector Double
cs,Vector Double
ds) = ((Double, Double) -> Double -> Double)
-> Double -> Vector (Double, Double) -> Vector Double
forall a b.
(Unbox a, Unbox b) =>
(a -> b -> b) -> b -> Vector a -> Vector b
V.prescanr' (\(Double
c,Double
d) Double
x' -> Double
d Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
cDouble -> Double -> Double
forall a. Num a => a -> a -> a
*Double
x') Double
0
                      (Vector (Double, Double) -> Vector Double)
-> Vector (Double, Double) -> Vector Double
forall a b. (a -> b) -> a -> b
$ ((Double, Double)
 -> ((Double, Double), (Double, Double)) -> (Double, Double))
-> (Double, Double)
-> Vector ((Double, Double), (Double, Double))
-> Vector (Double, Double)
forall a b.
(Unbox a, Unbox b) =>
(a -> b -> a) -> a -> Vector b -> Vector a
V.prescanl' (Double, Double)
-> ((Double, Double), (Double, Double)) -> (Double, Double)
forall {b}. Fractional b => (b, b) -> ((b, b), (b, b)) -> (b, b)
modify (Double
0,Double
0)
                      (Vector ((Double, Double), (Double, Double))
 -> Vector (Double, Double))
-> Vector ((Double, Double), (Double, Double))
-> Vector (Double, Double)
forall a b. (a -> b) -> a -> b
$ Vector (Double, Double)
-> Vector (Double, Double)
-> Vector ((Double, Double), (Double, Double))
forall a b.
(Unbox a, Unbox b) =>
Vector a -> Vector b -> Vector (a, b)
V.zip (Vector Double -> Vector Double -> Vector (Double, Double)
forall a b.
(Unbox a, Unbox b) =>
Vector a -> Vector b -> Vector (a, b)
V.zip Vector Double
as Vector Double
bs) (Vector Double -> Vector Double -> Vector (Double, Double)
forall a b.
(Unbox a, Unbox b) =>
Vector a -> Vector b -> Vector (a, b)
V.zip Vector Double
cs Vector Double
ds)
    where
      modify :: (b, b) -> ((b, b), (b, b)) -> (b, b)
modify (b
c',b
d') ((b
a,b
b),(b
c,b
d)) = 
                   let id :: b
id = b
1 b -> b -> b
forall a. Fractional a => a -> a -> a
/ (b
b b -> b -> b
forall a. Num a => a -> a -> a
- b
c'b -> b -> b
forall a. Num a => a -> a -> a
*b
a)
                   in
                   b
id b -> (b, b) -> (b, b)
forall a b. a -> b -> b
`seq` (b
cb -> b -> b
forall a. Num a => a -> a -> a
*b
id, (b
db -> b -> b
forall a. Num a => a -> a -> a
-b
d'b -> b -> b
forall a. Num a => a -> a -> a
*b
a)b -> b -> b
forall a. Num a => a -> a -> a
*b
id)