Sampling the Banana-shaped distribution

You can define a banana-shaped function and wrap it using pm.Potential. For example, defining a Rosenbrock function:

import numpy as np
import pymc3 as pm
import theano.tensor as tt
import theano

def pot1(z):
    z = z.T
    a = 1.
    b = 100.
    return (a-z[0])**2 + b*(z[1] - z[0]**2)**2

z = tt.matrix('z')
z.tag.test_value = pm.floatX([[0., 0.]])
pot1f = theano.function([z], pot1(z))

Visualize the function by evaluating it on a grid:

import matplotlib.pylab as plt
from mpl_toolkits import mplot3d
xlim=(-2,2)
ylim=(-1,3)
grid = pm.floatX(np.mgrid[xlim[0]:xlim[1]:100j,ylim[0]:ylim[1]:100j])
grid_2d = grid.reshape(2, -1).T
Z = pot1f(grid_2d)
fig = plt.figure()
ax = fig.gca(projection='3d')
surf = ax.plot_surface(grid[0], grid[1], Z.reshape(100,100),cmap='viridis',
                       linewidth=0, antialiased=False)
plt.show()

Sample from the above function in PyMC3, and visualized the samples:

def cust_logp(z):
    return -pot1(z)

with pm.Model() as pot1m:
    pm.DensityDist('pot1', logp=cust_logp, shape=(2,))
    trace1 = pm.sample(1000, step=pm.NUTS())
    trace2 = pm.sample(1000, step=pm.Metropolis())

_, ax = plt.subplots(1,2,figsize=(10,5))
tr1 = trace1['pot1']
ax[0].plot(tr1[:,0], tr1[:,1], 'ro--',alpha=.1)
tr2 = trace2['pot1']
ax[1].plot(tr2[:,0], tr2[:,1], 'bo--',alpha=.1)
plt.tight_layout()

There are also similar example in the doc of Normalizing flow: http://docs.pymc.io/notebooks/normalizing_flows_overview.html#Simulated-data-example

3 Likes