# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.6.0
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Discussion about a unique solution of vb linear regression with a laplace prior.
# + In the context of VB linear regression, penalty term of the laplace prior is a little bit different from l1-normalization
# + In this notebook, we describe discuss the following things:
#     1. how different the penalty term by simplifying the problem
#     2. In the simple case, the objective function is a convex function, thus we can optimize the approximated posterior globally.

# %matplotlib widget

from scipy.stats import norm
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d

# ## Comparison among the penalty terms:
# + First, we compare the penalty term of the laplace prior in the VB with L1, VB normal.
# + The penalty term of the laplace $-E[\log p_l(w)]$ is written as:
#
# $$
# -E[\log p(w|\beta)] = \frac{1}{\beta} \sum_{j=1}^M \mu_j + 2\sigma_j \phi(-\tau_j) - 2\mu_j \Phi(-\tau_j),
# $$
# where $\mu_j, \sigma_j$ is mean and variance of the j-th element of the vb posterior approximated by normal distribution,  
# $\tau_j = \frac{\mu_j}{\sigma_j}$, $\phi(\cdot)$ and $\Phi(\cdot)$ represents density function and cumulative function of the standard normal distribution.
#
# + The penalty term of the VB normal $-E[\log p_n(w)]$ is written as:
#
# $$
# -E[\log p_n(w)] = \frac{\beta}{2} \sum_{j=1}^{M} (\mu_j^2 + \sigma_j^2).
# $$
#
# To visualize the term, we consider the case $M=1$.

# +
mu = np.arange(-50, 50, step = 0.1)
sigma = np.arange(0.0001, 50, step = 0.1)

MU, SIGMA = np.meshgrid(mu, sigma)
# -

vb_laplace_term = lambda mu, sigma, beta = 1: (mu + 2 * sigma * norm.pdf(-mu/sigma) - 2 * mu * norm.cdf(-mu/sigma)) / beta
vb_normal_term = lambda mu, sigma, beta = 1: beta/2 * (mu**2 + sigma**2)
l1_term = lambda mu, sigma, beta = 1: (np.abs(mu)) / beta

z_laplace = vb_laplace_term(MU, SIGMA)
z_normal = vb_normal_term(MU, SIGMA, beta = 0.1)
z_l1 = l1_term(MU, SIGMA)

# +
# Figureと3DAxesを追加
fig = plt.figure(figsize = (8, 6))
ax = fig.add_subplot(111, projection="3d")

# 3DAxesの設定
ax.set_xlabel("mu", size = 15)
ax.set_ylabel("sigma", size = 15)
ax.set_zlabel("penalty_term", size = 15)

# 曲面を描画
ax.plot_surface(MU, SIGMA, z_l1, cmap = "summer")
ax.plot_surface(MU, SIGMA, z_laplace, cmap = "autumn")
# ax.plot_surface(MU, SIGMA, z_normal, cmap = "winter")
plt.show()
# -

# ### Summary
# + The value of laplace penalty term depends on the variance.
#     + The more the value is, the more the behavior differs from the l1-normalization.

# ## Visualization of the objective function
# + In the case of $M=1$, the objective function is the bivariate function, thus we can visualize the form:
#
# $$
# F(\mu, \sigma) = \|y - X\mu\|_2^2/2 - \log \sigma + \frac{1}{2} \sum_{i=1}^n x_i^2 \sigma^2 + \frac{1}{\beta} \{ \mu + 2\sigma \phi(-\tau) - 2\mu \Phi(-\tau) \},
# $$
# where X and y is generated by the same problem of linear regression.

n = 10
data_seed = 20201107

np.random.seed(data_seed)

true_w = 1

X = np.random.normal(size = n)
y = np.random.normal(size = n) + X * true_w

vb_laplace_obj_func = lambda mu, sigma, X, y, beta = 1: sum([(y[i] - X[i] * mu)**2 for i in range(len(X))])/2 - np.log(sigma) + (X**2).sum() * sigma**2 + (mu + 2 * sigma * norm.pdf(-mu/sigma) - 2 * mu * norm.cdf(-mu/sigma)) / beta

# +
mu = np.arange(-5, 5, step = 0.1)
sigma = np.arange(0.0001, 5, step = 0.1)

MU, SIGMA = np.meshgrid(mu, sigma)
# -

z_obj_vals = vb_laplace_obj_func(MU, SIGMA, X, y)

min_ind = np.where(z_obj_vals == z_obj_vals.min())

MU[min_ind]

SIGMA[min_ind]



# +
# Figureと3DAxesを追加
fig = plt.figure(figsize = (8, 6))
ax = fig.add_subplot(111)

ax.set_xlabel("mu", size = 15)
ax.set_ylabel("sigma", size = 15)

plt.pcolormesh(MU, SIGMA, z_obj_vals, cmap = 'summer')
 
# plt.gca().set_aspect('equal')
plt.show()

# +
# Figureと3DAxesを追加
fig = plt.figure(figsize = (8, 6))
ax = fig.add_subplot(111, projection="3d")

# 3DAxesの設定
ax.set_xlabel("mu", size = 15)
ax.set_ylabel("sigma", size = 15)
ax.set_zlabel("obj_val", size = 15)

# 曲面を描画
ax.plot_surface(MU, SIGMA, z_obj_vals, cmap = "autumn")
# ax.plot_surface(MU, SIGMA, z_normal, cmap = "winter")
plt.show()
# -
# ### Summary
# + We could capture the convexity of the objective function of the regression.
#     + Therefore, we can give a global minimum through a gradient method.



