(Warning: This article is a bit more technical than most of my stuff. It assumes prior knowledge of monads and monad transformers.)

Martin Erwig and Steve Kollmansberger wrote PFP, a really sweet Haskell library for computing with probabilities. To borrow their example:

die :: Dist Int
die = uniform [1..6]

If we roll a die, we get the expected distribution of results:

> die
1  16.7%
2  16.7%
3  16.7%
4  16.7%
5  16.7%
6  16.7%

If you haven't seen PFP before, I strongly encourage you to check it out. You can use it to solve all sorts of probability puzzles.

Anyway, I discovered an interesting way to implement PFP using monad transformers. Here's what it looks like:

type Dist = PerhapsT ([])

uniform = weighted . map (\x -> (x, 1))

In other words, Dist can be written by adding some semantics to the standard list monad.

Perhaps: A less specific version of Maybe

First, let's define a simple probability type:

newtype Prob = P Float
  deriving (Eq, Ord, Num)

instance Show Prob where
  show (P p) = show intPart ++ "." ++ show fracPart ++ "%"
    where digits = round (1000 * p)
          intPart = digits `div` 10
          fracPart = digits `mod` 10

Thanks to the deriving (Num) declaration, we can treat Prob like any other numeric type.

We can now define Perhaps, which represents a value with an associated probability:

data Perhaps a = Perhaps a Prob
  deriving (Show)

Now, this is just a generalization of Haskell's built-in Maybe type, which treats a value as either present (probability 1) or absent (probability 0). All we've added is a range of possibilities in between: Perhaps x 0.5 represents a 50% chance of having a value.

Note that there's one small trick here: When the probability of a value is 0, we may not actually know it! But because Haskell is a lazy language, we can write:

Perhaps undefined 0

We'll need a convenient way to test for this case, to make sure we don't try to use any undefined values:

neverHappens (Perhaps _ 0) = True
neverHappens _             = False

So Perhaps is just like Maybe. As it turns out, they're both monads, and they both have an associated monad transformer.

First, a monad

Like Haskell's standard Maybe type, Perhaps is a monad.

instance Functor Perhaps where
  fmap f (Perhaps x p) = Perhaps (f x) p

instance Monad Perhaps where
  return x = Perhaps x 1
  ph >>= f  | neverHappens ph  = never
            | otherwise        = Perhaps x (p1 * p2)
    where (Perhaps (Perhaps x p1) p2) = fmap f ph

We can also define a slightly extended interface:

class Monad m => MonadPerhaps m where
  perhaps :: a -> Prob -> m a
  never :: m a

instance MonadPerhaps Perhaps where
  never = Perhaps undefined 0
  perhaps = Perhaps

Next, a monad transformer

What if we're already using a monad, and we want to layer Perhaps-style semantics over the top? Well, we can always use a monad transformer. The following code is presented without further comment:

newtype PerhapsT m a = PerhapsT { runPerhapsT :: m (Perhaps a) }

instance MonadTrans PerhapsT where
  lift x = PerhapsT (liftM return x)

instance Monad m => Functor (PerhapsT m) where
  fmap = liftM

instance Monad m => Monad (PerhapsT m) where
  return = lift . return
  m >>= f = PerhapsT bound
    where bound = do
            ph <- runPerhapsT m
            case ph of
              (Perhaps x1 p1)  | p1 == 0    -> return never
                               | otherwise  -> do
                (Perhaps x2 p2) <- runPerhapsT (f x1)
                return (Perhaps x2 (p1 * p2))

Update: I removed the instance MonadPerhaps (PerhapsT m), because not all derived monads (including the one described here) will actually want to provide access to never and perhaps.

Recreating PFP

From here, everything is simple. We can implement PFP by applying our monad transformer to the standard list monad:

type Dist = PerhapsT ([])

uniform = weighted . map (\x -> (x, 1))

weighted :: [(a, Float)] -> Dist a
weighted [] =
  error "Empty probability distributuion"
weighted xws = PerhapsT (map weight xws)
  where weight (x,w) = Perhaps x (P (w / sum))
        sum = foldl' (\w1 (_,w2) -> w1+w2) 0 xws

What does this refactoring buy us? Well, it turns out that we can reuse PerhapsT elsewhere...

Part 2: Random sampling