4 files changed, 170 insertions, 2 deletions

diff --git a/packages/base/src/Internal/Algorithms.hs b/packages/base/src/Internal/Algorithms.hs
index 23a5e13..99c9e34 100644
--- a/packages/base/src/Internal/Algorithms.hs
+++ b/packages/base/src/Internal/Algorithms.hs
@@ -62,6 +62,7 @@ class (Numeric t,
     linearSolve' :: Matrix t -> Matrix t -> Matrix t
     cholSolve'   :: Matrix t -> Matrix t -> Matrix t
     triSolve'   :: UpLo -> Matrix t -> Matrix t -> Matrix t
+    triDiagSolve' :: Vector t -> Vector t -> Vector t -> Matrix t -> Matrix t
     ldlPacked'   :: Matrix t -> (Matrix t, [Int])
     ldlSolve'    :: (Matrix t, [Int]) -> Matrix t -> Matrix t
     linearSolveSVD' :: Matrix t -> Matrix t -> Matrix t
@@ -88,6 +89,7 @@ instance Field Double where
     mbLinearSolve' = mbLinearSolveR
     cholSolve' = cholSolveR
     triSolve' = triSolveR
+    triDiagSolve' = triDiagSolveR
     linearSolveLS' = linearSolveLSR
     linearSolveSVD' = linearSolveSVDR Nothing
     eig' = eigR
@@ -118,6 +120,7 @@ instance Field (Complex Double) where
     mbLinearSolve' = mbLinearSolveC
     cholSolve' = cholSolveC
     triSolve' = triSolveC
+    triDiagSolve' = triDiagSolveC
     linearSolveLS' = linearSolveLSC
     linearSolveSVD' = linearSolveSVDC Nothing
     eig' = eigC
@@ -356,10 +359,79 @@ cholSolve
     -> Matrix t -- ^ solution
 cholSolve = {-# SCC "cholSolve" #-} cholSolve'
 
--- | Solve a triangular linear system.
-triSolve :: Field t => UpLo -> Matrix t -> Matrix t -> Matrix t
+-- | Solve a triangular linear system. If `Upper` is specified then
+-- all elements below the diagonal are ignored; if `Lower` is
+-- specified then all elements above the diagonal are ignored.
+triSolve
+  :: Field t
+  => UpLo     -- ^ `Lower` or `Upper`
+  -> Matrix t -- ^ coefficient matrix
+  -> Matrix t -- ^ right hand sides
+  -> Matrix t -- ^ solution
 triSolve = {-# SCC "triSolve" #-} triSolve'
 
+-- | Solve a tridiagonal linear system. Suppose you wish to solve \(Ax = b\) where
+--
+-- \[
+-- A =
+-- \begin{bmatrix}
+--    1.0 & 4.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0
+-- \\ 3.0 & 1.0 & 4.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0
+-- \\ 0.0 & 3.0 & 1.0 & 4.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0
+-- \\ 0.0 & 0.0 & 3.0 & 1.0 & 4.0 & 0.0 & 0.0 & 0.0 & 0.0
+-- \\ 0.0 & 0.0 & 0.0 & 3.0 & 1.0 & 4.0 & 0.0 & 0.0 & 0.0
+-- \\ 0.0 & 0.0 & 0.0 & 0.0 & 3.0 & 1.0 & 4.0 & 0.0 & 0.0
+-- \\ 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 3.0 & 1.0 & 4.0 & 0.0
+-- \\ 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 3.0 & 1.0 & 4.0
+-- \\ 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 3.0 & 1.0
+-- \end{bmatrix}
+-- \quad
+-- b =
+-- \begin{bmatrix}
+--    1.0 &  1.0 &  1.0
+-- \\ 1.0 & -1.0 &  2.0
+-- \\ 1.0 &  1.0 &  3.0
+-- \\ 1.0 & -1.0 &  4.0
+-- \\ 1.0 &  1.0 &  5.0
+-- \\ 1.0 & -1.0 &  6.0
+-- \\ 1.0 &  1.0 &  7.0
+-- \\ 1.0 & -1.0 &  8.0
+-- \\ 1.0 &  1.0 &  9.0
+-- \end{bmatrix}
+-- \]
+--
+-- then
+--
+-- @
+-- dL =  fromList [3.0, 3.0, 3.0, 3.0, 3.0, 3.0, 3.0, 3.0]
+-- d  =  fromList [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
+-- dU =  fromList [4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0]
+--
+-- b = (9><3)
+--     [
+--       1.0,   1.0,   1.0,
+--       1.0,  -1.0,   2.0,
+--       1.0,   1.0,   3.0,
+--       1.0,  -1.0,   4.0,
+--       1.0,   1.0,   5.0,
+--       1.0,  -1.0,   6.0,
+--       1.0,   1.0,   7.0,
+--       1.0,  -1.0,   8.0,
+--       1.0,   1.0,   9.0
+--     ]
+--
+-- x = triDiagSolve dL d dU b
+-- @
+--
+triDiagSolve
+  :: Field t
+  => Vector t -- ^ lower diagonal: \(n - 1\) elements
+  -> Vector t -- ^ diagonal: \(n\) elements
+  -> Vector t -- ^ upper diagonal: \(n - 1\) elements
+  -> Matrix t -- ^ right hand sides
+  -> Matrix t -- ^ solution
+triDiagSolve = {-# SCC "triDiagSolve" #-} triDiagSolve'
+
 -- | Minimum norm solution of a general linear least squares problem Ax=B using the SVD. Admits rank-deficient systems but it is slower than 'linearSolveLS'. The effective rank of A is determined by treating as zero those singular valures which are less than 'eps' times the largest singular value.
 linearSolveSVD :: Field t => Matrix t -> Matrix t -> Matrix t
 linearSolveSVD = {-# SCC "linearSolveSVD" #-} linearSolveSVD'
diff --git a/packages/base/src/Internal/C/lapack-aux.c b/packages/base/src/Internal/C/lapack-aux.c
index 4a8129c..5018a45 100644
--- a/packages/base/src/Internal/C/lapack-aux.c
+++ b/packages/base/src/Internal/C/lapack-aux.c
@@ -668,6 +668,78 @@ int triSolveC_l_l(KOCMAT(a),OCMAT(b)) {
     OK
 }
 
+//////// tridiagonal real linear system ////////////
+
+int dgttrf_(integer *n,
+            doublereal *dl, doublereal *d, doublereal *du, doublereal *du2,
+            integer *ipiv,
+            integer *info);
+
+int dgttrs_(char *trans, integer *n, integer *nrhs,
+            doublereal *dl, doublereal *d, doublereal *du, doublereal *du2,
+            integer *ipiv, doublereal *b, integer *ldb,
+            integer *info);
+
+int triDiagSolveR_l(DVEC(dl), DVEC(d), DVEC(du), ODMAT(b)) {
+    integer n = dn;
+    integer nhrs = bc;
+    REQUIRES(n >= 1 && dln == dn - 1 && dun == dn - 1 && br == n, BAD_SIZE);
+    DEBUGMSG("triDiagSolveR_l");
+    integer res;
+    integer* ipiv = (integer*)malloc(n*sizeof(integer));
+    double* du2  = (double*)malloc((n - 2)*sizeof(double));
+    dgttrf_ (&n,
+             dlp, dp, dup, du2,
+             ipiv,
+             &res);
+    CHECK(res,res);
+    dgttrs_ ("N",
+             &n,&nhrs,
+             dlp, dp, dup, du2,
+             ipiv, bp, &n,
+             &res);
+    CHECK(res,res);
+    free(ipiv);
+    free(du2);
+    OK
+}
+
+//////// tridiagonal complex linear system ////////////
+
+int zgttrf_(integer *n,
+            doublecomplex *dl, doublecomplex *d, doublecomplex *du, doublecomplex *du2,
+            integer *ipiv,
+            integer *info);
+
+int zgttrs_(char *trans, integer *n, integer *nrhs,
+            doublecomplex *dl, doublecomplex *d, doublecomplex *du, doublecomplex *du2,
+            integer *ipiv, doublecomplex *b, integer *ldb,
+            integer *info);
+
+int triDiagSolveC_l(CVEC(dl), CVEC(d), CVEC(du), OCMAT(b)) {
+    integer n = dn;
+    integer nhrs = bc;
+    REQUIRES(n >= 1 && dln == dn - 1 && dun == dn - 1 && br == n, BAD_SIZE);
+    DEBUGMSG("triDiagSolveC_l");
+    integer res;
+    integer* ipiv = (integer*)malloc(n*sizeof(integer));
+    doublecomplex* du2 = (doublecomplex*)malloc((n - 2)*sizeof(doublecomplex));
+    zgttrf_ (&n,
+             dlp, dp, dup, du2,
+             ipiv,
+             &res);
+    CHECK(res,res);
+    zgttrs_ ("N",
+             &n,&nhrs,
+             dlp, dp, dup, du2,
+             ipiv, bp, &n,
+             &res);
+    CHECK(res,res);
+    free(ipiv);
+    free(du2);
+    OK
+}
+
 //////////////////// least squares real linear system ////////////
 
 int dgels_(char *trans, integer *m, integer *n, integer *
diff --git a/packages/base/src/Internal/LAPACK.hs b/packages/base/src/Internal/LAPACK.hs
index b4dd5cf..e306454 100644
--- a/packages/base/src/Internal/LAPACK.hs
+++ b/packages/base/src/Internal/LAPACK.hs
@@ -436,6 +436,28 @@ triSolveC :: UpLo -> Matrix (Complex Double) -> Matrix (Complex Double) -> Matri
 triSolveC Lower a b = linearSolveTRAux2 id ztrtrs_l "triSolveC" (fmat a) b
 triSolveC Upper a b = linearSolveTRAux2 id ztrtrs_u "triSolveC" (fmat a) b
 
+--------------------------------------------------------------------------------
+foreign import ccall unsafe "triDiagSolveR_l" dgttrs  :: R :> R :> R :> R ::> Ok
+foreign import ccall unsafe "triDiagSolveC_l" zgttrs  :: C :> C :> C :> C ::> Ok
+
+linearSolveGTAux2 g f st dl d du b
+    | ndl  == nd - 1 &&
+      ndu  == nd - 1 &&
+      nd   == r = unsafePerformIO . g $ do
+        s <- copy ColumnMajor b
+        (dl # d # du #! s) f #| st
+        return s
+    | otherwise = error $ st ++ " of nonsquare matrix"
+  where
+    ndl  = dim dl
+    nd   = dim d
+    ndu  = dim du
+    r    = rows b
+
+-- | Solves a tridiagonal system of linear equations.
+triDiagSolveR dl d du b = linearSolveGTAux2 id dgttrs "triDiagSolveR" dl d du b
+triDiagSolveC dl d du b = linearSolveGTAux2 id zgttrs "triDiagSolveC" dl d du b
+
 -----------------------------------------------------------------------------------
 
 foreign import ccall unsafe "linearSolveLSR_l"   dgels ::           R ::> R ::> Ok
diff --git a/packages/base/src/Numeric/LinearAlgebra.hs b/packages/base/src/Numeric/LinearAlgebra.hs
index 869330c..fd100e0 100644
--- a/packages/base/src/Numeric/LinearAlgebra.hs
+++ b/packages/base/src/Numeric/LinearAlgebra.hs
@@ -97,6 +97,8 @@ module Numeric.LinearAlgebra (
     -- ** Triangular
     UpLo(..),
     triSolve,
+    -- ** Tridiagonal
+    triDiagSolve,
     -- ** Sparse
     cgSolve,
     cgSolve',