%
% test for CG
%
clear all;
close all;

tol = 1e-10;
maxit = 100;

n = 32;
A = gallery('poisson',n); % 2-D poisson equation
xt= randn(n*n,1);   % true solution
b = A*xt;           % right-hand side
x0 = zeros(n*n,1);

[x, flag, relres, iter, resvec] = mycg(A,x0,b,tol,maxit);

semilogy(resvec, '+:');
xlim([1,120]);
title('CG for 2-D Poisson equation with n=32')
fprintf('||x-xt||/||xt|| = %.2e\n', norm(x-xt)/norm(xt));

%%%%% CG method %%%%%
function [x, flag, relres, iter, resvec] = mycg(A, x, b, tol, maxit)
% 输入参数:
%   A: 系数矩阵
%   x: 迭代初值
%   b: 右端项
%   tol: 停机准则
%   maxit: 最大迭代步数
%
% 输出参数:
%   x: 近似解
%   flag: 算法成功标识
%   relres: 相对残量范数
%   iter: 迭代步数
%   resvec: 每次迭代的相对残量
%
flag = 0;
iter = 0;
resvec = 1;

[dim,dim] = size(A);
bnrm2 = norm( b );
if  ( bnrm2 == 0.0 ), bnrm2 = 1.0; end

r = b - A*x;
relres = norm( r ) / bnrm2;
if ( relres < tol ) return, end

for iter = 1 : maxit
    rho = (r'*r);
    if ( iter > 1 )
        beta = rho / rho_1;
        p = r + beta*p;
    else
        p = r;
    end    
    q = A*p;
    alpha = rho / (p'*q );
    x = x + alpha * p;
    r = r - alpha*q;
    relres = norm( r ) / bnrm2;
    resvec = [resvec; relres];
    if ( relres <= tol ), break, end
    rho_1 = rho;
end
if ( relres > tol ) flag = 1; end
end