poisson_neumann_3d.f90

program poisson_neumann_3d
  use m_a3_all
  use m_gaussians

  implicit none

  integer, parameter :: box_size = 8
  integer, parameter :: n_boxes_base = 1
  integer, parameter :: n_iterations = 10
  integer, parameter :: n_var_cell = 4
  integer, parameter :: i_phi = 1
  integer, parameter :: i_rhs = 2
  integer, parameter :: i_tmp = 3
  integer, parameter :: i_err = 4

  type(a3_t)        :: tree
  type(ref_info_t)   :: refine_info
  integer            :: mg_iter
  integer            :: ix_list(3, n_boxes_base)
  real(dp)           :: dr
  character(len=100) :: fname
  type(mg3_t)       :: mg
  integer            :: count_rate,t_start,t_end

  print *, "Running poisson_neumann_3d"
  print *, "Number of threads", af_get_max_threads()

  ! The cell spacing at the coarsest grid level
  dr = 1.0_dp / box_size

  ! Initialize tree
  call a3_init(tree, & ! Tree to initialize
       box_size, &     ! A box contains box_size**DIM cells
       n_var_cell, &   ! Number of cell-centered variables
       0, &            ! Number of face-centered variables
       dr, &           ! Distance between cells on base level
       coarsen_to=2, & ! Add coarsened levels for multigrid
       cc_names=["phi", "rhs", "tmp", "err"]) ! Variable names

  ! Set up geometry. These indices are used to define the coordinates of a box,
  ! by default the box at [1,1] touches the origin (x,y) = (0,0)
  ix_list(:, 1) = [1, 1, 1]         ! Set index of box 1

  ! Create the base mesh, using the box indices and their neighbor information
  call a3_set_base(tree, 1, ix_list)

  call system_clock(t_start, count_rate)
  do
     ! For each box, set the initial conditions
     call a3_loop_box(tree, set_initial_condition)

     ! This updates the refinement of the tree, by at most one level per call.
     ! The second argument is a subroutine that is called for each box that can
     ! be refined or derefined, and it should set refinement flags. Information
     ! about the changes in refinement are returned in the third argument.
     call a3_adjust_refinement(tree, refine_routine, refine_info, 0)

     ! If no new boxes have been added, exit the loop
     if (refine_info%n_add == 0) exit
  end do
  call system_clock(t_end, count_rate)

  write(*,"(A,Es10.3,A)") " Wall-clock time generating AMR grid: ", &
       (t_end-t_start) / real(count_rate,dp), " seconds"

  call a3_print_info(tree)

  ! Initialize the multigrid options. This performs some basics checks and sets
  ! default values where necessary.
  mg%i_phi        = i_phi       ! Solution variable
  mg%i_rhs        = i_rhs       ! Right-hand side variable
  mg%i_tmp        = i_tmp       ! Variable for temporary space
  mg%sides_bc     => sides_bc   ! Method for boundary conditions

  ! This routine does not initialize the multigrid fields boxes%i_phi,
  ! boxes%i_rhs and boxes%i_tmp. These fileds will be initialized at the
  ! first call of mg3_fas_fmg
  call mg3_init_mg(mg)

  print *, "Multigrid iteration | max residual | max error"
  call system_clock(t_start, count_rate)

  do mg_iter = 1, n_iterations
     call mg3_fas_fmg(tree, mg, set_residual=.true., have_guess=(mg_iter>1))
     ! call mg3_fas_vcycle(tree, mg, tree%highest_lvl, set_residual=.true.)

     call a3_loop_box(tree, set_error)

     write(fname, "(A,I0)") "poisson_neumann_3d_", mg_iter
     call a3_write_silo(tree, trim(fname), dir="output")
  end do
  call system_clock(t_end, count_rate)

  write(*, "(A,I0,A,E10.3,A)") &
       " Wall-clock time after ", n_iterations, &
       " iterations: ", (t_end-t_start) / real(count_rate, dp), &
       " seconds"

contains

  ! Return the refinement flags for box
  subroutine refine_routine(box, cell_flags)
    type(box3_t), intent(in) :: box
    integer, intent(out)     :: cell_flags(box%n_cell, box%n_cell, box%n_cell)

    ! Refine around one corner
    if (box%lvl <= 4 .and. all(box%r_min < 0.25_dp)) then
       cell_flags(:, :, :) = af_do_ref
    else
       cell_flags(:, :, :) = af_keep_ref
    end if
  end subroutine refine_routine

  ! This routine sets the initial conditions for each box
  subroutine set_initial_condition(box)
    type(box3_t), intent(inout) :: box

    box%cc(:, :, :, i_rhs) = 0.0_dp
  end subroutine set_initial_condition

  ! Set the error compared to the analytic solution
  subroutine set_error(box)
    type(box3_t), intent(inout) :: box
    integer                     :: i, j, k, nc
    real(dp)                    :: rr(3)

    nc = box%n_cell
    do k = 1, nc; do j = 1, nc; do i = 1, nc
       rr = a3_r_cc(box, [i, j, k])
       box%cc(i, j, k, i_err) = box%cc(i, j, k, i_phi) - rr(1)
    end do; end do; end do
  end subroutine set_error

  ! This routine sets boundary conditions for a box, by filling its ghost cells
  ! with approriate values.
  subroutine sides_bc(box, nb, iv, bc_type)
    type(box3_t), intent(inout) :: box
    integer, intent(in)          :: nb      ! Direction for the boundary condition
    integer, intent(in)          :: iv      ! Index of variable
    integer, intent(out)         :: bc_type ! Type of boundary condition
    integer                      :: nc

    nc = box%n_cell

    ! Below the solution is specified in the approriate ghost cells
    select case (nb)
    case (a3_neighb_lowx)             ! Lower-x direction
       call a3_bc_dirichlet_zero(box, nb, iv, bc_type)
    case (a3_neighb_highx)             ! Higher-x direction
       bc_type = af_bc_dirichlet
       box%cc(nc+1, 1:nc, 1:nc, iv) = 1
    case default
       call a3_bc_neumann_zero(box, nb, iv, bc_type)
    end select
  end subroutine sides_bc

end program