| 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) |