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