From c7b6b27f4304416dcec67519b710d090850c3caa Mon Sep 17 00:00:00 2001 From: Michael Orlitzky 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.44.2