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 mrandomR lohi = Randomized $ randomR lohi
  58
  59 chooseCase :: Double -> [(Double, a)] -> a -> a
  60 chooseCase val ifCs elseR = case ifCs of
  61         [] -> elseR
  62         (cutoff, theR):ifCt -> if val < cutoff
  63                 then theR
  64                 else chooseCase (val - cutoff) ifCt elseR
  65
  66 -- An if-elsif...else-style construct where each "if" has a probability.
  67 withProb :: [(Double, Randomized a)] -> Randomized a -> Randomized a
  68 withProb ifCs elseR = do
  69         val <- mrandom
  70         chooseCase val ifCs elseR
  71
  72 -- Keep trying until we get what we want.
  73 filterRandomized :: (a -> Bool) -> Randomized a -> Randomized a
  74 filterRandomized f ra = do
  75         a <- ra
  76         if f a then return a else filterRandomized f ra
  77
  78 -- A randomized list of elements chosen independently from a distribution.
  79 -- Each element is under msplit for parallelizability.
  80 indReplicateRandom :: Int -> Randomized a -> Randomized [a]
  81 indReplicateRandom n ra = sequence $ replicate n $ msplit ra
  82
  83 -- An infinite randomized list of elements chosen independently from a
  84 -- distribution.  The list is under msplit to avoid an infinite loop when it is
  85 -- bound.
  86 indRepeatRandom :: Randomized a -> Randomized [a]
  87 indRepeatRandom ra = msplit $ sequence $ repeat $ msplit ra
  88
  89 -- Produces an array of elements chosen independently from a distribution.
  90 indRandomArray :: (IArray a e, Ix i) =>
  91         (i, i) -> Randomized e -> Randomized (a i e)
  92 indRandomArray bds re = do
  93         list <- indReplicateRandom (rangeSize bds) re
  94         return (listArray bds list)