jacobi_diis_solver Subroutine

    subroutine jacobi_diis_solver(n, rhs, x, eel, matvec, inv_diag, arg_tol, &
                                  arg_n_iter, arg_diis_max)
    
        use mod_constants, only: eps_rp
        use mod_memory, only: mallocate, mfree
        
        implicit none
    
        integer(ip), intent(in) :: n
        !! Size of the matrix
        real(rp), intent(in), optional :: arg_tol
        !! Optional convergence criterion in input, if not present
        !! OMMP_DEFAULT_SOLVER_TOL is used.
        real(rp) :: tol
        !! Convergence criterion, it is required that RMS norm < tol
        
        integer(ip), intent(in), optional :: arg_n_iter
        !! Optional maximum number of iterations for the solver, if not present
        !! OMMP_DEFAULT_SOLVER_ITER is used.
        integer(ip) :: n_iter
        !! Maximum number of iterations for the solver 
        
        integer(ip), intent(in), optional :: arg_diis_max
        !! Optional maximum number of points for diis extrapolation, if not present
        !! OMMP_DEFAULT_DIIS_MAX_POINTS is used.
        integer(ip) :: diis_max
        !! Maximum number of points for diis extrapolation, if zero or negative,
        !! diis extrapolation is not used.

        real(rp), dimension(n), intent(in) :: rhs
        !! Right hand side of the linear system
        real(rp), dimension(n), intent(inout) :: x
        !! In input, initial guess for the solver, in output the solution
        type(ommp_electrostatics_type), intent(in) :: eel
        !! Electrostatics data structure
        real(rp), dimension(n), intent(in) :: inv_diag
        !! Element-wise inverse of diagonal of LHS matrix
        external :: matvec
        !! Routine to perform matrix-vector product
        
        integer(ip) :: it, nmat
        real(rp) :: rms_norm, max_norm
        logical :: do_diis
        real(rp), allocatable :: x_new(:), y(:), x_diis(:,:), e_diis(:,:), bmat(:,:)
        character(len=OMMP_STR_CHAR_MAX) :: msg
        
        ! Optional arguments handling
        if(present(arg_tol)) then
            tol = arg_tol
        else
            tol = OMMP_DEFAULT_SOLVER_TOL
        end if
        
        if(present(arg_n_iter)) then
            n_iter = arg_n_iter
        else
            n_iter = OMMP_DEFAULT_SOLVER_ITER
        end if
        
        if(present(arg_diis_max)) then
            diis_max = arg_diis_max
        else
            diis_max = OMMP_DEFAULT_DIIS_MAX_POINTS
        end if

        do_diis =  (diis_max > 0)
        
        call ommp_message("Solving linear system with jacobi solver", OMMP_VERBOSE_LOW)
        write(msg, "(A, I4)") "Max iter:", n_iter
        call ommp_message(msg, OMMP_VERBOSE_LOW)
        write(msg, "(A, E8.1)") "Tolerance: ", tol
        call ommp_message(msg, OMMP_VERBOSE_LOW)
        if(do_diis) then
            write(msg, "(A, I4)") "DIIS is enabled with n = ", diis_max
        else
            write(msg, "(A)") "DIIS is disabled"
        endif
        call ommp_message(msg, OMMP_VERBOSE_LOW)
        
        ! Memory allocation
        call mallocate('jacobi_diis_solver [x_new]', n, x_new)
        call mallocate('jacobi_diis_solver [y]', n, y)
        if(do_diis) then
            call mallocate('jacobi_diis_solver [x_diis]', n, diis_max, x_diis)
            call mallocate('jacobi_diis_solver [e_diis]', n, diis_max, e_diis)
            call mallocate('jacobi_diis_solver [bmat]', diis_max+1, diis_max+1, bmat)
            nmat = 1
        endif
        
        ! if required, compute a guess
        rms_norm = dot_product(x, x)
        if(rms_norm < eps_rp) then
            call ommp_message("Input guess has zero norm, generating a guess&
                              & from preconditioner.", OMMP_VERBOSE_HIGH)
            x = inv_diag * rhs
        else
            call ommp_message("Using input guess as a starting point for&
                              & iterative solver.", OMMP_VERBOSE_HIGH)
        end if
        
        ! Jacobi iterations
        do it = 1, n_iter
            ! y = rhs - O x
            call matvec(eel, x, y, .false.)
            y = rhs - y

            ! x_new = D^-1 y
            x_new = inv_diag * y
            !call precnd(y, x_new)
            
            ! DIIS extrapolation
            if(do_diis) then
                x_diis(:,nmat) = x_new
                e_diis(:,nmat) = x_new - x
                call diis(n, nmat, diis_max, x_diis, e_diis, bmat, x_new)
            endif

            ! increment
            x = x_new - x
            ! compute norm
            call rmsvec(n, x, rms_norm, max_norm)
            ! update
            x = x_new
            
            write(msg, "('iter=',i4,' residual norm (rms, max): ', 2d14.4)") it, rms_norm, max_norm
            call ommp_message(msg, OMMP_VERBOSE_HIGH)

            ! Check convergence
            if(max_norm < tol) then
                call ommp_message("Required convergence threshold reached, &
                                  &exiting iterative solver.", OMMP_VERBOSE_HIGH)
                exit
            end if
        enddo
        
        call mfree('jacobi_diis_solver [x_new]', x_new)
        call mfree('jacobi_diis_solver [y]', y)
        if(do_diis) then
            call mfree('jacobi_diis_solver [x_diis]', x_diis)
            call mfree('jacobi_diis_solver [e_diis]', e_diis)
            call mfree('jacobi_diis_solver [bmat]', bmat)
        endif
      
        if(max_norm > tol) then
            call fatal_error("Iterative solver did not converged")
        end if

    end subroutine jacobi_diis_solver