poisson_lsf_test.f90

Test of the level-set functionality of the Poisson solver In this example, multiple irregular boundaries can be defined, for example a sphere or a heart shape. For the case of a sphere, an analytic solution is provided in 1d, 2d and 3d.

The example contains quite a few parameters that can be changed via the command line, using the config_fortran module documented at https://github.com/jannisteunissen/config_fortran

Options can for example be set like this: ./poisson_lsf_test_2d -max_refine_level=7 -shape=2

#include "../src/cpp_macros.h"
 
program poisson_lsf_test
  use m_af_all
  use m_config
  use omp_lib
 
  implicit none
 
  integer            :: box_size         = 8
  integer            :: n_iterations     = 10
  integer            :: max_refine_level = 4
  integer            :: min_refine_level = 2
  integer            :: i_phi
  integer            :: i_rhs
  integer            :: i_tmp
  integer            :: i_lsf
  integer            :: i_lsf_mask
  integer            :: i_error
  integer            :: i_field
  integer            :: i_field_norm
 
  logical :: cylindrical = .false.
  logical :: write_numpy = .false.
  integer :: refinement_type = 1
 
  ! Which shape to use
  integer             :: shape = 1
 
  real(dp) :: sharpness_t             = 4.0_dp
  real(dp) :: boundary_value          = 1.0_dp
  real(dp) :: solution_coeff          = 1.0_dp
  real(dp) :: solution_radius         = 0.25_dp
  real(dp) :: solution_r0(NDIM)       = 0.5_dp
  real(dp) :: solution_smallest_width = 1e-3_dp
 
  type(af_t)         :: tree
  type(ref_info_t)   :: ref_info
  integer            :: n, mg_iter, coord
  real(dp)           :: residu, max_error, max_field
  character(len=150) :: fname
  character(len=20)  :: output_suffix = ""
  character(len=20)  :: prolongation = "auto"
  type(mg_t)         :: mg
  type(CFG_t)        :: cfg
  real(dp)           :: error_squared, rmse
  real(dp)           :: mem_limit_gb = 8.0_dp
  real(dp)           :: t0, t_sum
  logical            :: write_output = .true.
 
  call af_add_cc_variable(tree, "phi", ix=i_phi)
  call af_add_cc_variable(tree, "rhs", ix=i_rhs)
  call af_add_cc_variable(tree, "tmp", ix=i_tmp)
  call af_add_cc_variable(tree, "lsf", ix=i_lsf)
  call af_add_cc_variable(tree, "lsf_mask", ix=i_lsf_mask)
  call af_add_cc_variable(tree, "error", ix=i_error)
  call af_add_cc_variable(tree, "field_norm", ix=i_field_norm)
  call af_add_fc_variable(tree, "field", ix=i_field)
 
  call af_set_cc_methods(tree, i_lsf, funcval=set_lsf)
  call af_set_cc_methods(tree, i_lsf_mask, af_bc_neumann_zero)
  call af_set_cc_methods(tree, i_field_norm, af_bc_neumann_zero)
 
  call cfg_update_from_arguments(cfg)
 
  call cfg_add_get(cfg, "cylindrical", cylindrical, &
       "Whether to use cylindrical coordinates")
  call cfg_add_get(cfg, "write_output", write_output, &
       "Whether to write Silo output")
  call cfg_add_get(cfg, "write_numpy", write_numpy, &
       "Whether to write Numpy output")
  call cfg_add_get(cfg, "output_suffix", output_suffix, &
       "Suffix to add to output")
  call cfg_add_get(cfg, "max_refine_level", max_refine_level, &
       "Maximum refinement level")
  call cfg_add_get(cfg, "min_refine_level", min_refine_level, &
       "Minimum refinement level")
  call cfg_add_get(cfg, "refinement_type", refinement_type, &
       "Type of refinement (1: uniform, 2: only boundary)")
  call cfg_add_get(cfg, "shape", shape, &
       "Index of the shape used")
  call cfg_add_get(cfg, "box_size", box_size, &
       "Size of grid boxes")
  call cfg_add_get(cfg, "n_iterations", n_iterations, &
       "Number of multigrid iterations")
  call cfg_add_get(cfg, "n_cycle_up", mg%n_cycle_up, &
       "Number of relaxation steps going up")
  call cfg_add_get(cfg, "n_cycle_down", mg%n_cycle_down, &
       "Number of relaxation steps going down")
  call cfg_add_get(cfg, "solution%sharpness", sharpness_t, &
       "Sharpness of solution (for some cases)")
  call cfg_add_get(cfg, "solution%r0", solution_r0, &
       "Origin of solution")
  call cfg_add_get(cfg, "solution%radius", solution_radius, &
       "Solution radius")
  call cfg_add_get(cfg, "solution%coeff", solution_coeff, &
       "Linear coefficient for solution")
  call cfg_add_get(cfg, "solution%smallest_width", solution_smallest_width, &
       "Smallest width to resolve in the solution")
  call cfg_add_get(cfg, "boundary_value", boundary_value, &
       "Value for Dirichlet boundary condition")
  call cfg_add_get(cfg, "mem_limit_gb", mem_limit_gb, &
       "Memory limit (GByte)")
  call cfg_add_get(cfg, "prolongation", prolongation, &
       "Prolongation method (linear, sparse, auto, custom)")
 
  call cfg_check(cfg)
  ! call CFG_write(cfg, "stdout")
 
  if (cylindrical) then
     coord = af_cyl
     ! Place solution on axis
     solution_r0(1) = 0.0_dp
  else
     coord = af_xyz
  end if
 
  tree%mg_i_lsf = i_lsf
  mg%i_phi = i_phi
  mg%i_rhs = i_rhs
  mg%i_tmp = i_tmp
  mg%sides_bc => bc_solution
  mg%lsf_boundary_value = boundary_value
  mg%lsf_dist => mg_lsf_dist_gss
  mg%lsf => get_lsf
  mg%lsf_length_scale = solution_smallest_width
 
  select case (prolongation)
  case ("linear")
     mg%prolongation_type = mg_prolong_linear
  case ("sparse")
     mg%prolongation_type = mg_prolong_sparse
  case ("auto")
     mg%prolongation_type = mg_prolong_auto
  case ("custom")
     mg%lsf_use_custom_prolongation = .true.
  case default
     error stop "Unknown prolongation method"
  end select
 
  ! Initialize tree
  call af_init(tree, & ! Tree to initialize
       box_size, &     ! A box contains box_size**DIM cells
       [dtimes(1.0_dp)], &
       [dtimes(box_size)], &
       coord=coord, &
       mem_limit_gb=mem_limit_gb)
 
  call af_refine_up_to_lvl(tree, min_refine_level)
  call mg_init(tree, mg)
 
  do n = 1, 100
     call mg_set_operators_tree(tree, mg)
     call af_adjust_refinement(tree, ref_routine, ref_info, ref_buffer=1)
     if (ref_info%n_add == 0) exit
  end do
 
  ! Set mask where lsf changes sign
  call af_loop_box(tree, set_lsf_mask)
  call af_gc_tree(tree, [i_lsf_mask])
 
  call af_print_info(tree)
 
  t_sum = 0
  write(*, "(A,A4,4A14)") "# ", "iter", "residu", "max_error", "rmse", "max_field"
 
  do mg_iter = 1, n_iterations
     t0 = omp_get_wtime()
     call mg_fas_fmg(tree, mg, .true., mg_iter>1)
     t_sum = t_sum + omp_get_wtime() - t0
 
     call mg_compute_phi_gradient(tree, mg, i_field, 1.0_dp, i_field_norm)
     call af_gc_tree(tree, [i_field_norm])
     call af_loop_box(tree, set_error)
 
     ! Determine the minimum and maximum residual and error
     call af_tree_maxabs_cc(tree, i_tmp, residu)
     call af_tree_maxabs_cc(tree, i_error, max_error)
     call af_tree_maxabs_cc(tree, i_field_norm, max_field)
     call af_tree_sum_cc(tree, i_error, error_squared, 2)
     rmse = sqrt(error_squared/af_total_volume(tree))
     write(*, "(A,i4,4E14.5)") "# ", mg_iter, residu, max_error, rmse, max_field
 
     write(fname, "(A,I0)") "output/poisson_lsf_test_" // dimname // &
          trim(output_suffix) // "_", mg_iter
     if (write_output) then
        call af_write_silo(tree, trim(fname))
     end if
 
     if (write_numpy) then
        call af_write_numpy(tree, trim(fname)//".npz", &
             ixs_cc=[i_phi, i_lsf_mask])
     end if
  end do
 
  write(*, "(A,E14.5)") " seconds_per_FMG_cycle", t_sum/n_iterations
 
  call af_stencil_print_info(tree)
 
contains
 
  ! Set refinement flags for box
  subroutine ref_routine(box, cell_flags)
    type(box_t), intent(in) :: box
    integer, intent(out)    :: cell_flags(DTIMES(box%n_cell))
    real(dp)                :: tmp(DTIMES(box%n_cell)), dr_min
    integer                 :: nc, IJK, ix_mask, n, cix(NDIM)
 
    nc = box%n_cell
    cell_flags(dtimes(:)) = af_keep_ref
 
    select case (refinement_type)
    case (1)
       ! Uniform refinement
       if (box%lvl < max_refine_level) then
          cell_flags(dtimes(:)) = af_do_ref
       end if
 
    case (2)
       ! Only refine at boundary
       ix_mask = af_stencil_index(box, mg_lsf_mask_key)
       if (ix_mask /= af_stencil_none .and. box%lvl < max_refine_level) then
          do n = 1, size(box%stencils(ix_mask)%sparse_ix, 2)
             cix = box%stencils(ix_mask)%sparse_ix(:, n)
             cell_flags(dindex(cix)) = af_do_ref
          end do
       end if
 
    case (3)
       ! 'Bad' refinement to test the method
       if (norm2(box%r_min - solution_r0) < solution_radius .and. &
            box%lvl < max_refine_level) then
          cell_flags(dtimes(:)) = af_do_ref
       end if
    case (4)
       ! Let refinement depend on radius
       dr_min = minval(af_lvl_dr(tree, max_refine_level))
 
       do kji_do(1, nc)
          tmp(ijk) = max(1.0_dp, norm2(af_r_cc(box, [ijk]) - solution_r0) &
               / solution_radius)
       end do; close_do
 
       if (minval(box%dr) > minval(tmp) * dr_min .and. box%lvl < max_refine_level) then
          cell_flags(dtimes(:)) = af_do_ref
       end if
    case default
       error stop "Invalid refinement_type"
    end select
  end subroutine ref_routine
 
  real(dp) function solution(r)
    real(dp), intent(in) :: r(NDIM)
    real(dp) :: distance, lsf
 
    select case (shape)
    case (1)
       ! Relative distance
       distance = norm2(r-solution_r0) / solution_radius
 
       ! Let values increase for distance -> infinity
       if (distance < 1.0_dp) then
          solution = boundary_value
       else if (ndim == 1) then
          solution = boundary_value + solution_coeff * (distance - 1.0_dp)
       else if (ndim == 2 .and. coord == af_xyz) then
          solution = boundary_value + solution_coeff * log(distance)
       else
          solution = boundary_value + solution_coeff * (1 - 1/distance)
       end if
    case (5)
       ! Triangle
       lsf = get_lsf(r)
       if (lsf <= 0.0_dp) then
          solution = boundary_value
       else
          solution = boundary_value * exp(-lsf)
       end if
    case default
       solution = 0.0_dp
    end select
  end function solution
 
  ! This routine sets the level set function
  subroutine set_lsf(box, iv)
    type(box_t), intent(inout) :: box
    integer, intent(in)        :: iv
    integer                    :: IJK, nc
    real(dp)                   :: rr(NDIM), norm_dr
 
    nc = box%n_cell
    norm_dr = norm2(box%dr)
 
    do kji_do(0,nc+1)
       rr = af_r_cc(box, [ijk])
       box%cc(ijk, iv) = get_lsf(rr)
    end do; close_do
  end subroutine set_lsf
 
  real(dp) function get_lsf(rr)
    real(dp), intent(in) :: rr(NDIM)
    real(dp)             :: qq(NDIM)
#if NDIM > 1
    real(dp)             :: dist1, dist2
#endif
 
    qq = (rr - solution_r0)/solution_radius
 
#if NDIM == 3
    ! Generalize shapes to 3D
    qq(1) = norm2(qq([1, 2]))
    qq(2) = qq(3)
    qq(3) = 0.0_dp
#endif
 
    select case (shape)
    case (1)
       get_lsf = norm2(qq) - 1.0_dp
#if NDIM > 1
    case (2)
       ! Heart centered on r0
       get_lsf = (qq(1)**2 + qq(2)**2 - 1)**3 - &
            qq(1)**2 * qq(2)**3
    case (3)
       ! Astroid
       get_lsf = abs(qq(1))**(2.0_dp/3) / 0.8_dp + &
            abs(qq(2))**(2.0_dp/3) / 1.5_dp - 0.8_dp
    case (4)
       ! Rhombus
       ! sharpness_t -> for sharpness of the top angle,
       ! larger sharpness_t equals more acute angle
       get_lsf = sharpness_t * abs(qq(1)) + abs(qq(2)) - 1.5_dp
    case (5)
       ! Triangle, uses signed distance from the triangle
       dist1 = gm_dist_line(rr, [solution_r0(1) - solution_r0(2)/sharpness_t, 0.0_dp], &
            solution_r0, 2)
       dist2 = gm_dist_line(rr, [solution_r0(1) + solution_r0(2)/sharpness_t, 0.0_dp], &
            solution_r0, 2)
 
       ! Determine sign of lsf function
       qq = rr - solution_r0
       get_lsf = sharpness_t * abs(qq(1)) + qq(2)
 
       ! Use sign in front of minimum distance
       get_lsf = sign(min(dist1, dist2), get_lsf)
    case (6)
       ! Heart v2
       get_lsf = qq(1)**2 + (qq(2) - abs(qq(1))**(2/3.0_dp))**2 - 1
    case (7)
       ! Spheroid
       get_lsf = sqrt(sharpness_t * qq(1)**2 + qq(2)**2) - 1.0_dp
#endif
    case default
       error stop "Invalid shape"
    end select
 
  end function get_lsf
 
  subroutine set_error(box)
    type(box_t), intent(inout) :: box
    integer                    :: IJK, nc
    real(dp)                   :: rr(NDIM)
 
    nc = box%n_cell
    do kji_do(0,nc+1)
       rr = af_r_cc(box, [ijk])
       box%cc(ijk, i_error) = box%cc(ijk, i_phi) - solution(rr)
    end do; close_do
  end subroutine set_error
 
  subroutine set_lsf_mask(box)
    type(box_t), intent(inout) :: box
    integer                    :: nc, ix
 
    nc = box%n_cell
    box%cc(dtimes(:), i_lsf_mask) = 0
 
    ix = af_stencil_index(box, mg%operator_key)
 
    if (ix == af_stencil_none) error stop "No stencil stored"
 
    if (allocated(box%stencils(ix)%f)) then
       where (abs(box%stencils(ix)%f) > 0)
          box%cc(dtimes(1:nc), i_lsf_mask) = 1
       end where
    end if
 
    ! This code can get the lsf mask that is used to check for roots
    ! ix_mask = af_stencil_index(box, mg_lsf_mask_key)
    ! if (ix_mask /= af_stencil_none) then
    !    do n = 1, size(box%stencils(ix_mask)%sparse_ix, 2)
    !       ix = box%stencils(ix_mask)%sparse_ix(:, n)
    !       box%cc(DINDEX(ix), i_lsf_mask) = 1
    !    end do
    ! end if
  end subroutine set_lsf_mask
 
  ! This routine sets boundary conditions for a box
  subroutine bc_solution(box, nb, iv, coords, bc_val, bc_type)
    type(box_t), intent(in) :: box
    integer, intent(in)     :: nb
    integer, intent(in)     :: iv
    real(dp), intent(in)    :: coords(NDIM, box%n_cell**(NDIM-1))
    real(dp), intent(out)   :: bc_val(box%n_cell**(NDIM-1))
    integer, intent(out)    :: bc_type
    integer                 :: n
 
    bc_type = af_bc_dirichlet
 
    do n = 1, box%n_cell**(ndim-1)
       bc_val(n) = solution(coords(:, n))
    end do
  end subroutine bc_solution
 
#if NDIM > 1
 
  subroutine gm_dist_vec_line(r, r0, r1, n_dim, dist_vec, frac)
    integer, intent(in)   :: n_dim
    real(dp), intent(in)  :: r(n_dim), r0(n_dim), r1(n_dim)
    real(dp), intent(out) :: dist_vec(n_dim)
    real(dp), intent(out) :: frac
    real(dp)              :: line_len2
 
    line_len2 = sum((r1 - r0)**2)
    frac = sum((r - r0) * (r1 - r0))
 
    if (frac <= 0.0_dp) then
       frac = 0.0_dp
       dist_vec = r - r0
    else if (frac >= line_len2) then
       frac = 1.0_dp
       dist_vec = r - r1
    else
       dist_vec = r - (r0 + frac/line_len2 * (r1 - r0))
       frac = sqrt(frac / line_len2)
    end if
  end subroutine gm_dist_vec_line
 
  function gm_dist_line(r, r0, r1, n_dim) result(dist)
    integer, intent(in)  :: n_dim
    real(dp), intent(in) :: r(n_dim), r0(n_dim), r1(n_dim)
    real(dp)             :: dist, dist_vec(n_dim), frac
    call gm_dist_vec_line(r, r0, r1, n_dim, dist_vec, frac)
    dist = norm2(dist_vec)
  end function gm_dist_line
#endif
 
end program poisson_lsf_test