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