import numpy as np
%pylab inline
import pylab as plt 

m_true = -0.9594
b_true = 4.294
f_true = 0.534

N = 50
x = np.sort(10*np.random.rand(N))
yerr = 0.1+0.5*np.random.rand(N)
y = m_true*x+b_true
y += np.abs(f_true*y) * np.random.randn(N)
y += yerr * np.random.randn(N)

plt.plot(x,y,"o")

A = np.vstack((np.ones_like(x), x)).T
C = np.diag(yerr * yerr)
cov = np.linalg.inv(np.dot(A.T, np.linalg.solve(C, A)))
b_ls, m_ls = np.dot(cov, np.dot(A.T, np.linalg.solve(C, y)))

y_ls=[b_ls+m_ls*i for i in x]
plot(x,y_ls)
plt.plot(x,y,"o")

import emcee

def lnprior(theta):
    m, b, lnf = theta
    if -5.0 < m < 0.5 and 0.0 < b < 10.0 and -10.0 < lnf < 1.0:
        return 0.0
    return -np.inf

def lnlike(theta, x, y, yerr):
    m, b, lnf = theta
    model = m * x + b
    inv_sigma2 = 1.0/(yerr**2 + model**2*np.exp(2*lnf))
    return -0.5*(np.sum((y-model)**2*inv_sigma2 - np.log(inv_sigma2)))

def lnprob(theta, x, y, yerr):
    return  lnprior(theta) + lnlike(theta, x, y, yerr)

import scipy.optimize as op
nll = lambda *args: -lnlike(*args)
result = op.minimize(nll, [m_true, b_true, np.log(f_true)], args=(x, y, yerr))
m_ml, b_ml, lnf_ml = result["x"]

ndim, nwalkers = 3, 100
pos = [result["x"] + 1e-4*np.random.randn(ndim) for i in range(nwalkers)]

sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, args=(x, y, yerr))

result=sampler.run_mcmc(pos, 1000)

samples = sampler.chain[:, 50:, :].reshape((-1, ndim))

samples.shape

xl = np.array([0, 10])
for m, b, lnf in samples[np.random.randint(len(samples), size=100)]:
    plt.plot(xl, m*xl+b, color="k", alpha=0.1)
plt.plot(xl, m_true*xl+b_true, color="r", lw=2, alpha=0.8)
plt.errorbar(x, y, yerr=yerr, fmt=".k")

import triangle
triangle.corner(samples, labels=["m", "b", "ln(f)"],
                      truths=[m_true, b_true, np.log(f_true)])

from emcee import PTSampler

mu1 = np.ones(2)
mu2 = -np.ones(2)

sigma1inv = np.diag([100.0, 100.0])
sigma2inv = np.diag([100.0, 100.0])

def logl(x):
    dx1 = x - mu1
    dx2 = x - mu2

    return np.logaddexp(-np.dot(dx1, np.dot(sigma1inv, dx1))/2.0,
                        -np.dot(dx2, np.dot(sigma2inv, dx2))/2.0)

#無情報事前分布
def logp(x):
    return 0.0

ntemps = 20
nwalkers = 100
ndim = 2

ptsampler=PTSampler(ntemps, nwalkers, ndim, logl, logp)

import time
#初期値
p0 = np.random.uniform(low=.0, high=1.0, size=(ntemps, nwalkers, ndim))

starttime=time.time()
#burn in
for p, lnprob, lnlike in ptsampler.sample(p0, iterations=1000):
    pass
sampler.reset()

for p, lnprob, lnlike in ptsampler.sample(p, lnprob0=lnprob,
                                           lnlike0=lnlike,
                                           iterations=10000, thin=10):
    pass
elapsed_time = time.time() - starttime

elapsed_time

ptsampler.chain.shape

ptsamples_low=ptsampler.chain[0,:,1000:,:].reshape((-1, ndim))
ptsamples_high=ptsampler.chain[-1,:,1000:,:].reshape((-1, ndim))

triangle.corner(ptsamples_low, labels=["mu1", "mu2"])

triangle.corner(ptsamples_high, labels=["mu1", "mu2"])

ptsampler_20=PTSampler(20, 20, ndim, logl, logp)

p0 = np.random.uniform(low=.0, high=1.0, size=(20, 20, ndim))

starttime=time.time()
#burn in
for p, lnprob, lnlike in ptsampler_20.sample(p0, iterations=1000):
    pass
sampler.reset()

for p, lnprob, lnlike in ptsampler_20.sample(p, lnprob0=lnprob,
                                           lnlike0=lnlike,
                                           iterations=10000, thin=10):
    pass
print time.time() - starttime

ptsamples_20_low=ptsampler_20.chain[0,:,1000:,:].reshape((-1, ndim))

triangle.corner(ptsamples_20_low, labels=["mu1", "mu2"])