arbitrary-dimensional matrices

Matrix.hs

module Matrix where

-- We want a datatype that can work like [a], [[a]], [[[a]]] etc.
-- We want the type system to ensure all the values of type a are at the same level, unlike in a rose tree:

import Data.List

data RoseTree a = Single a | Many [RoseTree a]
  deriving Show

mixed :: RoseTree Int
mixed = Many [Single 1, Many [Single 2, Single 3]]

data Matrix a = Element a | Dimension (Matrix [a])
  deriving Show

scalar :: Matrix Int
scalar = Element 1

vector :: Matrix Int
vector = Dimension (Element [1,2,3])

matrix :: Matrix Int
matrix = Dimension (Dimension (Element [[1,2,3]
                                       ,[4,5,6]]))

threeDee :: Matrix Int
threeDee = Dimension (Dimension (Dimension (Element [[[1,2,3]
                                                     ,[4,5,6]]
                                                    ,[[7,8,9]
                                                     ,[10,11,12]]])))

unDimension :: Matrix a -> Matrix [a]
unDimension (Dimension m) = m

catMatrices :: [Matrix a] -> Matrix a
catMatrices es@(Element _:_) = Dimension (Element [e | Element e <- es])
catMatrices ms@(Dimension _:_) = Dimension (catMatrices [m | Dimension m <- ms])

-- given a sparse n-dimensional matrix, sorted,
-- turn the outermost index into dense vector
unsparse1 :: [([Int],a)] -> [[([Int],a)]]
unsparse1 pairs = go 0 pairs
  where go _ [] = []
        go i pairs = let (now,rest) = break (\((j:_),_) -> i/=j) pairs
                     in strip now : go (i+1) rest
        strip = map (\((_:is),x) -> (is,x))

unsparse :: [([Int],a)] -> Matrix a
unsparse [([],x)] = Element x
unsparse pairs = catMatrices (map unsparse (unsparse1 pairs))

example0 = unsparse [([],1)]
--       = Element 1

example1 = unsparse [([0],1),([1],2),([2],3)]
--       = Dimension (Element [1,2,3])

example2 = unsparse [([0,0],1),([0,1],2),([1,0],3),([1,1],4)]
--       = Dimension (Dimension (Element [[1,2],[3,4]]))