program/RandomizedMonad.hs

   1 module RandomizedMonad (
   2         Randomized,
   3         msplit,
   4         runRandom1, runRandom, runRandomStd, runRandomNewStd,
   5         mrandomR, mrandom,
   6         withProb,
   7         filterRandomized,
   8         indReplicateRandom, indRepeatRandom, indRandomArray
   9 ) where
  10 import System.Random
  11 import Data.Array.IArray
  12 import Data.Ix
  13
  14 -- Needs -XRank2Types
  15 newtype Randomized a = Randomized (forall g. RandomGen g => (g -> (a, g)))
  16
  17 -- This implementation threads a single RandomGen through the whole process in
  18 -- order to satisfy the monad laws.
  19 instance Monad Randomized where
  20         ma >>= amb = Randomized (\g -> let
  21                         Randomized fa = ma
  22                         (a, g2) = fa g
  23                         Randomized fb = amb a
  24                         in fb g2
  25                 )
  26         return x = Randomized (\g -> (x, g))
  27
  28 -- Splits the generator and runs the argument on the left generator while
  29 -- threading the right generator on.  C.f. unsaveInterleaveIO.  Use this to
  30 -- make a sub-calculation parallelizable and evolvable without breaking
  31 -- same-seed reproducibility of the whole calculation.
  32 msplit :: Randomized a -> Randomized a
  33 msplit (Randomized fa) = Randomized
  34         (\g -> let (g1, g2) = split g in (fst (fa g1), g2))
  35
  36 runRandom1 :: RandomGen g => g -> Randomized a -> (a, g)
  37 runRandom1 g (Randomized fa) = fa g
  38
  39 runRandom :: RandomGen g => g -> Randomized a -> a
  40 runRandom g (Randomized fa) = fst (fa g)
  41
  42 -- Conveniences
  43 runRandomStd :: Randomized a -> IO a
  44 runRandomStd ra = do
  45         g <- getStdGen
  46         return $ runRandom g ra
  47
  48 runRandomNewStd :: Randomized a -> IO a
  49 runRandomNewStd ra = do
  50         g <- newStdGen
  51         return $ runRandom g ra
  52
  53 -- Monadic versions of random and randomR (to generate primitive-ish values)
  54 mrandom :: Random a => Randomized a
  55 mrandom = Randomized random
  56 mrandomR :: Random a => (a, a) -> Randomized a
  57 -- Eta-expand this one to keep GHC 6.6.1 on birdy happy.
  58 mrandomR lohi = Randomized (\g -> randomR lohi g)
  59
  60 chooseCase :: Double -> [(Double, a)] -> a -> a
  61 chooseCase val ifCs elseR = case ifCs of
  62         [] -> elseR
  63         (cutoff, theR):ifCt -> if val < cutoff
  64                 then theR
  65                 else chooseCase (val - cutoff) ifCt elseR
  66
  67 -- An if-elsif...else-style construct where each "if" has a probability.
  68 withProb :: [(Double, Randomized a)] -> Randomized a -> Randomized a
  69 withProb ifCs elseR = do
  70         val <- mrandom
  71         chooseCase val ifCs elseR
  72
  73 -- Keep trying until we get what we want.
  74 filterRandomized :: (a -> Bool) -> Randomized a -> Randomized a
  75 filterRandomized f ra = do
  76         a <- ra
  77         if f a then return a else filterRandomized f ra
  78
  79 -- A randomized list of elements chosen independently from a distribution.
  80 -- Each element is under msplit for parallelizability.
  81 indReplicateRandom :: Int -> Randomized a -> Randomized [a]
  82 indReplicateRandom n ra = sequence $ replicate n $ msplit ra
  83
  84 -- An infinite randomized list of elements chosen independently from a
  85 -- distribution.  The list is under msplit to avoid an infinite loop when it is
  86 -- bound.
  87 indRepeatRandom :: Randomized a -> Randomized [a]
  88 indRepeatRandom ra = msplit $ sequence $ repeat $ msplit ra
  89
  90 -- Produces an array of elements chosen independently from a distribution.
  91 indRandomArray :: (IArray a e, Ix i) =>
  92         (i, i) -> Randomized e -> Randomized (a i e)
  93 indRandomArray bds re = do
  94         list <- indReplicateRandom (rangeSize bds) re
  95         return (listArray bds list)