]> gitweb.michael.orlitzky.com - numerical-analysis.git/blob - src/Linear/Matrix.hs
Add the imap2 function to Linear.Matrix.
[numerical-analysis.git] / src / Linear / Matrix.hs
1 {-# LANGUAGE ExistentialQuantification #-}
2 {-# LANGUAGE FlexibleContexts #-}
3 {-# LANGUAGE FlexibleInstances #-}
4 {-# LANGUAGE MultiParamTypeClasses #-}
5 {-# LANGUAGE NoMonomorphismRestriction #-}
6 {-# LANGUAGE ScopedTypeVariables #-}
7 {-# LANGUAGE TypeFamilies #-}
8 {-# LANGUAGE RebindableSyntax #-}
9
10 -- | Boxed matrices; that is, boxed m-vectors of boxed n-vectors. We
11 -- assume that the underlying representation is
12 -- Data.Vector.Fixed.Boxed.Vec for simplicity. It was tried in
13 -- generality and failed.
14 --
15 module Linear.Matrix
16 where
17
18 import Data.List (intercalate)
19
20 import Data.Vector.Fixed (
21 (!),
22 N1,
23 N2,
24 N3,
25 N4,
26 N5,
27 S,
28 Z,
29 generate,
30 mk1,
31 mk2,
32 mk3,
33 mk4,
34 mk5
35 )
36 import qualified Data.Vector.Fixed as V (
37 and,
38 fromList,
39 head,
40 ifoldl,
41 imap,
42 map,
43 maximum,
44 replicate,
45 toList,
46 zipWith )
47 import Data.Vector.Fixed.Cont ( Arity, arity )
48 import Linear.Vector ( Vec, delete, element_sum )
49 import Normed ( Normed(..) )
50
51 import NumericPrelude hiding ( (*), abs )
52 import qualified NumericPrelude as NP ( (*) )
53 import qualified Algebra.Absolute as Absolute ( C )
54 import Algebra.Absolute ( abs )
55 import qualified Algebra.Additive as Additive ( C )
56 import qualified Algebra.Algebraic as Algebraic ( C )
57 import Algebra.Algebraic ( root )
58 import qualified Algebra.Ring as Ring ( C )
59 import qualified Algebra.Module as Module ( C )
60 import qualified Algebra.RealRing as RealRing ( C )
61 import qualified Algebra.ToRational as ToRational ( C )
62 import qualified Algebra.Transcendental as Transcendental ( C )
63 import qualified Prelude as P ( map )
64
65 -- | Our main matrix type.
66 data Mat m n a = (Arity m, Arity n) => Mat (Vec m (Vec n a))
67
68 -- Type synonyms for n-by-n matrices.
69 type Mat1 a = Mat N1 N1 a
70 type Mat2 a = Mat N2 N2 a
71 type Mat3 a = Mat N3 N3 a
72 type Mat4 a = Mat N4 N4 a
73 type Mat5 a = Mat N5 N5 a
74
75 -- | Type synonym for row vectors expressed as 1-by-n matrices.
76 type Row n a = Mat N1 n a
77
78 -- Type synonyms for 1-by-n row "vectors".
79 type Row1 a = Row N1 a
80 type Row2 a = Row N2 a
81 type Row3 a = Row N3 a
82 type Row4 a = Row N4 a
83 type Row5 a = Row N5 a
84
85 -- | Type synonym for column vectors expressed as n-by-1 matrices.
86 type Col n a = Mat n N1 a
87
88 -- Type synonyms for n-by-1 column "vectors".
89 type Col1 a = Col N1 a
90 type Col2 a = Col N2 a
91 type Col3 a = Col N3 a
92 type Col4 a = Col N4 a
93 type Col5 a = Col N5 a
94
95 -- We need a big column for Gaussian quadrature.
96 type N10 = S (S (S (S (S N5))))
97 type Col10 a = Col N10 a
98
99
100 instance (Eq a) => Eq (Mat m n a) where
101 -- | Compare a row at a time.
102 --
103 -- Examples:
104 --
105 -- >>> let m1 = fromList [[1,2],[3,4]] :: Mat2 Int
106 -- >>> let m2 = fromList [[1,2],[3,4]] :: Mat2 Int
107 -- >>> let m3 = fromList [[5,6],[7,8]] :: Mat2 Int
108 -- >>> m1 == m2
109 -- True
110 -- >>> m1 == m3
111 -- False
112 --
113 (Mat rows1) == (Mat rows2) =
114 V.and $ V.zipWith comp rows1 rows2
115 where
116 -- Compare a row, one column at a time.
117 comp row1 row2 = V.and (V.zipWith (==) row1 row2)
118
119
120 instance (Show a) => Show (Mat m n a) where
121 -- | Display matrices and vectors as ordinary tuples. This is poor
122 -- practice, but these results are primarily displayed
123 -- interactively and convenience trumps correctness (said the guy
124 -- who insists his vector lengths be statically checked at
125 -- compile-time).
126 --
127 -- Examples:
128 --
129 -- >>> let m = fromList [[1,2],[3,4]] :: Mat2 Int
130 -- >>> show m
131 -- ((1,2),(3,4))
132 --
133 show (Mat rows) =
134 "(" ++ (intercalate "," (V.toList row_strings)) ++ ")"
135 where
136 row_strings = V.map show_vector rows
137 show_vector v1 =
138 "(" ++ (intercalate "," element_strings) ++ ")"
139 where
140 v1l = V.toList v1
141 element_strings = P.map show v1l
142
143
144 -- | Convert a matrix to a nested list.
145 toList :: Mat m n a -> [[a]]
146 toList (Mat rows) = map V.toList (V.toList rows)
147
148
149 -- | Create a matrix from a nested list.
150 fromList :: (Arity m, Arity n) => [[a]] -> Mat m n a
151 fromList vs = Mat (V.fromList $ map V.fromList vs)
152
153
154 -- | Unsafe indexing. Much faster than the safe indexing.
155 (!!!) :: (Arity m, Arity n) => Mat m n a -> (Int, Int) -> a
156 (!!!) (Mat rows) (i, j) = (rows ! i) ! j
157
158
159 -- | Safe indexing.
160 --
161 -- Examples:
162 --
163 -- >>> let m = fromList [[1,2],[3,4]] :: Mat2 Int
164 -- >>> m !!? (-1,-1)
165 -- Nothing
166 -- >>> m !!? (-1,0)
167 -- Nothing
168 -- >>> m !!? (-1,1)
169 -- Nothing
170 -- >>> m !!? (0,-1)
171 -- Nothing
172 -- >>> m !!? (0,0)
173 -- Just 1
174 -- >>> m !!? (0,1)
175 -- Just 2
176 -- >>> m !!? (1,-1)
177 -- Nothing
178 -- >>> m !!? (1,0)
179 -- Just 3
180 -- >>> m !!? (1,1)
181 -- Just 4
182 -- >>> m !!? (2,-1)
183 -- Nothing
184 -- >>> m !!? (2,0)
185 -- Nothing
186 -- >>> m !!? (2,1)
187 -- Nothing
188 -- >>> m !!? (2,2)
189 -- Nothing
190 --
191 (!!?) :: (Arity m, Arity n) => Mat m n a -> (Int, Int) -> Maybe a
192 (!!?) matrix idx =
193 ifoldl2 f Nothing matrix
194 where
195 f k l found cur = if (k,l) == idx then (Just cur) else found
196
197
198 -- | The number of rows in the matrix.
199 nrows :: forall m n a. (Arity m) => Mat m n a -> Int
200 nrows _ = arity (undefined :: m)
201
202
203 -- | The number of columns in the first row of the
204 -- matrix. Implementation stolen from Data.Vector.Fixed.length.
205 ncols :: forall m n a. (Arity n) => Mat m n a -> Int
206 ncols _ = arity (undefined :: n)
207
208
209 -- | Return the @i@th row of @m@ as a matrix. Unsafe.
210 row :: (Arity m, Arity n) => Mat m n a -> Int -> Row n a
211 row m i =
212 construct lambda
213 where
214 lambda _ j = m !!! (i, j)
215
216
217 -- | Return the @j@th column of @m@ as a matrix. Unsafe.
218 column :: (Arity m, Arity n) => Mat m n a -> Int -> Col m a
219 column m j =
220 construct lambda
221 where
222 lambda i _ = m !!! (i, j)
223
224
225 -- | Transpose @m@; switch it's columns and its rows. This is a dirty
226 -- implementation, but I don't see a better way.
227 --
228 -- TODO: Don't cheat with fromList.
229 --
230 -- Examples:
231 --
232 -- >>> let m = fromList [[1,2], [3,4]] :: Mat2 Int
233 -- >>> transpose m
234 -- ((1,3),(2,4))
235 --
236 transpose :: (Arity m, Arity n) => Mat m n a -> Mat n m a
237 transpose matrix =
238 construct lambda
239 where
240 lambda i j = matrix !!! (j,i)
241
242
243 -- | Is @m@ symmetric?
244 --
245 -- Examples:
246 --
247 -- >>> let m1 = fromList [[1,2], [2,1]] :: Mat2 Int
248 -- >>> symmetric m1
249 -- True
250 --
251 -- >>> let m2 = fromList [[1,2], [3,1]] :: Mat2 Int
252 -- >>> symmetric m2
253 -- False
254 --
255 symmetric :: (Eq a, Arity m) => Mat m m a -> Bool
256 symmetric m =
257 m == (transpose m)
258
259
260 -- | Construct a new matrix from a function @lambda@. The function
261 -- @lambda@ should take two parameters i,j corresponding to the
262 -- entries in the matrix. The i,j entry of the resulting matrix will
263 -- have the value returned by lambda i j.
264 --
265 -- Examples:
266 --
267 -- >>> let lambda i j = i + j
268 -- >>> construct lambda :: Mat3 Int
269 -- ((0,1,2),(1,2,3),(2,3,4))
270 --
271 construct :: forall m n a. (Arity m, Arity n)
272 => (Int -> Int -> a) -> Mat m n a
273 construct lambda = Mat $ generate make_row
274 where
275 make_row :: Int -> Vec n a
276 make_row i = generate (lambda i)
277
278
279 -- | Create an identity matrix with the right dimensions.
280 --
281 -- Examples:
282 --
283 -- >>> identity_matrix :: Mat3 Int
284 -- ((1,0,0),(0,1,0),(0,0,1))
285 -- >>> identity_matrix :: Mat3 Double
286 -- ((1.0,0.0,0.0),(0.0,1.0,0.0),(0.0,0.0,1.0))
287 --
288 identity_matrix :: (Arity m, Ring.C a) => Mat m m a
289 identity_matrix =
290 construct (\i j -> if i == j then (fromInteger 1) else (fromInteger 0))
291
292
293 -- | Given a positive-definite matrix @m@, computes the
294 -- upper-triangular matrix @r@ with (transpose r)*r == m and all
295 -- values on the diagonal of @r@ positive.
296 --
297 -- Examples:
298 --
299 -- >>> let m1 = fromList [[20,-1], [-1,20]] :: Mat2 Double
300 -- >>> cholesky m1
301 -- ((4.47213595499958,-0.22360679774997896),(0.0,4.466542286825459))
302 -- >>> (transpose (cholesky m1)) * (cholesky m1)
303 -- ((20.000000000000004,-1.0),(-1.0,20.0))
304 --
305 cholesky :: forall m n a. (Algebraic.C a, Arity m, Arity n)
306 => (Mat m n a) -> (Mat m n a)
307 cholesky m = construct r
308 where
309 r :: Int -> Int -> a
310 r i j | i == j = sqrt(m !!! (i,j) - sum [(r k i)^2 | k <- [0..i-1]])
311 | i < j =
312 (((m !!! (i,j)) - sum [(r k i) NP.* (r k j) | k <- [0..i-1]]))/(r i i)
313 | otherwise = 0
314
315
316 -- | Returns True if the given matrix is upper-triangular, and False
317 -- otherwise. The parameter @epsilon@ lets the caller choose a
318 -- tolerance.
319 --
320 -- Examples:
321 --
322 -- >>> let m = fromList [[1,1],[1e-12,1]] :: Mat2 Double
323 -- >>> is_upper_triangular m
324 -- False
325 -- >>> is_upper_triangular' 1e-10 m
326 -- True
327 --
328 -- TODO:
329 --
330 -- 1. Don't cheat with lists.
331 --
332 is_upper_triangular' :: (Ord a, Ring.C a, Absolute.C a, Arity m, Arity n)
333 => a -- ^ The tolerance @epsilon@.
334 -> Mat m n a
335 -> Bool
336 is_upper_triangular' epsilon m =
337 and $ concat results
338 where
339 results = [[ test i j | i <- [0..(nrows m)-1]] | j <- [0..(ncols m)-1] ]
340
341 test :: Int -> Int -> Bool
342 test i j
343 | i <= j = True
344 -- use "less than or equal to" so zero is a valid epsilon
345 | otherwise = abs (m !!! (i,j)) <= epsilon
346
347
348 -- | Returns True if the given matrix is upper-triangular, and False
349 -- otherwise. A specialized version of 'is_upper_triangular\'' with
350 -- @epsilon = 0@.
351 --
352 -- Examples:
353 --
354 -- >>> let m = fromList [[1,0],[1,1]] :: Mat2 Int
355 -- >>> is_upper_triangular m
356 -- False
357 --
358 -- >>> let m = fromList [[1,2],[0,3]] :: Mat2 Int
359 -- >>> is_upper_triangular m
360 -- True
361 --
362 -- TODO:
363 --
364 -- 1. The Ord constraint is too strong here, Eq would suffice.
365 --
366 is_upper_triangular :: (Ord a, Ring.C a, Absolute.C a, Arity m, Arity n)
367 => Mat m n a -> Bool
368 is_upper_triangular = is_upper_triangular' 0
369
370
371 -- | Returns True if the given matrix is lower-triangular, and False
372 -- otherwise. This is a specialized version of 'is_lower_triangular\''
373 -- with @epsilon = 0@.
374 --
375 -- Examples:
376 --
377 -- >>> let m = fromList [[1,0],[1,1]] :: Mat2 Int
378 -- >>> is_lower_triangular m
379 -- True
380 --
381 -- >>> let m = fromList [[1,2],[0,3]] :: Mat2 Int
382 -- >>> is_lower_triangular m
383 -- False
384 --
385 is_lower_triangular :: (Ord a,
386 Ring.C a,
387 Absolute.C a,
388 Arity m,
389 Arity n)
390 => Mat m n a
391 -> Bool
392 is_lower_triangular = is_upper_triangular . transpose
393
394
395 -- | Returns True if the given matrix is lower-triangular, and False
396 -- otherwise. The parameter @epsilon@ lets the caller choose a
397 -- tolerance.
398 --
399 -- Examples:
400 --
401 -- >>> let m = fromList [[1,1e-12],[1,1]] :: Mat2 Double
402 -- >>> is_lower_triangular m
403 -- False
404 -- >>> is_lower_triangular' 1e-12 m
405 -- True
406 --
407 is_lower_triangular' :: (Ord a,
408 Ring.C a,
409 Absolute.C a,
410 Arity m,
411 Arity n)
412 => a -- ^ The tolerance @epsilon@.
413 -> Mat m n a
414 -> Bool
415 is_lower_triangular' epsilon = (is_upper_triangular' epsilon) . transpose
416
417
418 -- | Returns True if the given matrix is triangular, and False
419 -- otherwise.
420 --
421 -- Examples:
422 --
423 -- >>> let m = fromList [[1,0],[1,1]] :: Mat2 Int
424 -- >>> is_triangular m
425 -- True
426 --
427 -- >>> let m = fromList [[1,2],[0,3]] :: Mat2 Int
428 -- >>> is_triangular m
429 -- True
430 --
431 -- >>> let m = fromList [[1,2],[3,4]] :: Mat2 Int
432 -- >>> is_triangular m
433 -- False
434 --
435 is_triangular :: (Ord a,
436 Ring.C a,
437 Absolute.C a,
438 Arity m,
439 Arity n)
440 => Mat m n a
441 -> Bool
442 is_triangular m = is_upper_triangular m || is_lower_triangular m
443
444
445 -- | Return the (i,j)th minor of m.
446 --
447 -- Examples:
448 --
449 -- >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int
450 -- >>> minor m 0 0 :: Mat2 Int
451 -- ((5,6),(8,9))
452 -- >>> minor m 1 1 :: Mat2 Int
453 -- ((1,3),(7,9))
454 --
455 minor :: (m ~ S r,
456 n ~ S t,
457 Arity r,
458 Arity t)
459 => Mat m n a
460 -> Int
461 -> Int
462 -> Mat r t a
463 minor (Mat rows) i j = m
464 where
465 rows' = delete rows i
466 m = Mat $ V.map ((flip delete) j) rows'
467
468
469 class (Eq a, Ring.C a) => Determined p a where
470 determinant :: (p a) -> a
471
472 instance (Eq a, Ring.C a) => Determined (Mat (S Z) (S Z)) a where
473 determinant (Mat rows) = (V.head . V.head) rows
474
475 instance (Ord a,
476 Ring.C a,
477 Absolute.C a,
478 Arity n,
479 Determined (Mat (S n) (S n)) a)
480 => Determined (Mat (S (S n)) (S (S n))) a where
481 -- | The recursive definition with a special-case for triangular matrices.
482 --
483 -- Examples:
484 --
485 -- >>> let m = fromList [[1,2],[3,4]] :: Mat2 Int
486 -- >>> determinant m
487 -- -1
488 --
489 determinant m
490 | is_triangular m = product [ m !!! (i,i) | i <- [0..(nrows m)-1] ]
491 | otherwise = determinant_recursive
492 where
493 m' i j = m !!! (i,j)
494
495 det_minor i j = determinant (minor m i j)
496
497 determinant_recursive =
498 sum [ (-1)^(toInteger j) NP.* (m' 0 j) NP.* (det_minor 0 j)
499 | j <- [0..(ncols m)-1] ]
500
501
502
503 -- | Matrix multiplication.
504 --
505 -- Examples:
506 --
507 -- >>> let m1 = fromList [[1,2,3], [4,5,6]] :: Mat N2 N3 Int
508 -- >>> let m2 = fromList [[1,2],[3,4],[5,6]] :: Mat N3 N2 Int
509 -- >>> m1 * m2
510 -- ((22,28),(49,64))
511 --
512 infixl 7 *
513 (*) :: (Ring.C a, Arity m, Arity n, Arity p)
514 => Mat m n a
515 -> Mat n p a
516 -> Mat m p a
517 (*) m1 m2 = construct lambda
518 where
519 lambda i j =
520 sum [(m1 !!! (i,k)) NP.* (m2 !!! (k,j)) | k <- [0..(ncols m1)-1] ]
521
522
523
524 instance (Ring.C a, Arity m, Arity n) => Additive.C (Mat m n a) where
525
526 (Mat rows1) + (Mat rows2) =
527 Mat $ V.zipWith (V.zipWith (+)) rows1 rows2
528
529 (Mat rows1) - (Mat rows2) =
530 Mat $ V.zipWith (V.zipWith (-)) rows1 rows2
531
532 zero = Mat (V.replicate $ V.replicate (fromInteger 0))
533
534
535 instance (Ring.C a, Arity m, Arity n, m ~ n) => Ring.C (Mat m n a) where
536 -- The first * is ring multiplication, the second is matrix
537 -- multiplication.
538 m1 * m2 = m1 * m2
539
540
541 instance (Ring.C a, Arity m, Arity n) => Module.C a (Mat m n a) where
542 -- We can multiply a matrix by a scalar of the same type as its
543 -- elements.
544 x *> (Mat rows) = Mat $ V.map (V.map (NP.* x)) rows
545
546
547 instance (Algebraic.C a,
548 ToRational.C a,
549 Arity m)
550 => Normed (Mat (S m) N1 a) where
551 -- | Generic p-norms for vectors in R^n that are represented as nx1
552 -- matrices.
553 --
554 -- Examples:
555 --
556 -- >>> let v1 = vec2d (3,4)
557 -- >>> norm_p 1 v1
558 -- 7.0
559 -- >>> norm_p 2 v1
560 -- 5.0
561 --
562 norm_p p (Mat rows) =
563 (root p') $ sum [fromRational' (toRational x)^p' | x <- xs]
564 where
565 p' = toInteger p
566 xs = concat $ V.toList $ V.map V.toList rows
567
568 -- | The infinity norm.
569 --
570 -- Examples:
571 --
572 -- >>> let v1 = vec3d (1,5,2)
573 -- >>> norm_infty v1
574 -- 5
575 --
576 norm_infty (Mat rows) =
577 fromRational' $ toRational $ V.maximum $ V.map V.maximum rows
578
579
580 -- | Compute the Frobenius norm of a matrix. This essentially treats
581 -- the matrix as one long vector containing all of its entries (in
582 -- any order, it doesn't matter).
583 --
584 -- Examples:
585 --
586 -- >>> let m = fromList [[1, 2, 3],[4,5,6],[7,8,9]] :: Mat3 Double
587 -- >>> frobenius_norm m == sqrt 285
588 -- True
589 --
590 -- >>> let m = fromList [[1, -1, 1],[-1,1,-1],[1,-1,1]] :: Mat3 Double
591 -- >>> frobenius_norm m == 3
592 -- True
593 --
594 frobenius_norm :: (Algebraic.C a, Ring.C a) => Mat m n a -> a
595 frobenius_norm (Mat rows) =
596 sqrt $ element_sum $ V.map row_sum rows
597 where
598 -- | Square and add up the entries of a row.
599 row_sum = element_sum . V.map (^2)
600
601
602 -- Vector helpers. We want it to be easy to create low-dimension
603 -- column vectors, which are nx1 matrices.
604
605 -- | Convenient constructor for 2D vectors.
606 --
607 -- Examples:
608 --
609 -- >>> import Roots.Simple
610 -- >>> let fst m = m !!! (0,0)
611 -- >>> let snd m = m !!! (1,0)
612 -- >>> let h = 0.5 :: Double
613 -- >>> let g1 m = 1.0 + h NP.* exp(-((fst m)^2))/(1.0 + (snd m)^2)
614 -- >>> let g2 m = 0.5 + h NP.* atan((fst m)^2 + (snd m)^2)
615 -- >>> let g u = vec2d ((g1 u), (g2 u))
616 -- >>> let u0 = vec2d (1.0, 1.0)
617 -- >>> let eps = 1/(10^9)
618 -- >>> fixed_point g eps u0
619 -- ((1.0728549599342185),(1.0820591495686167))
620 --
621 vec1d :: (a) -> Col1 a
622 vec1d (x) = Mat (mk1 (mk1 x))
623
624 vec2d :: (a,a) -> Col2 a
625 vec2d (x,y) = Mat (mk2 (mk1 x) (mk1 y))
626
627 vec3d :: (a,a,a) -> Col3 a
628 vec3d (x,y,z) = Mat (mk3 (mk1 x) (mk1 y) (mk1 z))
629
630 vec4d :: (a,a,a,a) -> Col4 a
631 vec4d (w,x,y,z) = Mat (mk4 (mk1 w) (mk1 x) (mk1 y) (mk1 z))
632
633 vec5d :: (a,a,a,a,a) -> Col5 a
634 vec5d (v,w,x,y,z) = Mat (mk5 (mk1 v) (mk1 w) (mk1 x) (mk1 y) (mk1 z))
635
636 -- Since we commandeered multiplication, we need to create 1x1
637 -- matrices in order to multiply things.
638 scalar :: a -> Mat1 a
639 scalar x = Mat (mk1 (mk1 x))
640
641 dot :: (RealRing.C a, n ~ N1, m ~ S t, Arity t)
642 => Mat m n a
643 -> Mat m n a
644 -> a
645 v1 `dot` v2 = ((transpose v1) * v2) !!! (0, 0)
646
647
648 -- | The angle between @v1@ and @v2@ in Euclidean space.
649 --
650 -- Examples:
651 --
652 -- >>> let v1 = vec2d (1.0, 0.0)
653 -- >>> let v2 = vec2d (0.0, 1.0)
654 -- >>> angle v1 v2 == pi/2.0
655 -- True
656 --
657 angle :: (Transcendental.C a,
658 RealRing.C a,
659 n ~ N1,
660 m ~ S t,
661 Arity t,
662 ToRational.C a)
663 => Mat m n a
664 -> Mat m n a
665 -> a
666 angle v1 v2 =
667 acos theta
668 where
669 theta = (recip norms) NP.* (v1 `dot` v2)
670 norms = (norm v1) NP.* (norm v2)
671
672
673 -- | Retrieve the diagonal elements of the given matrix as a \"column
674 -- vector,\" i.e. a m-by-1 matrix. We require the matrix to be
675 -- square to avoid ambiguity in the return type which would ideally
676 -- have dimension min(m,n) supposing an m-by-n matrix.
677 --
678 -- Examples:
679 --
680 -- >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int
681 -- >>> diagonal m
682 -- ((1),(5),(9))
683 --
684 diagonal :: (Arity m) => Mat m m a -> Col m a
685 diagonal matrix =
686 construct lambda
687 where
688 lambda i _ = matrix !!! (i,i)
689
690
691 -- | Given a square @matrix@, return a new matrix of the same size
692 -- containing only the on-diagonal entries of @matrix@. The
693 -- off-diagonal entries are set to zero.
694 --
695 -- Examples:
696 --
697 -- >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int
698 -- >>> diagonal_part m
699 -- ((1,0,0),(0,5,0),(0,0,9))
700 --
701 diagonal_part :: (Arity m, Ring.C a)
702 => Mat m m a
703 -> Mat m m a
704 diagonal_part matrix =
705 construct lambda
706 where
707 lambda i j = if i == j then matrix !!! (i,j) else 0
708
709
710 -- | Given a square @matrix@, return a new matrix of the same size
711 -- containing only the on-diagonal and below-diagonal entries of
712 -- @matrix@. The above-diagonal entries are set to zero.
713 --
714 -- Examples:
715 --
716 -- >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int
717 -- >>> lt_part m
718 -- ((1,0,0),(4,5,0),(7,8,9))
719 --
720 lt_part :: (Arity m, Ring.C a)
721 => Mat m m a
722 -> Mat m m a
723 lt_part matrix =
724 construct lambda
725 where
726 lambda i j = if i >= j then matrix !!! (i,j) else 0
727
728
729 -- | Given a square @matrix@, return a new matrix of the same size
730 -- containing only the below-diagonal entries of @matrix@. The on-
731 -- and above-diagonal entries are set to zero.
732 --
733 -- Examples:
734 --
735 -- >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int
736 -- >>> lt_part_strict m
737 -- ((0,0,0),(4,0,0),(7,8,0))
738 --
739 lt_part_strict :: (Arity m, Ring.C a)
740 => Mat m m a
741 -> Mat m m a
742 lt_part_strict matrix =
743 construct lambda
744 where
745 lambda i j = if i > j then matrix !!! (i,j) else 0
746
747
748 -- | Given a square @matrix@, return a new matrix of the same size
749 -- containing only the on-diagonal and above-diagonal entries of
750 -- @matrix@. The below-diagonal entries are set to zero.
751 --
752 -- Examples:
753 --
754 -- >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int
755 -- >>> ut_part m
756 -- ((1,2,3),(0,5,6),(0,0,9))
757 --
758 ut_part :: (Arity m, Ring.C a)
759 => Mat m m a
760 -> Mat m m a
761 ut_part = transpose . lt_part . transpose
762
763
764 -- | Given a square @matrix@, return a new matrix of the same size
765 -- containing only the above-diagonal entries of @matrix@. The on-
766 -- and below-diagonal entries are set to zero.
767 --
768 -- Examples:
769 --
770 -- >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int
771 -- >>> ut_part_strict m
772 -- ((0,2,3),(0,0,6),(0,0,0))
773 --
774 ut_part_strict :: (Arity m, Ring.C a)
775 => Mat m m a
776 -> Mat m m a
777 ut_part_strict = transpose . lt_part_strict . transpose
778
779
780 -- | Compute the trace of a square matrix, the sum of the elements
781 -- which lie on its diagonal. We require the matrix to be
782 -- square to avoid ambiguity in the return type which would ideally
783 -- have dimension min(m,n) supposing an m-by-n matrix.
784 --
785 -- Examples:
786 --
787 -- >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int
788 -- >>> trace m
789 -- 15
790 --
791 trace :: (Arity m, Ring.C a) => Mat m m a -> a
792 trace matrix =
793 let (Mat rows) = diagonal matrix
794 in
795 element_sum $ V.map V.head rows
796
797
798 -- | Zip together two column matrices.
799 --
800 -- Examples:
801 --
802 -- >>> let m1 = fromList [[1],[1],[1]] :: Col3 Int
803 -- >>> let m2 = fromList [[1],[2],[3]] :: Col3 Int
804 -- >>> colzip m1 m2
805 -- (((1,1)),((1,2)),((1,3)))
806 --
807 colzip :: Arity m => Col m a -> Col m a -> Col m (a,a)
808 colzip c1 c2 =
809 construct lambda
810 where
811 lambda i j = (c1 !!! (i,j), c2 !!! (i,j))
812
813
814 -- | Zip together two column matrices using the supplied function.
815 --
816 -- Examples:
817 --
818 -- >>> let c1 = fromList [[1],[2],[3]] :: Col3 Integer
819 -- >>> let c2 = fromList [[4],[5],[6]] :: Col3 Integer
820 -- >>> colzipwith (^) c1 c2
821 -- ((1),(32),(729))
822 --
823 colzipwith :: Arity m
824 => (a -> a -> b)
825 -> Col m a
826 -> Col m a
827 -> Col m b
828 colzipwith f c1 c2 =
829 construct lambda
830 where
831 lambda i j = f (c1 !!! (i,j)) (c2 !!! (i,j))
832
833
834 -- | Map a function over a matrix of any dimensions.
835 --
836 -- Examples:
837 --
838 -- >>> let m = fromList [[1,2],[3,4]] :: Mat2 Int
839 -- >>> map2 (^2) m
840 -- ((1,4),(9,16))
841 --
842 map2 :: (a -> b) -> Mat m n a -> Mat m n b
843 map2 f (Mat rows) =
844 Mat $ V.map g rows
845 where
846 g = V.map f
847
848
849 -- | Fold over the entire matrix passing the coordinates @i@ and @j@
850 -- (of the row/column) to the accumulation function.
851 --
852 -- Examples:
853 --
854 -- >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int
855 -- >>> ifoldl2 (\i j cur _ -> cur + i + j) 0 m
856 -- 18
857 --
858 ifoldl2 :: forall a b m n.
859 (Int -> Int -> b -> a -> b)
860 -> b
861 -> Mat m n a
862 -> b
863 ifoldl2 f initial (Mat rows) =
864 V.ifoldl row_function initial rows
865 where
866 -- | The order that we need this in (so that @g idx@ makes sense)
867 -- is a little funny. So that we don't need to pass weird
868 -- functions into ifoldl2, we swap the second and third
869 -- arguments of @f@ calling the result @g@.
870 g :: Int -> b -> Int -> a -> b
871 g w x y = f w y x
872
873 row_function :: b -> Int -> Vec n a -> b
874 row_function rowinit idx r = V.ifoldl (g idx) rowinit r
875
876
877 -- | Map a function over a matrix of any dimensions, passing the
878 -- coordinates @i@ and @j@ to the function @f@.
879 --
880 -- Examples:
881 --
882 -- >>> let m = fromList [[1,2],[3,4]] :: Mat2 Int
883 -- >>> imap2 (\i j _ -> i+j) m
884 -- ((0,1),(1,2))
885 --
886 imap2 :: (Int -> Int -> a -> b) -> Mat m n a -> Mat m n b
887 imap2 f (Mat rows) =
888 Mat $ V.imap g rows
889 where
890 g i = V.imap (f i)