From c7b6b27f4304416dcec67519b710d090850c3caa Mon Sep 17 00:00:00 2001
From: Michael Orlitzky <michael@orlitzky.com>
Date: Mon, 3 Feb 2014 01:03:14 -0500
Subject: [PATCH] Small cleanups in Linear.Vector. Add 'diagonal' and 'trace'
 functions to Linear.Matrix.

---
 src/Linear/Matrix.hs | 38 ++++++++++++++++++++++++++++++++++++--
 src/Linear/Vector.hs | 14 ++++++++------
 2 files changed, 44 insertions(+), 8 deletions(-)

diff --git a/src/Linear/Matrix.hs b/src/Linear/Matrix.hs
index ea6bc57..5562e92 100644
--- 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
diff --git a/src/Linear/Vector.hs b/src/Linear/Vector.hs
index 9d43bfa..56bc2af 100644
--- a/src/Linear/Vector.hs
+++ b/src/Linear/Vector.hs
@@ -2,6 +2,7 @@
 {-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE NoImplicitPrelude #-}
+{-# LANGUAGE NoMonomorphismRestriction #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TypeFamilies #-}
 
@@ -22,7 +23,7 @@ import Data.Vector.Fixed (
   Vector(..),
   fromList,
   toList )
-import qualified Data.Vector.Fixed as V (
+import Data.Vector.Fixed (
   (!),
   foldl,
   length )
@@ -32,7 +33,8 @@ import Data.Vector.Fixed.Boxed (
   Vec3,
   Vec4,
   Vec5 )
-import NumericPrelude hiding ( abs )
+import NumericPrelude hiding ( abs, length, foldl )
+
 
 type Vec1 = Vec N1
 
@@ -51,8 +53,8 @@ type Vec1 = Vec N1
 --
 (!?) :: (Vector v a) => v a -> Int -> Maybe a
 (!?) v1 idx
-  | idx < 0 || idx >= V.length v1 = Nothing
-  | otherwise                     = Just $ v1 V.! idx
+  | idx < 0 || idx >= length v1 = Nothing
+  | otherwise                   = Just $ v1 ! idx
 
 
 -- | Remove an element of the given vector.
@@ -77,7 +79,7 @@ delete v1 idx =
     rhalf' = tail rhalf
 
 
--- | We provide our own sum because V.sum relies on a Num instance
+-- | We provide our own sum because sum relies on a Num instance
 --   from the Prelude that we don't have.
 --
 --   Examples:
@@ -88,4 +90,4 @@ delete v1 idx =
 --   6
 --
 element_sum :: (Additive.C a, Ring.C a, Vector v a) => v a -> a
-element_sum = V.foldl (+) (fromInteger 0)
+element_sum = foldl (+) (fromInteger 0)
-- 
2.43.2