# some preliminary stuff
%pylab inline
# import scipy.stats for the pdf of the Gaussian
import scipy.stats
# seed for the synthetic data
random.seed(615)
def get_synth_data(N, beta):
    """N : int, how many points
    beta : float, precision parameter for the target variable"""
    X_N = uniform(-1,1, N)
    X_N = column_stack((ones(N), X_N))
    # for this example the true function is f(x,a) = a_0 + a_1*x
    # where a_0 = -.3 and a_1 = .5 are the parameters that we
    # are going to estimate
    a = array([-.3, .5])
    t = dot(X_N, a)
    return X_N, t + normal(loc=0.0, scale=sqrt(1./beta), size=N)
# for this example we assume that we know beta,
# the precision parameter for the gaussian noise of the target variable
beta_ = 25.0
# for this example we assume that we know alpha,
# the precision parameter for the prior of w
alpha = 2.0
N = 20
# {X_n,t_n} n=1,..,N dataset 
X_N, t_N = get_synth_data(N, beta_)
# graphics stuff
x = linspace(-1, 1, 70)
y = linspace(-1, 1, 70)

# plot the prior on w
w0, w1 = np.meshgrid(x, y)
m_0 = [0, 0]
S_0 = 1/alpha * eye(2)
Z = bivariate_normal(w0, w1, sigmax=S_0[0,0], sigmay=S_0[1,1], mux=m_0[0], muy=m_0[1], sigmaxy=S_0[0,1])
title("Prior on w")
xlabel("w0"); ylabel("w1")
extent = (-1,1,-1,1)
imshow(Z, extent=extent, origin='lower')
show()

# plot some sample lines sampled from the prior on w
random.seed(1234)
N_lines = 6
W = multivariate_normal(m_0, S_0, N)
for w in W:
    plot(x, w[0] + w[1]*x)
ylim(-1,1)
title("Some sample lines, w sampled from the prior on w")
xlabel('x'); ylabel('y')
show()

# plot one point from the data set
n = 1
X_n = X_N[range(n),:]
t_n = t_N[range(n)]
xlim(-1,1); ylim(-1,1)
scatter(X_n[:,1], t_n)
title("First point in the data set")
xlabel('x'); ylabel('y')
show()

# plot the likelihood of the first point in the dataset
def likelihood(t_, x, w, beta):
    """t_ : float, target point
    x : array, location
    w : array, weight
    beta : float, precision parameter of the target variable
    Calculate the probability that the true function f(a,x)
    take the value t_ at the location x, given that a=w and
    beta is the precision parameter of the gaussian noise.
    """
    return scipy.stats.norm.pdf(t_,loc=dot(w,x), scale=sqrt(1./beta))
Z = zeros((len(y), len(x)))
for i, w1 in enumerate(y):
    for j, w0 in enumerate(x):
        Z[i, j] = likelihood(t_n[-1], X_n[-1,:], array([w0, w1]), beta_)
extent = (-1,1,-1,1)
imshow(Z, extent=extent, origin='lower')
plot(-0.3, 0.5, 'w+', markeredgewidth=2, markersize=12)
title("Likelihood of the first point in the data set, white cross = true value")
xlabel("w0"); ylabel("w1")
show()

# plot the posterior after one point has been observed
n = 1
X_n = X_N[range(n),:]
t_n = t_N[range(n)]
phi = X_n
S_n = inv(alpha*eye(2) + beta_ * dot(phi.T, phi))
m_n = beta_ * dot(S_n, dot(phi.T, t_n))
w0, w1 = meshgrid(x, y)
Z = bivariate_normal(w0, w1, sigmax=sqrt(S_n[0,0]), sigmay=sqrt(S_n[1,1]), mux=m_n[0], muy=m_n[1], sigmaxy=S_n[0,1])
title("Posterior on w after one point")
xlabel("w0"); ylabel("w1")
extent = (-1,1,-1,1)
imshow(Z, extent=extent, origin='lower')
plot(-0.3, 0.5, 'w+', markeredgewidth=2, markersize=12)
show()

# plot some lines with the parameters sampled
# from the posterior after one point
random.seed(1234)
N_points = 1
N_lines = 6
W = multivariate_normal(m_n, S_n, N_lines)
xlim(-1,1); ylim(-1,1)
for w in W:
    plot(x, w[0] + w[1]*x)
for x_i, t_i in zip(X_N[range(N_points),1], t_N[range(N_points)]):
    plot(x_i, t_i, 'o', mfc='none', markeredgewidth=2, markeredgecolor='b', markersize=9)
title("Some sampled lines, w of the lines sampled from the posterior, blue circle = point already observed")
xlabel('x'); ylabel('y')
show()

# plot first two points from the data set
n = 2
X_n = X_N[range(n),:]
t_n = t_N[range(n)]
xlim(-1,1); ylim(-1,1)
scatter(X_n[:,1], t_n)
title("First two points in the data set")
xlabel('x'); ylabel('y')
show()

# plot the likelihood of the second point in the dataset
Z = zeros((len(y), len(x)))
for i, w1 in enumerate(y):
    for j, w0 in enumerate(x):
        Z[i, j] = likelihood(t_n[-1], X_n[-1,:], array([w0, w1]), beta_)
extent = (-1,1,-1,1)
imshow(Z, extent=extent, origin='lower')
plot(-0.3, 0.5, 'w+', markeredgewidth=2, markersize=12)
title("Likelihood of the second point in the data set, white cross = true value")
xlabel("w0"); ylabel("w1")
show()

# plot the posterior after two points has been observed
n = 2
X_n = X_N[range(n),:]
t_n = t_N[range(n)]
phi = X_n
S_n = inv(alpha*eye(2) + beta_ * dot(phi.T, phi))
m_n = beta_ * dot(S_n, dot(phi.T, t_n))
w0, w1 = meshgrid(x, y)
Z = bivariate_normal(w0, w1, sigmax=sqrt(S_n[0,0]), sigmay=sqrt(S_n[1,1]), mux=m_n[0], muy=m_n[1], sigmaxy=S_n[0,1])
title("Posterior on w after two points")
xlabel("w0"); ylabel("w1")
extent = (-1,1,-1,1)
imshow(Z, extent=extent, origin='lower')
plot(-0.3, 0.5, 'w+', markeredgewidth=2, markersize=12)
show()

# plot some lines with the parameters sampled
# from the posterior after two points
random.seed(1234)
N_points = 2
N_lines = 6
W = multivariate_normal(m_n, S_n, N_lines)
xlim(-1,1); ylim(-1,1)
for w in W:
    plot(x, w[0] + w[1]*x)
for x_i, t_i in zip(X_N[range(N_points),1], t_N[range(N_points)]):
    plot(x_i, t_i, 'o', mfc='none', markeredgewidth=2, markeredgecolor='b', markersize=9)
title("Some sampled lines, w of the lines sampled from the posterior, blue circled = points already observed")
xlabel('x'); ylabel('y')
show()

# plot all the data set
n = N
X_n = X_N[range(n),:]
t_n = t_N[range(n)]
xlim(-1,1); ylim(-1,1)
scatter(X_n[:,1], t_n)
title("data set")
xlabel('x'); ylabel('y')
show()

# plot the likelihood of the last point in the dataset
Z = zeros((len(y), len(x)))
for i, w1 in enumerate(y):
    for j, w0 in enumerate(x):
        Z[i, j] = likelihood(t_n[-1], X_n[-1,:], array([w0, w1]), beta_)
extent = (-1,1,-1,1)
imshow(Z, extent=extent, origin='lower')
plot(-0.3, 0.5, 'w+', markeredgewidth=2, markersize=12)
title("Likelihood of the last point in the data set, white cross = true value")
xlabel("w0"); ylabel("w1")
show()

# plot the posterior after all the points in the data set has been observed
n = N
X_n = X_N[range(n),:]
t_n = t_N[range(n)]
phi = X_n
S_n = inv(alpha*eye(2) + beta_ * dot(phi.T, phi))
m_n = beta_ * dot(S_n, dot(phi.T, t_n))
w0, w1 = meshgrid(x, y)
Z = bivariate_normal(w0, w1, sigmax=sqrt(S_n[0,0]), sigmay=sqrt(S_n[1,1]), mux=m_n[0], muy=m_n[1], sigmaxy=S_n[0,1])
title("Posterior on w after all the points")
xlabel("w0"); ylabel("w1")
extent = (-1,1,-1,1)
imshow(Z, extent=extent, origin='lower')
plot(-0.3, 0.5, 'w+', markeredgewidth=2, markersize=12)
show()

# plot some lines with the parameters sampled
# from the posterior after all the points has been observed
random.seed(1234)
N_points = N
N_lines = 6
W = multivariate_normal(m_n, S_n, N_lines)
xlim(-1,1); ylim(-1,1)
for w in W:
    plot(x, w[0] + w[1]*x)
for x_i, t_i in zip(X_N[range(N_points),1], t_N[range(N_points)]):
    plot(x_i, t_i, 'o', mfc='none', markeredgewidth=2, markeredgecolor='b', markersize=9)
title("Some sampled lines, w of the lines sampled from the posterior, blue circled = data set")
xlabel('x'); ylabel('y')
show()