Commit | Line | Data |
---|---|---|
967c39ef MM |
1 | module RandomizedMonad ( |
2 | Randomized, | |
0df1b3e1 MM |
3 | msplit, |
4 | runRandom1, runRandom, runRandomStd, runRandomNewStd, | |
967c39ef | 5 | mrandomR, mrandom, |
8723ed6a | 6 | withProb, withWeight, |
066d7f53 | 7 | filterRandomized, |
0df1b3e1 | 8 | indReplicateRandom, indRepeatRandom, indRandomArray |
967c39ef | 9 | ) where |
070511a2 | 10 | import Control.Monad |
967c39ef | 11 | import System.Random |
066d7f53 MM |
12 | import Data.Array.IArray |
13 | import Data.Ix | |
967c39ef MM |
14 | |
15 | -- Needs -XRank2Types | |
0df1b3e1 | 16 | newtype Randomized a = Randomized (forall g. RandomGen g => (g -> (a, g))) |
967c39ef | 17 | |
0df1b3e1 MM |
18 | -- This implementation threads a single RandomGen through the whole process in |
19 | -- order to satisfy the monad laws. | |
070511a2 MM |
20 | |
21 | -- Migrate according to the guide at | |
22 | -- https://gitlab.haskell.org/ghc/ghc/-/wikis/migration/7.10#ghc-says-no-instance-for-applicative- | |
23 | -- ~ 2021-08-12 | |
24 | instance Functor Randomized where | |
25 | fmap = liftM | |
26 | ||
27 | instance Applicative Randomized where | |
28 | pure x = Randomized (\g -> (x, g)) | |
29 | (<*>) = ap | |
30 | ||
967c39ef MM |
31 | instance Monad Randomized where |
32 | ma >>= amb = Randomized (\g -> let | |
967c39ef | 33 | Randomized fa = ma |
0df1b3e1 | 34 | (a, g2) = fa g |
967c39ef MM |
35 | Randomized fb = amb a |
36 | in fb g2 | |
37 | ) | |
0df1b3e1 MM |
38 | |
39 | -- Splits the generator and runs the argument on the left generator while | |
40 | -- threading the right generator on. C.f. unsaveInterleaveIO. Use this to | |
41 | -- make a sub-calculation parallelizable and evolvable without breaking | |
42 | -- same-seed reproducibility of the whole calculation. | |
43 | msplit :: Randomized a -> Randomized a | |
44 | msplit (Randomized fa) = Randomized | |
45 | (\g -> let (g1, g2) = split g in (fst (fa g1), g2)) | |
46 | ||
47 | runRandom1 :: RandomGen g => g -> Randomized a -> (a, g) | |
48 | runRandom1 g (Randomized fa) = fa g | |
967c39ef MM |
49 | |
50 | runRandom :: RandomGen g => g -> Randomized a -> a | |
0df1b3e1 | 51 | runRandom g (Randomized fa) = fst (fa g) |
967c39ef MM |
52 | |
53 | -- Conveniences | |
54 | runRandomStd :: Randomized a -> IO a | |
55 | runRandomStd ra = do | |
56 | g <- getStdGen | |
57 | return $ runRandom g ra | |
58 | ||
59 | runRandomNewStd :: Randomized a -> IO a | |
60 | runRandomNewStd ra = do | |
61 | g <- newStdGen | |
62 | return $ runRandom g ra | |
63 | ||
64 | -- Monadic versions of random and randomR (to generate primitive-ish values) | |
967c39ef | 65 | mrandom :: Random a => Randomized a |
0df1b3e1 MM |
66 | mrandom = Randomized random |
67 | mrandomR :: Random a => (a, a) -> Randomized a | |
e42ffb75 MM |
68 | -- Eta-expand this one to keep GHC 6.6.1 on birdy happy. |
69 | mrandomR lohi = Randomized (\g -> randomR lohi g) | |
967c39ef MM |
70 | |
71 | chooseCase :: Double -> [(Double, a)] -> a -> a | |
72 | chooseCase val ifCs elseR = case ifCs of | |
73 | [] -> elseR | |
74 | (cutoff, theR):ifCt -> if val < cutoff | |
75 | then theR | |
76 | else chooseCase (val - cutoff) ifCt elseR | |
77 | ||
0df1b3e1 | 78 | -- An if-elsif...else-style construct where each "if" has a probability. |
967c39ef MM |
79 | withProb :: [(Double, Randomized a)] -> Randomized a -> Randomized a |
80 | withProb ifCs elseR = do | |
81 | val <- mrandom | |
82 | chooseCase val ifCs elseR | |
83 | ||
8723ed6a MM |
84 | -- Like withProb, but without an else case and with the "probabilities" scaled |
85 | -- so that they sum to 1. | |
86 | withWeight :: [(Double, Randomized a)] -> Randomized a | |
87 | withWeight ifCs = do | |
88 | val <- mrandomR (0, sum (map fst ifCs)) | |
89 | chooseCase val (tail ifCs) (snd (head ifCs)) | |
90 | ||
967c39ef MM |
91 | -- Keep trying until we get what we want. |
92 | filterRandomized :: (a -> Bool) -> Randomized a -> Randomized a | |
93 | filterRandomized f ra = do | |
94 | a <- ra | |
95 | if f a then return a else filterRandomized f ra | |
066d7f53 | 96 | |
0df1b3e1 MM |
97 | -- A randomized list of elements chosen independently from a distribution. |
98 | -- Each element is under msplit for parallelizability. | |
99 | indReplicateRandom :: Int -> Randomized a -> Randomized [a] | |
100 | indReplicateRandom n ra = sequence $ replicate n $ msplit ra | |
101 | ||
102 | -- An infinite randomized list of elements chosen independently from a | |
103 | -- distribution. The list is under msplit to avoid an infinite loop when it is | |
104 | -- bound. | |
105 | indRepeatRandom :: Randomized a -> Randomized [a] | |
106 | indRepeatRandom ra = msplit $ sequence $ repeat $ msplit ra | |
107 | ||
108 | -- Produces an array of elements chosen independently from a distribution. | |
066d7f53 MM |
109 | indRandomArray :: (IArray a e, Ix i) => |
110 | (i, i) -> Randomized e -> Randomized (a i e) | |
0df1b3e1 MM |
111 | indRandomArray bds re = do |
112 | list <- indReplicateRandom (rangeSize bds) re | |
066d7f53 | 113 | return (listArray bds list) |