Small cleanups in Linear.Vector.

[numerical-analysis.git] / src / Linear / Matrix.hs

diff --git a/src/Linear/Matrix.hs b/src/Linear/Matrix.hs
index c6f4a83c71af1c5231113ec555fa79f2dcfcc038..5562e92b7ef6101a754fbcc505eb2bf6bc22c37f 100644 (file)
--- a/src/Linear/Matrix.hs
+++ b/src/Linear/Matrix.hs
@@ -42,25 +42,24 @@ import qualified Data.Vector.Fixed as V (
   maximum,
   replicate,
   toList,
-  zipWith
-  )
-import Data.Vector.Fixed.Cont (Arity, arity)
-import Linear.Vector
-import Normed
+  zipWith )
+import Data.Vector.Fixed.Cont ( Arity, arity )
+import Linear.Vector ( Vec, delete, element_sum )
+import Normed ( Normed(..) )
 
 import NumericPrelude hiding ( (*), abs )
 import qualified NumericPrelude as NP ( (*) )
 import qualified Algebra.Absolute as Absolute ( C )
 import Algebra.Absolute ( abs )
-import qualified Algebra.Additive as Additive
-import qualified Algebra.Algebraic as Algebraic
-import Algebra.Algebraic (root)
-import qualified Algebra.Ring as Ring
-import qualified Algebra.Module as Module
-import qualified Algebra.RealRing as RealRing
-import qualified Algebra.ToRational as ToRational
-import qualified Algebra.Transcendental as Transcendental
-import qualified Prelude as P
+import qualified Algebra.Additive as Additive ( C )
+import qualified Algebra.Algebraic as Algebraic ( C )
+import Algebra.Algebraic ( root )
+import qualified Algebra.Ring as Ring ( C )
+import qualified Algebra.Module as Module ( C )
+import qualified Algebra.RealRing as RealRing ( C )
+import qualified Algebra.ToRational as ToRational ( C )
+import qualified Algebra.Transcendental as Transcendental ( C )
+import qualified Prelude as P ( map )
 
 data Mat m n a = (Arity m, Arity n) => Mat (Vec m (Vec n a))
 type Mat1 a = Mat N1 N1 a
@@ -608,6 +607,23 @@ angle v1 v2 =
    norms = (norm v1) NP.* (norm v2)
 
 
+-- | Retrieve the diagonal elements of the given matrix as a \"column
+--   vector,\" i.e. a m-by-1 matrix. We require the matrix to be
+--   square to avoid ambiguity in the return type which would ideally
+--   have dimension min(m,n) supposing an m-by-n matrix.
+--
+--   Examples:
+--
+--   >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int
+--   >>> diagonal m
+--   ((1),(5),(9))
+--
+diagonal :: (Arity m) => Mat m m a -> Mat m N1 a
+diagonal matrix =
+  construct lambda
+  where
+    lambda i _ = matrix !!! (i,i)
+
 
 -- | Given a square @matrix@, return a new matrix of the same size
 --   containing only the on-diagonal entries of @matrix@. The
@@ -696,3 +712,21 @@ ut_part_strict :: (Arity m, Ring.C a)
         => Mat m m a
         -> Mat m m a
 ut_part_strict = transpose . lt_part_strict . transpose
+
+
+-- | Compute the trace of a square matrix, the sum of the elements
+--   which lie on its diagonal. We require the matrix to be
+--   square to avoid ambiguity in the return type which would ideally
+--   have dimension min(m,n) supposing an m-by-n matrix.
+--
+--   Examples:
+--
+--   >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int
+--   >>> trace m
+--   15
+--
+trace :: (Arity m, Ring.C a) => Mat m m a -> a
+trace matrix =
+  let (Mat rows) = diagonal matrix
+  in
+    element_sum $ V.map V.head rows