anisotropic_diffusion_2d.f90

program anisotropic_diffusion_2d
  use m_a2_all

  implicit none

  integer, parameter  :: box_size   = 8
  integer, parameter  :: i_phi      = 1
  integer, parameter  :: i_phi_old  = 2
  integer, parameter  :: i_bx       = 3
  integer, parameter  :: i_by       = 4
  real(dp), parameter :: domain_len = 2 * acos(-1.0_dp)
  real(dp), parameter :: dr         = domain_len / box_size

  type(a2_t)        :: tree
  type(ref_info_t)   :: refine_info
  integer            :: ix_list(2, 1)
  integer            :: nb_list(a2_num_neighbors, 1)
  integer            :: time_steps, output_cnt
  integer            :: i, id, n, n_steps
  real(dp)           :: dt, time, end_time
  real(dp)           :: dt_adapt, dt_output
  real(dp)           :: velocity(2), dr_min(2)
  character(len=100) :: fname
  integer            :: count_rate, t_start, t_end

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

  ! Initialize tree
  call a2_init(tree, & ! Tree to initialize
       box_size, &     ! A box contains box_size**DIM cells
       n_var_cell=4, & ! Number of cell-centered variables
       n_var_face=1, & ! Number of face-centered variables
       dr=dr, &        ! Distance between cells on base level
       r_min=[-0.5_dp, -0.5_dp] * domain_len, &
       cc_names=["phi", "old", "Bx ", "By "]) ! Variable names

  call a2_set_cc_methods(tree, i_phi, a2_bc_neumann_zero, a2_gc_interp_lim)

  ! Set up geometry
  id             = 1
  ix_list(:, id) = 1            ! Set index of box
  nb_list(:, id) = id           ! Box is periodic, so its own neighbor

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

  output_cnt = 0
  time       = 0
  dt_adapt   = 0.25_dp
  dt_output  = 0.25_dp
  end_time   = 10_dp
  velocity(:) = 0.0_dp
  velocity(1) = 1.0_dp
  velocity(2) = -1.0_dp

  ! Set up the initial conditions
  call system_clock(t_start,count_rate)

  do
     call a2_loop_box(tree, set_initial_condition)
     call a2_gc_tree(tree, i_phi, a2_gc_interp_lim, a2_bc_neumann_zero)
     call a2_adjust_refinement(tree, &           ! tree
          refine_routine, & ! Refinement function
          refine_info, &    ! Information about refinement
          1)                ! Buffer width (in cells)
     if (refine_info%n_add == 0) exit
  end do

  call a2_gc_tree(tree, i_phi, a2_gc_interp_lim, a2_bc_neumann_zero)
  call a2_print_info(tree)

  time_steps = 0

  ! Starting simulation
  do
     time_steps = time_steps + 1
     dr_min  = a2_min_dr(tree)
     dt      = 0.25_dp * minval(dr_min)**2

     n_steps = ceiling(dt_adapt/dt)
     dt      = dt_adapt / n_steps

     if (output_cnt * dt_output <= time) then
        output_cnt = output_cnt + 1
        write(fname, "(A,I0)") "anisotropic_diffusion_2d_", output_cnt
        call a2_write_silo(tree, trim(fname), output_cnt, time, dir="output")
     end if

     if (time > end_time) exit

     do n = 1, n_steps
        ! Copy previous solution
        call a2_tree_copy_cc(tree, i_phi, i_phi_old)

        ! Two forward Euler steps over dt
        do i = 1, 2
           ! Call procedure get_fluxes for each id in tree, giving the list of boxes
           call a2_loop_boxes(tree, get_fluxes)

           ! Restrict fluxes from children to parents on refinement boundaries.
           call a2_consistent_fluxes(tree, [i_phi])

           ! Call procedure update_solution (see below) for each box in tree, with argument dt
           call a2_loop_box_arg(tree, update_solution, [dt])

           ! Restrict variables i_phi to all parent boxes
           call a2_restrict_tree(tree, i_phi)
        end do

        ! Take average of phi_old and phi
        call a2_loop_box(tree, average_phi)
        time = time + dt
     end do

     ! Fill ghost cells for variable i_phi
     call a2_gc_tree(tree, i_phi, a2_gc_interp_lim, a2_bc_neumann_zero)

     ! Adjust the refinement of a tree using refine_routine (see below) for grid
     ! refinement. On input, the tree should be balanced. On output, the tree is
     ! still balanced, and its refinement is updated (with at most one level per
     ! call).
     ! call a2_adjust_refinement(tree, refine_routine, refine_info, 1)

     ! Prolongation of i_phi values to new children (see below)
     ! call prolong_to_new_children(tree, refine_info)
  end do

  call system_clock(t_end,count_rate)
  write(*, "(A,I0,A,Es10.3,A)") &
       " Wall-clock time after ",time_steps, &
       " time steps: ", (t_end-t_start) / real(count_rate, dp), &
       " seconds"

  ! This call is not really necessary here, but cleaning up the data in a tree
  ! is important if your program continues with other tasks.
  call a2_destroy(tree)

contains

  subroutine refine_routine(box, cell_flags)
    type(box2_t), intent(in) :: box
    integer, intent(out)      :: cell_flags(box%n_cell, box%n_cell)
    ! real(dp)                  :: diff
    ! integer                   :: i, j, nc

    ! nc   = box%n_cell

    ! do j = 1, nc; do i = 1, nc
    !    diff = abs(box%cc(i+1, j, i_phi) + &
    !         box%cc(i-1, j, i_phi) + &
    !         box%cc(i, j+1, i_phi) + &
    !         box%cc(i, j-1, i_phi) - &
    !         4 * box%cc(i, j, i_phi)) * box%dr

    !    if (box%lvl < 2 .or. diff > 0.5e-3_dp .and. box%lvl < 5) then
    !       cell_flags(i, j) = af_do_ref
    !    else if (diff < 0.1_dp * 0.1e-3_dp) then
    !       cell_flags(i, j) = af_rm_ref
    !    else
    !       cell_flags(i, j) = af_keep_ref
    !    end if
    ! end do; end do

    if (box%lvl < 5) then
       cell_flags(:, :) = af_do_ref
    else
       cell_flags(:, :) = af_keep_ref
    end if
  end subroutine refine_routine

  subroutine set_initial_condition(box)
    type(box2_t), intent(inout) :: box
    integer                      :: i, j, nc
    real(dp)                     :: rr(2)

    nc = box%n_cell
    do j = 0, nc+1; do i = 0, nc+1
       rr = a2_r_cc(box, [i, j])
       box%cc(i, j, i_phi) = solution(rr)
       box%cc(i, j, i_bx:i_by) = get_bvec(rr)
    end do; end do
  end subroutine set_initial_condition

  function solution(rr) result(sol)
    real(dp), intent(in) :: rr(2)
    real(dp)             :: sol

    if (all(abs(rr - domain_len * [0.25_dp, 0.0_dp]) < 0.2_dp)) then
       sol = 1.0_dp
    else
       sol = 0.0_dp
    end if

  end function solution

  function get_bvec(rr) result(bvec)
    real(dp), intent(in) :: rr(2)
    real(dp) :: bvec(2), phi

    ! phi = atan2(rr(2), rr(1))
    ! bvec = [-sin(phi), cos(phi)]
    bvec = [1.0_dp, 1.0_dp] * sqrt(0.5_dp)
  end function get_bvec

  subroutine get_fluxes(boxes, id)
    use m_flux_schemes
    type(box2_t), intent(inout) :: boxes(:)
    integer, intent(in)         :: id
    integer                     :: nc, n, m
    real(dp)                    :: rr(2), bvec(2), grad(2), tmp(2), inv_dx

    nc     = boxes(id)%n_cell
    inv_dx = 1/boxes(id)%dr

    call a2_gc_box(boxes, id, i_phi, a2_gc_interp_lim, a2_bc_neumann_zero)

    do n = 1, nc+1
       do m = 1, nc
          rr = a2_rr_cc(boxes(id), [n-0.5_dp, m+0.0_dp])
          bvec = get_bvec(rr)

          grad(1) = boxes(id)%cc(n-1, m, i_phi) - boxes(id)%cc(n, m, i_phi)
          grad(2) = 0.25_dp * (&
               boxes(id)%cc(n-1, m-1, i_phi) - boxes(id)%cc(n-1, m+1, i_phi) + &
               boxes(id)%cc(n, m-1, i_phi) - boxes(id)%cc(n, m+1, i_phi))

          tmp = bvec(1) * inv_dx * bvec * grad
          boxes(id)%fc(n, m, 1, i_phi) = sum(tmp)

          if (boxes(id)%fc(n, m, 1, i_phi) > inv_dx * boxes(id)%cc(n-1, m, i_phi)) then
             boxes(id)%fc(n, m, 1, i_phi) = inv_dx * boxes(id)%cc(n-1, m, i_phi)
          else if (boxes(id)%fc(n, m, 1, i_phi) < -inv_dx * boxes(id)%cc(n, m, i_phi)) then
             boxes(id)%fc(n, m, 1, i_phi) = -inv_dx * boxes(id)%cc(n, m, i_phi)
          end if

          rr = a2_rr_cc(boxes(id), [m+0.0_dp, n-0.5_dp])
          bvec = get_bvec(rr)

          grad(2) = boxes(id)%cc(m, n-1, i_phi) - boxes(id)%cc(m, n, i_phi)
          grad(1) = 0.25_dp * (&
               boxes(id)%cc(m-1, n-1, i_phi) - boxes(id)%cc(m+1, n-1, i_phi) + &
               boxes(id)%cc(m-1, n, i_phi) - boxes(id)%cc(m+1, n, i_phi))

          tmp = bvec(2) * inv_dx * bvec * grad
          boxes(id)%fc(m, n, 2, i_phi) = sum(tmp)

          if (boxes(id)%fc(m, n, 2, i_phi) > inv_dx * boxes(id)%cc(m, n-1, i_phi)) then
             boxes(id)%fc(m, n, 2, i_phi) = inv_dx * boxes(id)%cc(m, n-1, i_phi)
          else if (boxes(id)%fc(m, n, 2, i_phi) < -inv_dx * boxes(id)%cc(m, n, i_phi)) then
             boxes(id)%fc(m, n, 2, i_phi) = -inv_dx * boxes(id)%cc(m, n, i_phi)
          end if
       end do
    end do

  end subroutine get_fluxes

  subroutine update_solution(box, dt)
    type(box2_t), intent(inout) :: box
    real(dp), intent(in)         :: dt(:)
    real(dp)                     :: inv_dr, fx(2), fy(2), tmp
    integer                      :: i, j, nc

    nc     = box%n_cell
    inv_dr = 1/box%dr

    do j = 1, nc
       do i = 1, nc
          fx = box%fc(i:i+1, j, 1, i_phi)
          fy = box%fc(i, j:j+1, 2, i_phi)

          tmp = dt(1) * inv_dr * (fx(1) - fx(2) + fy(1) - fy(2))
          ! if (tmp < -box%cc(i, j, i_phi)) tmp = -box%cc(i, j, i_phi)
          box%cc(i, j, i_phi) = box%cc(i, j, i_phi) + tmp
       end do
    end do
  end subroutine update_solution

  subroutine average_phi(box)
    type(box2_t), intent(inout) :: box

    box%cc(:, :, i_phi) = 0.5_dp * (box%cc(:, :, i_phi) + &
         box%cc(:, :, i_phi_old))
  end subroutine average_phi

end program anisotropic_diffusion_2d