%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np

def sample_data(a, b, sigma, n):
    x = np.sort( np.random.random(size = n) * 50 )
    y = a + b*x + np.random.normal(scale = sigma, size = n)
    return x,y

x, y = sample_data(5, 3, 8, 20)
_ = plt.scatter(x,y)

def log_likelihood(theta, x, y):
    a, b, sigma = theta
    return - np.sum((y - a - b*x)**2) / (2 * sigma**2) - x.size * np.log(sigma*sigma) * 0.5

log_likelihood((5, 3, 8), x, y)

import scipy.optimize as op
nll = lambda *args: -log_likelihood(*args) # Minimise the negative to get the maximum!
result = op.minimize(nll, [5, 3, 8], args=(x, y))
a_est, b_est, sigma_est = result["x"]
print(a_est, b_est, sigma_est, "-->", log_likelihood((a_est, b_est, sigma_est), x, y) )

def log_prior(theta):
    a, b, sigma = theta
    return - np.log(abs(sigma))

def log_prob(theta, x, y):
    return log_prior(theta) + log_likelihood(theta, x, y)

import emcee
ndim, nwalkers = 3, 50
sampler = emcee.EnsembleSampler(nwalkers, ndim, log_prob, args=[x, y])

starting_guesses = np.random.normal(0, 0.1, (nwalkers, ndim)) + [a_est, b_est, sigma_est]
sampler.run_mcmc(starting_guesses, 2000)
emcee_trace = sampler.chain[:, 1000:, :].reshape(-1, ndim).T

import triangle
figure = triangle.corner(emcee_trace.T, labels=["$a$", "$b$", "$\sigma$"], truths=[5, 3, 8])

def plot_MCMC_model(ax, xdata, ydata, trace):
    """Plot the linear model and 2sigma contours"""
    ax.plot(xdata, ydata, 'ok')

    alpha, beta = trace[:2]
    xfit = np.linspace(0, 50, 20)
    yfit = alpha[:, None] + beta[:, None] * xfit
    mu = yfit.mean(0)
    sig = 2 * yfit.std(0)

    ax.plot(xfit, mu, '-k')
    ax.fill_between(xfit, mu - sig, mu + sig, color='lightgray')

    ax.set_xlabel('x')
    ax.set_ylabel('y')
    
fig, axes = plt.subplots(figsize=(8,6))
axes.set_xlim([0,50])
plot_MCMC_model(axes, x, y, emcee_trace)

alpha, beta = emcee_trace[:2]
print("Line:", alpha.mean(), beta.mean())

def log_prior(theta):
    a, b, sigma, g = theta
    if g <= 0:
        return -np.inf
    beta_prior = -np.log(g * sigma**2) - np.sum( (a + b * x)**2 ) / ( 2 * g * sigma**2 )
    sigma_prior = - np.log(abs(sigma))
    g_prior = - 1.5 * np.log(g) - len(x) / (2 * g)
    return beta_prior + sigma_prior + g_prior

def log_prob(theta, x, y):
    lp = log_prior(theta)
    if not np.isfinite(lp):
        return -np.inf
    return lp + log_likelihood(theta[:3], x, y)

ndim, nwalkers = 4, 50
sampler = emcee.EnsembleSampler(nwalkers, ndim, log_prob, args=[x, y])

starting_guesses = np.random.normal(0, 0.1, (nwalkers, ndim)) + [a_est, b_est, sigma_est, 1]
sampler.run_mcmc(starting_guesses, 2000)
emcee_trace = sampler.chain[:, 1000:, :].reshape(-1, ndim).T

figure = triangle.corner(emcee_trace.T, labels=["$a$", "$b$", "$\sigma$", "$g$"], truths=[5, 3, 8, 1])

figure = triangle.corner(emcee_trace[:3,:].T, labels=["$a$", "$b$", "$\sigma$"], truths=[5, 3, 8])

fig, axes = plt.subplots(figsize=(8,6))
axes.set_xlim([0,50])
plot_MCMC_model(axes, x, y, emcee_trace[:3,:])

alpha, beta = emcee_trace[:2]
print("Line:", alpha.mean(), beta.mean())

# alpha = 1, sigma = 1
y = np.arange(-0.99, 2, 0.01)
f = np.exp(-0.5 * np.log(y+1) * np.log(y+1)) / (y + 1) / np.sqrt(2 * np.pi)
plt.plot(y,f)

def exact(T0, r, ti):
    Ti = T0 * np.exp(-r * ti)
    return Ti

def simulate(T0, r, ti, variance):
    Ti = T0 * np.exp(-r * ti)
    Ti += np.random.normal(0, np.sqrt(variance), Ti.size)
    return Ti

simulate(91, 0.26, np.arange(0,7), 0.00001)

simulate(91, 0.26, np.arange(0,7), 5)

def naive_regress(yi, ti):
    """Returns (a,r)"""
    tmean = np.sum(ti) / ti.size
    ymean = np.sum(yi) / yi.size
    r = - np.sum((ti - tmean) * (yi - ymean)) / np.sum((ti - tmean)**2)
    a = ymean + r * tmean
    return a, r

ti = np.arange(0,7)
Ti = simulate(91, 0.26, ti, 5)
xi = np.log(Ti)
a_naive, r_naive = naive_regress(xi, ti)
print(a_naive, r_naive, "-->", np.exp(a_naive), r_naive)

fig, axes = plt.subplots(1, 2, figsize=(15, 5))
_ = axes[0].plot(ti, a_naive - r_naive * ti, lw = 2)
_ = axes[0].plot(ti, xi, lw = 3)
_ = axes[1].plot(ti, np.exp(a_naive - r_naive * ti), lw = 2)
_ = axes[1].plot(ti, Ti, lw = 3)

def log_likelihood(theta, ti, Ti):
    a, r, sigma = theta
    return - np.sum((Ti - np.exp(a - r * ti))**2) / (2 * sigma**2) - Ti.size * np.log(sigma*sigma) * 0.5

log_likelihood((np.log(91), 0.26, 5), ti, Ti)

import scipy.optimize as op
nll = lambda *args: -log_likelihood(*args) # Minimise the negative to get the maximum!
result = op.minimize(nll, [np.log(91), 0.26, 5], args=(ti, Ti))
a_est, r_est, sigma_est = result["x"]
print(np.exp(a_est), r_est, sigma_est, "-->", log_likelihood((a_est, r_est, sigma_est), ti, Ti) )

def log_prior(theta):
    a, b, sigma = theta
    if sigma <= 0:
        return -np.inf
    return - np.log(sigma)

def log_prob(theta, ti, Ti):
    lp = log_prior(theta)
    if not np.isfinite(lp):
        return -np.inf
    return lp + log_likelihood(theta, ti, Ti)

import emcee
ndim, nwalkers = 3, 50
sampler = emcee.EnsembleSampler(nwalkers, ndim, log_prob, args=[ti, Ti])

starting_guesses = np.random.normal(0, 0.1, (nwalkers, ndim)) + [a_est, r_est, sigma_est]
sampler.run_mcmc(starting_guesses, 2000)
emcee_trace = sampler.chain[:, 1000:, :].reshape(-1, ndim).T

import triangle
figure = triangle.corner(emcee_trace.T, labels=["$a$", "$r$", "$\sigma$"], truths=[np.log(91), 0.26, 5])

fig, axes = plt.subplots(figsize=(12, 5))
_ = axes.plot(ti, Ti, "ok")

x = np.linspace(0,6,30)
a, r = emcee_trace[:2]
y = np.exp(a[:,None] - r[:,None]*x)
mu = y.mean(0)
s = y.std(0)
_ = axes.plot(x, mu, color="black")
_ = axes.fill_between(x, mu - s, mu + s, color='lightgray')
_ = axes.plot(ti, np.exp(a_naive - r_naive * ti))
_ = axes.legend(["Data", "Bayesian", "naive"])

Ti = np.array( [91, 73, 51, 41, 31, 23, 19] )
ti = np.array( range(7) )
xi = np.log(Ti)
a_naive, r_naive = naive_regress(xi, ti)
print(a_naive, r_naive, "-->", np.exp(a_naive), r_naive)

nll = lambda *args: -log_likelihood(*args) # Minimise the negative to get the maximum!
result = op.minimize(nll, [np.log(91), 0.26, 5], args=(ti, Ti))
a_est, r_est, sigma_est = result["x"]
print(np.exp(a_est), r_est, sigma_est, "-->", log_likelihood((a_est, r_est, sigma_est), ti, Ti) )

ndim, nwalkers = 3, 50
sampler = emcee.EnsembleSampler(nwalkers, ndim, log_prob, args=[ti, Ti])

starting_guesses = np.random.normal(0, 0.1, (nwalkers, ndim)) + [a_est, r_est, sigma_est]
sampler.run_mcmc(starting_guesses, 2000)
emcee_trace = sampler.chain[:, 1000:, :].reshape(-1, ndim).T

figure = triangle.corner(emcee_trace.T, labels=["$a$", "$r$", "$\sigma$"])

fig, axes = plt.subplots(figsize=(12, 5))
_ = axes.plot(ti, Ti, "ok")

x = np.linspace(0,6,30)
a, r = emcee_trace[:2]
y = np.exp(a[:,None] - r[:,None]*x)
mu = y.mean(0)
s = y.std(0)
_ = axes.plot(x, mu, color="black")
_ = axes.fill_between(x, mu - s, mu + s, color='lightgray')
_ = axes.plot(ti, np.exp(a_naive - r_naive * ti))
_ = axes.legend(["Data", "Bayesian", "naive"])

a, r = emcee_trace[:2]
print("Naive regression r:", r_naive, "Bayesian model r:", r.mean())