multigrid_matrix_3d.f90

program multigrid_matrix_3d
  use m_a3_all

  implicit none

  integer, parameter :: box_size = 16
  integer, parameter :: n_boxes_base = 1
  integer, parameter :: n_iterations = 10
  integer, parameter :: n_var_cell = 5
  integer, parameter :: i_phi = 1
  integer, parameter :: i_rhs = 2
  integer, parameter :: i_tmp = 3
  integer, parameter :: i_mat_diag = 4
  integer, parameter :: i_ins = 5
  integer, parameter :: if_mat_nb = 1

  type(a3_t)        :: tree
  type(ref_info_t)   :: refine_info
  integer            :: mg_iter
  integer            :: ix_list(3, n_boxes_base)
  real(dp)           :: dr, residu(2), my_sum
  character(len=100) :: fname
  type(mg3_t)       :: mg
  integer :: lvl, i, id, p_id

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

  ! The cell spacing at the coarsest grid level
  dr = 4.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
       1, &            ! 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", "a00", "ins"]) ! 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)

  do
     call a3_loop_box(tree, set_initial_condition)
     call a3_adjust_refinement(tree, refine_routine, refine_info)
     if (refine_info%n_add == 0) exit
  end do

  ! Restrict matrix
  do lvl = tree%highest_lvl-1, lbound(tree%lvls, 1), -1
     do i = 1, size(tree%lvls(lvl)%parents)
        p_id = tree%lvls(lvl)%parents(i)
        call restrict_matrix(tree%boxes, p_id, i_mat_diag, if_mat_nb)
     end do
  end do

  ! 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     => bc_dirichlet_one

  mg%box_op => box_matrix
  mg%box_gsrb => box_gsrb

  call mg3_init_mg(mg)

  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, set_residual=.true.)

     call a3_tree_min_cc(tree, i_tmp, residu(1))
     call a3_tree_max_cc(tree, i_tmp, residu(2))
     call a3_tree_sum_cc(tree, i_phi, my_sum)
     write(*,"(I8,13x,2Es14.5)") mg_iter, maxval(abs(residu)), my_sum

     write(fname, "(A,I0)") "multigrid_matrix_3d_", mg_iter
     call a3_write_silo(tree, trim(fname), dir="output")
  end do

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)

    if (box%lvl < 3) 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
    integer                     :: i, j, k, nc
    real(dp)                    :: rr(3), inv_dr2, dr2

    nc = box%n_cell
    dr2 = box%dr**2
    inv_dr2 = 1/box%dr**2

    do k = 0, nc+1; do j = 0, nc+1; do i = 0, nc+1
       ! Get the coordinate of the cell center at i, j, k
       rr = a3_r_cc(box, [i, j, k])
       rr = rr - 2.0_dp

       box%cc(i, j, k, i_rhs) = 0.0_dp
       box%cc(i, j, k, i_phi) = 0.0_dp

       ! box%cc(i, j, k, i_ins) = (sum(rr**2)-1)**3 - rr(1)**2 * rr(2)**3 ! Heart
       ! box%cc(i, j, k, i_ins) = norm2(rr-0.5_dp)-0.25_dp ! Circle
       ! box%cc(i, j, k, i_ins) = sqrt(200 * rr(1)**2 + rr(2)**2) - 1.5_dp ! Sharp ellipse
       box%cc(i, j, k, i_ins) = sum(abs(rr)) - 0.5_dp ! Square
       ! box%cc(i, j, k, i_ins) = 1.0_dp ! No boundary

    end do; end do; end do

  end subroutine set_initial_condition

  subroutine box_matrix(box, i_out, mg)
    type(box3_t), intent(inout) :: box
    integer, intent(in)         :: i_out
    type(mg3_t), intent(in)     :: mg
    integer                     :: i, j, nc, i_phi
    real(dp)                    :: inv_dr_sq, coeffs(2*3+1), vals(2*3+1)
    integer                     :: k

    nc = box%n_cell
    inv_dr_sq = 1 / box%dr**2
    i_phi = mg%i_phi

    do k = 1, nc
       do j = 1, nc
          do i = 1, nc
             coeffs(1:2) = box%fc(i:i+1, j, k, 1, if_mat_nb)
             coeffs(3:4) = box%fc(i, j:j+1, k, 2, if_mat_nb)
             coeffs(5:6) = box%fc(i, j, k:k+1, 3, if_mat_nb)
             coeffs(7) = box%cc(i, j, k, i_mat_diag)

             vals(1) = box%cc(i-1, j, k, i_phi)
             vals(2) = box%cc(i+1, j, k, i_phi)
             vals(3) = box%cc(i, j-1, k, i_phi)
             vals(4) = box%cc(i, j+1, k, i_phi)
             vals(5) = box%cc(i, j, k-1, i_phi)
             vals(6) = box%cc(i, j, k+1, i_phi)
             vals(7) = box%cc(i, j, k, i_phi)

             box%cc(i, j, k, i_out) = sum(coeffs * vals)
          end do
       end do
    end do
  end subroutine box_matrix

  subroutine box_gsrb(box, redblack_cntr, mg)
    type(box3_t), intent(inout) :: box
    integer, intent(in)         :: redblack_cntr
    type(mg3_t), intent(in)     :: mg
    integer                     :: i, i0, j, nc, i_phi, i_rhs
    real(dp)                    :: inv_dr_sq, coeffs(2*3+1), vals(2*3)
    integer                     :: k

    nc = box%n_cell
    inv_dr_sq = 1 / box%dr**2
    i_phi = mg%i_phi
    i_rhs = mg%i_rhs

    do k = 1, nc
       do j = 1, nc
          i0 = 2 - iand(ieor(redblack_cntr, k+j), 1)
          do i = i0, nc, 2
             coeffs(1:2) = box%fc(i:i+1, j, k, 1, if_mat_nb)
             coeffs(3:4) = box%fc(i, j:j+1, k, 2, if_mat_nb)
             coeffs(5:6) = box%fc(i, j, k:k+1, 3, if_mat_nb)
             coeffs(7) = box%cc(i, j, k, i_mat_diag)

             vals(1) = box%cc(i-1, j, k, i_phi)
             vals(2) = box%cc(i+1, j, k, i_phi)
             vals(3) = box%cc(i, j-1, k, i_phi)
             vals(4) = box%cc(i, j+1, k, i_phi)
             vals(5) = box%cc(i, j, k-1, i_phi)
             vals(6) = box%cc(i, j, k+1, i_phi)

             box%cc(i, j, k, i_phi) = (box%cc(i, j, k, i_rhs) - &
                  sum(coeffs(1:6) * vals)) / coeffs(7)
          end do
       end do
    end do
  end subroutine box_gsrb

  subroutine restrict_matrix(boxes, p_id, iv_diag, ivf_nb)
    type(box3_t), intent(inout) :: boxes(:)
    integer, intent(in)          :: p_id
    integer, intent(in)          :: iv_diag
    integer, intent(in)          :: ivf_nb
    integer :: c_id, nc, n

    nc = boxes(p_id)%n_cell

    do n = 1, a3_num_children
       c_id = boxes(p_id)%children(n)
       if (c_id <= af_no_box) cycle
       call a3_restrict_box(boxes(c_id), boxes(p_id), iv_diag)
       call a3_restrict_box_face(boxes(c_id), boxes(p_id), ivf_nb)
    end do

    boxes(p_id)%cc(1:nc, 1:nc, 1:nc, iv_diag) = &
         boxes(p_id)%cc(1:nc, 1:nc, 1:nc, iv_diag) * 0.25_dp
    boxes(p_id)%fc(1:nc+1, 1:nc+1, 1:nc+1, 1:3, ivf_nb) = &
         boxes(p_id)%fc(1:nc+1, 1:nc+1, 1:nc+1, 1:3, ivf_nb) * 0.25_dp

  end subroutine restrict_matrix

    ! This fills ghost cells near physical boundaries using Neumann zero
  subroutine bc_dirichlet_one(box, nb, iv, bc_type)
    type(box3_t), intent(inout) :: box
    integer, intent(in)         :: nb, iv
    integer, intent(out)        :: bc_type

    bc_type = af_bc_dirichlet
    call a3_set_box_gc(box, nb, iv, 1.0_dp)
  end subroutine bc_dirichlet_one

end program multigrid_matrix_3d