Probability Monad in Swift

I read about the probability monad in The Frequentist Approach to Probability a couple of years ago and thought it was pretty neat. I decided to make one in Swift, as an exercise in learning the language, after having done the same in Clojure and Java 8.

If you don’t know what a monad is, don’t worry, it’s not important for this post. Just know that it lets you do neat stuff with probability distributions, programmatically speaking. Ok… on to the code.

Protocols for Probabilities

I initially tried to do this with protocols. I had two. One that let us randomly sample values from a probability distribution

and another one that allowed for parameterization, like for the Poisson distribution.

You may have noticed that there’s no protocol that allows you to map one distribution into another, which is what would make this into a monad. That’s because I had not yet figured out how to do it with structs or with classes. It’s easy to map one set of values drawn from a distribution into another set of values according to a function. But I really needed to create a new struct with a specific get() function. And then I remembered that functions were first class values in Swift!

Turns out you don’t need protocols or classes for this at all. You can do it all with a pretty simple struct!

Probability Distributions via Closures

With a single generic struct, we have everything we need for the probability monad. To convert one distribution into another, we need only pass in a function that maps elements of one distribution into elements of the other.

Let’s see what we can do if we start from the uniform distribution.

For starters, we can easily generate the true-false distribution by mapping the Uniform distribution with a function that generates the booleans from a double. From there, it’s straightforward to transform the true-false distribution into the Bernoulli distribution

By passing the appropriate transformation function or closure to the map function, one type of distribution can be converted into another.

BTW, to use the random number generators you’ll need to import Foundation. If you’re on a Mac, you could also import GameplayKit’s GKRandomSource. Or you can always use something from C.

Composing Distributions

If you want to compose distributions, the our struct needs an appropriate function: flatMap.

A standard example is the distribution you get from combining a pair of six-sided dice. We can start with a single die:

Next, we can use the flatMap function to compose the distributions of a pair of six-sided dice by passing in a function that provides the behavior we need.

Now that you’ve seen all the pieces, here’s the final form of the probability distribution struct:

Computing Statistics for a Distribution

You may have noticed a few functions for summary statistics (mean, etc) and probability computation. The most important function is prob, which lets you use predicates to ask questions of the distribution. There are a few basic examples of what you can do below.

If I can figure how to properly implement the given() function I’ll add that in a future post. I also want to be able to handle more interesting distributions, like that seen in a probabilistic graphical model.

All the code is on Github.


Leave a Reply

Fill in your details below or click an icon to log in: Logo

You are commenting using your account. Log Out / Change )

Twitter picture

You are commenting using your Twitter account. Log Out / Change )

Facebook photo

You are commenting using your Facebook account. Log Out / Change )

Google+ photo

You are commenting using your Google+ account. Log Out / Change )

Connecting to %s