program/RandomizedMonad.hs

   1 module RandomizedMonad (
   2         Randomized,
   3         msplit,
   4         runRandom1, runRandom, runRandomStd, runRandomNewStd,
   5         mrandomR, mrandom,
   6         withProb, withWeight,
   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 -- Like withProb, but without an else case and with the "probabilities" scaled
  74 -- so that they sum to 1.
  75 withWeight :: [(Double, Randomized a)] -> Randomized a
  76 withWeight ifCs = do
  77         val <- mrandomR (0, sum (map fst ifCs))
  78         chooseCase val (tail ifCs) (snd (head ifCs))
  79
  80 -- Keep trying until we get what we want.
  81 filterRandomized :: (a -> Bool) -> Randomized a -> Randomized a
  82 filterRandomized f ra = do
  83         a <- ra
  84         if f a then return a else filterRandomized f ra
  85
  86 -- A randomized list of elements chosen independently from a distribution.
  87 -- Each element is under msplit for parallelizability.
  88 indReplicateRandom :: Int -> Randomized a -> Randomized [a]
  89 indReplicateRandom n ra = sequence $ replicate n $ msplit ra
  90
  91 -- An infinite randomized list of elements chosen independently from a
  92 -- distribution.  The list is under msplit to avoid an infinite loop when it is
  93 -- bound.
  94 indRepeatRandom :: Randomized a -> Randomized [a]
  95 indRepeatRandom ra = msplit $ sequence $ repeat $ msplit ra
  96
  97 -- Produces an array of elements chosen independently from a distribution.
  98 indRandomArray :: (IArray a e, Ix i) =>
  99         (i, i) -> Randomized e -> Randomized (a i e)
 100 indRandomArray bds re = do
 101         list <- indReplicateRandom (rangeSize bds) re
 102         return (listArray bds list)