! $Id: pcg_mod.f90 179 2011-01-31 15:39:40Z fhrouet $
! Preconditioned Conjugate-Gradient

module pcg

  use utils


contains


  subroutine conjugate_gradient(A, b, x, prec, L, itmax, eps, trace)
    ! This routine is pretty much a one-to-one translation of the
    ! Preconditioned Conjugate-Gradient algorithm in:
    !
    ! "Templates for the Solution of Linear Systems: Building Blocks for Iterative Methods"
    ! by
    ! R. Barrett, M. Berry, T.F. Chan, J. Demmel, J.M. Donato, J. Dongarra,
    ! V. Eijkhout, R. Pozo, C. Romine, H. Van der Vorst


    implicit none
    type(csc_mat)    :: A        ! the input matrix
    type(csc_mat)    :: L        ! the incomplete L factor of A
    real(kind(1.d0)) :: x(:)     ! the starting guess and final solution
    real(kind(1.d0)) :: b(:)     ! the RHS
    integer          :: itmax    ! max number of iterations
    real(kind(1.d0)) :: eps      ! tolerance for stopping criterion
    logical          :: prec     ! whether to use or not the preconditioner
    integer          :: trace    ! whether to print messages or not

    real(kind(1.d0)), allocatable :: p(:), r(:), z(:), q(:)
    integer                       :: it
    real(kind(1.d0))              :: rho, rho_old, beta, alpha, rnrm, bnrm
    integer                       :: t_start, t_end, t_rate

    ! allocate the needed vectors
    allocate(p(A%n))
    allocate(q(A%n))
    allocate(r(A%n))
    allocate(z(A%n))

    write(*,*)'Starting Conjugate Gradient...'

    call system_clock(t_start)

    ! compute the residual for the initial guess
    r = b
    call csc_spmv(A, -1.d0, x, 1.d0, r)

    rho = 0
    ! compute the norm of b; this is needed for the stopping criterion
    bnrm = dnrm2(a%n, b, 1)

    do it=1, itmax
       ! compute the residual norm
       rnrm = dnrm2(a%n, r, 1)

       ! if the relative residual norm is small enough, exit
       if (rnrm/bnrm .lt. eps) exit
       if(trace .eq. 2) write(*,'("   It:",i5, "   rnrm/bnrm: ",es10.3)')it-1,rnrm/bnrm

       if(prec) then
          ! apply preconditioner
          call csc_trsm(L, r, z)
       else
          z = r
       end if

       rho_old = rho
       rho = dotp(r, z, a%n)

       if(it .eq. 1) then
          p = z
       else
          beta = rho/rho_old
          p = z + beta*p
       end if

       call csc_spmv(A, 1.d0, p, 0.d0, q)

       alpha = rho / dotp(p,q,a%n)
       x = x+alpha*p
       r = r - alpha*q
    end do
    if(trace .eq. 2) write(*,'("   It:",i5, "   rnrm/bnrm: ",es10.3)')it-1,rnrm/bnrm
    call system_clock(t_end, t_rate)


    ! verify residual
    r = b
    call csc_spmv(A, -1.d0, x, 1.d0, r)
    rnrm = dnrm2(a%n, r, 1)
    ! dump some information
    if(trace .ge. 1) then
       write(*,'("   Norm2 of the residual is : ",es10.3)'),rnrm
       write(*,'("   Number of iterations     : ",i5)')it-1
       write(*,'("   Total time (sec.)        : ",f11.4)')real(t_end-t_start)/real(t_rate)
       write(*,'("   Time per iteration (sec.): ",f11.4)')(real(t_end-t_start)/real(t_rate))/real(it-1)
    end if

    ! deallocate everything
    deallocate(p, q, r, z)
    return

  end subroutine conjugate_gradient

  !------------------------------------------------------------------------------

  function dotp(x, y, n)
    ! This routine simply does the dot product of two vectors
    implicit none
    real(kind(1.d0)) :: x(:), y(:), dotp
    integer :: i, n

    dotp = 0.d0
    do i=1, n
       dotp = dotp + x(i)*y(i)
    end do

    return
  end function dotp

  !------------------------------------------------------------------------------

  subroutine csc_spmv(A, alpha, x, beta, y)
    ! This routine doesn the matrix-vector product:
    ! y = alpha*Ax + beta*y
    ! for a symmetric sparse matrix in CSC format

    implicit none
    type(csc_mat) :: A
    real(kind(1.d0)) :: x(:), y(:)
    real(kind(1.d0)) :: alpha, beta

    integer :: i, j, p

    y = beta*y
    do j=1, A%n
       do p=a%jptr(j), a%jptr(j+1)-1
          i = a%irn(p)
          y(i) = y(i)+alpha*a%val(p)*x(j)
          if(i .ne. j) y(j) = y(j)+alpha*a%val(p)*x(i)
       end do
    end do

    return

  end subroutine csc_spmv

  !------------------------------------------------------------------------------

  subroutine csc_trsm(A, x, y)
    ! This routine does two triangular system solutions. It assumes that L
    ! is lower-triangular, in CSC format and that the coefficients
    ! along each column are sorted in order of increasing row index
    implicit none
    type(csc_mat) :: A
    real(kind(1.d0)) :: x(:), y(:)

    integer :: i, p, j

    ! copy x into y to make an in-place trsm
    y = x

    ! first solve L*y = y
    do j=1, A%n
       y(j) = y(j)/a%val(a%jptr(j))
       do p=a%jptr(j)+1, a%jptr(j+1)-1
          i = a%irn(p)
          y(i) = y(i) - a%val(p)*y(j)
       end do
    end do

    ! do th second solve L'y=y
    do j=A%n, 1, -1
       do p=a%jptr(j+1)-1, a%jptr(j)+1, -1
          i = a%irn(p)
          y(j) = y(j)-a%val(p)*y(i)
       end do
       y(j) = y(j)/a%val(a%jptr(j))
    end do

    return
  end subroutine csc_trsm

  !------------------------------------------------------------------------------

  double precision function dnrm2(n,x,incx)
    ! .. scalar arguments ..
    integer incx,n
    ! ..
    ! .. array arguments ..
    double precision x(*)
    ! ..

    ! purpose
    ! =======

    ! DNRM2 returns the euclidean norm of a vector via the function
    ! name, so that

    ! DNRM2 := sqrt( x'*x )


    ! -- This version written on 25-October-1982.
    ! Modified on 14-October-1993 to inline the call to DLASSQ.
    ! Sven Hammarling, Nag Ltd.


    ! .. Parameters ..
    real(kind(1.d0)), parameter :: one=1.d0, zero=0.d0

    ! ..
    ! .. Local Scalars ..
    real(kind(1.d0)) :: absxi,norm,scale,ssq
    integer :: ix
    ! ..
    ! .. Intrinsic Functions ..

    if (n.lt.1 .or. incx.lt.1) then
       norm = zero
    else if (n.eq.1) then
       norm = abs(x(1))
    else
       scale = zero
       ssq = one
       ! The following loop is equivalent to this call to the LAPACK
       ! auxiliary routine:
       ! call dlassq( n, x, incx, scale, ssq )

       do  ix = 1,1 + (n-1)*incx,incx
          if (x(ix).ne.zero) then
             absxi = abs(x(ix))
             if (scale.lt.absxi) then
                ssq = one + ssq* (scale/absxi)**2
                scale = absxi
             else
                ssq = ssq + (absxi/scale)**2
             end if
          end if
       end do
       norm = scale*sqrt(ssq)
    end if

    dnrm2 = norm
    return
  end function dnrm2


end module pcg
