numerical-0.0.0.0: src/Numerical/Array/Layout/Builder.hs
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables#-}
-- {-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE NoImplicitPrelude #-}
module Numerical.Array.Layout.Builder where
import Control.Monad.Primitive ( PrimMonad, PrimState )
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Generic.Mutable as VGM
import Numerical.Array.Layout.Base
import Numerical.Array.Layout.Dense as Dense
--import Numerical.Array.Layou.Sparse as Sparse
--import Numerical.Data.Vector.Pair
import Control.Monad.ST (runST)
import Data.Typeable
import qualified Data.Foldable as F
import Data.Traversable as T
import Control.Applicative as A
import Numerical.Data.Vector.Pair
import Numerical.Array.Layout.Sparse
import Data.Vector.Algorithms.Intro as IntroSort
import Data.List (group)
import Numerical.InternalUtils
import Prelude hiding (error)
data BatchInit v = BatchInit { batchInitSize :: !Int
,batchInitKV :: !(Either [v] (IntFun v)) }
deriving (Typeable)
materializeBatchMV :: (PrimMonad m, VGM.MVector mv a) => BatchInit a -> m (mv (PrimState m) a)
materializeBatchMV (BatchInit size (Left ls )) =
do
newMV <- VGM.new size
_ <- Prelude.mapM (\(ix ,val )-> VGM.unsafeWrite newMV ix val ) (zip [0..] $ take size ls)
return newMV
materializeBatchMV (BatchInit size (Right (IntFun f) )) =
do
newMV <- VGM.new size
_ <- Prelude.mapM (\ix -> do v <- (f ix) ; VGM.unsafeWrite newMV ix v ) $ take size [0..]
return newMV
--- not sure if this is EVER useful
newtype AnyMV mv e = AMV (forall s . mv s e )
instance (Show a)=> Show (BatchInit a) where
show (BatchInit size (Left ls) ) | size > 100 = "(BatchInit " ++show size ++
"-- only showing the first 100 elements\n"
++ "(Left "++(show $ take 100 ls ) ++ "))\n"
| otherwise ="(BatchInit " ++show size ++
" (Left "++(show ls ) ++ "))\n"
show (BatchInit size (Right (IntFun f)) ) | size > 100 = "(BatchInit " ++show size ++
"-- only showing the first 100 elements\n"
++ "(Left "++(show $ runST (Prelude.mapM f [0..100]) ) ++ "))\n"
| otherwise ="(BatchInit " ++show size
++ "(Left "++(show $ runST (Prelude.mapM f [0,1..size -1]) ) ++ "))\n"
newtype IntFun a = IntFun (forall m. (PrimMonad m)=> Int -> m a )
-- This may change substantially in a future release, but for now
-- acts like
deriving (Typeable)
instance Functor IntFun where
fmap f (IntFun g) = IntFun (\x-> g x >>= (\ y -> return (f y)) )
{-# INLINE fmap #-}
instance Functor BatchInit where
{-# INLINE fmap #-}
fmap = \f bival ->
case bival of
(BatchInit size (Left ls))->
BatchInit size (Left (Prelude.map f ls ))
(BatchInit size (Right gfun))->
BatchInit size (Right $ fmap f gfun )
-- batchInit size should be Word rather than Int for size, but Vector is lame like that
{-
ChoiceT from monad lib is tempting
as is one of the ListT done right
Bundle from Vector 0.11 and Stream from 0.10 are both alluring too
but all of them make things complicated,
punt for now
ALso: I may want/need to distinguish sparse vs dense builders
and put them into different classes, punting that for now
-}
fromListBI :: [a] -> BatchInit a
fromListBI ls = BatchInit (length ls) (Left ls)
fromVectorBI :: VG.Vector v e => v e -> BatchInit e
fromVectorBI v = BatchInit size
(Right
(IntFun $
\i -> if i >= size
then error $ " out of bounds index on IntFun of size: " ++ show i
else return $ v VG.! i
))
where
size = VG.length v
fromMVectorBI :: (VGM.MVector mv e ) => AnyMV mv e -> BatchInit e
fromMVectorBI (AMV v) = BatchInit size
(Right
(IntFun $
\i -> if i >= size
then error $ " out of bounds index on IntFun of size: " ++ show i
else v `VGM.read` i
))
where
size = VGM.length v
class Layout form (rank::Nat) => LayoutBuilder form (rank::Nat) | form -> rank where
buildFormatM :: (store~FormatStorageRep form,Buffer store Int ,Buffer store a,PrimMonad m)=>
Index rank -> proxy form -> a
-> Maybe (BatchInit (Index rank ,a))
->m (form, BufferMut store (PrimState m) a )
buildFormatPure:: forall store form rank proxy m a. (LayoutBuilder form (rank::Nat)
,store~FormatStorageRep form,Buffer store Int ,Buffer store a, Monad m ) =>
Index rank -> proxy form -> a -> Maybe (BatchInit (Index rank ,a))
->m (form, BufferPure store a )
buildFormatPure shape prox defaultValue builder =
do res@(!_,!_)<-return $! theComputation
return res
where
theComputation :: (form,BufferPure store a )
!theComputation = runST $
do (form,buf) <- buildFormatM shape prox defaultValue builder
pureBuff <- VG.unsafeFreeze buf
return (form, pureBuff)
{-
this is a funky api for both dense and sparse arrays general builder format.
given the target shape, logical dimensions,a default value (only used for dense arrays)
and the list of manifest values (mostly only used for sparse), build the format
descriptor and the suitably initialized and sized values buffer
this api is only meant for internal use for building new array values
TODO: compare using a catenable priority heap vs just doing fast sorting.
-}
{-
the dense instances ignore the builder structure, which does suggest that maybe
there shoudl be a dense builder layout class and a sparse layout class separately
-}
instance LayoutBuilder (Format Direct 'Contiguous ('S 'Z) rep) ('S 'Z) where
buildFormatM (size:* _) _ defaultValue _ =
do
buf<- VGM.replicate size defaultValue
return (FormatDirectContiguous size,buf)
-- really wish I didn't have to write the foldable and traversable constraints
-- seems like a code smell?!
instance (F.Foldable (Shape r),T.Traversable (Shape r) ,A.Applicative (Shape r))
=> LayoutBuilder (Format Row 'Contiguous r rep) r where
buildFormatM ix _ defaultValue _ =
do
buf<- VGM.replicate (F.foldl' (*) 0 ix) defaultValue
return (FormatRowContiguous ix,buf)
instance (F.Foldable (Shape r),T.Traversable (Shape r) ,A.Applicative (Shape r))
=> LayoutBuilder (Format Column 'Contiguous r rep) r where
buildFormatM ix _ defaultValue _ =
do
buf<- VGM.replicate (F.foldl' (*) 0 ix) defaultValue
return (FormatColumnContiguous ix,buf)
isStrictlyMonotonicV ::(VG.Vector v e)=> (e -> e->Ordering)-> v e -> Maybe Int
isStrictlyMonotonicV cmp v = go 0 (VG.length v)
where
go !i !len | i+1 >= len = Nothing
| (v VG.! i) `lt` (v VG.! (i+1))= go (i+1) len
| otherwise = Just i
lt a b = case cmp a b of
LT -> True
_ -> False
instance (Buffer rep Int)=>LayoutBuilder (Format DirectSparse 'Contiguous ('S 'Z) rep ) ('S 'Z) where
buildFormatM (size:* _) _ _ Nothing = do
mvI <- VGM.new 0
vI <- VG.unsafeFreeze mvI
mvV <- VGM.new 0
return $! (FormatDirectSparseContiguous size 0 vI, mvV)
buildFormatM (size:* _) _ _ (Just builder)= do
-- need to use let so type inference doesnt totally barf
mvt@(MVPair (MVLeaf ix) (MVLeaf val)) <- materializeBatchMV $ fmap ( \((ix:*_),v)-> (ix,v)) builder
-- if i swap to using this instead of ix <- mat.. ; val <- mat..
--i get CRAZY type errors
-- could this be a ghc bug?
--ix <- materializeBatchMV $ fmap fst builtTup
--val <- materializeBatchMV $ fmap snd builtTup
_<- IntroSort.sortBy (\x y -> compare (fst x) (fst y)) mvt
-- this lets me sort a pair of arrays!
vIx <- VG.unsafeFreeze ix
optFail <- return $ isStrictlyMonotonicV compare vIx
--_hoelly
case optFail of
Nothing -> return (FormatDirectSparseContiguous size 0 vIx, val)
Just ixWrong -> error $ "DirectSparse Index duplication at index "++ show (vIx VG.! ixWrong)
instance (Buffer rep Int) => LayoutBuilder (Format CompressedSparseRow 'Contiguous ('S ('S 'Z)) rep ) ('S ('S 'Z)) where
buildFormatM (x:* y :* _) _ _ Nothing= do
mvi <- VGM.new 0
vi <- VG.unsafeFreeze mvi
mvval <- VGM.new 0
return $
(FormatContiguousCompressedSparseRow
(FormatContiguousCompressedSparseInternal y x vi vi), mvval )
buildFormatM (x:* y :* _) proxyFormat _ (Just builder) = do
mvtup@(MVPair (MVPair (MVLeaf mvectYs) (MVLeaf mvectXs)) (MVLeaf mvectVals))<-
materializeBatchMV $ fmap (\((xix:* yix :* _),val)-> ((yix,xix),val) ) builder
_ <- IntroSort.sortBy (\((y1,x1),_) ((y2,x2),_) -> basicCompareIndex proxyFormat (x1:*y1 :* Nil) (x2:*y2:* Nil) )
mvtup
vectXs <- unsafeBufferFreeze mvectXs
vectYs <- unsafeBufferFreeze mvectYs
--- predicate check here wrt monotonicity of
--- compute runlength partial sums of where ys go
-- need to actually check
yRunsVect <- return $
VG.replicate (y+1) (0::Int) VG.// computeStarts (computeRunLengths vectYs) 0 y
--_ <- (error "computeRUnCount") vectYs yRunsMVect
--yRunsVect <- unsafeBufferFreeze yRunsMVect
let xyVect = (VPair (VLeaf vectXs) (VLeaf vectYs))
optFail <- return $
isStrictlyMonotonicV (\(x1,y1) (x2,y2)->basicCompareIndex proxyFormat (x1:*y1:*Nil) (x2:*y2:*Nil))
xyVect
case optFail of
Nothing -> return $
(FormatContiguousCompressedSparseRow
(FormatContiguousCompressedSparseInternal y x vectXs yRunsVect), mvectVals )
Just i ->
error $ "illegal duplication in CSR builder (x,y) coordinates "
++ show (xyVect VG.! i) ++ " and " ++ show (xyVect VG.! (i+1))
++ "starting at position " ++ show i
computeRunLengths :: (VG.Vector v e, Eq e)=> v e -> [(e,Int)]
computeRunLengths = \y -> fmap (\x ->(head x,length x)) $ group $ VG.toList y
{-# SPECIALIZE INLINE computeStarts :: [(Int,Int)]->Int->Int ->[(Int,Int)] #-}
computeStarts:: (Enum a, Ord a, Num b )=>[(a,b)]-> a -> a -> [(a,b)]
computeStarts [] start end | start <= end = fmap (\x -> (x ,0)) [start..end]
| otherwise = error "bad start end arguments to computeStarts"
computeStarts ls start end | start <= end = go start 0 ls
| otherwise = error "bad start end arguments to computeStarts"
where
--go :: a ->b->[(a,b)]-> [(a,b)]
go !posNext preSum [] | posNext <= end = fmap (\x -> (x,preSum)) [posNext .. end]
| otherwise = error "impossible go computeStarts "
go !posNext !preSum gls@((posAt,atSum):rest)
| posNext < posAt= (posNext,preSum): go (succ posNext) preSum gls
| posNext == posAt = (posNext,preSum) : go (succ posNext) (preSum + atSum) rest
| otherwise = error "bad position in prefix stream for computeStarts go, literally unpossible "
--computeStarts :: (Eq a, Num a)=> [(a,Int)]->Int -> [(a,Int)]
--computeStarts [] len = map (\x -> (x ,0)) [0..len]
--computeStarts ls len = go 0 0 ls
-- where
-- go preSum place [] | place > len = []
-- | place == len = [(place,preSum)]
-- | otherwise = map