blas-0.5: Data/Matrix/Dense/Internal.hs
{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses #-}
{-# OPTIONS_GHC -fglasgow-exts #-}
-----------------------------------------------------------------------------
-- |
-- Module : Data.Matrix.Dense.Internal
-- Copyright : Copyright (c) , Patrick Perry <patperry@stanford.edu>
-- License : BSD3
-- Maintainer : Patrick Perry <patperry@stanford.edu>
-- Stability : experimental
--
module Data.Matrix.Dense.Internal (
-- * Dense matrix data types
DMatrix(..),
IOMatrix,
Matrix,
module BLAS.Matrix.Base,
module BLAS.Tensor,
-- * Converting to and from foreign pointers
toForeignPtr,
fromForeignPtr,
-- * Creating new matrices
newMatrix,
newMatrix_,
newListMatrix,
newColsMatrix,
newRowsMatrix,
listMatrix,
-- * Special matrices
newIdentity,
setIdentity,
-- * Row and column views
row,
col,
rows,
cols,
-- * Diagonal views
diag,
-- * Matrix views
submatrix,
-- * Converting to/from vectors
maybeFromRow,
maybeFromCol,
maybeToVector,
-- * Lifting scalar and vector operations
liftV,
liftV2,
-- * Casting matrices
coerceMatrix,
-- * Unsafe operations
unsafeThaw,
unsafeFreeze,
unsafeNewMatrix,
unsafeWithElemPtr,
unsafeRow,
unsafeCol,
unsafeDiag,
unsafeSubmatrix,
) where
import Control.Monad ( forM_, zipWithM_ )
import Data.Ix ( inRange, range )
import Foreign
import System.IO.Unsafe
import Unsafe.Coerce
import Data.AEq
import Data.Vector.Dense.Internal hiding ( toForeignPtr, fromForeignPtr,
unsafeFreeze, unsafeThaw, storageOf, offsetOf, unsafeWithElemPtr )
import qualified Data.Vector.Dense.Internal as V
import qualified Data.Vector.Dense.Operations as V
import BLAS.Access
import BLAS.Internal ( inlinePerformIO, checkedRow, checkedCol, checkedDiag,
checkedSubmatrix, diagStart, diagLen )
import BLAS.Elem ( Elem, BLAS1 )
import qualified BLAS.Elem as E
import BLAS.Matrix.Base hiding ( Matrix )
import qualified BLAS.Matrix.Base as C
import BLAS.Tensor
import BLAS.Types
-- | The mutable dense matrix data type. It can either store elements in
-- column-major order, or provide a view into another matrix. The view
-- transposes and conjugates the underlying matrix.
data DMatrix t mn e =
DM { storageOf :: {-# UNPACK #-} !(ForeignPtr e) -- ^ a pointer to the storage region
, offsetOf :: {-# UNPACK #-} !Int -- ^ an offset (in elements, not bytes) to the first element in the matrix.
, size1 :: {-# UNPACK #-} !Int -- ^ the number of rows in the matrix
, size2 :: {-# UNPACK #-} !Int -- ^ the number of columns in the matrix
, ldaOf :: {-# UNPACK #-} !Int -- ^ the leading dimension size of the matrix
, isHerm :: {-# UNPACK #-} !Bool -- ^ indicates whether or not the matrix is transposed and conjugated
}
type Matrix = DMatrix Imm
type IOMatrix = DMatrix Mut
unsafeFreeze :: DMatrix t mn e -> Matrix mn e
unsafeFreeze = unsafeCoerce
unsafeThaw :: DMatrix t mn e -> IOMatrix mn e
unsafeThaw = unsafeCoerce
-- | Coerce the phantom shape type from one type to another.
coerceMatrix :: DMatrix t mn e -> DMatrix t kl e
coerceMatrix = unsafeCoerce
-- | @fromForeignPtr f o mn l h@ creates a matrix view of the data pointed to
-- by @f@ starting at offset @o@ and having shape @mn@ and lda @l@. If @h@
-- is @True@ the matrix is interpreted as transposed and conjugated.
fromForeignPtr :: ForeignPtr e -> Int -> (Int,Int) -> Int -> Bool -> DMatrix t (m,n) e
fromForeignPtr f o (m,n) l h = DM f o m n l h
-- | Convert a dense matrix to a pointer, offset, size, lda, herm.
toForeignPtr :: DMatrix t (m,n) e -> (ForeignPtr e, Int, (Int,Int), Int, Bool)
toForeignPtr a = (storageOf a, offsetOf a, (size1 a, size2 a), ldaOf a, isHerm a)
indexOf :: DMatrix t (m,n) e -> (Int,Int) -> Int
indexOf a (i,j) =
let (i',j') = case isHerm a of
True -> (j,i)
False -> (i,j)
o = offsetOf a
l = ldaOf a
in o + i' + j'*l
{-# INLINE indexOf #-}
-- | Create a new matrix of the given size and initialize the given elements to
-- the given values. All other elements get initialized to zero.
newMatrix :: (BLAS1 e) => (Int,Int) -> [((Int,Int), e)] -> IO (DMatrix t (m,n) e)
newMatrix = newMatrixHelp writeElem
-- | Same as 'newMatrix' but do not do any bounds-checking.
unsafeNewMatrix :: (BLAS1 e) => (Int,Int) -> [((Int,Int), e)] -> IO (DMatrix t (m,n) e)
unsafeNewMatrix = newMatrixHelp unsafeWriteElem
newMatrixHelp :: (BLAS1 e) =>
(IOMatrix (m,n) e -> (Int,Int) -> e -> IO ())
-> (Int,Int) -> [((Int,Int),e)] -> IO (DMatrix t (m,n) e)
newMatrixHelp set mn ijes = do
x <- newZero mn
io <- unsafeInterleaveIO $ mapM_ (uncurry $ set $ unsafeThaw x) ijes
return $ io `seq` x
-- | Create a new matrix of given shape, but do not initialize the elements.
newMatrix_ :: (Elem e) => (Int,Int) -> IO (DMatrix t (m,n) e)
newMatrix_ (m,n)
| m < 0 || n < 0 =
ioError $ userError $
"Tried to create a matrix with shape `" ++ show (m,n) ++ "'"
| otherwise = do
f <- mallocForeignPtrArray (m*n)
return $ fromForeignPtr f 0 (m,n) (max 1 m) False
-- | Create a new matrix with the given elements in column-major order.
newListMatrix :: (Elem e) => (Int,Int) -> [e] -> IO (DMatrix t (m,n) e)
newListMatrix (m,n) es = do
a <- newMatrix_ (m,n)
unsafeWithElemPtr a (0,0) $ flip pokeArray (take (m*n) es)
return a
-- | Create a new matrix with the given elements in row-major order.
listMatrix :: (Elem e) => (Int,Int) -> [e] -> Matrix (m,n) e
listMatrix mn es = unsafePerformIO $ newListMatrix mn es
{-# NOINLINE listMatrix #-}
-- | Create a new matrix of the given shape with ones along the diagonal,
-- and zeros everywhere else.
newIdentity :: (BLAS1 e) => (Int,Int) -> IO (DMatrix t (m,n) e)
newIdentity mn = do
a <- newMatrix_ mn
setIdentity (unsafeThaw a)
return a
-- | Set the diagonal to ones, and set everywhere else to zero.
setIdentity :: (BLAS1 e) => IOMatrix (m,n) e -> IO ()
setIdentity a = do
s <- getSize a
case s of
0 -> return ()
_ -> setZero a >>
setConstant 1 (diag a 0)
-- | Form a matrix from a list of column vectors.
newColsMatrix :: (BLAS1 e) => (Int,Int) -> [DVector t m e] -> IO (DMatrix r (m,n) e)
newColsMatrix (m,n) cs = do
a <- newZero (m,n)
forM_ (zip [0..(n-1)] cs) $ \(j,c) ->
V.copyVector (unsafeCol (unsafeThaw a) j) c
return a
-- | Form a matrix from a list of row vectors.
newRowsMatrix :: (BLAS1 e) => (Int,Int) -> [DVector t n e] -> IO (DMatrix r (m,n) e)
newRowsMatrix (m,n) rs = do
a <- newZero (m,n)
forM_ (zip [0..(m-1)] rs) $ \(i,r) ->
V.copyVector (unsafeRow (unsafeThaw a) i) r
return a
-- | Evaluate a function with a pointer to the raw storage for the element
-- at the given index. It may be necessary to conjugate or scale values before
-- reading or writing to or from the location.
unsafeWithElemPtr :: (Elem e) => DMatrix t (m,n) e -> (Int,Int) -> (Ptr e -> IO a) -> IO a
unsafeWithElemPtr a ij f =
withForeignPtr (storageOf a) $ \ptr ->
let ptr' = ptr `advancePtr` (indexOf a ij)
in f ptr'
-- | Get a vector view of the given row in a matrix.
row :: (Elem e) => DMatrix t (m,n) e -> Int -> DVector t n e
row a = checkedRow (shape a) (unsafeRow a)
-- | Get a list of vector views of the rows of the matrix.
rows :: (Elem e) => DMatrix t (m,n) e -> [DVector t n e]
rows a = [ unsafeRow a i | i <- [0..numRows a - 1] ]
-- | Get a list of vector views of the columns of the matrix.
cols :: (Elem e) => DMatrix t (m,n) e -> [DVector t m e]
cols a = [ unsafeCol a j | j <- [0..numCols a - 1] ]
-- | Get a vector view of the given column in a matrix.
col :: (Elem e) => DMatrix t (m,n) e -> Int -> DVector t m e
col a = checkedCol (shape a) (unsafeCol a)
-- | Same as 'row', but does not do any bounds checking.
unsafeRow :: (Elem e) => DMatrix t (m,n) e -> Int -> DVector t n e
unsafeRow a i
| isHerm a =
conj $ unsafeCol (herm a) i
| otherwise =
let f = storageOf a
o = indexOf a (i,0)
n = numCols a
s = ldaOf a
c = False
in V.fromForeignPtr f o n s c
-- | Same as 'col', but does not do any bounds checking.
unsafeCol :: (Elem e) => DMatrix t (m,n) e -> Int -> DVector t m e
unsafeCol a j
| isHerm a =
conj $ unsafeRow (herm a) j
| otherwise =
let f = storageOf a
o = indexOf a (0,j)
m = numRows a
s = 1
c = False
in V.fromForeignPtr f o m s c
-- | @diag a 0@ gets a vector view of the main diagonal of @a@. @diag a k@ for
-- @k@ positive gets a view of the @k@th superdiagonal. For @k@ negative, it
-- gets a view of the @(-k)@th subdiagonal.
diag :: (Elem e) => DMatrix t (m,n) e -> Int -> DVector t k e
diag a = checkedDiag (shape a) (unsafeDiag a)
-- | Same as 'diag', but does not do any bounds checking.
unsafeDiag :: (Elem e) => DMatrix t (m,n) e -> Int -> DVector t k e
unsafeDiag a i
| isHerm a =
conj $ unsafeDiag (herm a) (negate i)
| otherwise =
let f = storageOf a
o = indexOf a (diagStart i)
n = diagLen (shape a) i
s = ldaOf a + 1
c = False
in V.fromForeignPtr f o n s c
-- | @submatrix a ij mn@ returns a view of the submatrix of @a@ with element @(0,0)@
-- being element @ij@ in @a@, and having shape @mn@.
submatrix :: (Elem e) => DMatrix t (m,n) e -> (Int,Int) -> (Int,Int) -> DMatrix t (k,l) e
submatrix a = checkedSubmatrix (shape a) (unsafeSubmatrix a)
-- | Same as 'submatrix' but does not do any bounds checking.
unsafeSubmatrix :: (Elem e) => DMatrix t (m,n) e -> (Int,Int) -> (Int,Int) -> DMatrix t (k,l) e
unsafeSubmatrix a (i,j) (m,n)
| isHerm a =
herm $ unsafeSubmatrix (herm a) (j,i) (n,m)
| otherwise =
let f = storageOf a
o = indexOf a (i,j)
l = ldaOf a
in fromForeignPtr f o (m,n) l False
-- | Create a matrix view of a row vector. This will fail if the
-- vector is conjugated and the stride is not @1@.
maybeFromRow :: (Elem e) => DVector t m e -> Maybe (DMatrix t (one,m) e)
maybeFromRow x
| isConj x && strideOf x == 1 =
let f = V.storageOf x
o = V.offsetOf x
n = dim x
l = max 1 n
h = True
in Just $ fromForeignPtr f o (n,1) l h
| (not . isConj) x =
let f = V.storageOf x
o = V.offsetOf x
n = dim x
s = strideOf x
l = max 1 s
h = False
in Just $ fromForeignPtr f o (1,n) l h
| otherwise =
Nothing
-- | Possibly create a matrix view of a column vector. This will fail
-- if the stride of the vector is not @1@ and the vector is not conjugated.
maybeFromCol :: (Elem e) => DVector t n e -> Maybe (DMatrix t (n,one) e)
maybeFromCol x
| isConj x = maybeFromRow (conj x) >>= return . herm
| strideOf x == 1 =
let f = V.storageOf x
o = V.offsetOf x
m = dim x
l = max 1 m
in Just $ fromForeignPtr f o (m,1) l False
| otherwise =
Nothing
maybeToVector :: (Elem e) => DMatrix t (m,n) e -> Maybe (Order, DVector t k e)
maybeToVector a@(DM f o m n ld h)
| h =
maybeToVector (herm a) >>= (\(ord,x) -> return (flipOrder ord, conj x))
| ld == m =
Just $ (ColMajor, V.fromForeignPtr f o (m*n) 1 False)
| m == 1 =
Just $ (ColMajor, V.fromForeignPtr f o n ld False)
| otherwise =
Nothing
-- | Take a unary elementwise vector operation and apply it to the elements of a matrix.
liftV :: (Elem e) => (DVector t k e -> IO ()) -> DMatrix t (m,n) e -> IO ()
liftV f a =
case maybeToVector a of
Just (_,x) -> f x
_ ->
let xs = case isHerm a of
True -> rows (coerceMatrix a)
False -> cols (coerceMatrix a)
in mapM_ f xs
-- | Take a binary elementwise vector operation and apply it to the elements of a pair
-- of matrices.
liftV2 :: (Elem e) => (DVector s k e -> DVector t k e -> IO ())
-> DMatrix s (m,n) e -> DMatrix t (m,n) e -> IO ()
liftV2 f a b =
case (maybeToVector a, maybeToVector b) of
(Just (RowMajor,x), Just (RowMajor,y)) -> f x y
(Just (ColMajor,x), Just (ColMajor,y)) -> f x y
_ ->
let (xs,ys) = case isHerm a of
True -> (rows (coerceMatrix a), rows (coerceMatrix b))
False -> (cols (coerceMatrix a), cols (coerceMatrix b))
in zipWithM_ f xs ys
instance C.Matrix (DMatrix t) where
numRows a | isHerm a = size2 a
| otherwise = size1 a
numCols a | isHerm a = size1 a
| otherwise = size2 a
herm a = a{ isHerm=(not . isHerm) a }
instance (BLAS1 e) => ITensor (DMatrix Imm (m,n)) (Int,Int) e where
size a = (numRows a * numCols a)
unsafeAt a = inlinePerformIO . unsafeReadElem a
{-# INLINE unsafeAt #-}
indices a = [ (i,j) | j <- range (0,n-1), i <- range (0,m-1) ]
where (m,n) = shape a
{-# INLINE indices #-}
elems = inlinePerformIO . getElems
{-# INLINE elems #-}
assocs = inlinePerformIO . getAssocs
{-# INLINE assocs #-}
(//) = replaceHelp writeElem
unsafeReplace = replaceHelp unsafeWriteElem
amap f a = listMatrix (shape a) (map f $ elems a)
replaceHelp :: (BLAS1 e) =>
(IOMatrix (m,n) e -> (Int,Int) -> e -> IO ())
-> Matrix (m,n) e -> [((Int,Int), e)] -> Matrix (m,n) e
replaceHelp set x ies =
unsafeFreeze $ unsafePerformIO $ do
y <- newCopy (unsafeThaw x)
mapM_ (uncurry $ set y) ies
return y
{-# NOINLINE replaceHelp #-}
instance (BLAS1 e) => IDTensor (DMatrix Imm (m,n)) (Int,Int) e where
zero = unsafePerformIO . newZero
{-# NOINLINE zero #-}
constant mn = unsafePerformIO . newConstant mn
{-# NOINLINE constant #-}
azipWith f a b
| shape b /= mn =
error ("azipWith: matrix shapes differ; first has shape `"
++ show mn ++ "' and second has shape `"
++ show (shape b) ++ "'")
| otherwise =
listMatrix mn (zipWith f (elems a) (elems b))
where
mn = shape a
instance (BLAS1 e) => RTensor (DMatrix t (m,n)) (Int,Int) e IO where
getSize a = return (numRows a * numCols a)
newCopy a | isHerm a =
newCopy (herm a) >>= return . herm
| otherwise = do
a' <- newMatrix_ (shape a)
liftV2 V.copyVector (unsafeThaw a') a
return a'
unsafeReadElem a (i,j)
| isHerm a = unsafeReadElem (herm a) (j,i) >>= return . E.conj
| otherwise = withForeignPtr (storageOf a) $ \ptr ->
peekElemOff ptr (indexOf a (i,j))
{-# INLINE unsafeReadElem #-}
getIndices = return . indices . unsafeFreeze
{-# INLINE getIndices #-}
getElems a = return $ go (cols a)
where go cs | cs `seq` False = undefined
go [] = []
go (c:cs) =
let e = inlinePerformIO $ getElems c
es = go cs
in e ++ es
{-# NOINLINE getElems #-}
getAssocs a = return $ go (cols a) 0
where go cs j | cs `seq` j `seq` False = undefined
go [] _ = []
go (c:cs) j =
let ie = inlinePerformIO $ getAssocs c
ije = map (\(i,e) -> ((i,j),e)) ie
ijes = go cs (j+1)
in ije ++ ijes
{-# NOINLINE getAssocs #-}
instance (BLAS1 e) => RDTensor (DMatrix t (m,n)) (Int,Int) e IO where
newZero mn = do
a <- newMatrix_ mn
setZero (unsafeThaw a)
return a
newConstant mn e = do
a <- newMatrix_ mn
setConstant e (unsafeThaw a)
return a
instance (BLAS1 e) => MTensor (DMatrix Mut (m,n)) (Int,Int) e IO where
setZero = liftV setZero
setConstant e = liftV (setConstant e)
canModifyElem a ij =
return $ inRange (bounds a) ij
{-# INLINE canModifyElem #-}
unsafeWriteElem a (i,j) e
| isHerm a = unsafeWriteElem (herm a) (j,i) $ E.conj e
| otherwise = withForeignPtr (storageOf a) $ \ptr ->
pokeElemOff ptr (indexOf a (i,j)) e
modifyWith f = liftV (modifyWith f)
instance (BLAS1 e, Show e) => Show (DMatrix Imm (m,n) e) where
show a | isHerm a =
"herm (" ++ show (herm a) ++ ")"
| otherwise =
"listMatrix " ++ show (shape a) ++ " " ++ show (elems a)
compareHelp :: (BLAS1 e) =>
(e -> e -> Bool) -> Matrix (m,n) e -> Matrix (m,n) e -> Bool
compareHelp cmp x y
| isHerm x && isHerm y =
compareHelp cmp (herm x) (herm y)
compareHelp cmp x y =
(shape x == shape y) && (and $ zipWith cmp (elems x) (elems y))
instance (BLAS1 e, Eq e) => Eq (DMatrix Imm (m,n) e) where
(==) = compareHelp (==)
instance (BLAS1 e, AEq e) => AEq (DMatrix Imm (m,n) e) where
(===) = compareHelp (===)
(~==) = compareHelp (~==)