#14 Linear least square fit using python

import os
import matplotlib.pyplot as plt
import numpy as np
import scipy.linalg as la
os.chdir('/home/vg/Downloads/projects/ex14')
os.getcwd()
#define model parameters
x=np.linspace(-1,1,200)
a,b,c=1,2,3
y_exact=a+b*x+c*x**2
#simulate noisy data
m=200
X=1-2*np.random.rand(m)
Y=a+b*X+c*X**2+np.random.randn(m)
#linear least square fit
A=np.vstack([X**0,X**1,X**2])
sol,r,rank,sv=la.lstsq(A.T,Y)
y_fit=sol[0]+sol[1]*x+sol[2]*x**2
fig,ax=plt.subplots(figsize=(12,4)) 
ax.plot(X,Y,'go',alpha=0.5,label='simulated data')
ax.plot(x,y_exact,'k',lw=2,label='true value $y=1+2x+3x^2$')
ax.plot(x,y_fit,'b',lw=2,label='least square fit')
#1st order polynomial
A=np.vstack([X**n for n in range(2)])
sol,r,rank,sv=la.lstsq(A.T,Y)
y_fit1=sum([s*x**n for n,s in enumerate(sol)])
#15th order polynomial
A=np.vstack([X**n for n in range(16)])
sol,r,rank,sv=la.lstsq(A.T,Y)
y_fit15=sum([s*x**n for n,s in enumerate(sol)])
ax.plot(x,y_fit1,'r',lw=2,label='least square fit (1st order)')
ax.plot(x,y_fit15,'m',lw=2,label='least square fit (15th order)')
ax.set_xlabel(r"$x$",fontsize=18)
ax.set_ylabel(r"$y$",fontsize=18)
ax. Legend(loc=2)
plt.savefig("example14.png", dpi=100)
plt.show()
plt.close()