jalla-0.1.0: Numeric/Jalla/Vector.hs
{-# LANGUAGE MultiParamTypeClasses, FunctionalDependencies,
FlexibleInstances, GeneralizedNewtypeDeriving, FlexibleContexts,
TypeSynonymInstances #-}
-----------------------------------------------------------------------------
--
-- Module : Math.Vector
-- Copyright : 2011 by Christian Gosch
-- License : BSD3
--
-- Maintainer : Christian Gosch <werbung@goschs.de>
-- Stability : Experimental
-- Portability : GHC only
--
-- |
--
-----------------------------------------------------------------------------
module Numeric.Jalla.Vector (
-- * Classes
-- ** Vectors
GVector(..),
CVector(..),
-- ** Vector/vector operations
VectorVector(..),
-- ** Vector/scalar operations
VectorScalar(..),
-- ** Indexable
module Numeric.Jalla.Indexable,
-- * Data Types
Vector(..),
-- * Construction, conversion, modification
-- ** Monadic, efficient vector modification
VMM,
getVector,
createVector,
modifyVector,
module Numeric.Jalla.IMM,
-- ** Vector maps
vectorAdd,
vectorMap,
vectorBinMap,
-- * Conversion From And To Lists
listVector,
vectorList,
-- * IO Functions
copyVector,
-- * Unsafe Functions
unsafeCopyVector,
unsafeVectorAdd,
-- * Re-exported
CFloat,
CDouble,
Complex
) where
import Numeric.Jalla.Foreign.BLAS
import Numeric.Jalla.Foreign.BlasOps
import Numeric.Jalla.Foreign.LAPACKE
import Numeric.Jalla.Foreign.LapackeOps
import Numeric.Jalla.Internal
import Numeric.Jalla.IMM
import Numeric.Jalla.Indexable
import Numeric.Jalla.Types
import Foreign.C.Types
import Foreign.Marshal.Array
import Foreign hiding (unsafePerformIO)
import System.IO.Unsafe (unsafePerformIO)
import Data.Ix
import Data.Complex
import Control.Monad.State
import Data.Convertible
class (Indexable (vec e) Index e, Field1 e) => GVector vec e where
vector :: Index -> vec e
vectorLength :: vec e -> Index
--(-|) :: vec e -> vec e -> e
-- (!) :: vec e -> Index -> e
class (BlasOps e, GVector vec e, Show (vec e)) => CVector vec e where
vectorAlloc :: Index -> IO (vec e)
withCVector :: vec e -> (Ptr e -> IO a) -> IO a
inc :: vec e -> Index
infixl 7 ||*
infixl 6 ||+,||-
{-| Vector/vector operations. -}
class (CVector vec e) => VectorVector vec e where
-- | Vector addition
(||+) :: vec e -> vec e -> vec e
v1 ||+ v2 = modifyVector v1 $ vectorAdd 1 v2
-- | Vector subtraction
(||-) :: vec e -> vec e -> vec e
v1 ||- v2 = modifyVector v1 $ vectorAdd (-1) v2
-- | Dot product
(||*) :: vec e -> vec e -> e
v1 ||* v2 = innerProduct v1 v2
innerProduct :: (BlasOps e, CVector vec e) => vec e -> vec e -> e
innerProduct v1 v2 | n == n2 = unsafePerformIO $
withCVector v1 $ \p1 ->
withCVector v2 $ \p2 ->
dot n p1 (inc v1) p2 (inc v2)
where
n = vectorLength v1
n2 = vectorLength v2
innerProduct _ _ | otherwise = error "innerProduct: vectors must have same length."
{-# SPECIALIZE NOINLINE innerProduct :: Vector CFloat -> Vector CFloat -> CFloat #-}
{-# SPECIALIZE NOINLINE innerProduct :: Vector CDouble -> Vector CDouble -> CDouble #-}
{-# SPECIALIZE NOINLINE innerProduct :: Vector (Complex CFloat) -> Vector (Complex CFloat) -> Complex CFloat #-}
{-# SPECIALIZE NOINLINE innerProduct :: Vector (Complex CDouble) -> Vector (Complex CDouble) -> Complex CDouble #-}
innerProductReal :: (BlasOpsReal e, CVector vec e) => vec e -> vec e -> CDouble
innerProductReal v1 v2 | n == n2 = realToFrac $ unsafePerformIO $
withCVector v1 $ \p1 ->
withCVector v2 $ \p2 ->
realdot n p1 (inc v1) p2 (inc v2)
where
n = vectorLength v1
n2 = vectorLength v2
innerProductReal _ _ | otherwise = error "innerProduct: vectors must have same length."
innerProductC :: (RealFloat e, BlasOpsComplex e, CVector vec (Complex e)) =>
vec (Complex e)
-> vec (Complex e)
-> Complex e
innerProductC v1 v2 | n == n2 = unsafePerformIO $
withCVector v1 $ \p1 ->
withCVector v2 $ \p2 ->
with (0 :+ 0) $ \pret ->
dotu_sub n p1 (inc v1) p2 (inc v2) pret >> peek pret
where
n = vectorLength v1
n2 = vectorLength v2
innerProductC _ _ | otherwise = error "innerProduct: vectors must have same length."
infixl 7 |.*,|./
infixl 6 |.+,|.-
{-| Vector manipulations by a scalar. -}
class (CVector vec e) => VectorScalar vec e where
(|.*) :: vec e -> e -> vec e
a |.* b = vectorMap (*b) a
(|./) :: vec e -> e -> vec e
a |./ b = vectorMap (/b) a
(|.+) :: vec e -> e -> vec e
a |.+ b = vectorMap (+b) a
(|.-) :: vec e -> e -> vec e
a |.- b = vectorMap ((-)b) a
{-| Vector is the 'CVector' type that is used in Jalla.
Somehow Haddock does not want to create documentation for the class instances
of 'Vector', I try to figure it out. -}
data BlasOps e => Vector e = Vector {vecP :: !(ForeignPtr e),
vecInc :: !Index,
vecLength :: !Index}
vectorAlloc' :: (BlasOps e) => Index -> IO (Vector e)
vectorAlloc' n = mallocForeignPtrArray n >>= \fp -> return $ Vector fp 1 n
instance (BlasOps e) => GVector Vector e where
vector n = unsafePerformIO $ vectorAlloc n
vectorLength = vecLength
-- v1 -| v2 = innerProduct v1 v2
{-| 'CVector' instance for 'Vector'. -}
instance (BlasOps e) => CVector Vector e where
vectorAlloc = vectorAlloc'
withCVector = withForeignPtr . vecP
inc = vecInc
instance BlasOps e => Indexable (Vector e) Index e where
v ! i = unsafePerformIO $ unsafeGetElem v i
instance BlasOps e => VectorVector Vector e
instance BlasOps e => VectorScalar Vector e
vectorList :: GVector vec e => vec e -> [e]
vectorList v = map (v !) [0..n-1] where n = vectorLength v
listVector :: (CVector vec e) => [e] -> vec e
listVector es = createVector n $ setElems ies
where n = length es
ies = zip [0..n-1] es
instance (BlasOps e, Eq e) => Eq (Vector e) where
a /= b | vectorLength a == vectorLength b = or [x /= y | (x,y) <- zip (vectorList a) (vectorList b)]
| otherwise = False
instance (BlasOps e, Show e) => Show (Vector e) where
show v = "listVector " ++ show (vectorList v)
{-| /Num/ instance for a /Vector/.
The operations are all /element-wise/. There may be the occasional error
by wrongly assuming that /(*)/ returns the inner product, which it doesn't.
This instance is basically only provided to get the + and - operators.
Note that this will /not/ work with 'sum', since
that assumes it can start with a "0". -}
instance (BlasOps e, Num e) => Num (Vector e) where
a + b = modifyVector a $ vectorAdd 1 b
a - b = modifyVector a $ vectorAdd (-1) b
a * b = vectorBinMap (*) a b
negate = vectorMap (* (-1))
abs = vectorMap abs
signum = vectorMap signum
fromInteger i = createVector 1 $ setElem 0 (fromIntegral i)
instance (BlasOps e, Num e, Fractional e) => Fractional (Vector e) where
a / b = vectorBinMap (/) a b
recip = vectorMap recip
fromRational r = createVector 1 $ setElem 0 (fromRational r)
{-| An instance of 'Vector' for 'Floating', for convenience.
Some of these don't make much sense in some situations,
but having the trigonometric functions and the like around can be pretty handy.
The functions work element-wise. -}
instance (BlasOps e, Num e, Fractional e) => Floating (Vector e) where
-- | Returns a 1-vector with /pi/ in it.
pi = createVector 1 $ setElem 0 pi
exp = vectorMap exp
sqrt = vectorMap sqrt
log = vectorMap log
-- | Takes the element-wise power.
a ** b = vectorBinMap (**) a b
-- | Computes 'logBase' the /element-wise/. It may be more useful to simply use /vectorMap (logBase b) v/.
logBase = vectorBinMap logBase
sin = vectorMap sin
tan = vectorMap tan
cos = vectorMap cos
asin = vectorMap asin
atan = vectorMap atan
acos = vectorMap acos
sinh = vectorMap sinh
tanh = vectorMap tanh
cosh = vectorMap cosh
asinh = vectorMap asinh
atanh = vectorMap atanh
acosh = vectorMap acosh
{-| Maps a unary function over the elements of a vector and returns the resulting vector. -}
vectorMap :: (CVector vec1 e1, CVector vec2 e2) => (e1 -> e2) -> vec1 e1 -> vec2 e2
vectorMap f v1 = unsafePerformIO $
vectorAlloc n >>= \v2 -> unsafeVectorMap f v1 v2 >> return v2
where n = vectorLength v1
{-# NOINLINE vectorMap #-}
{-# SPECIALIZE NOINLINE vectorMap :: (CFloat -> CFloat) -> Vector CFloat -> Vector CFloat #-}
{-# SPECIALIZE NOINLINE vectorMap :: (CDouble -> CDouble) -> Vector CDouble -> Vector CDouble #-}
{-# SPECIALIZE NOINLINE vectorMap :: (Complex CFloat -> Complex CFloat) -> Vector (Complex CFloat) -> Vector (Complex CFloat) #-}
{-# SPECIALIZE NOINLINE vectorMap :: (Complex CDouble -> Complex CDouble) -> Vector (Complex CDouble) -> Vector (Complex CDouble) #-}
{-| Maps a binary function to the elements of two vectors and returns the resulting vector. -}
vectorBinMap :: (CVector vec1 e1, CVector vec2 e2, CVector vec3 e3) =>
(e1 -> e2 -> e3) -- ^ The function /f/ to map.
-> vec1 e1 -- ^ The first input vector /v1/ for /f/.
-> vec2 e2 -- ^ The second input vector /v2/ for /f/.
-> vec3 e3 -- ^ The result vector. It will have length min(l1,l2), where l1,l2 are the lengths of /v1/ and /v2/.
vectorBinMap f v1 v2 = unsafePerformIO $
vectorAlloc n >>= \v3 -> unsafeVectorBinMap f v1 v2 v3 >> return v3
where n = min (vectorLength v1) (vectorLength v2)
{-| Make a copy of the input vector. Using the cblas_*copy functions. -}
copyVector :: (BlasOps e, CVector vec e, CVector vec2 e) => vec e -> IO (vec2 e)
copyVector v = vectorAlloc n >>= \ret ->
withCVector v $ \p ->
withCVector ret $ \pret ->
copy n p (inc v) pret (inc ret) >> return ret
where n = vectorLength v
--------------------------
-- Monadic vector manipulations
type VMMMonad vec e a = StateT (vec e) IO a
newtype VMM s vec e a = VMM { unVMM :: VMMMonad vec e a } deriving Monad
runVMM :: CVector vec e => vec e -> VMM s vec e a -> IO a
runVMM v action = evalStateT action' v
where
action' = unVMM action
instance (BlasOps e, CVector vec e) => IMM (VMM s vec e) Index (vec e) e where
-- create = createVector
-- modify = modifyVector
-- getO = getVector
setElem = setElem'
setElems = setElems'
fill = fill'
getElem = getElem'
createVector :: CVector vec e => Index -> VMM s vec e a -> vec e
createVector n action = unsafePerformIO $
vectorAlloc n >>= \mv -> runVMM mv (action >> (VMM get))
getVector :: CVector vec e => VMM s vec e (vec e)
getVector = VMM get
modifyVector :: CVector vec e => vec e -> VMM s vec e a -> vec e
modifyVector v action = unsafePerformIO $
copyVector v >>= \nv -> runVMM nv (action >> (VMM get))
where
n = vectorLength v
{-| Adds alpha * v to the current vector. -}
vectorAdd :: CVector vec e => e -> vec e -> VMM s vec e ()
vectorAdd alpha x = VMM $ (get >>= \v -> liftIO $ unsafeVectorAdd alpha x v)
{-| unsafeSetElem may fail gracefully,
therefore this method may or may not set the element, depending on a successful range check. -}
setElem' :: CVector vec e => Index -> e -> VMM s vec e ()
setElem' i e = VMM $ (get >>= \v -> liftIO $ unsafeSetElem v i e >> return ())
setElems' :: CVector vec e => [(Index,e)] -> VMM s vec e ()
setElems' ies = VMM $ (get >>= \v -> liftIO $ mapM_ (\(i,e) -> unsafeSetElem v i e) ies)
{-| Note: getElem' returns a Maybe. -}
getElem' :: CVector vec e => Index -> VMM s vec e e
getElem' i = VMM $ get >>= \v -> liftIO (unsafeGetElem v i)
fill' :: CVector vec e => e -> VMM s vec e ()
fill' e = VMM $ get >>= \v -> liftIO (unsafeFillVector v e)
---------------------------------------------------------------------------------------
-- Unsafe functions.
---------------------------------------------------------------------------------------
unsafeVectorMap :: (CVector vec1 e1, CVector vec2 e2) => (e1 -> e2) -> vec1 e1 -> vec2 e2 -> IO ()
unsafeVectorMap f v1 v2 =
withCVector v1 $ \v1p ->
withCVector v2 $ \v2p ->
unsafePtrMapInc i1 i2 f v1p v2p n
where
i1 = inc v1
i2 = inc v2
n = min (vectorLength v1) (vectorLength v2)
{-# SPECIALIZE INLINE unsafeVectorMap :: (CFloat -> CFloat) -> Vector CFloat -> Vector CFloat -> IO () #-}
{-# SPECIALIZE INLINE unsafeVectorMap :: (CDouble -> CDouble) -> Vector CDouble -> Vector CDouble -> IO () #-}
{-# SPECIALIZE INLINE unsafeVectorMap :: (Complex CFloat -> Complex CFloat) -> Vector (Complex CFloat) -> Vector (Complex CFloat) -> IO () #-}
{-# SPECIALIZE INLINE unsafeVectorMap :: (Complex CDouble -> Complex CDouble) -> Vector (Complex CDouble) -> Vector (Complex CDouble) -> IO () #-}
{-| Copies from one vector to the other, in-place and therefore unsafely.
Uses the BLAS 'copy' function. /min (vectorLength src) (vectorlength dest)/
elements are copied from the first to the second vector. -}
unsafeCopyVector :: (CVector vec e, CVector vec2 e) =>
vec e -- ^ The source vector.
-> vec2 e -- ^ The destination vector.
-> IO ()
unsafeCopyVector src dest =
withCVector src $ \srcp ->
withCVector dest $ \destp ->
copy n srcp (inc src) destp (inc dest)
where n = min (vectorLength src) (vectorLength dest)
unsafeVectorBinMap :: (CVector vec1 e1, CVector vec2 e2, CVector vec3 e3) =>
(e1 -> e2 -> e3)
-> vec1 e1
-> vec2 e2
-> vec3 e3
-> IO ()
unsafeVectorBinMap f v1 v2 v3 =
withCVector v1 $ \v1p ->
withCVector v2 $ \v2p ->
withCVector v3 $ \v3p ->
unsafe2PtrMapInc i1 i2 i3 f v1p v2p v3p n
where
i1 = inc v1
i2 = inc v2
i3 = inc v3
n = minimum [(vectorLength v1), (vectorLength v2), vectorLength v3]
{-| Computes v2 <- alpha * v1 + v2. The result is stored in the memory of v2, therefore this is
unsafe and low level, only for internal use. -}
unsafeVectorAdd :: (BlasOps e, CVector vec e) =>
e -- ^ alpha
-> vec e -- ^ Vector 1
-> vec e -- ^ Vector 2, will be changed in place!
-> IO ()
unsafeVectorAdd alpha v1 v2 | n == n2 =
withCVector v1 $ \p1 ->
withCVector v2 $ \p2 ->
axpy n alpha p1 (inc v1) p2 (inc v2)
where
n = vectorLength v1
n2 = vectorLength v2
unsafeVectorAdd _ v1 v2 | otherwise = error $ "unsafeVectorAdd: Vector lengths must match, when adding " ++ show v1 ++ "\nand\n" ++ show v2
{-# NOINLINE unsafeVectorAdd #-}
{-# SPECIALIZE NOINLINE unsafeVectorAdd :: CFloat -> Vector CFloat -> Vector CFloat -> IO () #-}
{-# SPECIALIZE NOINLINE unsafeVectorAdd :: CDouble -> Vector CDouble -> Vector CDouble -> IO () #-}
{-# SPECIALIZE NOINLINE unsafeVectorAdd :: Complex CFloat -> Vector (Complex CFloat) -> Vector (Complex CFloat) -> IO () #-}
{-# SPECIALIZE NOINLINE unsafeVectorAdd :: Complex CDouble -> Vector (Complex CDouble) -> Vector (Complex CDouble) -> IO () #-}
unsafeSetElem :: (BlasOps e, CVector vec e) => vec e -> Index -> e -> IO ()
unsafeSetElem v i e | i >= 0 && i < vectorLength v = withCVector v $
\p -> let p' = p `plusPtr` (i * sizeOf e * (inc v)) in poke p' e
unsafeSetElem _ _ _ | otherwise = error "unsafeSetElem: out of bounds."
unsafeGetElem :: (BlasOps e, CVector vec e) => vec e -> Index -> IO e
unsafeGetElem v i | i >= 0 && i < vectorLength v = withCVector v $ \p -> do
e1 <- peek p
let p' = p `plusPtr` (i * sizeOf e1 * (inc v))
peek p'
unsafeGetElem _ _ | otherwise = error "unsafeGetElem: out of bounds."
unsafeFillVector :: (BlasOps e, CVector vec e) => vec e -> e -> IO ()
unsafeFillVector v e =
withCVector v $ \p ->
unsafePtrMap1Inc i (const e) p n
where
-- Using /f/ instead of the unsafePtrMap1Inc should be /slightly/ more efficient,
-- but it is a good idea to have such dirty functions in a central place.
-- f _ _ 0 = return ()
-- f i p n = poke p e >> f i (advancePtr p i) (n - 1) where {p' = advancePtr p i; n' = n - 1 }
i = inc v
n = vectorLength v