module System.Random.MWC.Monad (
Rand
, liftR
, RandIO
, asRandIO
, RandST
, asRandST
, Seed
, runRand
, runWithSeed
, runWithVector
, runWithSystemRandom
, toRand
, uniform
, uniformR
, standard
, normal
, equiprobable
, choices
, save
) where
import Control.Applicative
import Control.Monad (ap)
import Control.Monad.ST (ST)
import Control.Monad.Primitive (PrimMonad, PrimState)
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Generic as G
import Data.Ord
import Data.Word (Word32)
import Debug.Trace
import qualified System.Random.MWC as MWC
import System.Random.MWC (Gen,Variate,Seed)
newtype Rand m a = Rand {
runRand :: Gen (PrimState m) -> m a
}
liftR :: PrimMonad m => m a -> Rand m a
liftR m = Rand $ const m
type RandST s a = Rand (ST s) a
asRandST :: RandST s a -> RandST s a
asRandST = id
type RandIO a = Rand IO a
asRandIO :: RandIO a -> RandIO a
asRandIO = id
runWithCreate :: PrimMonad m => Rand m a -> m a
runWithCreate m = runRand m =<< MWC.create
runWithVector :: (G.Vector v Word32, PrimMonad m) => Rand m a -> v Word32 -> m a
runWithVector m v = runRand m =<< MWC.initialize v
runWithSeed :: PrimMonad m => Seed -> Rand m a -> m a
runWithSeed seed m = runRand m =<< MWC.restore seed
runWithSystemRandom :: PrimMonad m => Rand m a -> IO a
runWithSystemRandom = MWC.withSystemRandom . runRand
instance PrimMonad m => Functor (Rand m) where
fmap f (Rand rnd) = Rand $ \g -> (return . f) =<< (rnd g)
instance PrimMonad m => Monad (Rand m) where
return x = Rand (\_ -> return x)
(Rand rnd) >>= f = Rand $ \g -> (\x -> runRand (f x) g) =<< rnd g
instance (PrimMonad m) => Applicative (Rand m) where
pure = return
(<*>) = ap
toRand :: PrimMonad m => (Gen (PrimState m) -> m a) -> Rand m a
toRand = Rand
uniform :: (PrimMonad m, Variate a) => Rand m a
uniform = Rand MWC.uniform
uniformR :: (PrimMonad m, Variate a) => (a,a) -> Rand m a
uniformR rng = Rand (MWC.uniformR rng)
standard :: PrimMonad m => Rand m Double
standard = Rand MWC.normal
normal :: PrimMonad m =>
Double
-> Double
-> Rand m Double
normal m s = (+ m) . (* s) <$> standard
equiprobable :: PrimMonad m => [Rand m a] -> Rand m a
equiprobable [] = error "System.Random.MWC.Monad.equiprobable: list must be nonempty"
equiprobable xs = worker (V.fromList xs)
where
worker v = do
p <- uniform
v V.! truncate ((p 2^^(53)) * fromIntegral (V.length v) :: Double)
choices :: PrimMonad m => [(Double,Rand m a)] -> Rand m a
choices xs
| null xs' = error "System.Random.MWC.Monad.choices: list must contain at least one nonnegative weight"
| otherwise = traceShow (ps, V.length vect)
$ worker vect ps
where
xs' = filter ((>0) . fst) xs
vect = V.fromList (map snd xs')
ps = U.init . U.scanl' (+) 0 $ U.map (/ U.sum q) q
where q = U.fromList (map fst xs')
worker vect probs = do
p <- uniform
let i = binary probs p
vect V.! (i 1)
binary :: (Ord a, U.Unbox a) => U.Vector a -> a -> Int
binary v x = binaryRange v x 0 (U.length v)
binaryRange :: (Ord a, U.Unbox a) => U.Vector a -> a -> Int -> Int -> Int
binaryRange v x = loop
where
loop i j | i >= j = j
| otherwise = case compare (U.unsafeIndex v k) x of
LT -> loop (k+1) j
EQ -> k
GT -> loop i k
where k = (i + j) `div` 2
save :: PrimMonad m => Rand m Seed
save = Rand MWC.save