solvers.f90

! solvers.f90 --- provides solvers ie. preconditioned CG, Chebyshev, Schwarz
!
! Copyright  (C)  2010  Florian Wilhelm <Florian.Wilhelm@kit.edu>
!
! Version: 1.0
! Keywords: scales ppm solver cg chebyshev schwarz
! Author: Florian Wilhelm <Florian.Wilhelm@kit.edu>
! Maintainer: Florian Wilhelm <Florian.Wilhelm@kit.edu>
! URL: https://www.dkrz.de/redmine/projects/show/scales-ppm
!
! Redistribution and use in source and binary forms, with or without
! modification, are  permitted provided that the following conditions are
! met:
!
! Redistributions of source code must retain the above copyright notice,
! this list of conditions and the following disclaimer.
!
! Redistributions in binary form must reproduce the above copyright
! notice, this list of conditions and the following disclaimer in the
! documentation and/or other materials provided with the distribution.
!
! Neither the name of the DKRZ GmbH nor the names of its contributors
! may be used to endorse or promote products derived from this software
! without specific prior written permission.
!
! THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
! IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
! TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
! PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER
! OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
! EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
! PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
! PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
! LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
! NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
! SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
!
! Commentary:
!
! Code:
!

MODULE solvers
    USE solver_public
    USE solver_internal
    USE linear_algebra
    USE mo_kind, ONLY: wp, dp
    USE ppm_extents, ONLY: extent_start, extent_end, extent_size, extent_type => extent
    USE ppm_base, ONLY: abort_ppm

    IMPLICIT NONE

    PRIVATE

    !-------------------------------------------------------------------
    ! Module Function & Procedure Interface
    !-------------------------------------------------------------------
    PUBLIC :: CG_SOLVER, CHEBYSHEV_SOLVER
    PUBLIC :: cg_method, precond_cg_method, chebyshev_method, precond_chebyshev_method, schwarz_method

    CONTAINS

    ! Preconditioned CG-Method with A*x=b and startvalue x
    FUNCTION precond_cg_method(A, b, x, ext_x, precond, exchange, tol_opt, maxiter_opt) RESULT(kiter)
        REAL(wp), INTENT(IN) :: b(:,:)                  ! right-hand-side b
        REAL(wp), INTENT(INOUT) :: x(:,:)               ! startvalue and result
        TYPE(extent_type), INTENT(IN) :: ext_x(:)       ! extent of x
        REAL(wp), OPTIONAL, INTENT(IN) :: tol_opt       ! tolerance for residual
        INTEGER, OPTIONAL, INTENT(IN) :: maxiter_opt    ! maximum iterations
        INTEGER :: i, j, jb, ib, il, jl, ie, je, inner_size, kiter, maxiter
        REAL(wp) :: delta, delta0, delta_old, dAd, alpha, beta, tol

        ! Auxiliary fields for CG-Method
        REAL(wp), ALLOCATABLE :: d(:,:), r(:,:), Ad(:,:), s(:,:)

        INTERFACE
            ! Matrix-Vector multiplication of the linear system to solve
            SUBROUTINE A(field, res_field)
                USE mo_kind, ONLY: wp, dp
                REAL(wp), INTENT(IN) :: field(:,:)
                REAL(wp), INTENT(OUT) :: res_field(:,:)
            END SUBROUTINE
            ! Preconditioner
            SUBROUTINE precond(r)
                USE mo_kind, ONLY: wp, dp
                REAL(wp), INTENT(INOUT) :: r(:,:)
            END SUBROUTINE
            ! Function to exchange boundaries if necessary
            SUBROUTINE exchange(a0, text)
                USE mo_kind, ONLY: wp, dp
                REAL(wp), INTENT(INOUT) :: a0(:,:)
                CHARACTER (LEN=*), INTENT(IN), OPTIONAL :: text
            END SUBROUTINE exchange
        END INTERFACE

        ! Check and set optional arguments
        IF ( PRESENT(maxiter_opt) ) THEN
            maxiter = maxiter_opt
        ELSE
            maxiter = config%maxiter
        ENDIF
        IF ( PRESENT(tol_opt) ) THEN
            tol = tol_opt**2 ! **2, because we compare later the square of a norm
        ELSE
            tol = config%tol**2
        ENDIF

        ! Define some variables for the ranges of the fields
        ie = SIZE(x,1)
        je = SIZE(x,2)
        ib = extent_start(ext_x(1))
        jb = extent_start(ext_x(2))
        il = extent_end(ext_x(1))
        jl = extent_end(ext_x(2))
        inner_size = extent_size(ext_x)

        ALLOCATE(d(ie,je), r(ie,je), Ad(ie,je), s(ie,je))

        ! Initialize boundary/halos of d
        CALL clear_halos(d, ext_x)

        ! Initialise residual and direction
        CALL A(x, Ad)
        r(ib:il,jb:jl) = b(ib:il,jb:jl) - Ad(ib:il,jb:jl)

        ! Apply Preconditioner
        d(ib:il,jb:jl) = r(ib:il,jb:jl)
        CALL precond(d)

        ! Calculate scalar product of residual and direction
        delta = arr_dotproduct(r(ib:il,jb:jl), d(ib:il,jb:jl))

        delta0 = delta
        kiter = 0

        DO WHILE (kiter < maxiter .AND. delta > tol*delta0)
            IF (config%do_exchange) CALL exchange(d,'precond_cg_method 1')
            kiter = kiter + 1

            ! Initialise auxiliary variables
            CALL A(d, Ad)

            dAd = arr_dotproduct( d(ib:il,jb:jl), Ad(ib:il,jb:jl) )

            alpha = delta/dAd

            ! One step of preconditioned CG-method
#ifndef BLAS
            x(ib:il,jb:jl) = x(ib:il,jb:jl) + alpha*d(ib:il,jb:jl)
            r(ib:il,jb:jl) = r(ib:il,jb:jl) - alpha*Ad(ib:il,jb:jl)
#else
            CALL DAXPY(inner_size,alpha,d(ib:il,jb:jl),1,x(ib:il,jb:jl),1)
            CALL DAXPY(inner_size,-alpha,Ad(ib:il,jb:jl),1,r(ib:il,jb:jl),1)
#endif

            ! Apply Preconditioner
            s(ib:il,jb:jl) = r(ib:il,jb:jl)
            CALL precond(s)

            ! Calculate scalar product of residual and direction
            delta_old = delta
            delta = arr_dotproduct( r(ib:il,jb:jl), s(ib:il,jb:jl) )

            beta = delta/delta_old

            !write(*,*) 'PCG iteration: ', kiter, ' current delta: ', delta

            ! Compute new direction d
#ifndef BLAS
            d(ib:il,jb:jl) = s(ib:il,jb:jl) + beta*d(ib:il,jb:jl)
#else
            CALL DSCAL(inner_size,beta,d(ib:il,jb:jl),1)
            CALL DAXPY(inner_size,1._wp,s(ib:il,jb:jl),1,d(ib:il,jb:jl),1)
#endif

        ENDDO

        DEALLOCATE(d, r, s, Ad)

        ! Communicate x (only d was used the whole time)
        IF (config%do_exchange) CALL exchange(x, 'precond_cg_method 2')

    END FUNCTION precond_cg_method

    ! CG-Method with A*x=b and startvalue x
    FUNCTION cg_method(A, b, x, ext_x, exchange, tol_opt, maxiter_opt) RESULT(kiter)
        USE preconditioners, ONLY: identity

        REAL(wp), INTENT(IN) :: b(:,:)                  ! right-hand-side b
        REAL(wp), INTENT(INOUT) :: x(:,:)               ! startvalue and result
        TYPE(extent_type), INTENT(IN) :: ext_x(:)       ! extent of x
        REAL(wp), OPTIONAL, INTENT(IN) :: tol_opt       ! tolerance for residual
        INTEGER, OPTIONAL, INTENT(IN) :: maxiter_opt    ! maximum iterations
        INTEGER :: kiter                                ! number of iterations

        INTERFACE
            ! Matrix-Vector multiplication of the linear system to solve
            SUBROUTINE A(field, res_field)
                USE mo_kind, ONLY: wp, dp
                REAL(wp), INTENT(IN) :: field(:,:)
                REAL(wp), INTENT(OUT) :: res_field(:,:)
            END SUBROUTINE
            ! Function to exchange boundaries if necessary
            SUBROUTINE exchange(a0, text)
                USE mo_kind, ONLY: wp, dp
                REAL(wp), INTENT(INOUT) :: a0(:,:)
                CHARACTER (LEN=*), INTENT(IN), OPTIONAL :: text
            END SUBROUTINE exchange
        END INTERFACE

        IF ( PRESENT(tol_opt) ) THEN
            IF ( PRESENT(maxiter_opt) ) THEN
                kiter = precond_cg_method(A, b, x, ext_x, identity, exchange, tol_opt, maxiter_opt)
            ELSE
                kiter = precond_cg_method(A, b, x, ext_x, identity, exchange, tol_opt)
            ENDIF
        ELSE
            kiter = precond_cg_method(A, b, x, ext_x, identity, exchange)
        ENDIF

    END FUNCTION cg_method

    ! Additive Schwarz method which uses a local solver for each partition
    FUNCTION schwarz_method(A, b, x, ext_x, local_solver, exchange, tol_opt, maxiter_opt) RESULT(kiter)
        REAL(wp), INTENT(IN) :: b(:,:)                  ! right-hand-side b
        REAL(wp), INTENT(INOUT) :: x(:,:)               ! startvalue and result
        TYPE(extent_type), INTENT(IN) :: ext_x(:)       ! extent of x
        REAL(wp), OPTIONAL, INTENT(IN) :: tol_opt       ! tolerance for residual
        INTEGER, OPTIONAL, INTENT(IN) :: maxiter_opt    ! maximum iterations
        INTEGER :: kiter, maxiter, iiter, ie, je, ierror
        !INTEGER :: my_rank
        REAL(wp) :: tol, numer, denom, local_relres

        INTERFACE
            ! Matrix-Vector multiplication of the linear system to solve
            SUBROUTINE A(field, res_field)
                USE mo_kind, ONLY: wp, dp
                REAL(wp), INTENT(IN) :: field(:,:)
                REAL(wp), INTENT(OUT) :: res_field(:,:)
            END SUBROUTINE
            ! The solver used by Schwarz method
            FUNCTION local_solver(A, b, x, ext_x, exchange, tol_opt, maxiter_opt) RESULT(kiter)
                USE mo_kind, ONLY: wp, dp
                USE ppm_extents, ONLY: extent_type => extent

                REAL(wp), INTENT(IN) :: b(:,:)                  ! right-hand-side b
                REAL(wp), INTENT(INOUT) :: x(:,:)               ! startvalue and result
                TYPE(extent_type), INTENT(IN) :: ext_x(:)       ! extent of x
                REAL(wp), OPTIONAL, INTENT(IN) :: tol_opt       ! tolerance for residual
                INTEGER, OPTIONAL, INTENT(IN) :: maxiter_opt    ! maximum iterations
                INTEGER :: kiter                                ! needed iterations
                INTERFACE
                    ! Matrix-Vector multiplication of the linear system to solve
                    SUBROUTINE A(field, res_field)
                        USE mo_kind, ONLY: wp, dp
                        REAL(wp), INTENT(IN) :: field(:,:)
                        REAL(wp), INTENT(OUT) :: res_field(:,:)
                    END SUBROUTINE
                    ! Function to exchange boundaries if necessary
                    SUBROUTINE exchange(a0, text)
                        USE mo_kind, ONLY: wp, dp
                        REAL(wp), INTENT(INOUT) :: a0(:,:)
                        CHARACTER (LEN=*), INTENT(IN), OPTIONAL :: text
                    END SUBROUTINE exchange
                END INTERFACE
            END FUNCTION
            ! Function to exchange boundaries if necessary
            SUBROUTINE exchange(a0, text)
                USE mo_kind, ONLY: wp, dp
                REAL(wp), INTENT(INOUT) :: a0(:,:)
                CHARACTER (LEN=*), INTENT(IN), OPTIONAL :: text
            END SUBROUTINE exchange
        END INTERFACE

        IF ( PRESENT(maxiter_opt) ) THEN
            maxiter = maxiter_opt
        ELSE
            maxiter = config%schwarz_maxiter
        ENDIF
        IF ( PRESENT(tol_opt) ) THEN
            tol = tol_opt
        ELSE
            tol = config%schwarz_tol
        ENDIF

        kiter = 0
        ie = SIZE(b,1)
        je = SIZE(b,2)

        denom = calc_abs_res(A, b, x, ext_x, .TRUE.)
        numer = denom

        !        CALL MPI_Comm_rank(mpi_comm_world, my_rank, ierror)

        DO WHILE (kiter < maxiter .AND. numer > tol*denom)
            kiter = kiter + 1

            ! Solve local problem according to solver as configured
            iiter = local_solver(A, b, x, ext_x, exchange, config%tol, config%maxiter)

            CALL exchange(x, 'schwarz_method 1')

!~             local_relres = calc_rel_res(A, b, x, ext_x, .FALSE.)
!~             WRITE(*,*) "CPU: ", my_rank, "kiter: ", kiter, "iiter: ", iiter, "local_relres: ", local_relres
!~             IF( my_rank == 0 ) WRITE(*,*) "Global relres", kiter, ":", numer/denom

            IF( kiter >= config%schwarz_miniter .AND. MOD(kiter, config%schwarz_checkrate) == 0 ) THEN
                numer = calc_abs_res(A, b, x, ext_x, .TRUE.)
            ENDIF
        END DO

    END FUNCTION schwarz_method

    ! Non-preconditioned Chebyshev method
    FUNCTION chebyshev_method(A, b, x, ext_x, lambda_min, lambda_max, exchange, tol_opt, maxiter_opt) RESULT(kiter)
        USE preconditioners, ONLY: identity

        REAL(wp), INTENT(IN) :: b(:,:)                  ! right-hand-side b
        REAL(wp), INTENT(INOUT) :: x(:,:)               ! startvalue and result
        TYPE(extent_type), INTENT(IN) :: ext_x(:)       ! extent of x
        REAL(wp), INTENT(IN) :: lambda_min, lambda_max  ! smallest and largest absolute eigenvalue
        REAL(wp), OPTIONAL, INTENT(IN) :: tol_opt       ! tolerance for residual
        INTEGER, OPTIONAL, INTENT(IN) :: maxiter_opt    ! maximum iterations
        INTEGER :: kiter                                ! number of iterations

        INTERFACE
            ! Matrix-Vector multiplication of the linear system to solve
            SUBROUTINE A(field, res_field)
                USE mo_kind, ONLY: wp, dp
                REAL(wp), INTENT(IN) :: field(:,:)
                REAL(wp), INTENT(OUT) :: res_field(:,:)
            END SUBROUTINE
            ! Function to exchange boundaries if necessary
            SUBROUTINE exchange(a0, text)
                USE mo_kind, ONLY: wp, dp
                REAL(wp), INTENT(INOUT) :: a0(:,:)
                CHARACTER (LEN=*), INTENT(IN), OPTIONAL :: text
            END SUBROUTINE exchange
        END INTERFACE

        IF ( PRESENT(tol_opt) ) THEN
            IF ( PRESENT(maxiter_opt) ) THEN
                kiter = precond_chebyshev_method(A, b, x, ext_x, lambda_min, lambda_max, identity, exchange, tol_opt, maxiter_opt)
            ELSE
                kiter = precond_chebyshev_method(A, b, x, ext_x, lambda_min, lambda_max, identity, exchange, tol_opt)
            ENDIF
        ELSE
            kiter = precond_chebyshev_method(A, b, x, ext_x, lambda_min, lambda_max, identity, exchange)
        ENDIF

    END FUNCTION chebyshev_method

    ! Preconditioned Chebyshev acceleration according to Saad "Iterative Methods for Sparse Linear Systems", Algorithm 12.1
    FUNCTION precond_chebyshev_method(A, b, x, ext_x, lambda_min, lambda_max, precond, exchange, tol_opt, maxiter_opt) RESULT(kiter)
        REAL(wp), INTENT(IN) :: b(:,:)                  ! right-hand-side b
        REAL(wp), INTENT(INOUT) :: x(:,:)               ! startvalue and result
        TYPE(extent_type), INTENT(IN) :: ext_x(:)       ! extent of x
        REAL(wp), INTENT(IN) :: lambda_min, lambda_max  ! smallest and largest absolute eigenvalue
        REAL(wp), OPTIONAL, INTENT(IN) :: tol_opt       ! tolerance for residual
        INTEGER, OPTIONAL, INTENT(IN) :: maxiter_opt    ! maximum iterations
        INTEGER :: i, j, jb, ib, il, jl, ie, je, kiter, maxiter
        REAL(wp) :: rho_old, rho_new, theta, delta, sigma, tol, denom, numer
        ! Auxiliary fields for CG-Method
        REAL(wp), ALLOCATABLE :: r(:,:), Ax(:,:), d(:,:)

        INTERFACE
            ! Matrix-Vector multiplication of the linear system to solve
            SUBROUTINE A(field, res_field)
                USE mo_kind, ONLY: wp, dp
                REAL(wp), INTENT(IN) :: field(:,:)
                REAL(wp), INTENT(OUT) :: res_field(:,:)
            END SUBROUTINE
            ! Preconditioner
            SUBROUTINE precond(r)
                USE mo_kind, ONLY: wp, dp
                REAL(wp), INTENT(INOUT) :: r(:,:)
            END SUBROUTINE
            ! Function to exchange boundaries if necessary
            SUBROUTINE exchange(a0, text)
                USE mo_kind, ONLY: wp, dp
                REAL(wp), INTENT(INOUT) :: a0(:,:)
                CHARACTER (LEN=*), INTENT(IN), OPTIONAL :: text
            END SUBROUTINE exchange
        END INTERFACE

        ! Check and set optional arguments
        IF ( PRESENT(maxiter_opt) ) THEN
            maxiter = maxiter_opt
        ELSE
            maxiter = config%maxiter
        ENDIF
        IF ( PRESENT(tol_opt) ) THEN
            tol = tol_opt**2 ! **2, because we compare later the square of a norm
        ELSE
            tol = config%tol**2
        ENDIF

        ie = SIZE(x,1)
        je = SIZE(x,2)
        ib = extent_start(ext_x(1))
        jb = extent_start(ext_x(2))
        il = extent_end(ext_x(1))
        jl = extent_end(ext_x(2))

        theta = (lambda_max + lambda_min) / 2._wp
        delta = (lambda_max - lambda_min) / 2._wp

        ! Check for trivial solution
        IF (delta == 0._wp) THEN
            IF (lambda_max == 0._wp) THEN
                x(ib:il,jb:jl) = 0._wp
            ELSE
                x(ib:il,jb:jl) = b(ib:il,jb:jl)/lambda_max
            ENDIF
            RETURN
        ENDIF

        ALLOCATE(r(ie,je), Ax(ie,je), d(ie,je))

        ! Residual and direction for first iteration
        CALL A(x, Ax)
        r(ib:il,jb:jl) = b(ib:il,jb:jl) - Ax(ib:il,jb:jl)
        CALL precond(r)
        d(ib:il,jb:jl) = 1._wp/theta*r(ib:il,jb:jl)

        ! Denominator for relativ residual
        denom = arr_dotproduct(r(ib:il,jb:jl))
        numer = denom

        kiter = 0
        sigma = theta/delta
        rho_old = 1._wp/sigma

        DO WHILE (kiter < maxiter .AND. numer > tol*denom)
            kiter = kiter + 1

            ! x_new = x_old + d_old
            x(ib:il,jb:jl) = x(ib:il,jb:jl) + d(ib:il,jb:jl)
            IF (config%do_exchange) CALL exchange(x, 'chebyshev 1')

            ! r_new = r_old - A*d_old = b - A*x_new
            CALL A(x, Ax)
            r(ib:il,jb:jl) = b(ib:il,jb:jl) - Ax(ib:il,jb:jl) ! == r - Ad in Saad Algorithm
            CALL precond(r)

            ! update rho
            rho_new = 1._wp/(2._wp*sigma - rho_old)

            ! d_new = rho_new*rho_old * d_old + (2*rho_new)/delta * r_new
            d(ib:il,jb:jl) = rho_new*rho_old*d(ib:il,jb:jl) + (2*rho_old)/delta*r(ib:il,jb:jl)
            rho_old = rho_new

            ! calculate current residual
            IF( kiter >= config%cheby_miniter .AND. MOD(kiter, config%cheby_checkrate) == 0 ) THEN
                numer = arr_dotproduct(r(ib:il,jb:jl))
            ENDIF

            !WRITE(*,*) 'CHEBY: kiter, ', kiter, 'rel_res', numer/denom
        ENDDO

        DEALLOCATE(r, Ax, d)

    END FUNCTION precond_chebyshev_method

END MODULE solvers