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 ea6bc5772422c620daa3057c0177a946442fc690..5562e92b7ef6101a754fbcc505eb2bf6bc22c37f 100644 (file)
--- a/src/Linear/Matrix.hs
+++ b/src/Linear/Matrix.hs
@@ -42,8 +42,7 @@ import qualified Data.Vector.Fixed as V (
   maximum,
   replicate,
   toList,
-  zipWith
-  )
+  zipWith )
 import Data.Vector.Fixed.Cont ( Arity, arity )
 import Linear.Vector ( Vec, delete, element_sum )
 import Normed ( Normed(..) )
@@ -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