% Test scripts that illustrate various aspects of steepest descent
% iterationm and of conjugate gradients.
%
% Feb 1, 2022, P.G. Martinsson, UT-Austin

function main_cg

%DRIVER1
%DRIVER2
%DRIVER3
DRIVER4

return

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Generate a minimization problem, and plot it.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function DRIVER1

rng(0)
[V,~]   = qr(randn(2));
A       = V*diag([1,5])*V';
x_exact = [1.5;1];
b       = A*x_exact;

[XX1,XX2] = meshgrid(x_exact(1) + linspace(-2,2,31),x_exact(2) + linspace(-2,2,31));
FF = 0.5*A(1,1)*XX1.*XX1 + A(1,2)*XX1.*XX2 + 0.5*A(2,2)*XX2.*XX2 - b(1)*XX1 - b(2)*XX2;
figure(1)
hold off
contour(XX1,XX2,FF) 
hold on
plot(x_exact(1),x_exact(2),'r.','MarkerSize',25)
axis equal
hold off

keyboard

return

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Execute steepest descent iteration to solve the minimization problem.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function DRIVER2

rng(0)
[V,~]   = qr(randn(2));
A       = V*diag([1,5])*V';
x_exact = [1.5;1];
b       = A*x_exact;

nstep = 10;

YY = zeros(2,nstep);
RR = [b,zeros(2,nstep-1)];
for istep = 2:nstep
  r           = RR(:,istep-1);
  t           = (r'*r)/(r'*A*r);
  YY(:,istep) = YY(:,istep-1) + t*r;
  RR(:,istep) = b - A*YY(:,istep);
  fprintf(1,'Steepest descent error at step %2i = %10.3e\n',istep,max(abs(RR(:,istep))))
end

[XX1,XX2] = meshgrid(x_exact(1) + linspace(-2,2,31),x_exact(2) + linspace(-2,2,31));
FF = 0.5*A(1,1)*XX1.*XX1 + A(1,2)*XX1.*XX2 + 0.5*A(2,2)*XX2.*XX2 - b(1)*XX1 - b(2)*XX2;
figure(1)
hold off
contour(XX1,XX2,FF)
hold on
plot(x_exact(1),x_exact(2),'r.','MarkerSize',25)
plot(YY(1,:),YY(2,:),'k.-')
hold off
axis equal

keyboard

return

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Execute both steepest descent and CG.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function DRIVER3

rng(0)
[V,~]   = qr(randn(2));
A       = V*diag([1,5])*V';
x_exact = [1.5;1];
b       = A*x_exact;

nstep = 5;

YY = zeros(2,nstep);
RR = [b,zeros(2,nstep-1)];
for istep = 2:nstep
  r           = RR(:,istep-1);
  t           = (r'*r)/(r'*A*r);
  YY(:,istep) = YY(:,istep-1) + t*r;
  RR(:,istep) = b - A*YY(:,istep);
  fprintf(1,'Steepest descent error at step %2i = %10.3e\n',istep,max(abs(RR(:,istep))))
end

ZZ = zeros(2,nstep);
RR = [b,zeros(2,nstep-1)];
PP = [b,zeros(2,nstep-1)];
for istep = 2:nstep
  alpha       = (RR(:,istep-1)'*RR(:,istep-1))/(PP(:,istep-1)'*A*PP(:,istep-1));
  ZZ(:,istep) = ZZ(:,istep-1) + alpha*PP(:,istep-1);
  RR(:,istep) = RR(:,istep-1) - alpha*A*PP(:,istep-1);
  beta        = (RR(:,istep)'*RR(:,istep))/(RR(:,istep-1)'*RR(:,istep-1));
  PP(:,istep) = RR(:,istep) + beta*PP(:,istep-1);
  fprintf(1,'CG error at step %2i = %10.3e\n',istep,max(abs(RR(:,istep))))
end

[XX1,XX2] = meshgrid(x_exact(1) + linspace(-2,2,31),x_exact(2) + linspace(-2,2,31));
FF = 0.5*A(1,1)*XX1.*XX1 + A(1,2)*XX1.*XX2 + 0.5*A(2,2)*XX2.*XX2 - b(1)*XX1 - b(2)*XX2;
figure(1)
hold off
contour(XX1,XX2,FF)
hold on
plot(x_exact(1),x_exact(2),'r.','MarkerSize',25)
plot(YY(1,:),YY(2,:),'k.-')
plot(ZZ(1,:),ZZ(2,:),'g.-')
hold off
axis equal

keyboard

return

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Compare steepest descent and CG for a larger problem.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function DRIVER4

rng(10)
n       = 100;
C       = randn(n);
%A       = C*C' + eye(n);    % Ill conditioned example.
A       = C*C' + n*eye(n);  % Well conditioned example.
b       = randn(n,1);
x_exact = A\b;

nstep = 40;

YY = zeros(n,nstep);
RR = [b,zeros(n,nstep-1)];
for istep = 2:nstep
  r           = RR(:,istep-1);
  t           = (r'*r)/(r'*A*r);
  YY(:,istep) = YY(:,istep-1) + t*r;
  RR(:,istep) = b - A*YY(:,istep);
  fprintf(1,'Steepest descent error at step %2i = %10.3e\n',istep,max(abs(RR(:,istep))))
end

ZZ = zeros(n,nstep);
RR = [b,zeros(n,nstep-1)];
PP = [b,zeros(n,nstep-1)];
for istep = 2:nstep
  alpha       = (RR(:,istep-1)'*RR(:,istep-1))/(PP(:,istep-1)'*A*PP(:,istep-1));
  ZZ(:,istep) = ZZ(:,istep-1) + alpha*PP(:,istep-1);
  RR(:,istep) = RR(:,istep-1) - alpha*A*PP(:,istep-1);
  beta        = (RR(:,istep)'*RR(:,istep))/(RR(:,istep-1)'*RR(:,istep-1));
  PP(:,istep) = RR(:,istep) + beta*PP(:,istep-1);
  fprintf(1,'CG error at step %2i = %10.3e\n',istep,max(abs(RR(:,istep))))
end

semilogy(1:nstep,max(abs(YY - x_exact*ones(1,nstep))),'b',...
         1:nstep,max(abs(ZZ - x_exact*ones(1,nstep))),'r',...
         'LineWidth',2)
legend('Steepest descent','CG')
title(sprintf('Condition number of A = %10.3e\n',max(svd(A))/min(svd(A))))
    
keyboard

return

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%