m_a3_multigrid.f90

module m_a3_multigrid
  use m_a3_types

  implicit none
  private

  ! The mg module supports different multigrid operators, and uses these tags to
  ! identify boxes / operators
  integer, parameter, public :: mg_normal_box = 1
  integer, parameter, public :: mg_lsf_box    = 2
  integer, parameter, public :: mg_ceps_box   = 3
  integer, parameter, public :: mg_veps_box   = 4

  ! Labels for the different steps of a multigrid cycle
  integer, parameter :: mg_cycle_down = 1
  integer, parameter :: mg_cycle_base = 2
  integer, parameter :: mg_cycle_up   = 3

  type, public :: mg3_t
     integer :: i_phi        = -1
     integer :: i_rhs        = -1
     integer :: i_tmp        = -1

     integer :: i_eps        = -1
     integer :: i_lsf        = -1
     integer :: i_bval       = -1

     integer :: n_cycle_down = -1
     integer :: n_cycle_up   = -1
     integer :: n_cycle_base = -1

     logical :: initialized = .false. 
     logical :: use_corners = .false. 
     logical :: subtract_mean = .false. 

     procedure(a3_subr_bc), pointer, nopass   :: sides_bc => null()

     procedure(a3_subr_rb), pointer, nopass   :: sides_rb => null()

     procedure(mg3_box_op), pointer, nopass   :: box_op => null()

     procedure(mg3_box_gsrb), pointer, nopass :: box_gsrb => null()

     procedure(mg3_box_corr), pointer, nopass :: box_corr => null()

     procedure(mg3_box_rstr), pointer, nopass :: box_rstr => null()
  end type mg3_t

  abstract interface

     subroutine mg3_box_op(box, i_out, mg)
       import
       type(box3_t), intent(inout) :: box
       type(mg3_t), intent(in)     :: mg
       integer, intent(in)         :: i_out
     end subroutine mg3_box_op

     subroutine mg3_box_gsrb(box, redblack_cntr, mg)
       import
       type(box3_t), intent(inout) :: box
       type(mg3_t), intent(in)     :: mg
       integer, intent(in)         :: redblack_cntr
     end subroutine mg3_box_gsrb

     subroutine mg3_box_corr(box_p, box_c, mg)
       import
       type(box3_t), intent(inout) :: box_c
       type(box3_t), intent(in)    :: box_p
       type(mg3_t), intent(in)     :: mg
     end subroutine mg3_box_corr

     subroutine mg3_box_rstr(box_c, box_p, iv, mg)
       import
       type(box3_t), intent(in)    :: box_c
       type(box3_t), intent(inout) :: box_p
       integer, intent(in)         :: iv
       type(mg3_t), intent(in)     :: mg
     end subroutine mg3_box_rstr
  end interface

  public :: mg3_init_mg
  public :: mg3_fas_fmg
  public :: mg3_fas_vcycle
  public :: mg3_box_op
  public :: mg3_box_gsrb
  public :: mg3_box_corr
  public :: mg3_box_rstr

  ! Automatic selection of operators
  public :: mg3_set_box_tag
  public :: mg3_auto_op
  public :: mg3_auto_gsrb
  public :: mg3_auto_corr
  public :: mg3_auto_rstr

  ! Methods for normal Laplacian
  public :: mg3_box_lpl
  public :: mg3_box_gsrb_lpl
  public :: mg3_box_corr_lpl
  public :: mg3_box_rstr_lpl
  public :: mg3_sides_rb

  ! Methods for Laplacian with jump in coefficient between boxes
  public :: mg3_box_lpld
  public :: mg3_box_gsrb_lpld
  public :: mg3_box_corr_lpld

  ! Methods for normal Laplacian with internal boundary conditions (LSF)
  public :: mg3_box_lpllsf
  public :: mg3_box_gsrb_lpllsf
  public :: mg3_box_corr_lpllsf
  public :: mg3_box_rstr_lpllsf

contains

  subroutine mg3_init_mg(mg)
    type(mg3_t), intent(inout) :: mg

    if (mg%i_phi < 0)                  stop "mg3_init_mg: i_phi not set"
    if (mg%i_tmp < 0)                  stop "mg3_init_mg: i_tmp not set"
    if (mg%i_rhs < 0)                  stop "mg3_init_mg: i_rhs not set"
    if (mg%i_lsf * mg%i_bval < 0) &
         stop "mg3_init_mg: you have to set both i_lsf and i_bval"

    if (.not. associated(mg%sides_bc)) stop "mg3_init_mg: sides_bc not set"

    ! Check whether these are set, otherwise use default
    if (mg%n_cycle_down < 0)           mg%n_cycle_down = 2
    if (mg%n_cycle_up < 0)             mg%n_cycle_up = 2
    if (mg%n_cycle_base < 0)           mg%n_cycle_base = 4

    ! Check whether methods are set, otherwise use default (for laplacian)
    if (.not. associated(mg%sides_rb)) mg%sides_rb => mg3_sides_rb
    if (.not. associated(mg%box_op))   mg%box_op => mg3_box_lpl
    if (.not. associated(mg%box_gsrb)) mg%box_gsrb => mg3_box_gsrb_lpl
    if (.not. associated(mg%box_corr)) mg%box_corr => mg3_box_corr_lpl
    if (.not. associated(mg%box_rstr)) mg%box_rstr => mg3_box_rstr_lpl

    mg%initialized = .true.
  end subroutine mg3_init_mg

  subroutine check_mg(mg)
    type(mg3_t), intent(in) :: mg
    if (.not. mg%initialized) stop "check_mg: you haven't called mg3_init_mg"
  end subroutine check_mg

  subroutine mg3_fas_fmg(tree, mg, set_residual, have_guess)
    use m_a3_ghostcell, only: a3_gc_ids
    use m_a3_utils, only: a3_boxes_copy_cc
    type(a3_t), intent(inout)       :: tree
    type(mg3_t), intent(in)         :: mg
    logical, intent(in)             :: set_residual
    logical, intent(in)             :: have_guess
    integer                         :: lvl, min_lvl

    call check_mg(mg)           ! Check whether mg options are set

    min_lvl = lbound(tree%lvls, 1)

    if (have_guess) then
       do lvl = tree%highest_lvl,  min_lvl+1, -1
          ! Set rhs on coarse grid and restrict phi
          call set_coarse_phi_rhs(tree, lvl, mg)
       end do
    else
       call init_phi_rhs(tree, mg)
    end if

    do lvl = min_lvl, tree%highest_lvl
       ! Store phi_old in tmp
       call a3_boxes_copy_cc(tree%boxes, tree%lvls(lvl)%ids, &
            mg%i_phi, mg%i_tmp)

       if (lvl > min_lvl) then
          ! Correct solution at this lvl using lvl-1 data
          ! phi = phi + prolong(phi_coarse - phi_old_coarse)
          call correct_children(tree%boxes, tree%lvls(lvl-1)%parents, mg)

          ! Update ghost cells
          call a3_gc_ids(tree, tree%lvls(lvl)%ids, mg%i_phi, &
               mg%sides_rb, mg%sides_bc)
       end if

       ! Perform V-cycle, only set residual on last iteration
       call mg3_fas_vcycle(tree, mg, &
            set_residual .and. lvl == tree%highest_lvl, lvl)
    end do
  end subroutine mg3_fas_fmg

  subroutine mg3_fas_vcycle(tree, mg, set_residual, highest_lvl)
    use m_a3_ghostcell, only: a3_gc_ids
    use m_a3_utils, only: a3_tree_sum_cc
    type(a3_t), intent(inout)    :: tree
    type(mg3_t), intent(in)      :: mg
    logical, intent(in)           :: set_residual
    integer, intent(in), optional :: highest_lvl
    integer                       :: lvl, min_lvl, i, id, max_lvl
    real(dp)                      :: sum_phi, mean_phi

    call check_mg(mg)           ! Check whether mg options are set
    min_lvl = lbound(tree%lvls, 1)
    max_lvl = tree%highest_lvl
    if (present(highest_lvl)) max_lvl = highest_lvl

    do lvl = max_lvl,  min_lvl+1, -1
       ! Downwards relaxation
       call gsrb_boxes(tree, tree%lvls(lvl)%ids, mg, mg_cycle_down)

       ! Set rhs on coarse grid, restrict phi, and copy i_phi to i_tmp for the
       ! correction later
       call update_coarse(tree, lvl, mg)
    end do

    lvl = min_lvl
    call gsrb_boxes(tree, tree%lvls(lvl)%ids, mg, mg_cycle_base)

    ! Do the upwards part of the v-cycle in the tree
    do lvl = min_lvl+1, max_lvl
       ! Correct solution at this lvl using lvl-1 data
       ! phi = phi + prolong(phi_coarse - phi_old_coarse)
       call correct_children(tree%boxes, tree%lvls(lvl-1)%parents, mg)

       ! Have to fill ghost cells after correction
       call a3_gc_ids(tree, tree%lvls(lvl)%ids, mg%i_phi, &
            mg%sides_rb, mg%sides_bc)

       ! Upwards relaxation
       call gsrb_boxes(tree, tree%lvls(lvl)%ids, mg, mg_cycle_up)
    end do

    if (set_residual) then
       !$omp parallel private(lvl, i, id)
       do lvl = min_lvl, max_lvl
          !$omp do
          do i = 1, size(tree%lvls(lvl)%ids)
             id = tree%lvls(lvl)%ids(i)
             call residual_box(tree%boxes(id), mg)
          end do
          !$omp end do nowait
       end do
       !$omp end parallel
    end if

    ! Subtract the mean of phi, for example when doing a fully periodic
    ! simulation. Note that the integrated r.h.s. term should also be zero then.
    if (mg%subtract_mean) then
       call a3_tree_sum_cc(tree, mg%i_phi, sum_phi)
       mean_phi = sum_phi / a3_total_volume(tree)
       !$omp parallel private(lvl, i, id)
       do lvl = min_lvl, max_lvl
          !$omp do
          do i = 1, size(tree%lvls(lvl)%ids)
             id = tree%lvls(lvl)%ids(i)
             tree%boxes(id)%cc(:, :, :, mg%i_phi) = &
                  tree%boxes(id)%cc(:, :, :, mg%i_phi) - mean_phi
          end do
          !$omp end do nowait
       end do
       !$omp end parallel
    end if

  end subroutine mg3_fas_vcycle

  subroutine mg3_sides_rb(boxes, id, nb, iv)
    use m_a3_ghostcell, only: a3_gc_prolong_copy
    type(box3_t), intent(inout) :: boxes(:)
    integer, intent(in)          :: id
    integer, intent(in)          :: nb
    integer, intent(in)          :: iv
    integer                      :: nc, ix, dix, i, di, j, dj, co(3)
    integer                      :: hnc, p_id, p_nb_id
    real(dp)                     :: grad(3-1)
    integer                      :: k, dk
    real(dp)                     :: tmp(0:boxes(id)%n_cell/2+1, 0:boxes(id)%n_cell/2+1)
    real(dp)                     :: gc(boxes(id)%n_cell, boxes(id)%n_cell)

    nc = boxes(id)%n_cell
    hnc = nc/2

    p_id = boxes(id)%parent
    p_nb_id = boxes(p_id)%neighbors(nb)
    co = a3_get_child_offset(boxes(id))

    associate(box => boxes(p_nb_id))
      ! First fill a temporary array with data next to the fine grid
      select case (nb)
      case (a3_neighb_lowx)
         tmp = box%cc(nc, co(2):co(2)+hnc+1, co(3):co(3)+hnc+1, iv)
      case (a3_neighb_highx)
         tmp = box%cc(1, co(2):co(2)+hnc+1, co(3):co(3)+hnc+1, iv)
      case (a3_neighb_lowy)
         tmp = box%cc(co(1):co(1)+hnc+1, nc, co(3):co(3)+hnc+1, iv)
      case (a3_neighb_highy)
         tmp = box%cc(co(1):co(1)+hnc+1, 1, co(3):co(3)+hnc+1, iv)
      case (a3_neighb_lowz)
         tmp = box%cc(co(1):co(1)+hnc+1, co(2):co(2)+hnc+1, nc, iv)
      case (a3_neighb_highz)
         tmp = box%cc(co(1):co(1)+hnc+1, co(2):co(2)+hnc+1, 1, iv)
      end select
    end associate

    ! Now interpolate the coarse grid data to obtain values 'straight' next to
    ! the fine grid points
    do j = 1, hnc
       do i = 1, hnc
          grad(1)          = 0.125_dp * (tmp(i+1, j) - tmp(i-1, j))
          grad(2)          = 0.125_dp * (tmp(i, j+1) - tmp(i, j-1))
          gc(2*i-1, 2*j-1) = tmp(i, j) - grad(1) - grad(2)
          gc(2*i, 2*j-1)   = tmp(i, j) + grad(1) - grad(2)
          gc(2*i-1, 2*j)   = tmp(i, j) - grad(1) + grad(2)
          gc(2*i, 2*j)     = tmp(i, j) + grad(1) + grad(2)
       end do
    end do

    if (a3_neighb_low(nb)) then
       ix = 1
       dix = 1
    else
       ix = nc
       dix = -1
    end if

    select case (a3_neighb_dim(nb))
    case (1)
       i = ix
       di = dix
       do k = 1, nc
          dk = -1 + 2 * iand(k, 1)
          do j = 1, nc
             dj = -1 + 2 * iand(j, 1)
             boxes(id)%cc(i-di, j, k, iv) = &
                  0.5_dp * gc(j, k) + &
                  0.75_dp * boxes(id)%cc(i, j, k, iv) - &
                  0.25_dp * boxes(id)%cc(i+di, j, k, iv)
          end do
       end do
    case (2)
       j = ix
       dj = dix
       do k = 1, nc
          dk = -1 + 2 * iand(k, 1)
          do i = 1, nc
             di = -1 + 2 * iand(i, 1)
             boxes(id)%cc(i, j-dj, k, iv) = &
                  0.5_dp * gc(i, k) + &
                  0.75_dp * boxes(id)%cc(i, j, k, iv) - &
                  0.25_dp * boxes(id)%cc(i, j+dj, k, iv)
          end do
       end do
    case (3)
       k = ix
       dk = dix
       do j = 1, nc
          dj = -1 + 2 * iand(j, 1)
          do i = 1, nc
             di = -1 + 2 * iand(i, 1)
             boxes(id)%cc(i, j, k-dk, iv) = &
                  0.5_dp * gc(i, j) + &
                  0.75_dp * boxes(id)%cc(i, j, k, iv) - &
                  0.25_dp * boxes(id)%cc(i, j, k+dk, iv)
          end do
       end do
    end select

  end subroutine mg3_sides_rb

  subroutine mg3_sides_rb_old(boxes, id, nb, iv)
    use m_a3_ghostcell, only: a3_gc_prolong_copy
    type(box3_t), intent(inout) :: boxes(:)
    integer, intent(in)         :: id
    integer, intent(in)         :: nb
    integer, intent(in)         :: iv
    integer                     :: nc, ix, dix, i, di, j, dj
    integer                     :: k, dk

    nc = boxes(id)%n_cell

    if (a3_neighb_low(nb)) then
       ix = 1
       dix = 1
    else
       ix = nc
       dix = -1
    end if

    call a3_gc_prolong_copy(boxes, id, nb, iv)

    select case (a3_neighb_dim(nb))
    case (1)
       i = ix
       di = dix
       do k = 1, nc
          dk = -1 + 2 * iand(k, 1)
          do j = 1, nc
             dj = -1 + 2 * iand(j, 1)
             ! Trilinear extrapolation (using 8 points)
             ! boxes(id)%cc(i-di, j, k, iv) = &
             !      0.5_dp * boxes(id)%cc(i-di, j, k, iv) + 0.0625_dp * (&
             !      27 * boxes(id)%cc(i, j, k, iv) &
             !      - 9 * boxes(id)%cc(i+di, j, k, iv) &
             !      - 9 * boxes(id)%cc(i, j+dj, k, iv) &
             !      - 9 * boxes(id)%cc(i, j, k+dk, iv) &
             !      + 3 * boxes(id)%cc(i+di, j+dj, k, iv) &
             !      + 3 * boxes(id)%cc(i+di, j, k+dk, iv) &
             !      + 3 * boxes(id)%cc(i, j+dj, k+dk, iv) &
             !      - 1 * boxes(id)%cc(i+di, j+dj, k+dk, iv))

             ! Extrapolation using 2 points
             boxes(id)%cc(i-di, j, k, iv) = &
                  0.5_dp * boxes(id)%cc(i-di, j, k, iv) + &
                  0.75_dp * boxes(id)%cc(i, j, k, iv) - &
                  0.25_dp * boxes(id)%cc(i+di, j+dj, k+dk, iv)
          end do
       end do
    case (2)
       j = ix
       dj = dix
       do k = 1, nc
          dk = -1 + 2 * iand(k, 1)
          do i = 1, nc
             di = -1 + 2 * iand(i, 1)

             ! boxes(id)%cc(i, j-dj, k, iv) = &
             !      0.5_dp * boxes(id)%cc(i, j-dj, k, iv) + 0.0625_dp * (&
             !      27 * boxes(id)%cc(i, j, k, iv) &
             !      - 9 * boxes(id)%cc(i+di, j, k, iv) &
             !      - 9 * boxes(id)%cc(i, j+dj, k, iv) &
             !      - 9 * boxes(id)%cc(i, j, k+dk, iv) &
             !      + 3 * boxes(id)%cc(i+di, j+dj, k, iv) &
             !      + 3 * boxes(id)%cc(i+di, j, k+dk, iv) &
             !      + 3 * boxes(id)%cc(i, j+dj, k+dk, iv) &
             !      - 1 * boxes(id)%cc(i+di, j+dj, k+dk, iv))

             boxes(id)%cc(i, j-dj, k, iv) = &
                  0.5_dp * boxes(id)%cc(i, j-dj, k, iv) + &
                  0.75_dp * boxes(id)%cc(i, j, k, iv) - &
                  0.25_dp * boxes(id)%cc(i+di, j+dj, k+dk, iv)
          end do
       end do
    case (3)
       k = ix
       dk = dix
       do j = 1, nc
          dj = -1 + 2 * iand(j, 1)
          do i = 1, nc
             di = -1 + 2 * iand(i, 1)

             ! boxes(id)%cc(i, j, k-dk, iv) = &
             !      0.5_dp * boxes(id)%cc(i, j, k-dk, iv) + 0.0625_dp * (&
             !      27 * boxes(id)%cc(i, j, k, iv) &
             !      - 9 * boxes(id)%cc(i+di, j, k, iv) &
             !      - 9 * boxes(id)%cc(i, j+dj, k, iv) &
             !      - 9 * boxes(id)%cc(i, j, k+dk, iv) &
             !      + 3 * boxes(id)%cc(i+di, j+dj, k, iv) &
             !      + 3 * boxes(id)%cc(i+di, j, k+dk, iv) &
             !      + 3 * boxes(id)%cc(i, j+dj, k+dk, iv) &
             !      - 1 * boxes(id)%cc(i+di, j+dj, k+dk, iv))

             boxes(id)%cc(i, j, k-dk, iv) = &
                  0.5_dp * boxes(id)%cc(i, j, k-dk, iv) + &
                  0.75_dp * boxes(id)%cc(i, j, k, iv) - &
                  0.25_dp * boxes(id)%cc(i+di, j+dj, k+dk, iv)
          end do
       end do
    end select

  end subroutine mg3_sides_rb_old

  ! Sets phi = phi + prolong(phi_coarse - phi_old_coarse)
  subroutine correct_children(boxes, ids, mg)
    type(box3_t), intent(inout) :: boxes(:)
    integer, intent(in)         :: ids(:)
    type(mg3_t), intent(in)     :: mg
    integer                     :: i, id, nc, i_c, c_id

    !$omp parallel do private(id, nc, i_c, c_id)
    do i = 1, size(ids)
       id = ids(i)
       nc = boxes(id)%n_cell

       ! Store the correction in i_tmp
       boxes(id)%cc(:, :, :, mg%i_tmp) = boxes(id)%cc(:, :, :, mg%i_phi) - &
            boxes(id)%cc(:, :, :, mg%i_tmp)

       do i_c = 1, a3_num_children
          c_id = boxes(id)%children(i_c)
          if (c_id == af_no_box) cycle
          call mg%box_corr(boxes(id), boxes(c_id), mg)
       end do
    end do
    !$omp end parallel do
  end subroutine correct_children

  subroutine gsrb_boxes(tree, ids, mg, type_cycle)
    use m_a3_ghostcell, only: a3_gc_box
    type(a3_t), intent(inout) :: tree
    type(mg3_t), intent(in)   :: mg
    integer, intent(in)        :: ids(:)
    integer, intent(in)        :: type_cycle
    integer                    :: n, i, n_cycle
    logical                    :: use_corners

    select case (type_cycle)
    case (mg_cycle_down)
       n_cycle = mg%n_cycle_down
    case (mg_cycle_up)
       n_cycle = mg%n_cycle_up
    case (mg_cycle_base)
       n_cycle = mg%n_cycle_base
    case default
       error stop "gsrb_boxes: invalid cycle type"
    end select

    !$omp parallel private(n, i)
    do n = 1, 2 * n_cycle
       !$omp do
       do i = 1, size(ids)
          call mg%box_gsrb(tree%boxes(ids(i)), n, mg)
       end do
       !$omp end do

       use_corners = mg%use_corners .or. &
            (type_cycle /= mg_cycle_down .and. n == 2 * n_cycle)

       !$omp do
       do i = 1, size(ids)
          call a3_gc_box(tree, ids(i), mg%i_phi, mg%sides_rb, &
               mg%sides_bc, use_corners)
       end do
       !$omp end do
    end do
    !$omp end parallel
  end subroutine gsrb_boxes

  ! Set rhs on coarse grid, restrict phi, and copy i_phi to i_tmp for the
  ! correction later
  subroutine update_coarse(tree, lvl, mg)
    use m_a3_utils, only: a3_n_cell, a3_box_add_cc, a3_box_copy_cc
    use m_a3_ghostcell, only: a3_gc_ids
    type(a3_t), intent(inout) :: tree
    integer, intent(in)        :: lvl
    type(mg3_t), intent(in)   :: mg
    integer                    :: i, id, p_id, nc
    real(dp), allocatable :: tmp(:,:,:)

    id = tree%lvls(lvl)%ids(1)
    nc = a3_n_cell(tree, lvl)
    allocate(tmp(1:nc, 1:nc, 1:nc))

    ! Restrict phi and the residual
    !$omp parallel do private(id, p_id, tmp)
    do i = 1, size(tree%lvls(lvl)%ids)
       id = tree%lvls(lvl)%ids(i)
       p_id = tree%boxes(id)%parent

       ! Copy the data currently in i_tmp, and restore it later (i_tmp holds the
       ! previous state of i_phi)
       tmp = tree%boxes(id)%cc(1:nc, 1:nc, 1:nc, mg%i_tmp)
       call residual_box(tree%boxes(id), mg)
       call mg%box_rstr(tree%boxes(id), tree%boxes(p_id), mg%i_tmp, mg)
       call mg%box_rstr(tree%boxes(id), tree%boxes(p_id), mg%i_phi, mg)
       tree%boxes(id)%cc(1:nc, 1:nc, 1:nc, mg%i_tmp) = tmp
    end do
    !$omp end parallel do

    call a3_gc_ids(tree, tree%lvls(lvl-1)%ids, mg%i_phi, &
         mg%sides_rb, mg%sides_bc)

    ! Set rhs_c = laplacian(phi_c) + restrict(res) where it is refined, and
    ! store current coarse phi in tmp.

    !$omp parallel do private(id)
    do i = 1, size(tree%lvls(lvl-1)%parents)
       id = tree%lvls(lvl-1)%parents(i)

       ! Set rhs = L phi
       call mg%box_op(tree%boxes(id), mg%i_rhs, mg)

       ! Add tmp (the fine grid residual) to rhs
       call a3_box_add_cc(tree%boxes(id), mg%i_tmp, mg%i_rhs)

       ! Story a copy of phi in tmp
       call a3_box_copy_cc(tree%boxes(id), mg%i_phi, mg%i_tmp)
    end do
    !$omp end parallel do
  end subroutine update_coarse

  subroutine set_coarse_phi_rhs(tree, lvl, mg)
    use m_a3_ghostcell, only: a3_gc_ids
    use m_a3_utils, only: a3_box_add_cc
    type(a3_t), intent(inout) :: tree
    integer, intent(in)        :: lvl
    type(mg3_t), intent(in)   :: mg
    integer                    :: i, id, p_id

    ! Fill ghost cells here to be sure
    if (lvl == tree%highest_lvl) then
       call a3_gc_ids(tree, tree%lvls(lvl)%ids, mg%i_phi, &
            mg%sides_rb, mg%sides_bc)
    end if

    !$omp parallel do private(id, p_id)
    do i = 1, size(tree%lvls(lvl)%ids)
       id = tree%lvls(lvl)%ids(i)
       p_id = tree%boxes(id)%parent

       call residual_box(tree%boxes(id), mg)
       call mg%box_rstr(tree%boxes(id), tree%boxes(p_id), mg%i_tmp, mg)
       call mg%box_rstr(tree%boxes(id), tree%boxes(p_id), mg%i_phi, mg)
    end do
    !$omp end parallel do

    call a3_gc_ids(tree, tree%lvls(lvl-1)%ids, mg%i_phi, &
         mg%sides_rb, mg%sides_bc)

    ! Set rhs_c = laplacian(phi_c) + restrict(res) where it is refined

    !$omp parallel do private(id)
    do i = 1, size(tree%lvls(lvl-1)%parents)
       id = tree%lvls(lvl-1)%parents(i)
       call mg%box_op(tree%boxes(id), mg%i_rhs, mg)
       call a3_box_add_cc(tree%boxes(id), mg%i_tmp, mg%i_rhs)
    end do
    !$omp end parallel do
  end subroutine set_coarse_phi_rhs

  subroutine init_phi_rhs(tree, mg)
    use m_a3_utils, only: a3_box_clear_cc
    type(a3_t), intent(inout) :: tree
    type(mg3_t), intent(in)   :: mg
    integer                    :: i, id, p_id, lvl, min_lvl

    ! Start from phi = 0 and restrict rhs
    min_lvl = lbound(tree%lvls, 1)

    !$omp parallel private(lvl, i, id, p_id)
    do lvl = tree%highest_lvl,  min_lvl+1, -1
       !$omp do
       do i = 1, size(tree%lvls(lvl)%ids)
          id = tree%lvls(lvl)%ids(i)
          call a3_box_clear_cc(tree%boxes(id), mg%i_phi)
          if (lvl > min_lvl) then
             p_id = tree%boxes(id)%parent
             call mg%box_rstr(tree%boxes(id), tree%boxes(p_id), mg%i_rhs, mg)
          end if
       end do
       !$omp end do
    end do
    !$omp end parallel
  end subroutine init_phi_rhs

  subroutine residual_box(box, mg)
    type(box3_t), intent(inout) :: box
    type(mg3_t), intent(in)     :: mg
    integer                     :: nc

    call mg%box_op(box, mg%i_tmp, mg)
    nc = box%n_cell
    box%cc(1:nc, 1:nc, 1:nc, mg%i_tmp) = box%cc(1:nc, 1:nc, 1:nc, mg%i_rhs) &
         - box%cc(1:nc, 1:nc, 1:nc, mg%i_tmp)
  end subroutine residual_box

  subroutine mg3_auto_gsrb(box, redblack_cntr, mg)
    type(box3_t), intent(inout) :: box
    integer, intent(in)         :: redblack_cntr
    type(mg3_t), intent(in)     :: mg

    if (box%tag == af_init_tag) call mg3_set_box_tag(box, mg)

    select case(box%tag)
    case (mg_normal_box)
       call mg3_box_gsrb_lpl(box, redblack_cntr, mg)
    case (mg_lsf_box)
       call mg3_box_gsrb_lpllsf(box, redblack_cntr, mg)
    case (mg_veps_box, mg_ceps_box)
       call mg3_box_gsrb_lpld(box, redblack_cntr, mg)
    end select
  end subroutine mg3_auto_gsrb

  subroutine mg3_auto_op(box, i_out, mg)
    type(box3_t), intent(inout) :: box
    integer, intent(in)         :: i_out
    type(mg3_t), intent(in)     :: mg

    if (box%tag == af_init_tag) call mg3_set_box_tag(box, mg)

    select case(box%tag)
    case (mg_normal_box)
       call mg3_box_lpl(box, i_out, mg)
    case (mg_lsf_box)
       call mg3_box_lpllsf(box, i_out, mg)
    case (mg_veps_box, mg_ceps_box)
       call mg3_box_lpld(box, i_out, mg)
    end select
  end subroutine mg3_auto_op

  subroutine mg3_auto_rstr(box_c, box_p, iv, mg)
    type(box3_t), intent(in)    :: box_c
    type(box3_t), intent(inout) :: box_p
    integer, intent(in)         :: iv
    type(mg3_t), intent(in)     :: mg

    ! We can only restrict after gsrb, so tag should always be set
    if (box_c%tag == af_init_tag) stop "mg3_auto_rstr: box_c tag not set"

    select case(box_c%tag)
    case (mg_normal_box, mg_veps_box, mg_ceps_box)
       call mg3_box_rstr_lpl(box_c, box_p, iv, mg)
    case (mg_lsf_box)
       call mg3_box_rstr_lpllsf(box_c, box_p, iv, mg)
    end select
  end subroutine mg3_auto_rstr

  subroutine mg3_auto_corr(box_p, box_c, mg)
    type(box3_t), intent(inout) :: box_c
    type(box3_t), intent(in)    :: box_p
    type(mg3_t), intent(in)     :: mg

    if (box_c%tag == af_init_tag) call mg3_set_box_tag(box_c, mg)

    select case(box_c%tag)
    case (mg_normal_box)
       call mg3_box_corr_lpl(box_p, box_c, mg)
    case (mg_lsf_box)
       call mg3_box_corr_lpllsf(box_p, box_c, mg)
    case (mg_veps_box, mg_ceps_box)
       call mg3_box_corr_lpld(box_p, box_c, mg)
    end select
  end subroutine mg3_auto_corr

  subroutine mg3_set_box_tag(box, mg)
    type(box3_t), intent(inout) :: box
    type(mg3_t), intent(in)     :: mg
    real(dp)                     :: a, b
    logical                      :: is_lsf, is_deps, is_eps

    is_lsf = .false.
    is_eps = .false.
    is_deps = .false.

    if (mg%i_lsf /= -1) then
       is_lsf = minval(box%cc(:, :, :, mg%i_lsf)) * &
            maxval(box%cc(:, :, :, mg%i_lsf)) < 0
    end if

    if (mg%i_eps /= -1) then
       a = minval(box%cc(:, :, :, mg%i_eps))
       b = maxval(box%cc(:, :, :, mg%i_eps))
       is_deps = (b > a)
       if (.not. is_deps) is_eps = (a < 1 .or. a > 1)
    end if

    if (count([is_lsf, is_eps, is_deps]) > 1) &
         stop "mg3_set_box_tag: Cannot set lsf and eps tag for same box"

    box%tag = mg_normal_box
    if (is_lsf) box%tag = mg_lsf_box
    if (is_eps) box%tag = mg_ceps_box
    if (is_deps) box%tag = mg_veps_box
  end subroutine mg3_set_box_tag

  subroutine mg3_box_corr_lpl(box_p, box_c, mg)
    use m_a3_prolong
    type(box3_t), intent(inout) :: box_c
    type(box3_t), intent(in)    :: box_p
    type(mg3_t), intent(in)     :: mg

    call a3_prolong_linear(box_p, box_c, mg%i_tmp, mg%i_phi, add=.true.)
  end subroutine mg3_box_corr_lpl

  subroutine mg3_box_gsrb_lpl(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)                    :: dx2
    integer                     :: k
    real(dp), parameter         :: sixth = 1/6.0_dp

    dx2   = box%dr**2
    nc    = box%n_cell
    i_phi = mg%i_phi
    i_rhs = mg%i_rhs

    ! The parity of redblack_cntr determines which cells we use. If
    ! redblack_cntr is even, we use the even cells and vice versa.
    do k = 1, nc
       do j = 1, nc
          i0 = 2 - iand(ieor(redblack_cntr, k+j), 1)
          do i = i0, nc, 2
             box%cc(i, j, k, i_phi) = sixth * ( &
                  box%cc(i+1, j, k, i_phi) + box%cc(i-1, j, k, i_phi) + &
                  box%cc(i, j+1, k, i_phi) + box%cc(i, j-1, k, i_phi) + &
                  box%cc(i, j, k+1, i_phi) + box%cc(i, j, k-1, i_phi) - &
                  dx2 * box%cc(i, j, k, i_rhs))
          end do
       end do
    end do
  end subroutine mg3_box_gsrb_lpl

  subroutine mg3_box_lpl(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
    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
             box%cc(i, j, k, i_out) = inv_dr_sq * (box%cc(i-1, j, k, i_phi) + &
                  box%cc(i+1, j, k, i_phi) + box%cc(i, j-1, k, i_phi) + &
                  box%cc(i, j+1, k, i_phi) + box%cc(i, j, k-1, i_phi) + &
                  box%cc(i, j, k+1, i_phi) - 6 * box%cc(i, j, k, i_phi))
          end do
       end do
    end do
  end subroutine mg3_box_lpl

  subroutine mg3_box_rstr_lpl(box_c, box_p, iv, mg)
    use m_a3_restrict, only: a3_restrict_box
    type(box3_t), intent(in)      :: box_c
    type(box3_t), intent(inout)   :: box_p
    integer, intent(in)           :: iv
    type(mg3_t), intent(in)       :: mg

    ! Don't use geometry for restriction, since this is inconsistent with the
    ! filling of ghost cells near refinement boundaries
    call a3_restrict_box(box_c, box_p, iv, use_geometry=.false.)
  end subroutine mg3_box_rstr_lpl

  subroutine mg3_box_gsrb_lpld(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_eps, i_rhs
    real(dp)                    :: dx2, u(2*3), a0, a(2*3), c(2*3)
    integer                     :: k

    dx2   = box%dr**2
    nc    = box%n_cell
    i_phi = mg%i_phi
    i_eps = mg%i_eps
    i_rhs = mg%i_rhs

    ! The parity of redblack_cntr determines which cells we use. If
    ! redblack_cntr is even, we use the even cells and vice versa.
    do k = 1, nc
       do j = 1, nc
          i0 = 2 - iand(ieor(redblack_cntr, k+j), 1)
          do i = i0, nc, 2
             a0 = box%cc(i, j, k, i_eps) ! value of eps at i,j
             u(1:2) = box%cc(i-1:i+1:2, j, k, i_phi) ! values at neighbors
             a(1:2) = box%cc(i-1:i+1:2, j, k, i_eps)
             u(3:4) = box%cc(i, j-1:j+1:2, k, i_phi)
             a(3:4) = box%cc(i, j-1:j+1:2, k, i_eps)
             u(5:6) = box%cc(i, j, k-1:k+1:2, i_phi)
             a(5:6) = box%cc(i, j, k-1:k+1:2, i_eps)
             c(:) = 2 * a0 * a(:) / (a0 + a(:))

             box%cc(i, j, k, i_phi) = &
                  (sum(c(:) * u(:)) - dx2 * box%cc(i, j, k, i_rhs)) / sum(c(:))
          end do
       end do
    end do
  end subroutine mg3_box_gsrb_lpld

  subroutine mg3_box_lpld(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, i_eps
    real(dp)                    :: inv_dr_sq, a0, u0, u(2*3), a(2*3)
    integer                     :: k

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

    do k = 1, nc
       do j = 1, nc
          do i = 1, nc
             u0 = box%cc(i, j, k, i_phi)
             a0 = box%cc(i, j, k, i_eps)
             u(1:2) = box%cc(i-1:i+1:2, j, k, i_phi)
             u(3:4) = box%cc(i, j-1:j+1:2, k, i_phi)
             u(5:6) = box%cc(i, j, k-1:k+1:2, i_phi)
             a(1:2) = box%cc(i-1:i+1:2, j, k, i_eps)
             a(3:4) = box%cc(i, j-1:j+1:2, k, i_eps)
             a(5:6) = box%cc(i, j, k-1:k+1:2, i_eps)

             box%cc(i, j, k, i_out) = inv_dr_sq * 2 * &
                  sum(a0*a(:)/(a0 + a(:)) * (u(:) - u0))
          end do
       end do
    end do

  end subroutine mg3_box_lpld

  subroutine mg3_box_corr_lpld(box_p, box_c, mg)
    type(box3_t), intent(inout)  :: box_c
    type(box3_t), intent(in)     :: box_p
    type(mg3_t), intent(in)      :: mg
    integer                      :: ix_offset(3), i_phi, i_corr, i_eps
    integer                      :: nc, i, j, i_c1, i_c2, j_c1, j_c2
    real(dp)                     :: u0, u(3), a0, a(3)
    integer                      :: k, k_c1, k_c2
    real(dp), parameter          :: third = 1/3.0_dp

    nc = box_c%n_cell
    ix_offset = a3_get_child_offset(box_c)
    i_phi = mg%i_phi
    i_corr = mg%i_tmp
    i_eps = mg%i_eps

    ! In these loops, we calculate the closest coarse index (_c1), and the
    ! one-but-closest (_c2). The fine cell lies in between.
    do k = 1, nc
       k_c1 = ix_offset(3) + ishft(k+1, -1) ! (k+1)/2
       k_c2 = k_c1 + 1 - 2 * iand(k, 1)     ! even: +1, odd: -1
       do j = 1, nc
          j_c1 = ix_offset(2) + ishft(j+1, -1) ! (j+1)/2
          j_c2 = j_c1 + 1 - 2 * iand(j, 1)     ! even: +1, odd: -1
          do i = 1, nc
             i_c1 = ix_offset(1) + ishft(i+1, -1) ! (i+1)/2
             i_c2 = i_c1 + 1 - 2 * iand(i, 1)     ! even: +1, odd: -1

             u0 = box_p%cc(i_c1, j_c1, k_c1, i_corr)
             u(1) = box_p%cc(i_c2, j_c1, k_c1, i_corr)
             u(2) = box_p%cc(i_c1, j_c2, k_c1, i_corr)
             u(3) = box_p%cc(i_c1, j_c1, k_c2, i_corr)
             a0 = box_p%cc(i_c1, j_c1, k_c1, i_eps)
             a(1) = box_p%cc(i_c2, j_c1, k_c1, i_eps)
             a(2) = box_p%cc(i_c1, j_c2, k_c1, i_eps)
             a(3) = box_p%cc(i_c1, j_c1, k_c2, i_eps)

             ! Get value of phi at coarse cell faces, and average
             box_c%cc(i, j, k, i_phi) = box_c%cc(i, j, k, i_phi) + third * &
                  sum((a0*u0 + a(:) * (1.5_dp * u(:) - 0.5_dp * u0)) / &
                  (a0 + a(:)))
          end do
       end do
    end do
  end subroutine mg3_box_corr_lpld

  ! Below: multigrid operators for internal boundary conditions. A level set
  ! function defines the location of the interface(s).

  subroutine lsf_dist_val(lsf_val_bval_a, lsf_val_bval_b, dist, val)

    real(dp), intent(in)  :: lsf_val_bval_a(3)
    real(dp), intent(in)  :: lsf_val_bval_b(3)
    real(dp), intent(out) :: dist
    real(dp), intent(out) :: val
    real(dp)              :: lsf_a, lsf_b, bval_a, bval_b

    lsf_a = lsf_val_bval_a(1)
    lsf_b = lsf_val_bval_b(1)

    if (lsf_a * lsf_b < 0) then
       ! There is a boundary between the points
       dist = lsf_a / (lsf_a - lsf_b)
       bval_a = lsf_val_bval_a(3)
       bval_b = lsf_val_bval_b(3)

       ! Interpolate between boundary values
       val  = bval_a * (1-dist) + bval_b * dist
    else
       ! Simply use the value at b
       dist = 1
       val  = lsf_val_bval_b(2)
    end if
  end subroutine lsf_dist_val

  subroutine mg3_box_corr_lpllsf(box_p, box_c, mg)
    type(box3_t), intent(inout) :: box_c
    type(box3_t), intent(in)    :: box_p
    type(mg3_t), intent(in)     :: mg
    integer                     :: i_phi, i_corr, i_lsf, ix_offset(3)
    integer                     :: nc, i, j, i_c1, i_c2, j_c1, j_c2
    real(dp)                    :: v_a(3), v_b(3), val(3+1), dist(3+1), c(3+1)
    integer                     :: k, k_c1, k_c2

    nc        = box_c%n_cell
    ix_offset = a3_get_child_offset(box_c)
    i_phi     = mg%i_phi
    i_corr    = mg%i_tmp
    i_lsf     = mg%i_lsf

    ! In these loops, we calculate the closest coarse index (_c1), and the
    ! one-but-closest (_c2). The fine cell lies in between.
    do k = 1, nc
       k_c1 = ix_offset(3) + ishft(k+1, -1) ! (k+1)/2
       k_c2 = k_c1 + 1 - 2 * iand(k, 1)     ! even: +1, odd: -1
       do j = 1, nc
          j_c1 = ix_offset(2) + ishft(j+1, -1) ! (j+1)/2
          j_c2 = j_c1 + 1 - 2 * iand(j, 1)     ! even: +1, odd: -1
          do i = 1, nc
             i_c1 = ix_offset(1) + ishft(i+1, -1) ! (i+1)/2
             i_c2 = i_c1 + 1 - 2 * iand(i, 1)     ! even: +1, odd: -1

             v_a(1:2) = box_c%cc(i, j, k, [i_lsf, i_corr])
             v_a(3) = 0.0_dp       ! Boundary value for correction is 0
             v_b(3) = 0.0_dp       ! Idem
             v_b(1:2) = box_p%cc(i_c1, j_c1, k_c1, [i_lsf, i_corr])
             call lsf_dist_val(v_a, v_b, dist(1), val(1))
             v_b(1:2) = box_p%cc(i_c2, j_c1, k_c1, [i_lsf, i_corr])
             call lsf_dist_val(v_a, v_b, dist(2), val(2))
             v_b(1:2) = box_p%cc(i_c1, j_c2, k_c1, [i_lsf, i_corr])
             call lsf_dist_val(v_a, v_b, dist(3), val(3))
             v_b(1:2) = box_p%cc(i_c1, j_c1, k_c2, [i_lsf, i_corr])
             call lsf_dist_val(v_a, v_b, dist(4), val(4))

             ! This expresses general interpolation between 4 points (on the lines
             ! between the fine and the 4 coarse values).
             c(1) = dist(2) * dist(3) * dist(4)
             c(2) = dist(1) * dist(3) * dist(4)
             c(3) = dist(1) * dist(2) * dist(4)
             c(4) = dist(1) * dist(2) * dist(3)
             box_c%cc(i, j, k, i_phi) = box_c%cc(i, j, k, i_phi) + &
                  sum(c * val)/sum(c)
          end do
       end do
    end do
  end subroutine mg3_box_corr_lpllsf

  subroutine mg3_box_gsrb_lpllsf(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, i_lsf, i_bval
    real(dp)                    :: dx2, dd(2*3), val(2*3), v_a(3), v_b(3)
    integer                     :: k

    dx2    = box%dr**2
    nc     = box%n_cell
    i_phi  = mg%i_phi
    i_rhs  = mg%i_rhs
    i_lsf  = mg%i_lsf
    i_bval = mg%i_bval

    ! The parity of redblack_cntr determines which cells we use. If
    ! redblack_cntr is even, we use the even cells and vice versa.
    do k = 1, nc
       do j = 1, nc
          i0 = 2 - iand(ieor(redblack_cntr, k+j), 1)
          do i = i0, nc, 2
             v_a = box%cc(i, j, k, [i_lsf, i_phi, i_bval])
             v_b = box%cc(i-1, j, k, [i_lsf, i_phi, i_bval])
             call lsf_dist_val(v_a, v_b, dd(1), val(1))
             v_b = box%cc(i+1, j, k, [i_lsf, i_phi, i_bval])
             call lsf_dist_val(v_a, v_b, dd(2), val(2))
             v_b = box%cc(i, j-1, k, [i_lsf, i_phi, i_bval])
             call lsf_dist_val(v_a, v_b, dd(3), val(3))
             v_b = box%cc(i, j+1, k, [i_lsf, i_phi, i_bval])
             call lsf_dist_val(v_a, v_b, dd(4), val(4))
             v_b = box%cc(i, j, k-1, [i_lsf, i_phi, i_bval])
             call lsf_dist_val(v_a, v_b, dd(5), val(5))
             v_b = box%cc(i, j, k+1, [i_lsf, i_phi, i_bval])
             call lsf_dist_val(v_a, v_b, dd(6), val(6))

             ! Solve for generalized Laplacian (see routine mg3_box_lpllsf)
             box%cc(i, j, k, i_phi) = 1 / (1/(dd(1)*dd(2)) + &
                  1/(dd(3)*dd(4)) + 1/(dd(5)*dd(6))) * ( &
                  (dd(2) * val(1) + dd(1) * val(2)) / &
                  ((dd(1) + dd(2)) * dd(1) * dd(2)) + &
                  (dd(4) * val(3) + dd(3) * val(4)) / &
                  ((dd(3) + dd(4)) * dd(3) * dd(4)) + &
                  (dd(6) * val(5) + dd(5) * val(6)) / &
                  ((dd(5) + dd(6)) * dd(5) * dd(6)) - &
                  0.5_dp * dx2 * box%cc(i, j, k, i_rhs))

          end do
       end do
    end do
  end subroutine mg3_box_gsrb_lpllsf

  subroutine mg3_box_lpllsf(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, i_lsf, i_bval
    real(dp)                    :: inv_dr_sq, dd(2*3), val(2*3)
    real(dp)                    :: f0, v_a(3), v_b(3)
    integer                     :: k

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

    do k = 1, nc
       do j = 1, nc
          do i = 1, nc
             v_a = box%cc(i, j, k, [i_lsf, i_phi, i_bval])
             v_b = box%cc(i-1, j, k, [i_lsf, i_phi, i_bval])
             call lsf_dist_val(v_a, v_b, dd(1), val(1))
             v_b = box%cc(i+1, j, k, [i_lsf, i_phi, i_bval])
             call lsf_dist_val(v_a, v_b, dd(2), val(2))
             v_b = box%cc(i, j-1, k, [i_lsf, i_phi, i_bval])
             call lsf_dist_val(v_a, v_b, dd(3), val(3))
             v_b = box%cc(i, j+1, k, [i_lsf, i_phi, i_bval])
             call lsf_dist_val(v_a, v_b, dd(4), val(4))
             v_b = box%cc(i, j, k-1, [i_lsf, i_phi, i_bval])
             call lsf_dist_val(v_a, v_b, dd(5), val(5))
             v_b = box%cc(i, j, k+1, [i_lsf, i_phi, i_bval])
             call lsf_dist_val(v_a, v_b, dd(6), val(6))

             ! Generalized Laplacian for neighbors at distance dd * dx
             f0 = box%cc(i, j, k, i_phi)
             box%cc(i, j, k, i_out) = 2 * inv_dr_sq * ( &
                  (dd(2) * val(1) + dd(1) * val(2) - (dd(1)+dd(2)) * f0) / &
                  ((dd(1) + dd(2)) * dd(1) * dd(2)) + &
                  (dd(4) * val(3) + dd(3) * val(4) - (dd(3)+dd(4)) * f0) / &
                  ((dd(3) + dd(4)) * dd(3) * dd(4)) + &
                  (dd(6) * val(5) + dd(5) * val(6) - (dd(5)+dd(6)) * f0) / &
                  ((dd(5) + dd(6)) * dd(5) * dd(6)))
          end do
       end do
    end do
  end subroutine mg3_box_lpllsf

  subroutine mg3_box_rstr_lpllsf(box_c, box_p, iv, mg)
    type(box3_t), intent(in)      :: box_c
    type(box3_t), intent(inout)   :: box_p
    integer, intent(in)           :: iv
    type(mg3_t), intent(in)       :: mg
    integer                       :: i, j, i_f, j_f, i_c, j_c
    integer                       :: hnc, ix_offset(3), n_ch
    logical                       :: child_mask(2, 2, 2)
    integer                       :: k, k_f, k_c

    hnc       = ishft(box_c%n_cell, -1) ! n_cell / 2
    ix_offset = a3_get_child_offset(box_c)

    do k = 1, hnc
       k_c = ix_offset(3) + k
       k_f = 2 * k - 1
       do j = 1, hnc
          j_c = ix_offset(2) + j
          j_f = 2 * j - 1
          do i = 1, hnc
             i_c = ix_offset(1) + i
             i_f = 2 * i - 1

             child_mask = (box_p%cc(i_c, j_c, k_c, mg%i_lsf) * &
                  box_c%cc(i_f:i_f+1, j_f:j_f+1, k_f:k_f+1, mg%i_lsf) > 0)
             n_ch = count(child_mask)

             if (n_ch < a3_num_children .and. n_ch > 0) then
                box_p%cc(i_c, j_c, k_c, iv) = 1 / n_ch * &
                     sum(box_c%cc(i_f:i_f+1, j_f:j_f+1, k_f:k_f+1, iv), &
                     mask=child_mask)
             else                  ! Take average of children
                box_p%cc(i_c, j_c, k_c, iv) = 0.125_dp * &
                     sum(box_c%cc(i_f:i_f+1, j_f:j_f+1, k_f:k_f+1, iv))
             end if
          end do
       end do
    end do
  end subroutine mg3_box_rstr_lpllsf

end module m_a3_multigrid