# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Fitting a line to correlated Gaussian data

# %matplotlib inline
import numpy as np
import matplotlib.pyplot as pl
import oktopus
from oktopus import MultivariateGaussianPosterior, MultivariateGaussianLikelihood, UniformPrior, JointPrior
from oktopus.models import ExpSquaredKernel, WhiteNoiseKernel
from matplotlib import rc
rc('text', usetex=False)
font = {'size'   : 18,}
rc('font', **font)

x = np.linspace(0, 10, 100)
def mean(m, b):
    return m * x + b
def cov(k, l, s):
    return ExpSquaredKernel(x).evaluate(k, l) + WhiteNoiseKernel(len(x)).evaluate(s)
np.random.seed(2)
fake_data = np.random.multivariate_normal(mean(3, 10), cov(1, 1, 1))
pl.plot(x, fake_data)
pl.title("Fake Data")
pl.ylabel("y")
pl.xlabel("x")

fig, ax = pl.subplots(figsize=(8, 14))
im = ax.imshow(cov(1, 1, 1))
ax.set_title("True Covariance Matrix")
cbar = fig.colorbar(im, fraction=0.046, pad=0.04)

prior_m = UniformPrior(1, 5, 'm')
prior_b = UniformPrior(4, 15, 'b')
prior_k = UniformPrior(0, 2, 'k')
prior_l = UniformPrior(0, 2, 'l')
prior_s = UniformPrior(0, 2, 's')
prior = JointPrior(prior_m, prior_b, prior_k, prior_l, prior_s)
prior([2.5, 11, 1.3, 1.5, 0.7])

posterior = MultivariateGaussianPosterior(fake_data, mean, cov, 2, prior)

map_est = posterior.fit(x0=(2.5, 11, 1.3, 1.5, 0.7))
map_est

pl.plot(x, mean(*map_est.x[:2]))
pl.plot(x, fake_data)
pl.ylabel("y")
pl.xlabel("x")

fig, ax = pl.subplots(figsize=(8, 14))
pl.imshow(cov(*map_est.x[2:]))
ax.set_title("Estimated Covariance Matrix")
cbar = fig.colorbar(im, fraction=0.046, pad=0.04)

import emcee
ndim, nwalkers = 5, 100
p0 = [map_est.x + 1e-4*np.random.randn(ndim) for i in range(nwalkers)]

sampler = emcee.EnsembleSampler(nwalkers, ndim, lambda params: - posterior(params))
_ = sampler.run_mcmc(p0, 1000)

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

import corner
fig = corner.corner(samples, labels=["$m$", "$b$", "$k$", "$l$", "$s$"],
                    truths=[3, 10, 1, 1, 1])

# +
from matplotlib.ticker import MaxNLocator
label = ["$m$", "$b$", "$k$", "$l$", "$s$"]

fig, axes = pl.subplots(5, 1, sharex=True, figsize=(8, 14))
for i in range(0, 5, 1):
    axes[i].plot(sampler.chain[:, 50:, i].T, color="k", alpha=0.2)
    axes[i].yaxis.set_major_locator(MaxNLocator(5))
    axes[i].set_ylabel(label[i])

axes[4].set_xlabel("step number")
fig.tight_layout(h_pad=0.0)
# -



StarcoderdataJupyterScriptsDedupFiltered