diff options
Diffstat (limited to 'lib/Data')
-rw-r--r-- | lib/Data/Packed/Matrix.hs | 15 |
1 files changed, 8 insertions, 7 deletions
diff --git a/lib/Data/Packed/Matrix.hs b/lib/Data/Packed/Matrix.hs index c56bf3d..4fd61a1 100644 --- a/lib/Data/Packed/Matrix.hs +++ b/lib/Data/Packed/Matrix.hs | |||
@@ -84,19 +84,20 @@ diag v = ST.runSTMatrix $ do | |||
84 | 84 | ||
85 | {- | creates a rectangular diagonal matrix | 85 | {- | creates a rectangular diagonal matrix |
86 | 86 | ||
87 | @> diagRect (constant 5 3) 3 4 | 87 | @> diagRect (constant 5 3) 3 4 :: Matrix Double |
88 | (3><4) | 88 | (3><4) |
89 | [ 5.0, 0.0, 0.0, 0.0 | 89 | [ 5.0, 0.0, 0.0, 0.0 |
90 | , 0.0, 5.0, 0.0, 0.0 | 90 | , 0.0, 5.0, 0.0, 0.0 |
91 | , 0.0, 0.0, 5.0, 0.0 ]@ | 91 | , 0.0, 0.0, 5.0, 0.0 ]@ |
92 | -} | 92 | -} |
93 | diagRect :: (Element t, Num t) => Vector t -> Int -> Int -> Matrix t | 93 | diagRect :: (Element t, Num t) => Vector t -> Int -> Int -> Matrix t |
94 | diagRect s r c | 94 | diagRect v r c |
95 | | dim s < min r c = error "diagRect" | 95 | | dim v < min r c = error "diagRect called with dim v < min r c" |
96 | | r == c = diag s | 96 | | otherwise = ST.runSTMatrix $ do |
97 | | r < c = trans $ diagRect s c r | 97 | m <- ST.newMatrix 0 r c |
98 | | otherwise = joinVert [diag s , zeros (r-c,c)] | 98 | let d = min r c |
99 | where zeros (r',c') = reshape c' $ constant 0 (r'*c') | 99 | mapM_ (\k -> ST.writeMatrix m k k (v@>k)) [0..d-1] |
100 | return m | ||
100 | 101 | ||
101 | -- | extracts the diagonal from a rectangular matrix | 102 | -- | extracts the diagonal from a rectangular matrix |
102 | takeDiag :: (Element t) => Matrix t -> Vector t | 103 | takeDiag :: (Element t) => Matrix t -> Vector t |