#Simple Conjugate Gradient in 2d

#g(x) = (1/2) x'Ax - b'x + c

set.seed(247)
p<-2
A<-rWishart(1,p+1,diag(p))[,,1]
b<-rnorm(p)
cval<--3

xstar<-solve(A)%*%b

x1<-x2<-seq(-10,10,by=0.05)

fQ<-function(x1v,x2v){return(A[1,1]*x1v^2+2*A[1,2]*x1v*x2v+A[2,2]*x2v^2)}
fL<-function(x1v,x2v){return(b[1]*x1v+b[2]*x2v)}

Qterm<-outer(x1,x2,fQ)
Lterm<-outer(x1,x2,fL)
Gmat<-0.5*Qterm-Lterm+cval

par(pty="s")
nc<-20
image(x1,x2,Gmat,col=heat.colors(nc),breaks=quantile(Gmat,prob=c(0:nc)/nc))
cont.vec<-unique(round(quantile(Gmat,prob=c(0:nc)/nc)/10,0)*10)

contour(x1,x2,Gmat,add=T,levels=cont.vec,col='gray')
abline(v=xstar[1],h=xstar[2],lty=2)
#############################################################
#Recursion: Initial approach 
x.old<-c(-7.5,7.5)
r.old<-b - A %*% x.old
alpha.old<-(t(r.old) %*% r.old/(t(r.old) %*% A %*% r.old))[1,1]
sse.old<-sum(r.old^2)

step.vals<-c(x.old,r.old,alpha.old,sse.old)

abstol<-1e-16

contour(x1,x2,Gmat,levels=cont.vec)
points(x.old[1],x.old[2],pch=3,lwd=2,col="red")

while(sse.old > abstol){

	xcur<-x.old
	x.old<-x.old+alpha.old*r.old
	r.old<-b - A %*% x.old
	alpha.old<-(t(r.old) %*% r.old/(t(r.old) %*% A %*% r.old))[1,1]
	sse.old<-sum(r.old^2)
	steps.val<-rbind(step.vals,c(x.old,r.old,alpha.old,sse.old))

	lines(c(xcur[1],x.old[1]),c(xcur[2],x.old[2]),col="red",lwd=2)

}

#############################################################
#Recursion: Orthogonal approach 

x.old<-rnorm(p)*5
r.old<-b - A %*% x.old
d.old<-r.old
alpha.old<-(t(r.old) %*% r.old/(t(d.old) %*% A %*% d.old))[1,1]
sse.old<-sum(r.old^2)
contour(x1,x2,Gmat,levels=cont.vec)
points(x.old[1],x.old[2],pch=3,lwd=2,col="red")

step.vals<-c(x.old,r.old,d.old,alpha.old,sse.old)

print(c(0,sse.old))
for(istep in 1:10){

	xcur<-x.old
	rcur<-r.old
	dcur<-d.old
	x.old<-x.old+alpha.old*d.old
	r.old<-rcur - alpha.old * A %*% d.old
	d.old<-r.old+((t(r.old) %*% r.old)/(t(rcur) %*% rcur))[1,1] * dcur
	sse.old<-sum(r.old^2)
	alpha.old<-((t(r.old) %*% r.old)/(t(d.old) %*% A %*% d.old))[1,1]
	step.vals<-rbind(step.vals,c(x.old,r.old,d.old,alpha.old,sse.old))
	points(x.old[1],x.old[2],pch=3,lwd=2,col="red")
	lines(c(xcur[1],x.old[1]),c(xcur[2],x.old[2]),col="red",lwd=2)
	print(c(istep,sse.old))
}