space_plasma_2d.f90

program space_plasma_2d
  use m_a2_all

  implicit none

  ! Cell-centered variables
  integer, parameter :: i_cs_ion     = 1
  integer, parameter :: i_cs         = 2
  integer, parameter :: i_o          = 3
  integer, parameter :: i_o_min      = 4
  integer, parameter :: i_cs_ion_old = 5
  integer, parameter :: i_cs_old     = 6
  integer, parameter :: i_o_old      = 7
  integer, parameter :: i_o_min_old  = 8

  ! Domain
  integer, parameter  :: box_size           = 8
  real(dp), parameter :: domain_len         = 3000.0_dp
  real(dp), parameter :: dr                 = domain_len / box_size
  integer, parameter  :: min_refinement_lvl = 4
  integer, parameter  :: max_refinement_lvl = 7

  ! Time steps
  real(dp) :: dt_adapt  = 60.0_dp
  real(dp) :: dt_output = 60.0_dp
  real(dp) :: end_time  = 3600.0_dp

  ! Physical parameters
  real(dp), parameter :: dc_cs_ion_parallel = 10.0_dp
  real(dp), parameter :: dc_cs_ion_perpendicular = 2.0_dp
  real(dp), parameter :: dc_cs = 20.0_dp
  real(dp), parameter :: k_attachment = 1.0e-3_dp
  real(dp), parameter :: k_recombination = 1.0e-3_dp
  real(dp), parameter :: k_photoionization = 5.0e-4_dp
  real(dp), parameter :: initial_cs_density = 1.0_dp
  real(dp), parameter :: initial_ionization_fraction = 0.5_dp
  real(dp), parameter :: initial_o_density = 1e-7
  real(dp), parameter :: initial_sigma = 50.0_dp

  real(dp)           :: max_diffusion_coeff
  type(a2_t)         :: tree
  type(ref_info_t)   :: refine_info
  integer            :: ix_list(2, 1)
  integer            :: time_steps, output_cnt
  integer            :: i, n, n_steps
  real(dp)           :: dt, time
  real(dp)           :: dr_min(2)
  character(len=100) :: fname

  print *, "Running space_plasma_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=8, & ! Number of cell-centered variables
       n_var_face=2, & ! 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=[character(len=10) :: "Cs_ion", "Cs", "O", "O_min", &
       "Cs_ion_old", "Cs_old", "O_old", "O_min_old"]) ! Variable names

  ! Set default methods
  call a2_set_cc_methods(tree, i_cs_ion, a2_bc_neumann_zero, a2_gc_interp_lim)
  call a2_set_cc_methods(tree, i_cs, a2_bc_neumann_zero, a2_gc_interp_lim)
  call a2_set_cc_methods(tree, i_o, a2_bc_neumann_zero, a2_gc_interp_lim)
  call a2_set_cc_methods(tree, i_o_min, a2_bc_neumann_zero, a2_gc_interp_lim)

  ! Set index of box
  ix_list(:, 1) = [1, 1]

  ! Create the base mesh
  call a2_set_base(tree, 1, ix_list)

  output_cnt = 0
  time       = 0

  do
     call a2_loop_box(tree, set_initial_condition)
     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_cs_ion, a2_gc_interp_lim, a2_bc_neumann_zero)
  call a2_gc_tree(tree, i_cs, a2_gc_interp_lim, a2_bc_neumann_zero)

  call a2_print_info(tree)

  time_steps = 0

  max_diffusion_coeff = max(dc_cs_ion_parallel, dc_cs_ion_perpendicular, &
       dc_cs)

  ! Starting simulation
  do
     time_steps = time_steps + 1
     dr_min  = a2_min_dr(tree)

     ! Limit time step for diffusion
     dt      = 0.25_dp * minval(dr_min)**2 / max_diffusion_coeff

     ! Limit for reactions (using the maximum density, this can be improved)
     dt = min(dt, 0.5_dp/(k_attachment*initial_cs_density), &
          0.5_dp/(k_recombination*initial_cs_density), &
          0.5_dp/k_photoionization)

     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)") "space_plasma_2d_", output_cnt
        call a2_gc_tree(tree, i_cs_ion, a2_gc_interp_lim, a2_bc_neumann_zero)
        call a2_gc_tree(tree, i_cs, a2_gc_interp_lim, a2_bc_neumann_zero)
        call a2_gc_tree(tree, i_o, a2_gc_interp_lim, a2_bc_neumann_zero)
        call a2_gc_tree(tree, i_o_min, a2_gc_interp_lim, a2_bc_neumann_zero)

        call a2_write_silo(tree, trim(fname), output_cnt, time, &
             ixs_cc=[i_cs_ion, i_cs, i_o, i_o_min], dir="output")
     end if

     if (time > end_time) exit

     do n = 1, n_steps
        ! Copy previous solution
        call a2_tree_copy_cc(tree, i_cs_ion, i_cs_ion_old)
        call a2_tree_copy_cc(tree, i_cs, i_cs_old)
        call a2_tree_copy_cc(tree, i_o, i_o_old)
        call a2_tree_copy_cc(tree, i_o_min, i_o_min_old)

        ! Explicit trapezoidal rule. First take 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_cs_ion, i_cs])

           ! 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_Cs_ion to all parent boxes
           call a2_restrict_tree(tree, i_cs_ion)
           call a2_restrict_tree(tree, i_cs)
        end do

        ! Now take average of phi_old and phi (this is the explicit trapezoidal rule)
        call a2_loop_box(tree, average_solution)
        time = time + dt
     end do

     ! Fill ghost cells before refinement
     call a2_gc_tree(tree, i_cs_ion, a2_gc_interp_lim, a2_bc_neumann_zero)
     call a2_gc_tree(tree, i_cs, a2_gc_interp_lim, a2_bc_neumann_zero)

     call a2_adjust_refinement(tree, refine_routine, refine_info, 1)
  end do

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_cs_ion) + &
            box%cc(i-1, j, i_cs_ion) + &
            box%cc(i, j+1, i_cs_ion) + &
            box%cc(i, j-1, i_cs_ion) - &
            4 * box%cc(i, j, i_cs_ion)) * box%dr / &
            (initial_cs_density * initial_ionization_fraction)

       if (box%lvl < min_refinement_lvl .or. diff > 0.1_dp &
            .and. box%lvl < max_refinement_lvl) then
          cell_flags(i, j) = af_do_ref
       else if (diff < 0.01_dp) then
          cell_flags(i, j) = af_rm_ref
       else
          cell_flags(i, j) = af_keep_ref
       end if
    end do; end do
  end subroutine refine_routine

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

    nc = box%n_cell
    do j = 0, nc+1; do i = 0, nc+1
       distance = norm2(a2_r_cc(box, [i, j]))
       density = exp(-0.5_dp * distance**2/initial_sigma**2)

       box%cc(i, j, i_cs_ion) = initial_ionization_fraction * density
       box%cc(i, j, i_cs) = (1 - initial_ionization_fraction) * density
       box%cc(i, j, i_o) = initial_o_density
    end do; end do
  end subroutine set_initial_condition

  subroutine get_fluxes(boxes, id)
    use m_flux_schemes
    type(box2_t), intent(inout) :: boxes(:)
    integer, intent(in)         :: id
    integer                     :: nc
    real(dp)                    :: inv_dx
    real(dp), allocatable       :: diff_coeff(:, :, :)

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

    allocate(diff_coeff(1:nc+1, 1:nc+1, 2))

    call a2_gc_box(boxes, id, i_cs_ion, a2_gc_interp_lim, a2_bc_neumann_zero)
    call a2_gc_box(boxes, id, i_cs, a2_gc_interp_lim, a2_bc_neumann_zero)

    diff_coeff(:, :, 1) = dc_cs_ion_parallel
    diff_coeff(:, :, 2) = dc_cs_ion_perpendicular
    call flux_diff_2d(boxes(id)%cc(:, :, i_cs_ion), diff_coeff, inv_dx, nc, 1)
    boxes(id)%fc(:, :, :, i_cs_ion) = diff_coeff

    diff_coeff(:, :, :) = dc_cs
    call flux_diff_2d(boxes(id)%cc(:, :, i_cs), diff_coeff, inv_dx, nc, 1)
    boxes(id)%fc(:, :, :, i_cs) = diff_coeff

  end subroutine get_fluxes

  subroutine update_solution(box, dt)
    type(box2_t), intent(inout) :: box
    real(dp), intent(in)        :: dt(:)
    real(dp)                    :: inv_dr, recombination, photoionization
    real(dp)                    :: attachment, electron_density
    integer                     :: i, j, nc

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

    do j = 1, nc
       do i = 1, nc
          electron_density = box%cc(i, j, i_cs_ion) - box%cc(i, j, i_o_min)
          recombination = dt(1) * k_recombination * box%cc(i, j, i_cs_ion)**2
          photoionization = dt(1) * k_photoionization * box%cc(i, j, i_cs)
          box%cc(i, j, i_cs_ion) = box%cc(i, j, i_cs_ion) - recombination + photoionization
          box%cc(i, j, i_cs) = box%cc(i, j, i_cs) + recombination - photoionization

          attachment = dt(1) * k_attachment * box%cc(i, j, i_o) * electron_density
          box%cc(i, j, i_o) = box%cc(i, j, i_o) - attachment
          box%cc(i, j, i_o_min) = box%cc(i, j, i_o_min) + attachment

          ! Fluxes
          box%cc(i, j, i_cs_ion) = box%cc(i, j, i_cs_ion) + dt(1) * inv_dr * ( &
               box%fc(i, j, 1, i_cs_ion) - box%fc(i+1, j, 1, i_cs_ion) + &
               box%fc(i, j, 2, i_cs_ion) - box%fc(i, j+1, 2, i_cs_ion))
          box%cc(i, j, i_cs) = box%cc(i, j, i_cs) + dt(1) * inv_dr * ( &
               box%fc(i, j, 1, i_cs) - box%fc(i+1, j, 1, i_cs) + &
               box%fc(i, j, 2, i_cs) - box%fc(i, j+1, 2, i_cs))
       end do
    end do
  end subroutine update_solution

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

    box%cc(:, :, :) = 0.5_dp * (box%cc(:, :, :) + &
         box%cc(:, :, :))
  end subroutine average_solution

end program space_plasma_2d