! Solves Ax=b using the Conjugate gradient method (see wikipedia for the exact algorithm)
! For large A, use a compressed storage scheme and do the linear algebra using sparse BLAS (MKL)/SPARSKIT etc.
!
! stali@purdue.edu

program main
  implicit none

  real(8), allocatable :: A(:,:), x(:), b(:)
  real(8) :: tolerance
  integer :: max_iteration

  allocate(A(5,5),x(5),b(5))
  A = reshape((/ 2.0, -1.0,  0.0,  0.0,  0.0, &
                -1.0,  2.0, -1.0,  0.0,  0.0, &
                 0.0, -1.0,  2.0, -1.0,  0.0, &
                 0.0,  0.0, -1.0,  2.0, -1.0, &
                 0.0,  0.0,  0.0, -1.0,  2.0  /), shape=(/5,5/))

  b = (/ 0.0, 0.0, 0.0, 0.0, 100.0/)
  tolerance = 0.000000001; max_iteration = 20

  call cg (A,x,b,tolerance,max_iteration)

  print*, "x =", x
  deallocate(A,x,b)

contains

  !-----
  subroutine cg (A,x,b,tolerance,max_iteration)
    implicit none

    real(8) :: A(:,:), x(:), b(:), tolerance
    integer :: max_iteration
    real(8) :: r(size(x)), p(size(x)), Ap(size(x))
    real(8) :: alpha, beta, rtr
    integer :: k

    x = 0.0

    r = b-matmul(A,x)
    p = r
    k = 1
    do 
       Ap = matmul(A,p)
       rtr = dot_product(r,r)
       alpha = rtr / dot_product(p,Ap)
       x = x + alpha * p
       r = r - alpha * Ap
       if (sqrt(sum(r**2)) <= tolerance .or. k == max_iteration) exit
       beta = dot_product(r,r) / rtr
       p = r + beta * p
       k = k + 1
    end do

    if (k == max_iteration) then
       print*, "CG stops without converging at iteration", k
    else
       print*, "CG converged in", k, "iterations"
    end if
  end subroutine cg
  !-----

end program main