Sunday, August 19, 2007

Anatomy of a new monad

The monad described here was originally written in response to a posting by Chad Scherrer on Haskell-Cafe:

I need to build a function
buildSample :: [A] -> State StdGen [(A,B,C)]

given lookup functions
f :: A -> [B]
g :: A -> [C]

The idea is to first draw randomly form the [A], then apply each
lookup function and draw randomly from the result of each.

I suggested Chad try using ListT for the non-determinism, but this didn't work, so I decided to try solving it myself. I discovered that actually ListT doesn't work here: you need a completely custom monad. So here it is. Hopefully this will help other people who find they need a custom monad. This isn't meant to be yet another ab-initio monad tutorial: try reading the Wikibook or All About Monads if you need one.

The basic idea of the MonteCarlo monad is that each action in the monad returns a list of possible results, just like the list monad. However the list monad then takes all of these possible results forwards into the next step, potentially leading to combinatorial explosion if you can't prune the tree somehow. It was this explosion that was giving Chad problems when he scaled up his original solution. So instead the MonteCarlo monad picks one of the results at random and goes forwards with that.

Picking a random element means we need to thread a random number generator StdGen as a form of monadic state. So the monad type looks like this:
newtype MonteCarlo a = MonteCarlo {runMC :: StdGen -> (StdGen, [a])}
This is based on the state monad, except that the type parameter "s" is replaced by StdGen. I could have used the state monad here, but it would have meant working its "bind" and "return" into MonteCarlo, which would have been more trouble than it was worth.

Picking a random item from a list is going to be necessary, and as seen below it is actually needed more than once. So:

-- Internal function to pick a random element from a list
pickOne :: [a] -> StdGen -> (StdGen, a)
pickOne xs g1 = let (n, g2) = randomR (0, length xs - 1) g1 in (g2, xs !! n)
Monad instance
Now for the Monad instance declaration. In this we have to declare the bind function (>>=) and the return function.
instance Monad MonteCarlo where
MonteCarlo m >>= f = MonteCarlo $ \g1 ->
(g2, xs) = m g1
(g3, x) = pickOne xs g2
(g4, f') = case xs of
[] -> (g2, mzero)
[x1] -> (g2, f x1)
_ -> (g3, f x)
in runMC f' g4

return x = MonteCarlo $ \g -> (g, [x])
The return function shows the minimal structure for a Monte-Carlo action: wrapped up in the MonteCarlo type is a function that takes a StdGen state and returns the state (in this case unmodified) along with the potential results. If the action had used the generator then the result pair would have had the new generator state instead of the old one.

The bind (>>=) is a bit more complicated. The type for monadic bind is:
(>>=) :: m a -> (a -> m b) -> m b
In the MonteCarlo instance of Monad this becomes:
(>>=) :: MonteCarlo a -> (a -> MonteCarlo b) -> MonteCarlo b

The job of bind is to take its two arguments, both of which are functions under the covers, and compose them into a new function. Then it wraps that function up in the MonteCarlo type.

We unwrap the first argument using pattern matching to get the function inside "m". The second parameter "f" is a function that takes the result of the first and turns it into a new action.

The result of the bind operation has to be wrapped up in MonteCarlo. This is done by the line
MonteCarlo $ \g1 -> 
The "g1" lambda parameter is the input generator state for this action when it is run, and the rest of the definition is the way we compose up the functions. We need to call the first function "m" to get a result and a new generator state "g2", and then feed the result into the second argument "f".

The "let" expression threads the random generator manually through the first two random things we need to do:
  1. Get the possible results of the first argument as a list.
  2. Pick one of these at random.
However there is a twist because the result set could be empty. This is interpreted as failure, and the "MonadPlus" instance below will explain more about failing. In addition the result set could also be a single item, in which case there is no need to generate a random number to pick it. So the f-prime value is the result of applying "f" to whatever item was picked, or an empty list if the "m" action returned an empty list. We also avoid getting a new generator state if the random pick was not needed.

Finally we use "runMC" to strip the MonteCarlo type wrapper off the second result, because the whole thing is being wrapped by the MonteCarlo type constructor at the top level of bind.

And there you have it. In summary, when defining an instance of "bind" you have to construct a new function out of the two arguments by extracting the result of the first argument and passing it to the second. This then has to be wrapped up as whatever kind of function is sitting inside your monadic action type. The arguments to this hidden function are the monadic "state" and therefore have to be threaded through bind in whatever way suits the semantics of your monad.

MonadPlus instance
MonadPlus is used for monads that can fail or be added together in some way. If your underlying type is a Monoid then its a good bet that you can come up with a MonadPlus instance from the "mempty" and "mappend" functions of the Monoid. In this case the underlying type is a list, so the "mempty" is an empty list and the "mappend" is concatenation.
instance MonadPlus MonteCarlo where
mzero = MonteCarlo $ \g -> (g, [])
mplus (MonteCarlo m1) (MonteCarlo m2) = MonteCarlo $ \g ->
(g1, xs1) = m1 g
(g2, xs2) = m2 g1
in (g2, xs1 ++ xs2)

Here the "mzero" returns an empty list. The bind operation interprets an empty list as failure, and so indicates this to its caller by returning an empty list. Thus a single failure aborts the entire computation.

"mplus" meanwhile threads the random number state "g" through the two arguments, and then returns the results concatenated. The bind operation will then pick one of these results, so "mplus" has become a kind of "or" operation. However note that the alternation is done over the results, not the original arguments. If you say "a `mplus` b" and "a" returns ten times as many results as "b" then the odds are 10:1 that you will be getting one of the results of "a".

One important combinator introduces a list of items to pick one from.
-- | Convert a list of items into a Monte-Carlo action.
returnList :: [a] -> MonteCarlo a
returnList xs = MonteCarlo $ \g -> (g, xs)
Finally we need to be able to run a MonteCarlo computation. You could do this using runMC, but that gives you the full list of results for whatever the last action was, which is not what you want. Instead you want to pick one of them at random. So there is a runMonteCarlo function that looks an awful lot like the bind operation, and for similar reasons. This returns a Maybe type with "Just" for success and "Nothing" for failure.
-- | Run a Monte-Carlo simulation to generate a zero or one results.
runMonteCarlo :: MonteCarlo a -> StdGen -> Maybe a
runMonteCarlo (MonteCarlo m) g1 =
(g2, xs) = m g1
(_, x) = pickOne xs g2
in case xs of
[] -> Nothing
[x1] -> Just x1
_ -> Just x
Exercises for the reader:
  1. Add a new combinator allowing access to the random number generator.
  2. Each computation returns only a single result, but Monte Carlo methods depend on statistical analysis of many results. How should this be done?
  3. Can we introduce a "try" combinator to limit the scope of failure?


Neil Bartlett said...

Shouldn't your runMonteCarlo function return the generator along with the result, so that it can be threaded with other random functions? I.e.:

runMonteCarlo :: MonteCarlo a -> StdGen -> (StdGen, Maybe a)

Or perhaps:

runMonteCarlo :: RandomGen g => MonteCarlo a -> g -> (g, Maybe a)

Paul Johnson said...

Could do. I was following in the footsteps of the QuickCheck "Gen" monad, which doesn't.

The best thing would probably be to do as you suggest, but then have an "evalMonteCarlo" which throws away the state.

I did wonder about genericising the whole thing by RandomGen, but decided not to bother.

Anonymous said...

As your 2nd exercise for the reader indicates, a Monte Carlo monad would be interested in integrating by a large number of random samples. What you appear to have here is a Monad that would help with building uniformly distributed Markov chain simulations. Monte Carlo integration is often used in combination with a Markov chain to calculate the probabilities of the possible resulting states.


A possible improvement to the presented Monad would be to allow for non-uniformly distributed selection of the next state at each step.

Paul Johnson said...

anonymous: I've had a look at the pages you referenced. But I'm not that much of a statistician.

It sounds like the MonteCarlo monad should be renamed the Markov monad. Am I right?

A non-uniform selection of next states could be done, but it would require a different monad. The type would be something like:

newtype Markov a = Markov (StdGen -> (StdGen, [(Int, a)])

The integer components in the returned values would be the relative probabilities of each value. The original monad is then the special case in which all these values are the same.

Would this do the job?