gdritter repos treant / 61040cf
Some early, hacky, unfinished, ill-advised efforts towards a Gen-like library Getty Ritter 5 years ago
3 changed file(s) with 166 addition(s) and 0 deletion(s). Collapse all Expand all
1 dist
2 dist-*
3 *~
4 cabal-dev
5 *.o
6 *.hi
7 *.chi
8 *.chs.h
9 *.dyn_o
10 *.dyn_hi
11 .hpc
12 .hsenv
13 .cabal-sandbox/
14 cabal.sandbox.config
15 *.prof
16 *.aux
17 *.hp
18 *.eventlog
19 cabal.project.local
20 .ghc.environment.*
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
1 name: treant
2 version: 0.1.0.0
3 -- synopsis:
4 -- description:
5 license: BSD3
6 author: Getty Ritter <gettylefou@gmail.com>
7 maintainer: Getty Ritter <gettylefou@gmail.com>
8 copyright: @2018 Getty Ritter
9 -- category:
10 build-type: Simple
11 cabal-version: >=1.14
12
13 library
14 hs-source-dirs: src
15 ghc-options: -Wall
16 build-depends: base >=4.7 && <5
17 , random
18 , text
19 , containers
20 default-language: Haskell2010
21 default-extensions: ScopedTypeVariables
22 exposed-modules: Treant