from mpl_toolkits.mplot3d import proj3d
from numpy.linalg import inv
import matplotlib.pyplot as plt
import numpy as np
from numpy import matrix, linalg, ones, array

Q = np.eye(3)*.1 # error covariance matrix
#Q[0,0]=1

beta = matrix(ones((2,1))) # this is what we are trying estimate
W = matrix([[1,2],
            [2,3],
            [1,1]])

ntrials = 50 
epsilon = np.random.multivariate_normal((0,0,0),Q,ntrials).T 
y=W*beta+epsilon

K=inv(W.T*inv(Q)*W)*matrix(W.T)*inv(Q) 
b=K*y #estimated beta from data

fig = plt.figure()
fig.set_size_inches([6,6])

# some convenience definitions for plotting
bb = array(b)
bm = bb.mean(1)
yy = array(y)
ax = fig.add_subplot(111, projection='3d')

ax.plot3D(yy[0,:],yy[1,:],yy[2,:],'mo',label='y',alpha=0.3)
ax.plot3D([beta[0,0],0],[beta[1,0],0],[0,0],'r-',label=r'$\beta$')
ax.plot3D([bm[0],0],[bm[1],0],[0,0],'g-',lw=1,label=r'$\hat{\beta}_m$')
ax.plot3D(bb[0,:],bb[1,:],0*bb[1,:],'.g',alpha=0.5,lw=3,label=r'$\hat{\beta}$')
ax.legend(loc=0,fontsize=18)
plt.show()

from  matplotlib.patches import Ellipse

fig, ax = plt.subplots()
fig.set_size_inches((6,6))
ax.set_aspect(1)
ax.plot(bb[0,:],bb[1,:],'g.')
ax.plot([beta[0,0],0],[beta[1,0],0],'r--',label=r'$\beta$',lw=4.)
ax.plot([bm[0],0],[bm[1],0],'g-',lw=1,label=r'$\hat{\beta}_m$')
ax.legend(loc=0,fontsize=18)
ax.grid()

bm_cov = inv(W.T*Q*W)
U,S,V = linalg.svd(bm_cov) 

err = np.sqrt((matrix(bm))*(bm_cov)*(matrix(bm).T))
theta = np.arccos(U[0,1])/np.pi*180

ax.add_patch(Ellipse(bm,err*2/np.sqrt(S[0]),err*2/np.sqrt(S[1])
                       ,angle=theta,color='pink',alpha=0.5))

plt.show()