Basics of Probability Monads

One extremely simple way to model probability is with a list of the universe of all events that can happen, along with the chances of each.

newtype Prob a = Prob { runProb :: [(a, Rational)] }
  deriving Show

We’ll use Rational (from Data.Ratio) here because it helps us avoid floating point arithmetic issues and so on. Realistically, we’d want to restrict a with some kind of Event typeclass, and we’d want to bound our probabilities to the range \([0, 1]\) but this model is good enough to show how to build a monad out of a bunch of events-with-likelihoods.

It’s going to be useful to have a nicely-formatted way of printing out the chance of each event happening, so let’s do that too:

display :: Show a => Prob a -> IO ()
display (Prob dist) = do
  mapM_ putStrLn [show x ++ ": " ++ showNum p | (x, p) <- dist]
  where showNum p = show (Ratio.numerator p) ++ "/" ++ show (Ratio.denominator p)

At this point we can start building up a bunch of instances. Functor is reasonably straightforward: we just map the function over each event, almost like the fmap instance for lists.

instance Functor Prob where
  fmap f (Prob dist) = Prob $ do
    (x, p) <- dist
    return (f x, p)

One thing to note here is that our do block is working with lists, so we could also have used list comprehensions like below. However, I’ll favor do since I think it’s a smidge cleaner.

instance Functor Prob where
  fmap f (Prob dist) = Prob [(f x, p) | (x, p) <- dist]

Moving on to applicatives. pure is also somewhat common sense: we take the single event given, and give it all the probability mass (its chances of occurring are 100%).

instance Applicative Prob where
  pure x = Prob [(x, 1)]

A reasonable instance definition of <*> would take all the event-altering functions in a and apply them to all the events in b. We multiply probabilities because this reduces them in just the right way to maintain the correct overall sum total probability mass. If we take two coins, the chance of flipping heads on either is \[ 1/2 \] but the chance of getting heads on both is \[ 1/4 \]. You can also think of applicatives as doing something like the product of lists here, so it makes sense that we’d also want the product of probabilities.

  a <*> b = Prob $ do
    (f, p) <- runProb a
    (x, q) <- runProb b
    return (f x, p * q)

Now for the monad instance. return is just pure.

instance Monad Prob where
  return = pure

Bind is a bit more interesting. As with <*>, we’re taking the list of all events out of our prior Prob and as with fmap we apply f to each piece. We then take our probabilities and multiply, along with passing back the newly-created event (I called it y).

  old >>= f = Prob $ do
    (x, p) <- runProb old
    (y, q) <- runProb $ f x
    return (y, p * q)

What can we do with this kind of simple monad for managing probabilities? Let’s say we want to model a fair coin in a simple way. We can use this ADT:

data Coin = Heads | Tails
  deriving (Bounded, Enum, Show)

The instances for Bounded and Enum now let us assign an equal chance to each constructor of the type.

enumerateProb :: (Enum a, Bounded a) => Prob a
enumerateProb = equalChance [minBound..maxBound]

equalChance :: [a] -> Prob a
equalChance xs = Prob [(x, chance) | x <- xs]
  where chance = 1 % length xs
        a % b = fromIntegral a Ratio.% fromIntegral b

This makes it trivial to create more ADTs where every constructor has the same likelihood of showing up:

data RPS = Rock | Paper | Scissors
  deriving (Bounded, Enum, Show)

Another kind of thing we can do with Prob is basic calculations you’d expect to be able to do with any sort of probabilistic model. For example, we can get the expected value of some set of events that we can give Fractional value to.

expected :: (Fractional a) => Prob a -> a
expected = sum . map (\(x, p) -> x * fromRational p) . runProb

If I now say that I’ll give you $2 for a coin flip that’s heads if you give me $1 for every flip that’s tails, you can figure out that your maximum reasonable buy-in price per round should be 50 cents.

valueCoin Heads = 2
valueCoin Tails = -1

main = do
  let coin = enumerateProb :: Prob Coin
  display coin
  print . expected $ valueCoin <$> coin