src/FixedMatrix.hs

   1 {-# LANGUAGE ScopedTypeVariables #-}
   2 {-# LANGUAGE FlexibleContexts #-}
   3 {-# LANGUAGE FlexibleInstances #-}
   4 {-# LANGUAGE MultiParamTypeClasses #-}
   5 {-# LANGUAGE TypeFamilies #-}
   6
   7 module FixedMatrix
   8 where
   9
  10 import FixedVector as FV
  11 import qualified Data.Vector.Fixed as V
  12 import Data.Vector.Fixed.Internal (arity)
  13
  14 type Mat v w a = Vn v (Vn w a)
  15 type Mat2 a = Mat Vec2D Vec2D a
  16 type Mat3 a = Mat Vec3D Vec3D a
  17 type Mat4 a = Mat Vec4D Vec4D a
  18
  19 -- | Convert a matrix to a nested list.
  20 toList :: (V.Vector v (Vn w a), V.Vector w a) => Mat v w a -> [[a]]
  21 toList m = Prelude.map V.toList (V.toList m)
  22
  23 -- | Create a matrix from a nested list.
  24 fromList :: (V.Vector v (Vn w a), V.Vector w a) => [[a]] -> Mat v w a
  25 fromList vs = V.fromList $ Prelude.map V.fromList vs
  26
  27
  28 -- | Unsafe indexing.
  29 (!) :: (V.Vector v (Vn w a), V.Vector w a) => Mat v w a -> (Int, Int) -> a
  30 (!) m (i, j) = (row m i) V.! j
  31
  32 -- | Safe indexing.
  33 (!?) :: (V.Vector v (Vn w a), V.Vector w a) => Mat v w a
  34                                                -> (Int, Int)
  35                                                -> Maybe a
  36 (!?) m (i, j)
  37   | i < 0 || j < 0 = Nothing
  38   | i > V.length m = Nothing
  39   | otherwise = if j > V.length (row m j)
  40                 then Nothing
  41                 else Just $ (row m j) V.! j
  42
  43
  44 -- | The number of rows in the matrix.
  45 nrows :: forall v w a. (V.Vector v (Vn w a), V.Vector w a) => Mat v w a -> Int
  46 nrows = V.length
  47
  48 -- | The number of columns in the first row of the
  49 --   matrix. Implementation stolen from Data.Vector.Fixed.length.
  50 ncols :: forall v w a. (V.Vector v (Vn w a), V.Vector w a) => Mat v w a -> Int
  51 ncols _ = arity (undefined :: V.Dim w)
  52
  53 -- | Return the @i@th row of @m@. Unsafe.
  54 row :: (V.Vector v (Vn w a), V.Vector w a) => Mat v w a
  55                                            -> Int
  56                                            -> Vn w a
  57 row m i = m V.! i
  58
  59
  60 -- | Return the @j@th column of @m@. Unsafe.
  61 column :: (V.Vector v a, V.Vector v (Vn w a), V.Vector w a) => Mat v w a
  62                                                             -> Int
  63                                                             -> Vn v a
  64 column m j =
  65   V.map (element j) m
  66   where
  67     element = flip (V.!)
  68
  69
  70 -- | Transose @m@; switch it's columns and its rows. This is a dirty
  71 --   implementation.. it would be a little cleaner to use imap, but it
  72 --   doesn't seem to work.
  73 transpose :: (V.Vector v (Vn w a),
  74               V.Vector w (Vn v a),
  75               V.Vector v a,
  76               V.Vector w a)
  77              => Mat v w a
  78              -> Mat w v a
  79 transpose m = V.fromList column_list
  80   where
  81     column_list = [ column m i | i <- [0..(ncols m)-1] ]
  82
  83 -- | Is @m@ symmetric?
  84 symmetric :: (V.Vector v (Vn w a),
  85               V.Vector w a,
  86               v ~ w,
  87               V.Vector w Bool,
  88               Eq a)
  89              => Mat v w a
  90              -> Bool
  91 symmetric m =
  92   m == (transpose m)
  93
  94
  95 -- | Construct a new matrix from a function @lambda@. The function
  96 --   @lambda@ should take two parameters i,j corresponding to the
  97 --   entries in the matrix. The i,j entry of the resulting matrix will
  98 --   have the value returned by lambda i j.
  99 construct :: forall v w a.
 100              (V.Vector v (Vn w a),
 101               V.Vector w a)
 102              => (Int -> Int -> a)
 103              -> Mat v w a
 104 construct lambda = rows
 105   where
 106     -- The arity trick is used in Data.Vector.Fixed.length.
 107     imax = (arity (undefined :: V.Dim v)) - 1
 108     jmax = (arity (undefined :: V.Dim w)) - 1
 109     row' i = V.fromList [ lambda i j | j <- [0..jmax] ]
 110     rows = V.fromList [ row' i | i <- [0..imax] ]