Skip to content

HTTPS clone URL

Subversion checkout URL

You can clone with HTTPS or Subversion.

Download ZIP
tree: 42c63b1972
Fetching contributors…

Cannot retrieve contributors at this time

221 lines (194 sloc) 7.731 kb
-----------------------------------------------------------------------------
-- |
-- Module : Control.Monad.MC.Sample
-- Copyright : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability : experimental
--
module Control.Monad.MC.Sample (
-- * Sampling from lists
sample,
sampleWithWeights,
sampleSubset,
sampleSubset',
-- * Sampling @Int@s
sampleInt,
sampleIntWithWeights,
sampleIntSubset,
sampleIntSubset',
-- * Shuffling
shuffle,
shuffleInt,
shuffleInt',
) where
import Control.Monad
import Control.Monad.ST
import Control.Monad.MC.Base
import Control.Monad.MC.Walker
import Data.Vector.Unboxed( MVector, Unbox )
import qualified Data.Vector as BV
import qualified Data.Vector.Mutable as BMV
import qualified Data.Vector.Unboxed as V
import qualified Data.Vector.Generic.Mutable as MV
-- | @sample xs@ samples a value uniformly from the elements of @xs@. The
-- results are undefined if @length xs@ is zero.
sample :: (MonadMC m) => [a] -> m a
sample xs = let
n = length xs
in sampleHelp n xs $ sampleInt n
{-# INLINE sample #-}
-- | @sampleWithWeights wxs@ samples a value from the list with the given
-- weight.
sampleWithWeights :: (MonadMC m) => [(Double, a)] -> m a
sampleWithWeights wxs = let
(ws,xs) = unzip wxs
n = length xs
in sampleHelp n xs $ sampleIntWithWeights ws n
{-# INLINE sampleWithWeights #-}
-- | @sampleSubset k xs@ samples a subset of size @k@ from @xs@ by
-- sampling without replacement. The return value is a list of length @k@
-- with the elements in the subset in the order that they were sampled. Note
-- also that the elements are lazily generated.
sampleSubset :: (MonadMC m) => Int -> [a] -> m [a]
sampleSubset k xs = let
n = length xs
in sampleListHelp n xs $ sampleIntSubset k n
{-# INLINE sampleSubset #-}
-- | Strict version of 'sampleSubset'.
sampleSubset' :: (MonadMC m) => Int -> [a] -> m [a]
sampleSubset' k xs = do
s <- sampleSubset k xs
length s `seq` return s
{-# INLINE sampleSubset' #-}
sampleHelp :: (Monad m) => Int -> [a] -> m Int -> m a
sampleHelp _n xs f = let
arr = BV.fromList xs
in liftM (BV.unsafeIndex arr) f
{-# INLINE sampleHelp #-}
sampleHelpU :: (Unbox a, Monad m) => Int -> [a] -> m Int -> m a
sampleHelpU _n xs f = let
arr = V.fromList xs
in liftM (V.unsafeIndex arr) f
{-# INLINE sampleHelpU #-}
{-# RULES "sampleHelp/Double" forall n xs f.
sampleHelp n (xs :: [Double]) f = sampleHelpU n xs f #-}
{-# RULES "sampleHelp/Int" forall n xs f.
sampleHelp n (xs :: [Int]) f = sampleHelpU n xs f #-}
sampleListHelp :: (Monad m) => Int -> [a] -> m [Int] -> m [a]
sampleListHelp _n xs f = let
arr = BV.fromList xs
in liftM (map $ BV.unsafeIndex arr) f
{-# INLINE sampleListHelp #-}
sampleListHelpU :: (Unbox a, Monad m) => Int -> [a] -> m [Int] -> m [a]
sampleListHelpU _n xs f = let
arr = V.fromList xs
in liftM (map $ V.unsafeIndex arr) f
{-# RULES "sampleListHelp/Double" forall n xs f.
sampleListHelp n (xs :: [Double]) f = sampleListHelpU n xs f #-}
{-# RULES "sampleListHelp/Int" forall n xs f.
sampleListHelp n (xs :: [Int]) f = sampleListHelpU n xs f #-}
-- | @sampleInt n@ samples integers uniformly from @[ 0..n-1 ]@. It is an
-- error to call this function with a non-positive @n@.
sampleInt :: (MonadMC m) => Int -> m Int
sampleInt n | n < 1 = fail "invalid argument"
| otherwise = uniformInt n
{-# INLINE sampleInt #-}
-- | @sampleIntWithWeights ws n@ samples integers from @[ 0..n-1 ]@ with the
-- probability of choosing @i@ proportional to @ws !! i@. The list @ws@ must
-- have length equal to @n@. Also, the elements of @ws@ must be non-negative
-- with at least one nonzero entry.
sampleIntWithWeights :: (MonadMC m) => [Double] -> Int -> m Int
sampleIntWithWeights ws n =
let qjs = computeTable n ws
in liftM (indexTable qjs) (uniform 0 1)
{-# INLINE sampleIntWithWeights #-}
-- | @sampleIntSubset k n@ samples a subset of size @k@ by sampling without
-- replacement from the integers @{ 0, ..., n-1 }@. The return value is a
-- list of length @k@ with the elements in the subset in the order that they
-- were sampled. Note also that the elements are lazily generated.
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 <- MV.new n :: ST s (MVector s Int)
sequence_ [ MV.unsafeWrite ints i i | i <- [0 .. n-1] ]
sampleIntSubsetHelp ints us (n-1)
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 <- MV.unsafeRead ints u
MV.unsafeWrite ints u =<< MV.unsafeRead ints n'
is <- sampleIntSubsetHelp ints us (n'-1)
return (i:is)
{-# INLINE sampleIntSubset #-}
-- | Strict version of 'sampleIntSubset'.
sampleIntSubset' :: (MonadMC m) => Int -> Int -> m [Int]
sampleIntSubset' k n = do
s <- sampleIntSubset k n
length s `seq` return s
{-# INLINE sampleIntSubset' #-}
-- | @shuffle xs@ randomly permutes the list @xs@ and returns
-- the result. All permutations of the elements of @xs@ are equally
-- likely.
shuffle :: (MonadMC m) => [a] -> m [a]
shuffle xs = let
n = length xs
in shuffleInt n >>= \swaps -> (return . BV.toList . BV.create) $ do
marr <- MV.new n :: ST s (BMV.MVector s a)
zipWithM_ (MV.unsafeWrite marr) [0 .. n-1] xs
mapM_ (swap marr) swaps
return marr
where
swap marr (i,j) | i == j = return ()
| otherwise = do
x <- MV.unsafeRead marr i
y <- MV.unsafeRead marr j
MV.unsafeWrite marr i y
MV.unsafeWrite marr j x
{-# INLINE shuffle #-}
shuffleU :: (Unbox a, MonadMC m) => [a] -> m [a]
shuffleU xs = let
n = length xs
in shuffleInt n >>= \swaps -> (return . V.toList . V.create) $ do
marr <- MV.new n
zipWithM_ (MV.unsafeWrite marr) [0 .. n-1] xs
mapM_ (swap marr) swaps
return marr
where
swap marr (i,j) | i == j = return ()
| otherwise = do
x <- MV.unsafeRead marr i
y <- MV.unsafeRead marr j
MV.unsafeWrite marr i y
MV.unsafeWrite marr j x
{-# INLINE shuffleU #-}
{-# RULES "shuffle/Double" forall xs.
shuffle (xs :: [Double]) = shuffleU xs #-}
{-# RULES "shuffle/Int" forall xs.
shuffle (xs :: [Int]) = shuffleU xs #-}
-- | @shuffleInt n@ generates a sequence of swaps equivalent to a
-- uniformly-chosen random permutatation of the integers @{0, ..., n-1}@.
-- For an input of @n@, there are @n-1@ swaps, which are lazily generated.
shuffleInt :: (MonadMC m) => Int -> m [(Int,Int)]
shuffleInt n =
let shuffleIntHelp i | i <= 1 = return []
| otherwise = unsafeInterleaveMC $ do
j <- uniformInt i
ijs <- shuffleIntHelp (i-1)
return $ (i-1,j):ijs in
shuffleIntHelp n
{-# INLINE shuffleInt #-}
-- | Strict version of 'shuffleInt'.
shuffleInt' :: (MonadMC m) => Int -> m [(Int,Int)]
shuffleInt' n = do
ss <- shuffleInt n
length ss `seq` return ss
{-# INLINE shuffleInt' #-}
Jump to Line
Something went wrong with that request. Please try again.