MCMC in Python: PyMC to sample uniformly from a convex body

This post is a little tutorial on how to use PyMC to sample points uniformly at random from a convex body.  This computational challenge says: if you have a magic box which will tell you yes/no when you ask, “Is this point (in n-dimensions) in the convex set S”, can you come up with a random point which is nearly uniformly distributed over S?

MCMC has been the main approach to solving this problem, and it has been a great success for the polynomial-time dogma, starting with the work of Dyer, Frieze, and Kannan which established the running-time upper bound of .  The idea is this: you start with some point in S and try to move to a new, nearby point randomly.  “Randomly how?”, you wonder.  That is the art.

In PyMC, if you want to sample points uniformly from the unit square, all you need to do is define your stochastic variable, make a Markov chain, and sample from it. Here is how:

from pymc import *
X = Uniform('X', zeros(2), ones(2))
mc = MCMC([X])
mc.sample(10000)

If you have matplotlib set up right, you should be able to get a nice picture of this now

plot(X.trace()[:,0], X.trace()[:,1],',')

Not a very complicated convex body, but it’s a start. And it was easy! Five lines of code? No problem.

Say you want something to sample from some more complicated, though, like anything-that’s-not-a-rectangle. One way to do this is to use PyMC’s potential object. Below is a constraint that makes S into a ball of diameter one, centered at (0.5, 0.5). It just says that the log-probability for any point outside of this ball is

from pymc import *
X = Uniform('X', zeros(2), ones(2))
@potential
def s1(X=X):
    if ((X-0.5)**2).sum() >= 0.25:
        return -inf
    else:
        return 0.0

Did that give you an error? For a fraction of readers, it will; PyMC gets unhappy when it has a sample with probability zero. Try setting X to something in the ball first:

from pymc import *
X = Uniform('X', zeros(2), ones(2))
X.value = [0.5,0.5]
@potential
def s2(X=X):
    if ((X-0.5)**2).sum() >= 0.25:
        return -inf
    else:
        return 0.0
mc = MCMC([X, s2])
mc.sample(10000)
plot(X.trace()[:,0], X.trace()[:,1],',')

But, this is too easy to sample from.  It is a nice, plump convex body.  The hard stuff comes from the convex bodies with skinny bits.  If you live in high dimensions, these skinny bits show up everywhere.  For example, if you have a ball just like the one above, but in 20 dimensions instead of 2, check out what happens:

from pymc import *
X = Uniform('X', zeros(20), ones(20), value=0.5*ones(20))
@potential
def s3(X=X):
    if ((X-0.5)**2).sum() >= 0.25:
        return -inf
    else:
        return 0.0
mc = MCMC([X, s3])
mc.sample(10000)
plot(X.trace()[:,0], X.trace()[:,1],',')

That’s not looking nice anymore. The problem is this. If you don’t tell PyMC how you want the Markov transitions to go, it picks something for you. It defaults to a Metropolis walk, where each stochastic variable is drawn from its prior distribution. In this case, that means that for 10,000 samples, PyMC draws uniformly from the 20-dimensional hypercube, and then rejects it and stays at if falls outside of the ball we want. Since the volume of the ball is much smaller than the volume of the cube, roughly 75% of these samples are rejected. (In the picture above, I jittered the points to reveal the duplicate points.)

How can we make things nice again? Instead of sampling from the prior every time, we can take a small, random step away from the point that we are currently at.

from pymc import *
X = Uniform('X', zeros(20), ones(20), value=0.5*ones(20))
@potential
def s4(X=X):
    if sum((X-0.5)**2) >= 0.25:
        return -inf
    else:
        return 0.0
mc = MCMC([X, s4])
mc.use_step_method(Metropolis, X, dist='Normal', sig=0.05)
mc.sample(50*10000,thin=50)
plot(X.trace()[:,0], X.trace()[:,1],',')

The new trick is in line 10, where I asked for a particular step method. It appears that the dist parameter tells PyMC to use move to a new point normally distributed around the current point, and the sig parameter controls the variance of the normal distribution. Now things looks nice again. Maybe this isn’t actually a great example, because it takes longer for this chain to mix, so the pictures that look nice come from keeping 1 sample out of every 50, meaning the chain takes 50 times longer to run. If I really wanted samples from the 20-dimensional ball, rejection sampling is a better way to go. (But if I really wanted samples from the 20-dimensional ball, there are even better ways to draw them…)

Up next, how to implement the hit-and-run walk in PyMC. (Actually, not next…)