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 ->
    let
       (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 ->
     let
       (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 =
  let
     (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?