petsc_solver.F90

Go to the documentation of this file.

#include "petsc/finclude/petscvecdef.h"
#include "petsc/finclude/petscmatdef.h"
#include "petsc/finclude/petsckspdef.h"
#include "petsc/finclude/petscdmdadef.h"
module petsc_solver_mod
  use datadefn_mod, only : default_precision, string_length
  use state_mod, only : model_state_type
  use monc_component_mod, only : component_descriptor_type
  use grids_mod, only : local_grid_type, x_index, y_index, z_index
  use logging_mod, only : log_debug, log_info, log_error, log_log, log_master_log, log_get_logging_level, log_is_master
  use optionsdatabase_mod, only : options_get_integer, options_get_real, options_get_string
  use conversions_mod, only : conv_to_string
  use communication_types_mod, only : halo_communication_type, neighbour_description_type, field_data_wrapper_type
  use halo_communication_mod, only : copy_buffer_to_field, copy_field_to_buffer, perform_local_data_copy_for_field, &
       init_halo_communication, finalise_halo_communication, blocking_halo_swap, get_single_field_per_halo_cell
  use petscvec
  use petscmat
  use petscksp
  use petscsys
  use petscdmda
  implicit none

#ifndef TEST_MODE
  private
#endif

  ksp :: ksp
  dm :: da
  integer :: z_start, z_end, y_start, y_end, x_start, x_end
  type(halo_communication_type), save :: halo_swap_state

  real(kind=DEFAULT_PRECISION), dimension(:,:,:), allocatable :: p_source, prev_p
  real(kind=DEFAULT_PRECISION) :: cx, cy
  real(kind=DEFAULT_PRECISION), dimension(:), allocatable :: rdzn, rdz, rho, rhon, dz, dzn

  public petsc_solver_get_descriptor

  contains

  type(component_descriptor_type) function petsc_solver_get_descriptor()
    petsc_solver_get_descriptor%name="petsc_solver"
    petsc_solver_get_descriptor%version=0.1
    petsc_solver_get_descriptor%initialisation=>init_callback
    petsc_solver_get_descriptor%timestep=>timestep_callback
    petsc_solver_get_descriptor%finalisation=>finalisation_callback
  end function petsc_solver_get_descriptor

 subroutine finalisation_callback(current_state)
    type(model_state_type), target, intent(inout) :: current_state

    petscerrorcode :: ierr

    call petscfinalize(ierr)
  end subroutine finalisation_callback  

  subroutine init_callback(current_state)
    type(model_state_type), target, intent(inout) :: current_state

    petscerrorcode :: ierr
    ksptype :: ksp_type
    pc :: pc_instance
    pctype :: pc_type
    integer :: y_lengths(current_state%parallel%dim_sizes(y_index)), x_lengths(current_state%parallel%dim_sizes(x_index))
    integer new_communicator, my_transposed_rank, x_rank_val, y_rank_val

    call apply_dimension_bounds(current_state%global_grid%size(y_index), current_state%parallel%dim_sizes(y_index), y_lengths)
    call apply_dimension_bounds(current_state%global_grid%size(x_index), current_state%parallel%dim_sizes(x_index), x_lengths)

    allocate(p_source(current_state%local_grid%size(z_index)-1, current_state%local_grid%size(y_index), &
         current_state%local_grid%size(x_index)), prev_p(current_state%local_grid%size(z_index)-1, &
         current_state%local_grid%size(y_index), current_state%local_grid%size(x_index)))

    cx=current_state%global_grid%configuration%horizontal%cx
    cy=current_state%global_grid%configuration%horizontal%cy
    allocate(rdzn(current_state%local_grid%size(z_index)-1), rdz(current_state%local_grid%size(z_index)-1), &
         rho(current_state%local_grid%size(z_index)-1), rhon(current_state%local_grid%size(z_index)-1), &
         dz(current_state%local_grid%size(z_index)-1), dzn(current_state%local_grid%size(z_index)-1))
    rdzn=current_state%global_grid%configuration%vertical%rdzn(2:)
    rdz=current_state%global_grid%configuration%vertical%rdz(2:)
    rho=current_state%global_grid%configuration%vertical%rho(2:)
    rhon=current_state%global_grid%configuration%vertical%rhon(2:)
    dz=current_state%global_grid%configuration%vertical%dz(2:)
    dzn=current_state%global_grid%configuration%vertical%dzn(2:)

    z_start=current_state%local_grid%halo_size(z_index)+2
    z_end=current_state%local_grid%halo_size(z_index)+current_state%local_grid%size(z_index)
    y_start=current_state%local_grid%halo_size(y_index)+1
    y_end=current_state%local_grid%halo_size(y_index)+current_state%local_grid%size(y_index)
    x_start=current_state%local_grid%halo_size(x_index)+1
    x_end=current_state%local_grid%halo_size(x_index)+current_state%local_grid%size(x_index)

    y_rank_val = current_state%parallel%my_rank / current_state%parallel%dim_sizes(x_index)
    x_rank_val = mod(current_state%parallel%my_rank, current_state%parallel%dim_sizes(x_index))    
    my_transposed_rank = x_rank_val*current_state%parallel%dim_sizes(y_index) + y_rank_val

    call mpi_comm_split(current_state%parallel%monc_communicator, 1, my_transposed_rank, new_communicator, ierr)
    petsc_comm_world = new_communicator

    call petscinitialize(petsc_null_character, ierr)    
    call kspcreate(petsc_comm_world, ksp, ierr)
    call dmdacreate3d(petsc_comm_world, dm_boundary_none, dm_boundary_periodic, dm_boundary_periodic, dmda_stencil_star, &
         current_state%global_grid%size(z_index)-1, current_state%global_grid%size(y_index), &
         current_state%global_grid%size(x_index), 1, current_state%parallel%dim_sizes(y_index), &
         current_state%parallel%dim_sizes(x_index), 1, 1, (/ current_state%global_grid%size(z_index)-1 /), &
         y_lengths, x_lengths, da, ierr)
    call kspsetdm(ksp, da, ierr)
    call kspsetcomputerhs(ksp, compute_rhs, petsc_null_object, ierr)
    call kspsetcomputeoperators(ksp, compute_matrix, petsc_null_object, ierr)
    call kspsetcomputeinitialguess(ksp, compute_initial_guess, petsc_null_object, ierr)
    call petscoptionssetvalue("-options_left", "no", ierr)
    if (trim(options_get_string(current_state%options_database, "solver_type")) .ne. "auto") then
      call petscoptionssetvalue("-ksp_type", trim(options_get_string(current_state%options_database, "solver_type")), ierr)
    end if
    if (trim(options_get_string(current_state%options_database, "preconditioner_type")) .ne. "auto") then
      call petscoptionssetvalue("-pc_type", trim(options_get_string(current_state%options_database, "preconditioner_type")), ierr)
    end if
    if (trim(options_get_string(current_state%options_database, "norm_type")) .ne. "auto") then
      if (trim(options_get_string(current_state%options_database, "norm_type")) .eq. "preconditioned" .or. &
           trim(options_get_string(current_state%options_database, "norm_type")) .eq. "unpreconditioned" .or. &
           trim(options_get_string(current_state%options_database, "norm_type")) .eq. "natural" .or. &
           trim(options_get_string(current_state%options_database, "norm_type")) .eq. "none") then
        call petscoptionssetvalue("-ksp_norm_type", trim(options_get_string(current_state%options_database, "norm_type")), ierr)
      else
        call log_master_log(log_error, "Configured PETSc norm type of '"//&
             trim(options_get_string(current_state%options_database, "norm_type"))//"' not recognised")
      end if      
    end if    
    call petscoptionssetvalue("-ksp_rtol", conv_to_string(options_get_real(current_state%options_database, "tolerance")), ierr)
    if (options_get_integer(current_state%options_database, "max_iterations") .gt. 0) then
      call petscoptionssetvalue("-ksp_max_it", &
           conv_to_string(options_get_integer(current_state%options_database, "max_iterations")), ierr)
    end if
    call kspsetfromoptions(ksp, ierr)
    call kspsetinitialguessnonzero(ksp, petsc_true, ierr)

    if (log_is_master() .and. log_get_logging_level() == log_debug) then
      call kspmonitorset(ksp,ksp_monitor,petsc_null_object, petsc_null_function,ierr)
    end if    
    prev_p=0.0_default_precision
    call init_halo_communication(current_state, get_single_field_per_halo_cell, halo_swap_state, 1, .false.)
    call kspgettype(ksp, ksp_type, ierr)
    call kspgetpc(ksp, pc_instance, ierr)
    call pcgettype(pc_instance, pc_type, ierr)
    call log_master_log(log_info, "PETSc iterative solver initialised, using "//trim(pc_type)//" preconditioner with "//&
         trim(ksp_type)//" solver")
  end subroutine init_callback

  subroutine ksp_monitor(ksp, n, rnorm, dummy, ierr)
    ksp              ksp
    petscerrorcode ierr
    petscint n,dummy
    petscreal rnorm

    call log_log(log_debug, "Iteration="//trim(conv_to_string(n))//" Rnorm="//trim(conv_to_string(rnorm, &
         exponent_small_numbers=.true.)))
  end subroutine ksp_monitor  

  subroutine timestep_callback(current_state)
    type(model_state_type), target, intent(inout) :: current_state
    
    petscerrorcode :: ierr
    vec x
    petscscalar, pointer :: xx(:)
    integer :: its, i, j, k
    petscreal :: rnorm
    kspconvergedreason :: reason

    call complete_psrce_calculation(current_state, current_state%local_grid%halo_size(y_index), &
         current_state%local_grid%halo_size(x_index))    

    p_source=current_state%p%data(z_start:z_end, y_start:y_end, x_start:x_end)
    call kspsolve(ksp, petsc_null_object, petsc_null_object, ierr)
    call kspgetsolution(ksp, x, ierr)
    call vecgetarrayreadf90(x, xx, ierr)
    
    call copy_petsc_pointer_to_data(xx, z_start, z_end, y_start, y_end, x_start, x_end, current_state%p%data)
    call vecrestorearrayreadf90(x, xx, ierr)
    
    prev_p=current_state%p%data(z_start:z_end, y_start:y_end, x_start:x_end)

    if (log_is_master() .and. log_get_logging_level() == log_debug) then
      call kspgetiterationnumber(ksp, its, ierr)
      call kspgetresidualnorm(ksp, rnorm, ierr)
      call kspgetconvergedreason(ksp, reason, ierr)
      call log_log(log_debug, "Converged in "//trim(conv_to_string(its))//" rnorm="//trim(conv_to_string(&
           rnorm, exponent_small_numbers=.true.))// " ("//trim(determine_convegence_reason(reason))//")")
    end if
    call blocking_halo_swap(current_state, halo_swap_state, copy_p_to_halo_buffer, &
         perform_local_data_copy_for_p, copy_halo_buffer_to_p)
  end subroutine timestep_callback

  character(len=STRING_LENGTH) function determine_convegence_reason(r_code)
    integer, intent(in) :: r_code

    if (r_code == 1 .or. r_code == 2) then
      determine_convegence_reason="relative tolerance of residual norm"
    else if (r_code == 9 .or. r_code == 3) then
      determine_convegence_reason="residual norm less than absolute tolerance"
    else if (r_code == 4) then
      determine_convegence_reason="number of iterations exceeded maximum"
    else if (r_code == -3) then
      determine_convegence_reason="diverged as number of iterations exceeded maximum before any convergence criteria"
    else if (r_code == -4) then
      determine_convegence_reason="diverged as divergence tolerance exceeded"
    else if (r_code .gt. 0) then
      determine_convegence_reason="convergence code of "//conv_to_string(r_code)
    else if (r_code .lt. 0) then
      determine_convegence_reason="divergence code of "//conv_to_string(r_code)
    end if    
  end function determine_convegence_reason  

  subroutine compute_initial_guess(ksp, x, dummy, ierr)
    petscerrorcode :: ierr
    ksp :: ksp
    vec :: x
    integer :: dummy(*)

    dm dm
    petscscalar, pointer :: xx(:)
    petscint :: zs, ys, xs, zm, ym, xm

    call kspgetdm(ksp, dm, ierr)
    call dmdagetcorners(dm, zs, ys, xs, zm, ym, xm, ierr)
    call vecgetarrayf90(x, xx, ierr)
    call copy_data_to_petsc_pointer(xx, zs, ys, xs, zm, ym, xm, prev_p)
    call vecrestorearrayf90(x, xx, ierr)
  end subroutine compute_initial_guess  

  subroutine compute_rhs(ksp, b, dummy, ierr)
    petscerrorcode :: ierr
    ksp :: ksp
    vec :: b
    integer :: dummy(*)
    
    dm dm
    petscscalar, pointer :: xx(:)
    petscint :: zs, ys, xs, zm, ym, xm

    call kspgetdm(ksp, dm, ierr)
    call dmdagetcorners(dm, zs, ys, xs, zm, ym, xm, ierr)
    call vecgetarrayf90(b, xx, ierr)
    call copy_data_to_petsc_pointer(xx, zs, ys, xs, zm, ym, xm, p_source)
    call vecrestorearrayf90(b, xx, ierr)
  end subroutine compute_rhs

  subroutine copy_data_to_petsc_pointer(x, zs, ys, xs, zm, ym, xm, src_values)
    integer, intent(in) :: zs, ys, xs, zm, ym, xm
    petscscalar, intent(inout) :: x(zs:zs+zm-1, ys:ys+ym-1, xs:xs+xm-1)
    real(kind=DEFAULT_PRECISION), dimension(:,:,:), intent(in) :: src_values

    integer :: i, j, k

    do i=xs, xs+xm-1
      do j=ys, ys+ym-1
        do k=zs, zs+zm-1
          x(k, j, i)=src_values((k-zs)+1, (j-ys)+1, (i-xs)+1)
        end do
      end do
    end do
  end subroutine copy_data_to_petsc_pointer

  subroutine copy_petsc_pointer_to_data(x, z_start_idx, z_end_idx, y_start_idx, y_end_idx, x_start_idx, x_end_idx, target_values)
     integer, intent(in) :: z_start_idx, z_end_idx, y_start_idx, y_end_idx, x_start_idx, x_end_idx
     petscscalar :: x((z_end_idx-z_start_idx)+1, (y_end_idx-y_start_idx)+1, (x_end_idx-x_start_idx)+1)
     real(kind=DEFAULT_PRECISION), dimension(:,:,:), intent(inout) :: target_values

     target_values(z_start_idx:z_end_idx, y_start_idx:y_end_idx, x_start_idx:x_end_idx)=x
  end subroutine copy_petsc_pointer_to_data

  subroutine compute_matrix(ksp, A, B, dummy, ierr)
    petscerrorcode :: ierr
    ksp :: ksp
    mat :: a, b
    integer :: dummy(*)

    dm dm
    petscint :: zs, ys,xs, zm, ym, xm, mz, my, mx
    real(kind=DEFAULT_PRECISION) ::  v(7), Hx, Hy
    real(kind=DEFAULT_PRECISION), dimension(:), allocatable ::  HzU, HzD, d_sc  
    matstencil :: row(4),col(4,7)
    integer :: i, j, k, num

    call kspgetdm(ksp, dm, ierr)
    call dmdagetinfo(dm,petsc_null_integer, mz, my, mx, petsc_null_integer, petsc_null_integer, petsc_null_integer, &
     petsc_null_integer, petsc_null_integer, petsc_null_integer, petsc_null_integer, petsc_null_integer, petsc_null_integer, ierr)
    call dmdagetcorners(dm, zs, ys, xs, zm, ym, xm, ierr)

    allocate(hzu(0:mz-1), hzd(0:mz-1), d_sc(0:mz-1))

    do k=0, mz-1
      d_sc(k)=rdz(k+1) / rhon(k+1)
      if (k .gt. 0) then
        hzd(k)=rho(k) * rdzn(k+1)
      else
        hzd(k)=0.0_default_precision
      end if
      if (k .lt. mz-1) then
        hzu(k)=rho(k+1) * rdzn(k+2)
      else
        hzu(k)=0.0_default_precision
      end if
    end do   

    hx = cx * cx
    hy = cy * cy  
    do i=xs, xs+xm-1
      do j=ys,ys+ym-1
        do k=zs, zs+zm-1
          row(matstencil_i) = k
          row(matstencil_j) = j
          row(matstencil_k) = i

          v(1)=hy
          col(matstencil_i, 1)=k
          col(matstencil_j, 1)=j-1
          col(matstencil_k, 1)=i
          v(2)=hx
          col(matstencil_i, 2)=k
          col(matstencil_j, 2)=j
          col(matstencil_k, 2)=i-1
          v(3)=d_sc(k) * (-(hzu(k) + hzd(k))) - (2.0*(hx + hy))
          col(matstencil_i, 3)=k
          col(matstencil_j, 3)=j
          col(matstencil_k, 3)=i
          v(4)=hx
          col(matstencil_i, 4)=k
          col(matstencil_j, 4)=j 
          col(matstencil_k, 4)=i+1
          v(5)=hy
          col(matstencil_i, 5)=k
          col(matstencil_j, 5)=j+1
          col(matstencil_k, 5)=i
          num=6
          if (k .gt. zs) then
            v(num)=(d_sc(k)*hzd(k))
            col(matstencil_i, num)=k-1
            col(matstencil_j, num)=j 
            col(matstencil_k, num)=i
            num=num+1
          end if
          if (k .lt. mz-1) then
            v(num)=(d_sc(k)*hzu(k))
            col(matstencil_i, num)=k+1
            col(matstencil_j, num)=j
            col(matstencil_k, num)=i
            num=num+1
          end if
          call matsetvaluesstencil(b, 1, row, num-1, col, v, insert_values, ierr)
        end do
      end do
    end do
    call matassemblybegin(b, mat_final_assembly, ierr)
    call matassemblyend(b, mat_final_assembly, ierr)
    if ( a .ne. b) then
      call matassemblybegin(a, mat_final_assembly, ierr)
      call matassemblyend(a, mat_final_assembly, ierr)
    endif
    deallocate(hzu, hzd, d_sc)
  end subroutine compute_matrix

  subroutine apply_dimension_bounds(dim_size, dim_processes, length_array)
    integer, intent(in) :: dim_size, dim_processes
    integer, intent(out) :: length_array(:)

    integer :: i, dimension_division, dimension_extra

    dimension_division = dim_size / dim_processes
    length_array= dimension_division    
    
    dimension_extra = dim_size - (dimension_division * dim_processes)
    if (dimension_extra .gt. 0) then
      do i=1, dimension_extra
        length_array(i)=length_array(i)+1
      end do
    end if
  end subroutine apply_dimension_bounds

  subroutine complete_psrce_calculation(current_state, y_halo_size, x_halo_size)
    type(model_state_type), target, intent(inout) :: current_state
    integer, intent(in) :: y_halo_size, x_halo_size

    integer :: ierr, combined_handles(2), i, j, k

    combined_handles(1)=current_state%psrce_x_hs_recv_request
    combined_handles(2)=current_state%psrce_y_hs_recv_request
    call mpi_waitall(2, combined_handles, mpi_statuses_ignore, ierr)

    do j=current_state%local_grid%local_domain_start_index(y_index), current_state%local_grid%local_domain_end_index(y_index)
      do k=2,current_state%local_grid%size(z_index)
#ifdef U_ACTIVE
        current_state%p%data(k,j,x_halo_size+1)=current_state%p%data(k,j,x_halo_size+1)-&
               current_state%psrce_recv_buffer_x(k-1,j-x_halo_size)
#endif
#ifdef V_ACTIVE
        if (j .gt. y_halo_size+1) current_state%p%data(k, j, x_halo_size+1)=current_state%p%data(k, j, x_halo_size+1)-&
             current_state%global_grid%configuration%horizontal%cy * current_state%sv%data(k, j-1, x_halo_size+1)
#endif
      end do
    end do
      
#ifdef V_ACTIVE
    do i=current_state%local_grid%local_domain_start_index(x_index), current_state%local_grid%local_domain_end_index(x_index)
      do k=2,current_state%local_grid%size(z_index)
        current_state%p%data(k,y_halo_size+1,i)=current_state%p%data(k,y_halo_size+1,i)-&
             current_state%psrce_recv_buffer_y(k-1,i-y_halo_size)
      end do
    end do
#endif

    combined_handles(1)=current_state%psrce_x_hs_send_request
    combined_handles(2)=current_state%psrce_y_hs_send_request
    call mpi_waitall(2, combined_handles, mpi_statuses_ignore, ierr)
  end subroutine complete_psrce_calculation 

  subroutine copy_p_to_halo_buffer(current_state, neighbour_description, dim, source_index, &
       pid_location, current_page, source_data)
    type(model_state_type), intent(inout) :: current_state
    integer, intent(in) :: dim, pid_location, source_index
    integer, intent(inout) :: current_page(:)
    type(neighbour_description_type), intent(inout) :: neighbour_description
    type(field_data_wrapper_type), dimension(:), intent(in), optional :: source_data

    call copy_field_to_buffer(current_state%local_grid, neighbour_description%send_halo_buffer, current_state%p%data, &
         dim, source_index, current_page(pid_location))

    current_page(pid_location)=current_page(pid_location)+1
  end subroutine copy_p_to_halo_buffer

  !! @param dim The dimension we receive for
  !! @param target_index The target index for the dimension we are receiving for
  !! @param neighbour_location The location in the local neighbour data stores of this neighbour
  !! @param current_page The current, next, halo swap page to read from (all previous have been read and copied already)
  !! @param source_data Optional source data which is written into
  subroutine copy_halo_buffer_to_p(current_state, neighbour_description, dim, target_index, &
       neighbour_location, current_page, source_data)
    type(model_state_type), intent(inout) :: current_state
    integer, intent(in) :: dim, target_index, neighbour_location
    integer, intent(inout) :: current_page(:)
    type(neighbour_description_type), intent(inout) :: neighbour_description
    type(field_data_wrapper_type), dimension(:), intent(in), optional :: source_data

    call copy_buffer_to_field(current_state%local_grid, neighbour_description%recv_halo_buffer, current_state%p%data, &
         dim, target_index, current_page(neighbour_location))

    current_page(neighbour_location)=current_page(neighbour_location)+1
  end subroutine copy_halo_buffer_to_p

  subroutine perform_local_data_copy_for_p(current_state, halo_depth, involve_corners, source_data)
    type(model_state_type), intent(inout) :: current_state
    integer, intent(in) :: halo_depth
    logical, intent(in) :: involve_corners
    type(field_data_wrapper_type), dimension(:), intent(in), optional :: source_data

    call perform_local_data_copy_for_field(current_state%p%data, current_state%local_grid, &
         current_state%parallel%my_rank, halo_depth, involve_corners)
  end subroutine perform_local_data_copy_for_p
end module petsc_solver_mod