improved luSolve', tests

4 files changed, 118 insertions, 70 deletions

diff --git a/packages/base/src/Internal/Modular.hs b/packages/base/src/Internal/Modular.hs
index d158111..1cae1ac 100644
--- a/packages/base/src/Internal/Modular.hs
+++ b/packages/base/src/Internal/Modular.hs
@@ -34,7 +34,9 @@ import Internal.Container
 import Internal.Vectorized (prodI,sumI,prodL,sumL)
 import Internal.LAPACK (multiplyI, multiplyL)
 import Internal.Algorithms(luFact)
-import Internal.Util(Normed(..),Indexable(..),gaussElim, gaussElim_1, gaussElim_2,luST, magnit)
+import Internal.Util(Normed(..),Indexable(..),
+                     gaussElim, gaussElim_1, gaussElim_2,
+                     luST, luSolve', luPacked', magnit)
 import Internal.ST(mutable)
 import GHC.TypeLits
 import Data.Proxy(Proxy)
@@ -350,7 +352,11 @@ test = (ok, info)
     lg = (3><3) (repeat (3*10^(9::Int))) :: Matrix Z
     lgm = fromZ lg :: Matrix (Mod 10000000000 Z)
 
-    gen n = diagRect 1 (konst 5 n) n n :: Numeric t => Matrix t
+    gen  n = diagRect 1 (konst 5 n) n n :: Numeric t => Matrix t
+    
+    rgen n = gen n :: Matrix R
+    cgen n = complex (rgen n) + fliprl (complex (rgen n)) * scalar (0:+1) :: Matrix C
+    sgen n = single (cgen n)
     
     checkGen x = norm_Inf $ flatten $ invg x <> x - ident (rows x)
     
@@ -360,6 +366,11 @@ test = (ok, info)
       where
         (l,u,p,_ :: Int) = luFact $ mutable (luST okf) t
 
+    checkSolve aa = norm_Inf $ flatten (aa <> x - bb)
+       where
+         bb = flipud aa
+         x = luSolve' (luPacked' aa) bb
+
     info = do
         print v
         print m
@@ -383,9 +394,9 @@ test = (ok, info)
         print $ lgm <> lgm
         
         print (checkGen (gen 5 :: Matrix R))
-        print (checkGen (gen 5 :: Matrix C))
         print (checkGen (gen 5 :: Matrix Float))
-        print (checkGen (gen 5 :: Matrix (Complex Float)))
+        print (checkGen (cgen 5 :: Matrix C))
+        print (checkGen (sgen 5 :: Matrix (Complex Float)))
         print (invg (gen 5) :: Matrix (Mod 7 I))
         print (invg (gen 5) :: Matrix (Mod 7 Z))
         
@@ -394,11 +405,18 @@ test = (ok, info)
 
         print $ checkLU (magnit 0) (gen 5 :: Matrix R)
         print $ checkLU (magnit 0) (gen 5 :: Matrix Float)
-        print $ checkLU (magnit 0) (gen 5 :: Matrix C)
-        print $ checkLU (magnit 0) (gen 5 :: Matrix (Complex Float))
+        print $ checkLU (magnit 0) (cgen 5 :: Matrix C)
+        print $ checkLU (magnit 0) (sgen 5 :: Matrix (Complex Float))
         print $ checkLU (magnit 0) (gen 5 :: Matrix (Mod 7 I))
         print $ checkLU (magnit 0) (gen 5 :: Matrix (Mod 7 Z))
 
+        print $ checkSolve (gen 5 :: Matrix R)
+        print $ checkSolve (gen 5 :: Matrix Float)
+        print $ checkSolve (cgen 5 :: Matrix C)
+        print $ checkSolve (sgen 5 :: Matrix (Complex Float))
+        print $ checkSolve (gen 5 :: Matrix (Mod 7 I))
+        print $ checkSolve (gen 5 :: Matrix (Mod 7 Z))
+
 
     ok = and
       [ toInt (m #> v) == cmod 11 (toInt m #> toInt v )
@@ -407,16 +425,22 @@ test = (ok, info)
       , am <> gaussElim   am bm == bm
       , (checkGen (gen 5 :: Matrix R)) < 1E-15
       , (checkGen (gen 5 :: Matrix Float)) < 2E-7
-      , (checkGen (gen 5 :: Matrix C)) < 1E-15
-      , (checkGen (gen 5 :: Matrix (Complex Float))) < 2E-7
+      , (checkGen (cgen 5 :: Matrix C)) < 1E-15
+      , (checkGen (sgen 5 :: Matrix (Complex Float))) < 2E-7
       , (checkGen (gen 5 :: Matrix (Mod 7 I))) == 0
       , (checkGen (gen 5 :: Matrix (Mod 7 Z))) == 0
-      , (checkLU (magnit 1E-10) (gen 5 :: Matrix R)) < 1E-15
+      , (checkLU (magnit 1E-10) (gen 5 :: Matrix R)) < 2E-15
       , (checkLU (magnit 1E-5) (gen 5 :: Matrix Float)) < 1E-6
-      , (checkLU (magnit 1E-10) (gen 5 :: Matrix C)) < 1E-15
-      , (checkLU (magnit 1E-5) (gen 5 :: Matrix (Complex Float))) < 1E-6
+      , (checkLU (magnit 1E-10) (cgen 5 :: Matrix C)) < 5E-15
+      , (checkLU (magnit 1E-5) (sgen 5 :: Matrix (Complex Float))) < 1E-6
       , (checkLU (magnit 0) (gen 5 :: Matrix (Mod 7 I))) == 0
       , (checkLU (magnit 0) (gen 5 :: Matrix (Mod 7 Z))) == 0
+      , checkSolve (gen 5 :: Matrix R) < 2E-15
+      , checkSolve (gen 5 :: Matrix Float) < 1E-6
+      , checkSolve (cgen 5 :: Matrix C) < 4E-15
+      , checkSolve (sgen 5 :: Matrix (Complex Float)) < 1E-6
+      , checkSolve (gen 5 :: Matrix (Mod 7 I)) == 0
+      , checkSolve (gen 5 :: Matrix (Mod 7 Z)) == 0
       , prodElements (konst (9:: Mod 10 I) (12::Int)) == product (replicate 12 (9:: Mod 10 I))
       , gm <> gm == konst 0 (3,3)
       , lgm <> lgm == konst 0 (3,3)
diff --git a/packages/base/src/Internal/ST.hs b/packages/base/src/Internal/ST.hs
index 25e7f03..d1defda 100644
--- a/packages/base/src/Internal/ST.hs
+++ b/packages/base/src/Internal/ST.hs
@@ -228,10 +228,12 @@ extractMatrix (STMatrix m) rr rc = unsafeIOToST (extractR m 0 (idxs[i1,i2]) 0 (i
     (i1,i2) = getRowRange (rows m) rr
     (j1,j2) = getColRange (cols m) rc
 
+-- | r0 c0 height width
 data Slice s t = Slice (STMatrix s t) Int Int Int Int
 
 slice (Slice (STMatrix m) r0 c0 nr nc) = (m, idxs[r0,r0+nr-1,c0,c0+nc-1])
 
+gemmm :: Element t => t -> Slice s t -> t -> Slice s t -> Slice s t -> ST s ()
 gemmm beta (slice->(r,pr)) alpha (slice->(a,pa)) (slice->(b,pb)) = res
   where
     res = unsafeIOToST (gemm u v a b r)
diff --git a/packages/base/src/Internal/Util.hs b/packages/base/src/Internal/Util.hs
index d9777ae..079663d 100644
--- a/packages/base/src/Internal/Util.hs
+++ b/packages/base/src/Internal/Util.hs
@@ -54,7 +54,7 @@ module Internal.Util(
     -- ** 2D
     corr2, conv2, separable,
     block2x2,block3x3,view1,unView1,foldMatrix,
-    gaussElim_1, gaussElim_2, gaussElim, luST, luSolve'
+    gaussElim_1, gaussElim_2, gaussElim, luST, luSolve', luSolve'', luPacked'
 ) where
 
 import Internal.Vector
@@ -64,7 +64,7 @@ import Internal.Element
 import Internal.Container
 import Internal.Vectorized
 import Internal.IO
-import Internal.Algorithms hiding (Normed,linearSolve',luSolve')
+import Internal.Algorithms hiding (Normed,linearSolve',luSolve', luPacked')
 import Numeric.Matrix()
 import Numeric.Vector()
 import Internal.Random
@@ -686,6 +686,35 @@ luST ok (r,_) x = do
     v <- unsafeFreezeVector p
     return (toList v)
 
+{- | Experimental implementation of 'luPacked'
+     for any Fractional element type, including 'Mod' n 'I' and 'Mod' n 'Z'.
+
+>>> let m = ident 5 + (5><5) [0..] :: Matrix (Z ./. 17)
+(5><5)
+ [  1,  1,  2,  3,  4
+ ,  5,  7,  7,  8,  9
+ , 10, 11, 13, 13, 14
+ , 15, 16,  0,  2,  2
+ ,  3,  4,  5,  6,  8 ]
+
+>>> let (l,u,p,s) = luFact $ luPacked' m
+>>> l
+(5><5)
+ [  1,  0, 0,  0, 0
+ ,  6,  1, 0,  0, 0
+ , 12,  7, 1,  0, 0
+ ,  7, 10, 7,  1, 0
+ ,  8,  2, 6, 11, 1 ]
+>>> u
+(5><5)
+ [ 15, 16,  0,  2,  2
+ ,  0, 13,  7, 13, 14
+ ,  0,  0, 15,  0, 11
+ ,  0,  0,  0, 15, 15
+ ,  0,  0,  0,  0,  1 ]
+
+-}
+luPacked' x = mutable (luST (magnit 0)) x
 
 --------------------------------------------------------------------------------
 
@@ -693,35 +722,79 @@ rowRange m = [0..rows m -1]
 
 at k = Pos (idxs[k])
 
-backSust lup rhs = foldl' f (rhs?[]) (reverse ls)
+backSust' lup rhs = foldl' f (rhs?[]) (reverse ls)
   where
     ls  = [ (d k , u k , b k) | k <- rowRange lup ]
       where
         d k = lup ?? (at k, at k)
         u k = lup ?? (at k, Drop (k+1))
         b k = rhs ?? (at k, All)
-    
+
     f x (d,u,b) = (b - u<>x) / d
                        ===
                         x
 
 
-forwSust lup rhs = foldl' f (rhs?[]) ls
+forwSust' lup rhs = foldl' f (rhs?[]) ls
   where
     ls  = [ (l k , b k) | k <- rowRange lup ]
       where
         l k = lup ?? (at k, Take k)
         b k = rhs ?? (at k, All)
-    
+
     f x (l,b) =     x
                    ===
                 (b - l<>x)
 
 
-luSolve' (lup,p) b = backSust lup (forwSust lup pb)
+luSolve'' (lup,p) b = backSust' lup (forwSust' lup pb)
   where
     pb = b ?? (Pos (fixPerm' p), All)
 
+--------------------------------------------------------------------------------
+
+forwSust lup rhs = fst $ mutable f rhs
+  where
+    f (r,c) x = do
+        l <- unsafeThawMatrix lup
+        let go k = gemmm 1 (Slice x k 0 1 c) (-1) (Slice l k 0 1 k) (Slice x 0 0 k c)
+        mapM_ go [0..r-1]
+
+
+backSust lup rhs = fst $ mutable f rhs
+  where
+    f (r,c) m = do
+        l <- unsafeThawMatrix lup
+        let d k = recip (lup `atIndex` (k,k))
+            u k = Slice l k (k+1) 1 (r-1-k)
+            b k = Slice m k 0 1 c
+            x k = Slice m (k+1) 0 (r-1-k) c
+            scal k = rowOper (SCAL (d k) (Row k) AllCols) m
+
+            go k = gemmm 1 (b k) (-1) (u k) (x k) >> scal k
+        mapM_ go [r-1,r-2..0]
+
+
+{- | Experimental implementation of 'luSolve' for any Fractional element type, including 'Mod' n 'I' and 'Mod' n 'Z'.
+
+>>> let a = (2><2) [1,2,3,5] :: Matrix (Z ./. 13)
+(2><2)
+ [ 1, 2
+ , 3, 5 ]
+>>> b
+(2><3)
+ [ 5, 1, 3
+ , 8, 6, 3 ]
+
+>>> luSolve' (luPacked' a) b
+(2><3)
+ [ 4,  7, 4
+ , 7, 10, 6 ]
+
+-}
+luSolve' (lup,p) b = backSust lup (forwSust lup pb)
+  where
+    pb = b ?? (Pos (fixPerm' p), All)
 
 --------------------------------------------------------------------------------
 
diff --git a/packages/base/src/Numeric/LinearAlgebra.hs b/packages/base/src/Numeric/LinearAlgebra.hs
index e899445..0b8abbb 100644
--- a/packages/base/src/Numeric/LinearAlgebra.hs
+++ b/packages/base/src/Numeric/LinearAlgebra.hs
@@ -81,7 +81,6 @@ module Numeric.LinearAlgebra (
     cholSolve,
     cgSolve,
     cgSolve',
-    linearSolve',
 
     -- * Inverse and pseudoinverse
     inv, pinv, pinvTol,
@@ -123,7 +122,7 @@ module Numeric.LinearAlgebra (
     schur,
 
     -- * LU
-    lu, luPacked, luFact, luPacked',
+    lu, luPacked, luPacked', luFact,
 
     -- * Matrix functions
     expm,
@@ -166,7 +165,6 @@ import Internal.Random
 import Internal.Sparse((!#>))
 import Internal.CG
 import Internal.Conversion
-import Internal.ST(mutable)
 
 {- | infix synonym of 'mul'
 
@@ -241,53 +239,4 @@ nullspace m = nullspaceSVD (Left (1*eps)) m (rightSV m)
 -- | return an orthonormal basis of the range space of a matrix. See also 'orthSVD'.
 orth m = orthSVD (Left (1*eps)) m (leftSV m)
 
-{- | Experimental implementation of LU factorization
-     working on any Fractional element type, including 'Mod' n 'I' and 'Mod' n 'Z'.
-
->>> let m = ident 5 + (5><5) [0..] :: Matrix (Z ./. 17)
-(5><5)
- [  1,  1,  2,  3,  4
- ,  5,  7,  7,  8,  9
- , 10, 11, 13, 13, 14
- , 15, 16,  0,  2,  2
- ,  3,  4,  5,  6,  8 ]
-
->>> let (l,u,p,s) = luFact $ luPacked' m
->>> l
-(5><5)
- [  1,  0, 0,  0, 0
- ,  6,  1, 0,  0, 0
- , 12,  7, 1,  0, 0
- ,  7, 10, 7,  1, 0
- ,  8,  2, 6, 11, 1 ]
->>> u
-(5><5)
- [ 15, 16,  0,  2,  2
- ,  0, 13,  7, 13, 14
- ,  0,  0, 15,  0, 11
- ,  0,  0,  0, 15, 15
- ,  0,  0,  0,  0,  1 ]
-
--}
-luPacked' x = mutable (luST (magnit 0)) x
-
-{- | Experimental implementation of gaussian elimination
-     working on any Fractional element type, including 'Mod' n 'I' and 'Mod' n 'Z'.
-
->>> let a = (2><2) [1,2,3,5] :: Matrix (Z ./. 13)
-(2><2)
- [ 1, 2
- , 3, 5 ]
->>> b
-(2><3)
- [ 5, 1, 3
- , 8, 6, 3 ]
-
->>> let x = linearSolve' a b
-(2><3)
- [ 4,  7, 4
- , 7, 10, 6 ]
-
--}
-linearSolve' x y = gaussElim x y