module Data.Eigen.Util (
rowAdd
, colAdd
, scaleRow
, scaleCol
, fromList'
, hstack
, vstack
, delRow
, delRows
, delCol
, delCols
, kronecker
) where
import Data.Eigen.Matrix as E
import Data.Maybe (fromJust)
import Data.List as L
import Data.Vector.Storable as V
to2DList _ [] = []
to2DList n es = row : to2DList n rest
where
( row, rest ) = L.splitAt n es
fromList' :: E.Elem a b => Int -> [ a ] -> E.Matrix a b
fromList' n elems = E.fromList $ to2DList n elems
hstack :: E.Elem a b => [ E.Matrix a b ] -> E.Matrix a b
hstack mats = E.generate allrows allcols generateFunc
where
allcols = L.sum ncols
allrows = E.rows $ L.head mats
ncols = L.map E.cols mats
whichMat c = fromJust $ L.findIndex (>c) $ L.scanl1 (+) ncols
generateFunc i j = (mats!!matId) E.! (i, j (L.sum $ L.take matId ncols) )
where
matId = whichMat j
vstack :: E.Elem a b => [ E.Matrix a b ] -> E.Matrix a b
vstack mats = E.transpose $ hstack $ L.map E.transpose mats
kronecker mat1 mat2 = vstack $ L.map hstack result
where
result = to2DList c1 $ E.fold' ( \c e -> ((E.map (*e) mat2):c) ) [] mat1
[ (r1,c1), (r2,c2) ] = [ E.dims mat1, E.dims mat2 ]
(r, c) = ( r1*r2, c1*c2 )
rowAdd :: E.Elem a b => (a, Int) -> (a, Int) -> E.Matrix a b -> E.Matrix a b
rowAdd (k1,r1) (k2,r2) mat = E.imap (
\i j v -> if i == r1 then k1 * v + k2 * ( mat E.! (r2,j) ) else v ) mat
colAdd :: E.Elem a b => (a, Int) -> (a, Int ) -> E.Matrix a b -> E.Matrix a b
colAdd (k1, c1) (k2, c2) mat = E.imap (
\i j v -> if j == c1 then k1 * v + k2 * ( mat E.! (i,c2) ) else v ) mat
scaleRow :: E.Elem a b => Int -> a -> E.Matrix a b -> E.Matrix a b
scaleRow row c mat = E.imap ( \i j v -> if i == row then c * v else v ) mat
scaleCol :: E.Elem a b => Int -> a -> E.Matrix a b -> E.Matrix a b
scaleCol col c mat = E.imap ( \i j v -> if j == col then c * v else v ) mat
delRow :: E.Elem a b => Int -> E.Matrix a b -> E.Matrix a b
delRow r mat = E.fromList $ deleteAt r $ E.toList mat
deleteAt :: Int -> [a] -> [a]
deleteAt n ls = let (ys,zs) = L.splitAt n ls in ys L.++ (L.tail zs)
delCol :: E.Elem a b => Int -> E.Matrix a b -> E.Matrix a b
delCol c mat = E.fromList $ L.map (\row -> deleteAt c row) $ E.toList mat
delRows :: E.Elem a b => [ Int ] -> E.Matrix a b -> E.Matrix a b
delRows rows mat = delRows' (L.sort rows) mat
where
delRows' [] mat = mat
delRows' (r:rs) mat = delRows' (L.map (\e->e1) rs) (delRow r mat)
delCols :: E.Elem a b => [ Int ] -> E.Matrix a b -> E.Matrix a b
delCols cols mat = delCols' (L.sort cols) mat
where
delCols' [] mat = mat
delCols' (c:cs) mat = delCols' (L.map (\e->e1) cs) (delCol c mat)