Ticket #3005: Vectors2.hs

{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances, TypeFamilies, ScopedTypeVariables #-}

import TypeNats2
import qualified Data.List as List (transpose)

class Multiply a b where
  type Product a b
  times :: a -> b -> Product a b

data Matrix r c t = Matrix [[t]] deriving (Eq)

type Vector n t = Matrix n One t

instance (Show t) => Show (Matrix r c t) where
  show (Matrix rows) = unlines (map show rows)

-- creation functions

matrix :: (Natural r, Natural c) => r -> c -> [[t]] -> (Matrix r c t)
matrix r c rows | (validNumRows r rows) && (all (validNumCols c) rows) = Matrix rows
matrix _ _ _ = error "dimension mismatch"

vector :: (Natural n) => n -> [t] -> Vector n t
vector n ts = matrix n one (map (\t -> [t]) ts)

validNumRows r rows = (length rows) == (naturalToIntegral r)

validNumCols c row = (length row) == (naturalToIntegral c)

-- element accessors

unsafeAt :: (Natural r, Natural c) => Int -> Int -> Matrix r c t -> t
unsafeAt r c (Matrix rows) = rows !! r !! c

at :: (Natural r, Natural c, Natural rAt, Natural cAt, LEqNat rAt r, LEqNat cAt c) => rAt -> cAt -> Matrix r c t -> t
at r c m = unsafeAt (naturalToIntegral r) (naturalToIntegral c) m

rows :: forall r c t. (Natural r, Natural c) => (Matrix r c t) -> Integer
rows _ = naturalToIntegral (undefined :: r)

cols :: forall r c t. (Natural r, Natural c) => (Matrix r c t) -> Integer
cols _ = naturalToIntegral (undefined :: c)

size :: forall r c t. (Natural r, Natural c) => (Matrix r c t) -> (Integer,Integer)
size m = (rows m, cols m)

-- matrix is a functor

instance (Natural r, Natural c) => Functor (Matrix r c) where
  fmap f (Matrix rows) = Matrix (map (map f) rows)

-- matrix products

matrixPartialProductElem n m1 m2 i j = sum [(unsafeAt i r m1) * (unsafeAt r j m2) | r <- [0..n-1]]

matrixPartialProductRow n m1 m2 i cols = [matrixPartialProductElem n m1 m2 i j | j <- [0..cols-1]]

instance (Natural a, Natural b, Natural c, Num t) => Multiply (Matrix a b t) (Matrix b c t) where
  type Product (Matrix a b t) (Matrix b c t) = Matrix a c t
  times m1 m2 = let rows = naturalToIntegral (undefined :: a)
                    cols = naturalToIntegral (undefined :: c)
                    n = naturalToIntegral (undefined :: b)
                in
                    Matrix [matrixPartialProductRow n m1 m2 i cols | i <- [0..rows-1]]

instance (Natural r, Natural c, Num t) => Multiply t (Matrix r c t) where
  type Product t (Matrix r c t) = Matrix r c t
  times x m = fmap (* x) m

-- useful math

zipMatrixWith :: (Natural rows, Natural cols) => (a -> b -> c) -> (Matrix rows cols a) -> (Matrix rows cols b) -> (Matrix rows cols c)
zipMatrixWith f (Matrix m1) (Matrix m2) = Matrix (zipWith (\row1 row2 -> zipWith f row1 row2) m1 m2)

dot :: (Natural n, Num t) => Vector n t -> Vector n t -> Vector n t
dot = zipMatrixWith (*)

cross :: (Num t) => Vector Three t -> Vector Three t -> Vector Three t
cross v1 v2 = vector three [a2 * b3 - a3 * b2, a3 * b1 - a1 * b3, a1 * b2 - a2 * b1]
              where a1 = at zero zero v1
                    a2 = at one zero v1
                    a3 = at two zero v1
                    b1 = at zero zero v2
                    b2 = at one zero v2
                    b3 = at two zero v2

transpose :: (Natural a, Natural b) => (Matrix a b t) -> (Matrix b a t)
transpose (Matrix rows) = Matrix (List.transpose rows)

plus :: (Natural r, Natural c, Num t) => (Matrix r c t) -> (Matrix r c t) -> (Matrix r c t)
plus = zipMatrixWith (+)

minus :: (Natural r, Natural c, Num t) => (Matrix r c t) -> (Matrix r c t) -> (Matrix r c t)
minus = zipMatrixWith (-)

identity :: (Natural n, Num t) => n -> Matrix n n t
identity n = matrix n n [[identity' (i == j) | j <- [0..(naturalToIntegral n)-1]] | i <- [0..(naturalToIntegral n)-1]]

identity' :: (Num t) => Bool -> t
identity' True = 1
identity' False = 0

flatten :: [[a]] -> [a]
flatten = foldl (++) []

norm :: (Natural n, Floating t) => (Vector n t) -> t
norm (Matrix rows) = sqrt $ sum (map (\x -> x * x) (flatten rows))