
sample_size = 10000
cov = np.array([[1.,-0.65],[-0.65,2]])
cov_inv = np.linalg.inv(cov)
mu_ = np.array([1,-3])
samples = []
mu = np.random.normal(0,1,size=2)
for i in range(sample_size):
    X = np.random.multivariate_normal(mu_,cov,size=1)
    samples.append(X)
samples = np.concatenate(samples)

plt.figure(figsize=(6,6))
_ = plt.hist2d(samples[:,0],samples[:,1],bins=40,)
plt.axis("equal")
plt.xlabel("$x_1$")
_ =plt.ylabel("$x_2$")

"""
Explanation: In this notebook we will explore how to use Hamiltonian Monte Carlo for sampling from posterior parameters of a Multivariate Gaussian Distribution.
We start by generating samples from a bivariate Normal Distribution using a fixed mean and covariance, then we sample from posterior of distribution of each parameter given the other parameter is known. Finally we combine the two to sample from posterior distribution of joint parameter. Although we can apply HMC on both parameters, here we alternate between sampling from one parameter given the other (simulating Gibbs sampling).
Generating Data
We generate a set of data by sampling from a Multivariate Normal Distribution. Probability Density of a Multivariate Normal Distribution is defined as:
\begin{align}
 \frac{1}{\sqrt{2 \pi \det(\Sigma)}} \exp{\left(-(\frac{1}{2}(x-\mu)^T\Sigma^{-1}(x-\mu))\right)}
\end{align}
where:
$\mu: R^{N \times 1}$ vector, indicating the mean value.
$\Sigma: N^{N \times N}$ matrix, indicating the covariance matrix.
Here we generate samples from such a distribution with $N=2$.
End of explanation
"""


def logEnergy_mu(X,mu, cov):
    cov = np.reshape(cov, (X.shape[1],X.shape[1]))
    cov_inv = np.linalg.inv(cov)
    sum_ = 0   #0.5*(mu).T.dot(mu) uncomment this to add a normal prior.
    _, logdet = np.linalg.slogdet(cov)
    sum_ += X.shape[0]*logdet + ((X.T-mu).T.dot(cov_inv)*(X.T-mu).T).sum()
    return sum_


def grad_logEnergy_mu(X,mu, cov):
    cov = np.reshape(cov, (X.shape[1],X.shape[1]))
    sum_ = 0#-mu
    cov_inv = np.linalg.inv(cov)
    sum_ += -np.dot(cov_inv,X.T-mu).sum(axis=1, keepdims = True)
    return sum_


def p_sampler_mu():
    return np.random.normal(0,1, size= (X.shape[1],1))

U_mu = lambda mu: logEnergy_mu(samples, mu, cov)

U_der_mu = lambda mu: grad_logEnergy_mu(samples, mu, cov)

current_q = np.array([[-1.0],[-3.0]])
eps = 0.001
sample_gen = 2000
hist_val = []
hist_prob = []
hist_accept = []
for i in range(sample_gen):
    current_q, prob, accept = hmc.hmc(U_mu, U_der_mu, eps, L = 30, current_q = current_q, p_sampler=p_sampler_mu )
    hist_val.append(current_q)
    hist_prob.append(prob[0])
    hist_accept.append(accept)
    if i % 500 == 0:
        print(i)
        print(prob[0], accept)
        print(current_q)
hist_prob = np.concatenate(hist_prob)
hist_prob[hist_prob>1] = 1
hist_accept = np.array(hist_accept)
hist_val = np.concatenate(hist_val,axis=1).T

import matplotlib.gridspec as gridspec
fig = plt.figure(figsize=(6*3,6))
gs = gridspec.GridSpec(2, 1)
ax1 = fig.add_subplot(gs[0])
ax1.plot(np.abs(hist_val[:,0]-mu_[0].flatten()),alpha=0.7,label="$\mu_1$",linewidth=2)
ax1.plot(np.abs(hist_val[:,1]-mu_[1].flatten()),alpha=0.7,label="$\mu_2$",linewidth=2)
ax1.set_ylabel("absolute error")
ax1.set_xlabel("iteration")
ax1.set_yscale("log")
ax1.legend(loc=1,ncol=2)

ax2 = fig.add_subplot(gs[1],sharex=ax1)
ax2.plot(hist_prob,linewidth=2)
ax2.set_ylabel("P(accept)")
ax2.set_xlabel("iteration")

#plt.show()

"""
Explanation: Inference
As we said earlier, we will resort to MCMA using Hamiltonian Monte Carlo. To give a bit of context for this method following quote from the reference maybe useful:

Hamiltonian dynamics has a physical interpretation that can provide useful intuitions. In
two dimensions, we can visualize the dynamics as that of a frictionless puck that slides over a surface of varying height. The state of this system consists of the position of the puck, given by a 2D vector $q$, and the momentum of the puck (its mass times its velocity), given by a 2D vector $p$. The potential energy, $U(q)$, of the puck is proportional to the height of the surface at its current position, and its kinetic energy, $K(p)$, is equal to $|p|2/(2m)$, where m is the mass of the puck. On a level part of the surface, the puck moves at a constant velocity, equal to $p/m$. If it encounters a rising slope, the puck’s momentum allows it to continue, with its kinetic energy decreasing and its potential energy increasing, until the kinetic energy (and hence p) is zero, at which point it will slide back down (with kinetic energy increasing and potential energy decreasing).
In non-physical MCMC applications of Hamiltonian dynamics, the position will correspond
to the variables of interest. The potential energy will be minus the log of the probability
density for these variables. Momentum variables, one for each position variable, will
be introduced artificially.

Following the paper we learn that there are two quantities that we need for sampling using this method, 1) negative log likelihood $L$ and 2) derivative of negative of log likelihood with respect to parameter of interest, here $\mu$ and $\Sigma$, denoted by $\frac{\partial L}{\partial \mu}$ and $\frac{\partial L}{\partial \Sigma}$ respectively.
Mean
In this section we derive the equations for inference of the mean, assuming covariance matrix is known. Based on earlier explanation we can derive the two quantities of interest using following formulas:
\begin{align}
  L &= \sum_i^N  \frac{-1}{2}\log\det(\Sigma) -(\frac{1}{2}(x_i-\mu)^T\Sigma^{-1}(x_i-\mu))\
  \frac{\partial L}{\partial \mu} &= \sum_i^N  - \Sigma^{-1}(x_i-\mu)
\end{align}
For details of this derivation refer to Matrix cook book.
End of explanation
"""


def logEnergy_cov(X,mu,cov):
    cov = np.reshape(cov, (X.shape[1],X.shape[1]))
    sum_ = 0
    _, logdet = np.linalg.slogdet(cov)
    cov_inv = np.linalg.inv(cov)
    sum_ += X.shape[0]*logdet + ((X.T-mu).T.dot(cov_inv)*(X.T-mu).T).sum()
    return sum_
def grad_logEnergy_cov(X,mu,cov):
    cov = np.reshape(cov, (X.shape[1],X.shape[1]))
    sum_ = 0
    cov_inv = np.linalg.inv(cov)
    sum_ = X.shape[0]*cov_inv-cov_inv.dot(np.einsum('ki,kj->ij', X-mu.T, X-mu.T)).dot(cov_inv)
    return np.reshape(sum_,(sum_.size,1)) 


U_cov = lambda cov: logEnergy_cov(samples, mu_[:,np.newaxis], cov)
U_der_cov = lambda cov: grad_logEnergy_cov(samples, mu_[:,np.newaxis], cov)

def p_sampler_cov():
    scale = np.diag(np.ones(samples.shape[1]))
    mat = stats.invwishart.rvs(5,scale,1)
    return np.reshape(mat,(mat.size,1)) 

current_q = np.array([[0.1,.0],[0.0,0.1]]).reshape((4,1))
eps = 0.001
sample_gen = 2000
hist_val = []
hist_prob = []
hist_accept = []
for i in range(sample_gen):
    current_q, prob, accept = hmc.hmc(U_cov, U_der_cov
                                      , eps, L = 30, current_q = current_q,p_sampler= p_sampler_cov )
    hist_val.append(current_q)
    hist_prob.append(prob[0])
    hist_accept.append(accept)
    if i % 500 == 0:
        print(i)
        print(prob[0], accept)
        print(current_q)
        print(np.linalg.eigvals(np.reshape(current_q,(samples.shape[1],samples.shape[1]))))
hist_prob = np.concatenate(hist_prob)
hist_prob[hist_prob>1] = 1
hist_accept = np.array(hist_accept)
hist_val = np.concatenate(hist_val,axis=1).T

fig = plt.figure(figsize=(6*3,6))
gs = gridspec.GridSpec(2, 1)
ax1 = fig.add_subplot(gs[0])
ax1.plot(np.abs(hist_val[:,0]-cov[0,0].flatten()),alpha=0.7,label="$\Sigma_{11}$",linewidth=2)
ax1.plot(np.abs(hist_val[:,1]-cov[0,1].flatten()),alpha=0.7,label="$\Sigma_{12}$",linewidth=2)
ax1.plot(np.abs(hist_val[:,3]-cov[1,1].flatten()),alpha=0.7,label="$\Sigma_{22}$",linewidth=2)

ax1.set_ylabel("absolute error")
ax1.set_xlabel("iteration")
ax1.set_yscale("log")
ax1.legend(loc=1,ncol=3)

ax2 = fig.add_subplot(gs[1],sharex=ax1)
ax2.plot(hist_prob,linewidth=2)
ax2.set_ylabel("P(accept)")
ax2.set_xlabel("iteration")

"""
Explanation: Covariance
In this section we derive the necessary equations for inference of the covariance matrix, assuming the vector mean is known. Based on earlier explanation we can derive the two quantities of interest using following formulas:
\begin{align}
  L &= \frac{-1}{2}\log\det(\Sigma) -(\frac{1}{2}(x-\mu)^T\Sigma^{-1}(x-\mu))\
  \frac{\partial L}{\partial \Sigma} &= \frac{-1}{2}( \Sigma^{-1} - \Sigma^{-1}(x-\mathbf{\mu})(x-\mathbf{\mu})^T\Sigma^{-1})
\end{align}
For details of this derivation refer to Matrix cook book.
End of explanation
"""


current_mu = np.array([[-1.0],[3.0]])
current_cov = np.array([[0.1,.0],[0.0,1.0]]).reshape((4,1))
eps = 0.001
sample_gen = 2000
hist_val_mu = []
hist_prob_mu = []
hist_accept_mu = []
hist_val_cov = []
hist_prob_cov = []
hist_accept_cov = []

for i in range(sample_gen):
    U_mu = lambda mu: logEnergy_mu(samples, mu, current_cov)
    U_der_mu = lambda mu: grad_logEnergy_mu(samples, mu, current_cov)
    U_cov = lambda cov: logEnergy_cov(samples, current_mu, cov)
    U_der_cov = lambda cov: grad_logEnergy_cov(samples, current_mu, cov)
    
    current_mu, prob_mu, accept_mu = hmc.hmc(U_mu, U_der_mu, eps, L = 30, current_q = current_mu,p_sampler= p_sampler_mu )
    hist_val_mu.append(current_mu)
    hist_prob_mu.append(prob_mu[0])
    hist_accept_mu.append(accept_mu)
    current_cov, prob_cov, accept_cov = hmc.hmc(U_cov, U_der_cov, eps, L = 30, current_q = current_cov,p_sampler= p_sampler_cov )
    hist_val_cov.append(current_cov)
    hist_prob_cov.append(prob_cov[0])
    hist_accept_cov.append(accept_cov)
    if i % 500 == 0:
        print(i)
        print("mu:",prob_mu[0], accept_mu,current_mu)
        print("cov:",prob_cov[0], accept_cov,current_cov)
hist_prob_mu = np.concatenate(hist_prob_mu)
hist_prob_mu[hist_prob_mu>1] = 1
hist_accept_mu = np.array(hist_accept_mu)
hist_val_mu = np.concatenate(hist_val_mu,axis=1).T

hist_prob_cov = np.concatenate(hist_prob_cov)
hist_prob_cov[hist_prob_cov>1] = 1
hist_accept_cov = np.array(hist_accept_cov)
hist_val_cov = np.concatenate(hist_val_cov,axis=1).T

"""
Explanation: Putting everything together
Gibbs sampling style
In this section we use inference in previous sections to infere both parameters $\mu$ and $\Sigma$ simultaneously. In our first approach we alternate between parameters assuming th eother is fixed.
End of explanation
"""


fig = plt.figure(figsize=(6*3,6))
gs = gridspec.GridSpec(2, 2)
ax1 = fig.add_subplot(gs[0,0])
ax1.plot(np.abs(hist_val_cov[:,0]-cov[0,0].flatten()),alpha=0.7,label="$\Sigma_{11}$",linewidth=2)
ax1.plot(np.abs(hist_val_cov[:,1]-cov[0,1].flatten()),alpha=0.7,label="$\Sigma_{12}$",linewidth=2)
ax1.plot(np.abs(hist_val_cov[:,3]-cov[1,1].flatten()),alpha=0.7,label="$\Sigma_{22}$",linewidth=2)

ax1.set_ylabel("absolute error")
ax1.set_xlabel("iteration")
ax1.set_yscale("log")
ax1.legend(loc=1,ncol=2)

ax2 = fig.add_subplot(gs[1,0],sharex=ax1)
ax2.plot(hist_prob_cov,linewidth=2)
ax2.set_ylabel("P(accept)")
ax2.set_xlabel("iteration")

ax3 = fig.add_subplot(gs[0,1])
ax3.plot(np.abs(hist_val_mu[:,0]-mu_[0].flatten()),alpha=0.7,label="$\mu_{1}$",linewidth=2)
ax3.plot(np.abs(hist_val_mu[:,1]-mu_[1].flatten()),alpha=0.7,label="$\mu_{2}$",linewidth=2)

ax3.set_ylabel("absolute error")
ax3.set_xlabel("iteration")
ax3.set_yscale("log")
ax3.legend(loc=1,ncol=3)

ax4 = fig.add_subplot(gs[1,1],sharex=ax1)
ax4.plot(hist_prob_cov,linewidth=2)
ax4.set_ylabel("P(accept)")
ax4.set_xlabel("iteration")

"""
Explanation: As we can see below, inference still works however it takes longer (close to 1000 samples unlike 10 to 100 samples in previous cases) for MCMC method to pass the burnning phase and generate samples from the posterior.
End of explanation
"""


