RCodensity fun

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ConstraintKinds #-}
module Prob where
import Data.Ratio
import qualified Data.Map as M

newtype ProbData a = ProbD { runProbD :: M.Map a Rational } deriving (Show,Eq)

retProb :: Ord a => a -> ProbData a
retProb a = ProbD $ M.singleton a 1

bindProb :: Ord b => ProbData a -> (a -> ProbData b) -> ProbData b
bindProb m k = ProbD $ M.unionsWith (+) $ map (\(a,p) -> M.map (*p) (runProbD $ k a)) (M.toList $ runProbD m)

newtype RCodensity c m a = Cod { runCod :: forall r. c r => (a -> m r) -> m r }
instance Functor (RCodensity c m) where
    fmap f m = Cod $ \c -> runCod m (c . f)

instance Monad (RCodensity c m) where
    return a = Cod ($ a)
    m >>= k  = Cod $ \c -> runCod m $ \a -> runCod (k a) c

newtype Prob a = Prob { runProbM :: RCodensity Ord ProbData a } deriving (Monad, Functor)

liftProbD :: ProbData a -> Prob a
liftProbD p = Prob $ Cod $ \c -> bindProb p c

uniform :: Ord a => [a] -> Prob a
uniform xs = liftProbD $ ProbD $ M.fromList $ map (\a -> (a,p)) xs where
    p = 1 / fromIntegral (length xs)

runProb :: Ord a => Prob a -> ProbData a
runProb xs = runCod (runProbM xs) retProb

optimize :: Ord a => Prob a -> Prob a
optimize = liftProbD . runProb

d6 = uniform [1..6]

twod6 = optimize $ do
   x <- d6
   y <- d6
   return (x+y)