<jupyter_start><jupyter_text>TALENT Course 11 Learning from Data: Bayesian Methods and Machine Learning York, UK, June 10-28, 2019  <NAME>, Chalmers University of Technology, Sweden  MCMC sampling - part IIb<jupyter_code>%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt

import seaborn as sns
sns.set(context='talk')
import emcee
from emcee import PTSampler
import corner
from mpl_toolkits.mplot3d import axes3d
from matplotlib import cm<jupyter_output><empty_output><jupyter_text>Prior<jupyter_code># Use a Gaussian prior centered at (0,0) with variance sqrt(10)
mup = np.zeros(2)
sigp=np.sqrt(10.)
sigmap = np.diag([sigp**2, sigp**2])
sigmapinv = np.linalg.inv(sigmap)
normp = 1/np.sqrt(np.linalg.det(2*np.pi*sigmap))
# Define the log prior
def logp(x):
    dxp = x - mup
    return -np.dot(dxp, np.dot(sigmapinv, dxp))/2.0 \
        + np.log(normp)<jupyter_output><empty_output><jupyter_text>Likelihood<jupyter_code>def setup_modes(sig=0.2,ratio=3.):
    # Means of the two Gaussian modes
    mu1 = np.ones(2)
    mu2 = -np.ones(2)

    # Width in each dimension (zero correlation)
    _sig = sig
    sigma1 = np.diag([_sig**2, _sig**2])
    sigma2 = np.diag([_sig**2, _sig**2])
    sigma1inv = np.linalg.inv(sigma1)
    sigma2inv = np.linalg.inv(sigma2)

    # amplitudes of the two modes and the corresponding norm factor
    _amp1 = ratio / (1+ratio)
    _amp2 = 1. / (1+ratio)
    norm1 = _amp1 / np.sqrt(np.linalg.det(2*np.pi*sigma1))
    norm2 = _amp2 / np.sqrt(np.linalg.det(2*np.pi*sigma2))
    return (mu1,mu2,sigma1inv,sigma2inv,norm1,norm2)
params_modes = setup_modes()

# Define the log likelihood function
def logl(x,params=params_modes):
    (mu1,mu2,sigma1inv,sigma2inv,norm1,norm2)=params
    dx1 = x - mu1
    dx2 = x - mu2
    return np.logaddexp(-np.dot(dx1, np.dot(sigma1inv, dx1))/2.0 \
                        + np.log(norm1),\
                        -np.dot(dx2, np.dot(sigma2inv, dx2))/2.0 \
                        + np.log(norm2) )<jupyter_output><empty_output><jupyter_text>Posterior<jupyter_code>def log_posterior(x):
    return logp(x) + logl(x)

@np.vectorize
def posterior(y,x):
    xv=np.array([x,y])
    return np.exp(logl(xv)+logp(xv))<jupyter_output><empty_output><jupyter_text>MH Sampling and convergence<jupyter_code>print('emcee sampling (version: )', emcee.__version__)

ndim = 2  # number of parameters in the model
nwarmup = 1000
nsteps = 10000

# MH-Sampler setup
stepsize = .05
cov = stepsize * np.eye(ndim)
p0 = np.random.rand(ndim)

# initialize the sampler
sampler = emcee.MHSampler(cov, ndim, log_posterior)
# Sample the posterior distribution

# Warm-up
if nwarmup > 0:
    print(f'Performing {nwarmup} warnup iterations.')
    pos, prob, state = sampler.run_mcmc(p0, nwarmup)
    sampler.reset()
else:
    pos = p0
    
# Perform iterations, starting at the final position from the warmup.
print(f'MH sampler performing {nsteps} samples.')
%time sampler.run_mcmc(pos, nsteps)
print("done")

print(f"Mean acceptance fraction: {sampler.acceptance_fraction:.3f}")

samples = sampler.flatchain
    
# make a corner plot with the posterior distribution
fig = corner.corner(samples, quantiles=[0.16, 0.5, 0.84], labels=[r"$\theta_0$", r"$\theta_1$"],
                       show_titles=True, title_kwargs={"fontsize": 12})
fix, ax = plt.subplots(2,2,figsize=(12,5*ndim))
for irow in range(ndim):
    ax[irow,0].plot(np.arange(samples.shape[0]),samples[:,irow])
    ax[irow,0].set_ylabel(r'$\theta_{0}$'.format(irow))
    ax[irow,1].hist(samples[:,irow],orientation='horizontal',bins=30)

ax[0,1].set_title('Histogram')
ax[0,0].set_title('Trace Plot')

plt.tight_layout()
for irow in range(ndim):
    print(f"Standard Error of the Mean for theta_{irow}: {samples[:,irow].std()/np.sqrt(samples.shape[0]):.1e}")<jupyter_output>Standard Error of the Mean for theta_0: 2.0e-03
Standard Error of the Mean for theta_1: 2.0e-03<jupyter_text>Check for between chain variations<jupyter_code>no_of_chains=2
chains=[]

for ichain in range(no_of_chains):
    sampler.reset()
    p0 = np.random.rand(ndim)/10 + (-1)**ichain
    # Warm-up
    if nwarmup > 0:
        print(f'Chain {ichain} performing {nwarmup} warnup iterations.')
        pos, prob, state = sampler.run_mcmc(p0, nwarmup)
        sampler.reset()
    else:
        pos = p0
    
    # Perform iterations, starting at the final position from the warmup.
    print(f'MH sampler {ichain} performing {nsteps} samples.')
    sampler.run_mcmc(pos, nsteps)
    print("done")
    print(f"Mean acceptance fraction: {sampler.acceptance_fraction:.3f}")

    chains.append(sampler.flatchain)
chain1 = chains[0]
chain2 = chains[1]
fix, ax = plt.subplots(2,1,figsize=(12,10))
for icol in range(ndim):
    ax[icol].plot(np.arange(chain1.shape[0]),chain1[:,icol])
    ax[icol].plot(np.arange(chain2.shape[0]),chain2[:,icol],alpha=.7)
    ax[icol].set_ylabel(r'$\theta_{0}$'.format(icol))<jupyter_output><empty_output><jupyter_text>This is a multimodal distribution!<jupyter_code>fig = plt.figure()
ax = fig.gca()

# Make data.
X = np.arange(-4, 4, 0.05)
Y = np.arange(-4, 4, 0.05)
X, Y = np.meshgrid(X, Y)
Z=posterior(Y,X)

ax.set_xlim([-1.6,1.6])
ax.set_ylim([-1.6,1.6])
ax.contour(X, Y, Z, 10)
CS=ax.contourf(X, Y, Z, cmap=plt.cm.cubehelix_r);
cbar = plt.colorbar(CS)
ax.set_aspect('equal', 'box')<jupyter_output><empty_output><jupyter_text>PT Sampler<jupyter_code># Now we can construct a sampler object that will drive the PTMCMC; 
# arbitrarily, we choose to use 21 temperatures 
# (the default is for each temperature to increase by a factor of sqrt(2), 
# so the highest temperature will be T=1024, resulting in an effective 
# \sigma_T=32\sigma=3.2, which is about the separation of our modes). 

#ntemps = 21
#temps = np.array([np.sqrt(2)**i for i in range(ntemps)])

# Modified temperature ladder to improve the integration for evidence calculation
# need more low temperatures, i.e. finer grid near beta=1
ntemps_lo = 8
ntemps_hi = 21
temps_lo = np.array([2**(i/8.) for i in range(ntemps_lo)])
temps_hi = np.array([np.sqrt(2)**i for i in range(ntemps_hi)])
temps = np.concatenate((temps_lo,temps_hi[temps_hi>max(temps_lo)]))
ntemps=len(temps)

# Let us use 10 walkers in the ensemble at each temperature:
nwalkers = 10
ndim = 2

nburnin = 1000
niterations=1000
nthin = 10 # only record every nthin iteration

nthreads = 8
betas=1/temps
sampler=PTSampler(ntemps, nwalkers, ndim, logl, logp, threads=nthreads,betas=betas)

#Making the sampling multi-threaded is as simple as adding the threads=Nthreads 
# argument to PTSampler. We could have modified the temperature ladder using the 
# betas optional argument (which should be an array of \beta=1/T values). 

#First, we run the sampler for 1000 burn-in iterations:

#p0 = np.random.uniform(low=-1.0, high=1.0, size=(ntemps, nwalkers, ndim))
p0 = np.random.normal(loc=mup, scale=sigp, size=(ntemps, nwalkers, ndim))

print("Running burn-in phase")
for p, lnprob, lnlike in sampler.sample(p0, iterations=nburnin):
    pass
sampler.reset()

print("Running MCMC chains")
#Now we sample for nwalkers*niterations, recording every nthin-th sample:

for p, lnprob, lnlike in sampler.sample(p, lnprob0=lnprob,
                                        lnlike0=lnlike,
                                        iterations=niterations, thin=nthin):
    pass 
#The resulting samples (nwalkers*niterations/nthin of them)
# are stored as the sampler.chain property:

assert sampler.chain.shape == (ntemps, nwalkers, niterations/nthin, ndim)
# Chain has shape (ntemps, nwalkers, nsteps, ndim)
# Zero temperature mean:
mu0 = np.mean(np.mean(sampler.chain[0,...], axis=0), axis=0)
print("The zero temperature mean is: {}".format(mu0))

(mu1,mu2,sigma1inv,sigma2inv,norm1,norm2)=params_modes
print("To be compared with likelihood distribution: ")
print("... peak 1: {}, peak 2: {}".format(mu1, mu2))
mcmc_data0 = sampler.chain[0,...].reshape(-1,ndim)
figure = corner.corner(mcmc_data0)

# Extract the axes
axes = np.array(figure.axes).reshape((ndim, ndim))
# Plot the posterior for four different temperatures: beta[0],beta[4],beta[10],beta[19]
ntempered=np.array([ntemps-1,ntemps-11,8,0])
for nt in ntempered:
    mcmc_data_nt = sampler.chain[nt,...].reshape(-1,ndim)
    figure = corner.corner(mcmc_data_nt)
    # Extract the axes
    axes = np.array(figure.axes).reshape((ndim, ndim))
    axes[ndim-1,ndim-1].set_title('T=%i' %temps[nt])<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14144775941994533, 'fraction_numerical': 0.03233091643884464, 'mean_word_length': 4.342222222222222, 'pattern_counts': {'":': 1, '<': 33, '<?xml version=': 0, '>': 36, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '604453', 'license': 'CC0-1.0', 'n_tokens_mistral': 3370, 'n_tokens_neox': 3186, 'n_words': 734, 'path': 'topics/mcmc-sampling/MCMC-PT.ipynb', 'repo_name': 'asemposki/Bayes2019', 'size': '8424'}