module Control.Monad.MC.Sample (
sample,
sampleWithWeights,
sampleSubset,
sampleInt,
sampleIntWithWeights,
sampleIntSubset,
shuffle,
shuffleInt,
) where
import Control.Monad
import Control.Monad.ST
import Control.Monad.MC.Base
import Control.Monad.MC.Walker
import Data.Array.Base
import Data.Array.IArray
import Data.Array.ST
import Data.Array.Vector
sample :: (MonadMC m) => Int -> [a] -> m a
sample n xs =
sampleHelp n xs $ sampleInt n
sampleWithWeights :: (MonadMC m) => [Double] -> Int -> [a] -> m a
sampleWithWeights ws n xs =
sampleHelp n xs $ sampleIntWithWeights ws n
sampleSubset :: (MonadMC m) => Int -> Int -> [a] -> m [a]
sampleSubset k n xs =
sampleListHelp n xs $ sampleIntSubset k n
sampleHelp :: (Monad m) => Int -> [a] -> m Int -> m a
sampleHelp n (xs :: [a]) f = let
arr = listArray (0,n1) xs :: Array Int a
in liftM (unsafeAt arr) f
sampleHelpUA :: (UA a, Monad m) => Int -> [a] -> m Int -> m a
sampleHelpUA n xs f = let
arr = newU n (\marr -> zipWithM_ (writeMU marr) [0..n1] xs)
in liftM (indexU arr) f
sampleListHelp :: (Monad m) => Int -> [a] -> m [Int] -> m [a]
sampleListHelp n (xs :: [a]) f = let
arr = listArray (0,n1) xs :: Array Int a
in liftM (map $ unsafeAt arr) f
sampleListHelpUA :: (UA a, Monad m) => Int -> [a] -> m [Int] -> m [a]
sampleListHelpUA n xs f = let
arr = newU n (\marr -> zipWithM_ (writeMU marr) [0..n1] xs)
in liftM (map $ indexU arr) f
sampleInt :: (MonadMC m) => Int -> m Int
sampleInt n | n < 1 = fail "invalid argument"
| otherwise = uniformInt n
sampleIntWithWeights :: (MonadMC m) => [Double] -> Int -> m Int
sampleIntWithWeights ws n =
let qjs = computeTable n ws
in liftM (indexTable qjs) (uniform 0 1)
sampleIntSubset :: (MonadMC m) => Int -> Int -> m [Int]
sampleIntSubset k n | k < 0 = fail "negative subset size"
| k > n = fail "subset size is too big"
| otherwise = do
us <- randomIndices k n
return $ runST $ do
ints <- newMU n
sequence_ [ writeMU ints i i | i <- [0 .. n1] ]
sampleIntSubsetHelp ints us (n1)
where
randomIndices k' n' | k' == 0 = return []
| otherwise = unsafeInterleaveMC $ do
u <- uniformInt n'
us <- randomIndices (k'1) (n'1)
return (u:us)
sampleIntSubsetHelp _ [] _ = return []
sampleIntSubsetHelp ints (u:us) n' = unsafeInterleaveST $ do
i <- readMU ints u
writeMU ints u =<< readMU ints n'
is <- sampleIntSubsetHelp ints us (n'1)
return (i:is)
shuffle :: (MonadMC m) => Int -> [a] -> m [a]
shuffle n (xs :: [a]) =
shuffleInt n >>= \swaps -> (return . runST) $ do
marr <- newListArray (0,n1) xs :: ST s (STArray s Int a)
mapM_ (swap marr) swaps
getElems marr
where
swap marr (i,j) | i == j = return ()
| otherwise = do
x <- unsafeRead marr i
y <- unsafeRead marr j
unsafeWrite marr i y
unsafeWrite marr j x
shuffleUA :: (UA a, MonadMC m) => Int -> [a] -> m [a]
shuffleUA n (xs :: [a]) =
shuffleInt n >>= \swaps -> (return . runST) $ do
marr <- newMU n
zipWithM_ (writeMU marr) [0 .. n1] xs
mapM_ (swap marr) swaps
arr <- unsafeFreezeAllMU marr
return $ fromU arr
where
swap marr (i,j) | i == j = return ()
| otherwise = do
x <- readMU marr i
y <- readMU marr j
writeMU marr i y
writeMU marr j x
shuffleInt :: (MonadMC m) => Int -> m [(Int,Int)]
shuffleInt n =
let shuffleIntHelp i | i <= 1 = return []
| otherwise = unsafeInterleaveMC $ do
j <- uniformInt i
ijs <- shuffleIntHelp (i1)
return $ (i1,j):ijs in
shuffleIntHelp n