BUG NUTS and HamiltonianMC samplers broken for a simple problem #1428
Description
Hello pymc3 devs! Thanks for all of your help so far!
Here is a weird thing. I made some simple changes to the model I sent previously, introducing some extra structure in the prior.
Now the NUTS and HMC samplers are just completely failing.
BTW, I had to turn up the number of samples taken for emcee to get it to even start to converge for this problem. STAN has no trouble and samples the problem in ~0.5 seconds once the C++ compiles.
The plot:
The code:
#!/usr/bin/env python
import numpy as np
import matplotlib.pyplot as plt
import seaborn
import pystan
import emcee
import pymc3
def run_emcee(X1, X2, Y):
def loglike(x):
"""loglike in python"""
alpha = x[0]
beta1 = x[1]
beta2 = x[2]
sigma = x[3]
sigma_beta = x[4]
if sigma <= 0.0 or sigma > 1.0:
return -np.inf
if sigma_beta <= 0.0 or sigma_beta > 1.0:
return -np.inf
sigma_beta2 = sigma_beta * sigma_beta
log_sigma_beta = np.log(sigma_beta)
loglike = 0.0
loglike += -0.5 * alpha * alpha / 1
loglike += -0.5 * beta1 * beta1 / sigma_beta2
loglike -= log_sigma_beta
loglike += -0.5 * beta2 * beta2 / sigma_beta2
loglike -= log_sigma_beta
mu = alpha + beta1 * X1 + beta2 * X2
loglike += np.sum(-0.5 * ((Y - mu)**2) / sigma / sigma - np.log(sigma))
return loglike
ndim = 5
nwalkers = 80
p0 = [np.random.rand(ndim) for i in range(nwalkers)]
sampler = emcee.EnsembleSampler(nwalkers, ndim, loglike)
sampler.run_mcmc(p0, 10000)
_chains = sampler.chain[0:4, 5000:, -1]
chains = []
for i in range(4):
chains.append(_chains[i, :])
return chains
def run_stan(X1, X2, Y):
stan_code = """
data {
int<lower=0> num_obs;
vector[num_obs] y;
vector[num_obs] x1;
vector[num_obs] x2;
}
parameters {
real<lower=-10, upper=10> alpha;
vector[2] beta;
real<lower=0.0, upper=1.0> sigma;
real<lower=0.0, upper=1.0> sigma_beta;
}
transformed parameters {
vector[num_obs] mu;
mu = alpha + beta[1] * x1 + beta[2] * x2;
}
model {
alpha ~ normal(0, 1);
beta ~ normal(0, sigma_beta);
y ~ normal(mu, sigma);
}
"""
model = pystan.StanModel(model_code=stan_code)
mcmc_data = {'num_obs': len(Y), 'y': Y, 'x1': X1, 'x2': X2}
base_init = {'alpha': 0.1, 'beta': [0.2, 0.5], 'sigma': 0.3, 'sigma_beta': 0.2}
init = [base_init] * 4
fitted = model.sampling(warmup=500,
verbose=True,
thin=1,
init=init,
iter=1000,
n_jobs=4,
data=mcmc_data,
chains=4,
pars=['alpha', 'beta', 'sigma', 'sigma_beta'],
seed=46576)
return fitted
# Initialize random number generator
np.random.seed(123)
# True parameter values
alpha, sigma = 1, 0.5
beta = [-1, 2.5]
sigma_beta = 0.4
# Size of dataset
size = 100
# Predictor variable
X1 = np.random.randn(size)
X2 = np.random.randn(size) * 0.2
# Simulate outcome variable
Y = alpha + beta[0] * sigma_beta * X1 + beta[1] * sigma_beta * X2 + np.random.randn(size) * sigma
# build pymc3 model
verbose = 2
with pymc3.Model(verbose=verbose) as basic_model:
# Priors for unknown model parameters
sigma_beta = pymc3.Uniform('sigma_beta', lower=0.0, upper=1.0)
alpha = pymc3.Normal('alpha', mu=0, sd=1)
sigma = pymc3.Uniform('sigma', lower=0.0, upper=1.0)
beta = pymc3.Normal('beta', mu=0, sd=sigma_beta, shape=2)
mu = alpha + beta[0] * X1 + beta[1] * X2
# Likelihood (sampling distribution) of observations
Y_obs = pymc3.Normal('Y_obs', mu=mu, sd=sigma, observed=Y)
# now sample
chains = []
names = []
pymc3_samplers = [(pymc3.NUTS, 'pymc3 NUTS'),
(pymc3.HamiltonianMC, 'pymc3 HamiltonianMC'),
(pymc3.Metropolis, 'pymc3 Metropolis'),
(pymc3.Slice, 'pymc3 Slice'),]
for stepper, name in pymc3_samplers:
with basic_model:
trace = pymc3.sample(1000, njobs=4, progressbar=bool(verbose),
step=stepper(basic_model.cont_vars),
random_seed=[89, 50, 51, 34])
chains.append(trace.get_values('sigma_beta', combine=False, burn=500))
names.append(name)
# run STAN
if True:
fitted = run_stan(X1, X2, Y)
sigma_index = fitted.flatnames.index('sigma_beta')
_schains = fitted.extract('sigma_beta', permuted=False)[:, :, sigma_index]
schains = []
for i in range(4):
schains.append(_schains[:, i])
chains.append(schains)
names.append('STAN NUTS')
# run emcee
if True:
_echains = run_emcee(X1, X2, Y)
chains.append(_echains)
names.append('emcee')
# make the plot
fig, axs = plt.subplots(nrows=3, ncols=2, sharex=True, figsize=(10, 10))
axs = axs.ravel()
for ax, chain, name in zip(axs, chains, names):
try:
for c in chain:
seaborn.kdeplot(c, ax=ax, legend=False);
except:
for c in chain:
seaborn.distplot(c, bins=25, kde=False, rug=True, ax=ax);
ax.set_title(name)
axs[-2].set_xlabel('sigma beta')
axs[-1].set_xlabel('sigma beta')
fig.tight_layout()
plt.show()