Monads for Markov chains
Suppose you need to model a finite Markov chain in code. There are essentially two ways of doing that: one is to simply run a simulation of the Markov chain using a random number generator to obtain dice rolls and random cards from the decks, the other is to create a stochastic matrix containing the transition probabilities for each pair of states. In this post I will show how a single monadic description of the Markov chain dynamics can be used to obtain both a simulator and the transition matrix.
{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances, GeneralizedNewtypeDeriving #-} import Control.Arrow import Control.Monad import Control.Monad.State.Strict import Data.Array import Random
Let’s start with an example of Markov chain and how we would like to be able to implement in Haskell. Consider a simplified version of the familiar Monopoly game: there are 40 squares (numbered 0 to 39), you throw two 6-sided dice each turn, some special squares have particular effects (see below), if you get a double roll three times in a row, you go to jail. The special squares are: 30: go to jail 2, 17, 33: Community Chest 7, 22, 36: Chance Community Chest (CC) and Chance (CH) make you take a card from a deck and move to some other place depending on what’s written on the card. You will find the details on the code, so I won’t explain them here. This is of course a Markov chain, where the states can be represented by:
type Square = Int data GameState = GS { position :: Square, doubles :: Int } deriving (Eq, Ord, Show)
and a description of the game can be given in a monadic style like this:
sGO :: Square sGO = 0 sJAIL :: Square sJAIL = 10 finalize :: Square -> Game Square finalize n | n == 2 || n == 17 || n == 33 = cc n | n == 7 || n == 22 || n == 36 = ch n | n == 30 = return sJAIL | otherwise = return n cc :: Square -> Game Square cc n = do i <- choose (1 :: Int, 16) return $ case i of 1 -> sGO 2 -> sJAIL _ -> n ch :: Square -> Game Square ch n = do i <- choose (1 :: Int, 16) return $ case i of 1 -> sGO 2 -> sJAIL 3 -> 11 4 -> 24 5 -> 39 6 -> 5 7 -> nextR n 8 -> nextR n 9 -> nextU n 10 -> n - 3 _ -> n where nextR n = let n' = n + 5 in n' - (n' `mod` 5) nextU n | n >= 12 && n < 28 = 28 | otherwise = 12 roll :: Game (Int, Int) roll = let r1 = choose (1, 6) in liftM2 (,) r1 r1 markDouble :: Bool -> Game () markDouble True = modify $ \s -> s { doubles = doubles s + 1 } markDouble False = modify $ \s -> s { doubles = 0 } goTo :: Square -> Game () goTo n = let n' = n `mod` 40 in modify $ \s -> s { position = n' } game :: Game () game = do n <- liftM position get (a, b) <- roll markDouble (a == b) d <- liftM doubles get if d == 3 then do markDouble False goTo sJAIL else do let n' = n + a + b n'' <- finalize n' goTo n''
As you can see, Game is a state monad, with an additional function choose that gives us a random element of a range:
class MonadState s m => MonadMC s m where choose :: (Enum a) => (a, a) -> m a
This can be implemented very easily using the (strict) state monad and a random generator:
newtype MCSim s a = MCSim (State ([s], StdGen) a) deriving Monad instance MonadState s (MCSim s) where get = MCSim $ liftM (head . fst) get put x = MCSim . modify $ \(xs, g) -> (x : xs, g) instance MonadMC s (MCSim s) where choose (a, b) = MCSim $ do (xs, g) <- get let bnds = (fromEnum a, fromEnum b) let (y, g') = randomR bnds g put (xs, g') return . toEnum $ y -- type Game a = MCSim GameState a runSim :: StdGen -> Int -> s -> MCSim s () -> [s] runSim g n start m = fst $ execState m' ([start], g) where (MCSim m') = foldr (>>) (return ()) $ replicate n m
The runSim function runs the simulation and returns the list of visited states. This is already quite nice, but the best thing is that the same code can be used to create the transition matrix, just swapping in a new implementation of the Game type alias:
newtype MC s a = MC (s -> [(s, Double, a)]) instance Monad (MC s) where return x = MC $ \s -> return (s, 1.0, x) (MC m) >>= f = MC $ \s -> do (s', p, x) <- m s let (MC m') = f x (s'', q, y) <- m' s' return (s'', p * q, y) instance MonadState s (MC s) where get = MC $ \s -> return (s, 1.0, s) put x = MC $ \s -> return (x, 1.0, ()) instance MonadMC s (MC s) where choose (a, b) = let r = [a..b] p = recip . fromIntegral . length $ r in MC $ \s -> map (\x -> (s, p, x)) r type Game a = MC GameState a
The idea is that we keep track of all possible destination states for a given state, with associated conditional probabilities. For those familiar with Eric Kidd’s series on probability monads, this is basically:
type MC s a = StateT s (PerhapsT [] a)
Now, how to get a transition matrix from such a monad? Of course, we have to require that the states are indexable:
markov :: Ix s => MC s () -> (s, s) -> Array (s, s) Double markov (MC m) r = accumArray (+) 0.0 (double r) $ range r >>= transitions where mkAssoc s (s', p, _) = ((s, s'), p) transitions s = map (mkAssoc s) $ m s double (a, b) = ((a, a), (b, b))
So we iterate over all states and use the probability values contained in the monad to fill in the array cells corresponding to the selected state pair. To actually apply this to our Monopoly example, we need to make GameState indexable:
nextState :: GameState -> GameState nextState (GS p d) = if d == 2 then GS (p + 1) 0 else GS p (d + 1) instance Ix GameState where range (s1, s2) = takeWhile (<= s2) . iterate nextState $ s1 index (s1, s2) s = let poss = (position s1, position s2) in index poss (position s) * 3 + doubles s - doubles s1 inRange (s1, s2) s = s1 <= s && s <= s2 rangeSize (s1, s2) = index (s1, s2) s2 + 1
then finally we can try:
monopoly :: (GameState, GameState) monopoly = (GS 0 0, GS 39 2) initialState :: Array GameState Double initialState = let n = rangeSize monopoly p = recip $ fromIntegral n in listArray monopoly $ replicate n p statDistr :: Int -> [(GameState, Double)] statDistr n = let mat = markov game monopoly distributions = iterate (.* mat) initialState st = distributions !! n in assocs st
where .* is a simple vector-matrix multiplication function:
infixl 5 .* (.*) :: (Ix i, Num a) => Array i a -> Array (i, i) a -> Array i a (.*) x y = array resultBounds [(i, sum [x!k * y!(k,i) | k <- range (l,u)]) | i <- range (l'',u'') ] where (l, u) = bounds x ((l', l''), (u', u'')) = bounds y resultBounds | (l,u)==(l',u') = (l'', u'') | otherwise = error ".*: incompatible bounds"
Calling statDistr 100 will return an association list of states with corresponding probability in an approximation of the stationary distribution, computed by applying the power method to the transition matrix. The number 100 is a pure guess, I don’t know how to estimate the number of iterations necessary for convergence, but that is out of the scope of this post, anyway.
Comments