Minimalistic MCMC implementation

mcmc.py

import numpy as np

def mcmc(logfunction, x0, nsteps, sigma_p):
    samples = np.empty((nsteps, len(x0)))
    logL0 = logfunction(x0)
    naccepts = 0
    for i in range(nsteps):
        x1 = np.random.normal(x0, sigma_p)
        logL1 = logfunction(x1)
        if logL1 - logL0 > np.log(np.random.uniform()):
            x0, logL0 = x1, logL1
            naccepts += 1
        samples[i,:] = x0
    return samples, naccepts

if __name__ == '__main__':
    ctr = np.array([3.14, 0.])
    sigmas = np.array([0.01, 1])
    #def logfunction(x):
    #    return -0.5 * (((x - ctr)/sigmas)**2).sum()
    def rosenbrock(x):
        return -( (1 - x[0])**2 + 100 * (x[1] - x[0]**2)**2)
    logfunction = rosenbrock
    
    x0 = np.array([0.,0.])
    sigma_p = np.array([1.,1.])
    samples, _ = mcmc(logfunction, x0, 1000, sigma_p)
    x0 = samples[-1]
    sigma_p = (samples.std(axis=0)**2 + 0.01**2)**0.5
    samples, _ = mcmc(logfunction, x0, 1000, sigma_p)
    
    x0 = samples[-1]
    sigma_p = (samples.std(axis=0)**2 + 0.001**2)**0.5
    samples, _ = mcmc(logfunction, x0, 10000, sigma_p)
    
    import corner
    import matplotlib.pyplot as plt
    corner.corner(samples)
    plt.savefig('corner.pdf', bbox_inches='tight')
    plt.close()