| 1 |
{-# LANGUAGE TypeFamilies #-}
|
| 2 |
|
| 3 |
module Treant where
|
| 4 |
|
| 5 |
import qualified System.Random as Random
|
| 6 |
|
| 7 |
-- * The relevant types
|
| 8 |
|
| 9 |
newtype Size = Size { fromSize :: Int } deriving (Eq, Show)
|
| 10 |
|
| 11 |
data Gen r = Gen
|
| 12 |
{ fromGen :: Size -> Random.StdGen -> (r, Random.StdGen)
|
| 13 |
}
|
| 14 |
|
| 15 |
runGen :: Gen r -> IO r
|
| 16 |
runGen (Gen k) = do
|
| 17 |
gen <- Random.getStdGen
|
| 18 |
let (r, gen') = k (Size 0) gen
|
| 19 |
Random.setStdGen gen'
|
| 20 |
pure r
|
| 21 |
|
| 22 |
instance Functor Gen where
|
| 23 |
fmap f (Gen k) = Gen $ \ size gen ->
|
| 24 |
let (r, gen') = k size gen
|
| 25 |
in (f r, gen')
|
| 26 |
|
| 27 |
instance Applicative Gen where
|
| 28 |
pure r = Gen $ \_ gen -> (r, gen)
|
| 29 |
Gen f <*> Gen x = Gen $ \size gen ->
|
| 30 |
let (f', gen' ) = f size gen
|
| 31 |
(x', gen'') = x size gen'
|
| 32 |
in (f' x', gen'')
|
| 33 |
|
| 34 |
instance Monad Gen where
|
| 35 |
Gen x >>= f = Gen $ \size gen ->
|
| 36 |
let (x', gen') = x size gen
|
| 37 |
in fromGen (f x') size gen'
|
| 38 |
|
| 39 |
|
| 40 |
class Distribution dist where
|
| 41 |
type Item dist
|
| 42 |
generate :: dist -> Gen (Item dist)
|
| 43 |
|
| 44 |
-- * Typical mathy distributions
|
| 45 |
|
| 46 |
data Normal r = Normal
|
| 47 |
{ normalMu :: r
|
| 48 |
, normalSigma :: r
|
| 49 |
} deriving (Eq, Show)
|
| 50 |
|
| 51 |
instance Floating r => Distribution (Normal r) where
|
| 52 |
type Item (Normal r) = r
|
| 53 |
generate Normal {} = error "unimplemented"
|
| 54 |
|
| 55 |
data Cauchy r = Cauchy
|
| 56 |
{ cauchyX :: r
|
| 57 |
, cauchyGamma :: r
|
| 58 |
} deriving (Eq, Show)
|
| 59 |
|
| 60 |
instance Floating r => Distribution (Cauchy r) where
|
| 61 |
type Item (Cauchy r) = r
|
| 62 |
generate Cauchy {} = error "unimplemented"
|
| 63 |
|
| 64 |
-- * Categorical distribution
|
| 65 |
|
| 66 |
newtype Categorical a = Categorical
|
| 67 |
{ categoricalItems :: [a]
|
| 68 |
} deriving (Eq, Show)
|
| 69 |
|
| 70 |
instance Distribution (Categorical r) where
|
| 71 |
type Item (Categorical r) = r
|
| 72 |
generate (Categorical rs) = do
|
| 73 |
let len = length rs
|
| 74 |
index <- int
|
| 75 |
pure (rs !! (index `mod` len))
|
| 76 |
|
| 77 |
-- * Weighted distribution
|
| 78 |
|
| 79 |
newtype Weighted n a = Weighted
|
| 80 |
{ weightedItems :: [(n, a)]
|
| 81 |
} deriving (Eq, Show)
|
| 82 |
|
| 83 |
instance Integral n => Distribution (Weighted n r) where
|
| 84 |
type Item (Weighted n r) = r
|
| 85 |
generate (Weighted rs) = do
|
| 86 |
let bound = sum (map fst rs)
|
| 87 |
choice <- int
|
| 88 |
let idx = fromIntegral choice `mod` bound
|
| 89 |
findChoice _ [] = error "internal error"
|
| 90 |
findChoice n ((p, x):xs)
|
| 91 |
| n < p = x
|
| 92 |
| otherwise = findChoice (n-p) xs
|
| 93 |
pure (findChoice idx rs)
|
| 94 |
|
| 95 |
-- * Dirac distribution
|
| 96 |
|
| 97 |
newtype Dirac r = Dirac
|
| 98 |
{ diracConstant :: r
|
| 99 |
} deriving (Eq, Show)
|
| 100 |
|
| 101 |
instance Distribution (Dirac r) where
|
| 102 |
type Item (Dirac r) = r
|
| 103 |
generate dist = pure (diracConstant dist)
|
| 104 |
|
| 105 |
|
| 106 |
-- * Some typical generator functions
|
| 107 |
|
| 108 |
random :: Random.Random r => Gen r
|
| 109 |
random = Gen $ \size gen -> Random.random gen
|
| 110 |
|
| 111 |
int :: Gen Int
|
| 112 |
int = random
|
| 113 |
|
| 114 |
list :: Gen a -> Gen [a]
|
| 115 |
list = undefined
|
| 116 |
|
| 117 |
maybe :: Gen a -> Gen (Maybe a)
|
| 118 |
maybe = undefined
|
| 119 |
|
| 120 |
d :: Int -> Gen Int
|
| 121 |
d n = generate (Categorical [1..n])
|
| 122 |
|
| 123 |
d20 :: Gen Int
|
| 124 |
d20 = d 20
|