helmholtz_solver_petsc_s.f90 Source File


Source Code

submodule(helmholtz_solver_m) helmholtz_solver_petsc_s
    !! Helmholtz solver based on PETSC library
#include <petsc/finclude/petscksp.h>
    use MPI
    use PETSCKSP
    use screen_io_m, only : get_stderr
    use error_handling_m, only : handle_error
    use status_codes_m, only : PARALLAX_ERR_HELMHOLTZ
    use boundaries_perp_m, only : bnd_types_to_bnd_descrs
    implicit none

contains

    module subroutine create_petsc(self, multigrid, &
                                   bnd_type_core, bnd_type_wall, &
                                   bnd_type_dome, bnd_type_out, &
                                   co, lambda, xi)
        class(helmholtz_solver_petsc_t), intent(inout) :: self
        type(multigrid_t), intent(inout) :: multigrid
        integer, intent(in) :: bnd_type_core
        integer, intent(in) :: bnd_type_wall
        integer, intent(in) :: bnd_type_dome
        integer, intent(in) :: bnd_type_out
        real(FP), dimension(multigrid%get_np(1)), intent(in) :: co
        real(FP), dimension(multigrid%get_np_inner(1)), intent(in) :: lambda
        real(FP), dimension(multigrid%get_np_inner(1)), intent(in) :: xi

        type(mesh_cart_t), pointer :: mesh => null()
        type(csrmat_t), pointer :: csrpntr_wrk => null()

        integer :: lvl, lvl_rev
        
        mesh => multigrid%get_mesh_pointer(lvl=1)

        self%ndim = mesh%get_n_points()
        self%np_inner = multigrid%get_np_inner(1)
        self%nlvl = multigrid%get_nlvls()
        
        ! Get restriction and prolongation matrices from multigrid
        allocate(self%rcsr_zb_raw(self%nlvl-1))
        allocate(self%rcsr_petsc_mat(self%nlvl-1))
        
        allocate(self%pcsr_zb_raw(self%nlvl-1))
        allocate(self%pcsr_petsc_mat(self%nlvl-1))
        
        do lvl_rev = self%nlvl-1, 1, -1
            lvl = self%nlvl - lvl_rev
            ! Restriction matrix
            csrpntr_wrk => multigrid%get_mrcsr_pointer(lvl)
            call build_mat_csr_petsc_pair(csrpntr_wrk, &
                                          self%rcsr_zb_raw(lvl_rev), &
                                          self%rcsr_petsc_mat(lvl_rev))
            csrpntr_wrk => null()
            
            ! Prolongation matrix
            csrpntr_wrk => multigrid%get_mpcsr_pointer(lvl)            
            call build_mat_csr_petsc_pair(csrpntr_wrk, &
                                          self%pcsr_zb_raw(lvl_rev), &
                                          self%pcsr_petsc_mat(lvl_rev))
            csrpntr_wrk => null()
        enddo

        allocate(self%mgr_solver)
        call self%mgr_solver%create(multigrid, &
                                    bnd_type_core, bnd_type_wall, &
                                    bnd_type_dome, bnd_type_out, &
                                    co, lambda, xi)
                            
        ! Get elliptic matrices from mgr_solver
        allocate(self%hcsr_petsc_mat(self%nlvl))
        allocate(self%hcsr_zb_raw(self%nlvl))
        
        do lvl_rev = self%nlvl, 1, -1
            lvl = self%nlvl + 1 - lvl_rev
            call build_mat_csr_petsc_pair(self%mgr_solver%get_hcsr(lvl), &
                                          self%hcsr_zb_raw(lvl_rev), &
                                          self%hcsr_petsc_mat(lvl_rev))
        enddo
        
        mesh => null()
    
    end subroutine

    module subroutine update_petsc(self, co, lambda, xi, &
                                   bnd_type_core, bnd_type_wall, &
                                   bnd_type_dome, bnd_type_out)
        class(helmholtz_solver_petsc_t), intent(inout) :: self
        real(FP), dimension(self%ndim),     intent(in), target :: co
        real(FP), dimension(self%np_inner), intent(in), target :: lambda
        real(FP), dimension(self%np_inner), intent(in), target :: xi
        integer, intent(in), optional :: bnd_type_core
        integer, intent(in), optional :: bnd_type_wall
        integer, intent(in), optional :: bnd_type_dome
        integer, intent(in), optional :: bnd_type_out

        logical :: update_with_bnds
        integer :: lvl_rev, lvl, ierr

        update_with_bnds = present(bnd_type_core) &
                            .and. present(bnd_type_wall) &
                            .and. present(bnd_type_dome) &
                            .and. present(bnd_type_out)

        if (update_with_bnds) then
            call self%mgr_solver%update(co, lambda, xi, &
                                        bnd_type_core, bnd_type_wall, &
                                        bnd_type_dome, bnd_type_out)
        else
            call self%mgr_solver%update(co, lambda, xi)
        endif


        ! Get updated elliptic matrices from mgr_solver
        do lvl_rev = self%nlvl, 1, -1
            lvl = self%nlvl + 1 - lvl_rev

            call MatDestroy(self%hcsr_petsc_mat(lvl_rev), ierr)

            call build_mat_csr_petsc_pair(self%mgr_solver%get_hcsr(lvl), &
                                          self%hcsr_zb_raw(lvl_rev), &
                                          self%hcsr_petsc_mat(lvl_rev))
        enddo

        call KSPSetOperators(self%ksp,  self%hcsr_petsc_mat(self%nlvl), &
                             self%hcsr_petsc_mat(self%nlvl), ierr)

       
    end subroutine

    module subroutine init_petsc(self, rtol, atol, maxiter, pc_sel)
        class(helmholtz_solver_petsc_t), intent(inout) :: self
        real(FP), intent(in) :: rtol
        real(FP), intent(in) :: atol
        integer, intent(in) :: maxiter
        integer :: ierr, nrows, i, lvl_rev
        character(len=*), intent(in) :: pc_sel
        
        KSP :: local_ksp
        
        self%rtol = rtol
        self%atol = atol
        self%maxiter = maxiter
        
        nrows = self%hcsr_zb_raw(self%nlvl)%ndim
    
        ! Allocate PETSc vectors
        call VecCreate(PETSC_COMM_SELF, self%b, ierr)
        call VecCreate(PETSC_COMM_SELF, self%x, ierr)
        call VecSetSizes(self%b, PETSC_DECIDE, nrows, ierr)
        call VecSetSizes(self%x, PETSC_DECIDE, nrows, ierr)
        call VecSetFromOptions(self%b, ierr)
        call VecSetFromOptions(self%x, ierr)
        ! Allocate PETSc solver
        call KSPCreate(PETSC_COMM_SELF, self%ksp, ierr)
        call KSPSetFromOptions(self%ksp, ierr)
        call KSPSetOperators(self%ksp,  self%hcsr_petsc_mat(self%nlvl), &
                             self%hcsr_petsc_mat(self%nlvl), ierr)
        ! Define KSP type GMRES
        call KSPSetType(self%ksp, KSPGMRES, ierr)
 
        if (pc_sel == 'PCMG') then
            ! Define PC type MG
            call KSPGetPC(self%ksp, self%pc, ierr)       
            call PCSetType(self%pc, PCMG, ierr)
            call PCMGSetLevels(self%pc, self%nlvl, PETSC_COMM_SELF, ierr)  
          
            do lvl_rev = 1, self%nlvl-1
                call PCMGSetInterpolation(self%pc, lvl_rev, &
                                          self%pcsr_petsc_mat(lvl_rev), ierr) 
                call PCMGSetRestriction(self%pc, lvl_rev, &
                                        self%rcsr_petsc_mat(lvl_rev), ierr) 
                call PCMGGetSmoother(self%pc, lvl_rev, local_ksp, ierr)
                call KSPSetOperators(local_ksp, &
                                     self%hcsr_petsc_mat(lvl_rev+1), &
                                     self%hcsr_petsc_mat(lvl_rev+1), ierr)
            end do
        endif
        
        ! Set KSP default options
        call KSPSetTolerances(self%ksp, self%rtol, self%atol, & 
                              PETSC_DEFAULT_REAL, self%maxiter, ierr)
        call KSPSetNormType(self%ksp, KSP_NORM_UNPRECONDITIONED, ierr)
        ! Set KSP options from file (overwrites default options)
        call KSPSetUp(self%ksp, ierr)
        ! Work array        
        allocate(self%iwrk(nrows))
        do i = 1, nrows
            self%iwrk(i) = i - 1
        end do
        
    end subroutine

    module subroutine solve_petsc(self, rhs, sol, res, info)
        class(helmholtz_solver_petsc_t), intent(inout) :: self
        real(FP), dimension(self%ndim), intent(inout) :: rhs
        real(FP), dimension(self%ndim), intent(inout) :: sol
        real(FP), intent(out) :: res
        integer, intent(out) :: info
        
        integer :: ierr, nrows
        
        nrows = self%hcsr_zb_raw(self%nlvl)%ndim
        
        ! Assemble right hand side and solution/initial guess vector
        call VecSetValues(self%b, nrows, self%iwrk, rhs, &
                          INSERT_VALUES, ierr)
        call VecSetValues(self%x, nrows, self%iwrk, sol, &
                          INSERT_VALUES, ierr)
        call VecAssemblyBegin(self%b, ierr)
        call VecAssemblyEnd(self%b, ierr)
        call VecAssemblyBegin(self%x, ierr)
        call VecAssemblyEnd(self%x, ierr)
        
        ! Solve 
        call KSPSolve(self%ksp, self%b, self%x, ierr)
        
        ! Get solution
        call KSPGetSolution(self%ksp, self%x, ierr)
        call VecGetValues(self%x, nrows, self%iwrk, sol, ierr)
        
        ! Compute residuum
        call KSPGetResidualNorm(self%ksp, res, ierr)
        if (res < self%rtol) then
            call KSPGetIterationNumber(self%ksp, info, ierr)
        else
            info = -1
        endif
                
    end subroutine

    module subroutine destructor_petsc(self)
        type(helmholtz_solver_petsc_t), intent(inout) :: self
        
        integer :: lvl, lvl_f, ierr
        
        call KSPDestroy(self%ksp, ierr)
        call VecDestroy(self%x, ierr)
        call VecDestroy(self%b, ierr)
        deallocate(self%iwrk)
        
        do lvl_f = 1, self%nlvl - 1
            call MatDestroy(self%rcsr_petsc_mat(lvl_f), ierr)
            call MatDestroy(self%pcsr_petsc_mat(lvl_f), ierr)
        enddo
        
        do lvl = 1, self%nlvl
            call MatDestroy(self%hcsr_petsc_mat(lvl), ierr)
        enddo
        
        if (allocated(self%rcsr_petsc_mat)) deallocate(self%rcsr_petsc_mat)
        if (allocated(self%rcsr_zb_raw)) deallocate(self%rcsr_zb_raw)
        
        if (allocated(self%pcsr_petsc_mat)) deallocate(self%pcsr_petsc_mat)
        if (allocated(self%pcsr_zb_raw)) deallocate(self%pcsr_zb_raw)
       
        if (allocated(self%hcsr_petsc_mat)) deallocate(self%hcsr_petsc_mat)
        if (allocated(self%hcsr_zb_raw)) deallocate(self%hcsr_zb_raw)
        
        if (allocated(self%mgr_solver)) deallocate(self%mgr_solver)
        
    end subroutine
    
    subroutine build_mat_csr_petsc_pair(csrmat_in, csrmat_out, petscmat_out)
        !! Creates a zero based indexed copy of csrmat_in 
        !! and a PETSc matrix that is associated with it
        type(csrmat_t), intent(in) :: csrmat_in
        !! Matrix
        type(csrmat_t), intent(out) :: csrmat_out
        !! Sorted copy of csrmat_in
        Mat, intent(out) :: petscmat_out
        !! PETSc matrix associated with csrmat_out
        
        integer :: ierr
    
        ! Create a copy csrmat_wrk of csrmat_in
        call csrmat_in%create_copy(csrmat_out)
        
        ! Modify to zero based indexing
        csrmat_out%i = csrmat_out%i - 1
        csrmat_out%j = csrmat_out%j - 1

        call MatCreateSeqAIJWithArrays(PETSC_COMM_SELF, &
                                       csrmat_out%ndim, csrmat_out%ncol, &
                                       csrmat_out%i, csrmat_out%j,&
                                       csrmat_out%val, petscmat_out, ierr)
        call MatSetFromOptions(petscmat_out, ierr)

    end subroutine
    
end submodule