MPI-AMRVAC  2.2
The MPI - Adaptive Mesh Refinement - Versatile Advection Code
m_octree_mg_3d.t
Go to the documentation of this file.
1 
2 
3 ! Single module generated from the octree-mg sources.
4 ! This file can be easier to include in existing projects.
5 !
6 ! Notes:
7 ! 1. The module name is here extended by _1d, _2d or _3d
8 ! 2. The free space Poisson solver is not included here.
9 ! 3. It is best to make changes in the original repository at
10 ! https://github.com/jannisteunissen/octree-mg
11 ! 4. This file can be generated as follows:
12 ! cd octree-mg/single_module
13 ! ./to_single_module.sh
14 !
15 ! The modules can be compiled with:
16 ! mpif90 -c m_octree_mg_1d.f90 [other options]
17 ! mpif90 -c m_octree_mg_2d.f90 [other options]
18 ! mpif90 -c m_octree_mg_3d.f90 [other options]
19 ! mpif90 -c m_octree_mg.f90 -cpp -DNDIM=<1,2,3> [other options]
20 
21 
23  use mpi
24  implicit none
25  private
26 
27  !! File ../src/m_data_structures.f90
28 
29  !> Type of reals
30  integer, parameter, public :: dp = kind(0.0d0)
31 
32  !> Type for 64-bit integers
33  integer, parameter, public :: i8 = selected_int_kind(18)
34 
35  !> Indicates a standard Laplacian
36  integer, parameter, public :: mg_laplacian = 1
37 
38  !> Indicates a variable-coefficient Laplacian
39  integer, parameter, public :: mg_vlaplacian = 2
40 
41  !> Indicates a constant-coefficient Helmholtz equation
42  integer, parameter, public :: mg_helmholtz = 3
43 
44  !> Indicates a variable-coefficient Helmholtz equation
45  integer, parameter, public :: mg_vhelmholtz = 4
46 
47  !> Cartesian coordinate system
48  integer, parameter, public :: mg_cartesian = 1
49  !> Cylindrical coordinate system
50  integer, parameter, public :: mg_cylindrical = 2
51  !> Spherical coordinate system
52  integer, parameter, public :: mg_spherical = 3
53 
54  integer, parameter, public :: mg_smoother_gs = 1
55  integer, parameter, public :: mg_smoother_gsrb = 2
56  integer, parameter, public :: mg_smoother_jacobi = 3
57 
58  !> Problem dimension
59  integer, parameter, public :: mg_ndim = 3
60 
61  !> Number of predefined multigrid variables
62  integer, parameter, public :: mg_num_vars = 4
63  !> Maximum number of variables
64  integer, parameter, public :: mg_max_num_vars = 10
65  !> Index of solution
66  integer, parameter, public :: mg_iphi = 1
67  !> Index of right-hand side
68  integer, parameter, public :: mg_irhs = 2
69  !> Index of previous solution (used for correction)
70  integer, parameter, public :: mg_iold = 3
71  !> Index of residual
72  integer, parameter, public :: mg_ires = 4
73 
74  !> Index of the variable coefficient (at cell centers)
75  integer, parameter, public :: mg_iveps = 5
76 
77  !> Minimum allowed grid level
78  integer, parameter, public :: mg_lvl_lo = -20
79  !> Maximum allowed grid level
80  integer, parameter, public :: mg_lvl_hi = 20
81 
82  !> Value to indicate a Dirichlet boundary condition
83  integer, parameter, public :: mg_bc_dirichlet = -10
84 
85  !> Value to indicate a Neumann boundary condition
86  integer, parameter, public :: mg_bc_neumann = -11
87 
88  !> Value to indicate a continuous boundary condition
89  integer, parameter, public :: mg_bc_continuous = -12
90 
91  !> Special value that indicates there is no box
92  integer, parameter, public :: mg_no_box = 0
93  !> Special value that indicates there is a physical boundary
94  integer, parameter, public :: mg_physical_boundary = -1
95 
96  !> Maximum number of timers to use
97  integer, parameter, public :: mg_max_timers = 20
98 
99  ! Numbering of children (same location as **corners**)
100  integer, parameter, public :: mg_num_children = 8
101 
102  ! Index offset for each child
103  integer, parameter, public :: mg_child_dix(3, 8) = reshape( &
104  [0,0,0, 1,0,0, 0,1,0, 1,1,0, &
105  0,0,1, 1,0,1, 0,1,1, 1,1,1], [3,8])
106  ! Reverse child index in each direction
107  integer, parameter, public :: mg_child_rev(8, 3) = reshape( &
108  [2,1,4,3,6,5,8,7, 3,4,1,2,7,8,5,6, 5,6,7,8,1,2,3,4], [8,3])
109  ! Children adjacent to a neighbor
110  integer, parameter, public :: mg_child_adj_nb(4, 6) = reshape( &
111  [1,3,5,7, 2,4,6,8, 1,2,5,6, 3,4,7,8, 1,2,3,4, 5,6,7,8], [4,6])
112  ! Which children have a low index per dimension
113  logical, parameter, public :: mg_child_low(3, 8) = reshape([ &
114  .true., .true., .true., .false., .true., .true., &
115  .true., .false., .true., .false., .false., .true., &
116  .true., .true., .false., .false., .true., .false., &
117  .true., .false., .false., .false., .false., .false.], [3, 8])
118 
119  ! Neighbor topology information
120  integer, parameter, public :: mg_num_neighbors = 6
121  integer, parameter, public :: mg_neighb_lowx = 1
122  integer, parameter, public :: mg_neighb_highx = 2
123  integer, parameter, public :: mg_neighb_lowy = 3
124  integer, parameter, public :: mg_neighb_highy = 4
125  integer, parameter, public :: mg_neighb_lowz = 5
126  integer, parameter, public :: mg_neighb_highz = 6
127  ! Index offsets of neighbors
128  integer, parameter, public :: mg_neighb_dix(3, 6) = reshape( &
129  [-1,0,0, 1,0,0, 0,-1,0, 0,1,0, 0,0,-1, 0,0,1], [3,6])
130  ! Which neighbors have a lower index
131  logical, parameter, public :: mg_neighb_low(6) = &
132  [.true., .false., .true., .false., .true., .false.]
133  ! Opposite of nb_low, but now as -1,1 integers
134  integer, parameter, public :: mg_neighb_high_pm(6) = [-1, 1, -1, 1, -1, 1]
135  ! Reverse neighbors
136  integer, parameter, public :: mg_neighb_rev(6) = [2, 1, 4, 3, 6, 5]
137  ! Direction (dimension) for a neighbor
138  integer, parameter, public :: mg_neighb_dim(6) = [1, 1, 2, 2, 3, 3]
139 
140  !> Lists of blocks per refinement level
141  type, public :: mg_lvl_t
142  integer, allocatable :: leaves(:)
143  integer, allocatable :: parents(:)
144  integer, allocatable :: ref_bnds(:)
145  integer, allocatable :: ids(:)
146  integer, allocatable :: my_leaves(:)
147  integer, allocatable :: my_parents(:)
148  integer, allocatable :: my_ref_bnds(:)
149  integer, allocatable :: my_ids(:)
150  end type mg_lvl_t
151 
152  !> Box data structure
153  type, public :: mg_box_t
154  integer :: rank !< Which process owns this box
155  integer :: id !< Box id (index in boxes(:) array)
156  integer :: lvl !< Refinement level
157  integer :: ix(3) !< Spatial index
158  integer :: parent !< Id of parent
159  integer :: children(2**3) !< Ids of children
160  integer :: neighbors(2*3) !< Ids of neighbors
161  real(dp) :: r_min(3) !< Minimum coordinate
162  real(dp) :: dr(3) !< Grid spacing
163  !> Cell-centered data
164  real(dp), allocatable :: cc(:, :, :, :)
165  end type mg_box_t
166 
167  !> Buffer type (one is used for each pair of communicating processes)
168  type, public :: mg_buf_t
169  integer :: i_send !< Index in send array
170  integer :: i_recv
171  integer :: i_ix
172  integer, allocatable :: ix(:) !< Will be used to sort the data
173  real(dp), allocatable :: send(:)
174  real(dp), allocatable :: recv(:)
175  end type mg_buf_t
176 
177  type, public :: mg_comm_t
178  integer, allocatable :: n_send(:, :)
179  integer, allocatable :: n_recv(:, :)
180  end type mg_comm_t
181 
182  type, public :: mg_bc_t
183  integer :: bc_type = mg_bc_dirichlet !< Type of boundary condition
184  real(dp) :: bc_value = 0.0_dp !< Value (for e.g. Dirichlet or Neumann)
185  !> To set user-defined boundary conditions (overrides bc(:))
186  procedure(mg_subr_bc), pointer, nopass :: boundary_cond => null()
187  !> To set a user-defined refinement boundary method
188  procedure(mg_subr_rb), pointer, nopass :: refinement_bnd => null()
189  end type mg_bc_t
190 
191  type, public :: mg_timer_t
192  character(len=20) :: name
193  real(dp) :: t = 0.0_dp
194  real(dp) :: t0
195  end type mg_timer_t
196 
197  type, public :: mg_t
198  !> Whether the multigrid tree structure has been created
199  logical :: tree_created = .false.
200  !> Whether storage has been allocated for boxes
201  logical :: is_allocated = .false.
202  !> Number of extra cell-centered variable (e.g., for coefficients)
203  integer :: n_extra_vars = 0
204  !> MPI communicator
205  integer :: comm = -1
206  !> Number of MPI tasks
207  integer :: n_cpu = -1
208  !> MPI rank of this task
209  integer :: my_rank = -1
210  !> Size of boxes in cells, be equal for all dimensions
211  integer :: box_size = -1
212  !> Highest grid level in the tree
213  integer :: highest_lvl = -1
214  !> Lowest grid level in the tree
215  integer :: lowest_lvl = -1
216  !> First normal level of the quadtree/octree, at coarser levels parents
217  !> have only one child
218  integer :: first_normal_lvl = -1
219  !> Total number of boxes in the tree (among all processes)
220  integer :: n_boxes = 0
221  !> Size of boxes per level (differs for coarsest levels)
222  integer :: box_size_lvl(mg_lvl_lo:mg_lvl_hi)
223  !> Size of domain per level (if uniformly refined)
224  integer :: domain_size_lvl(3, mg_lvl_lo:mg_lvl_hi)
225  !> Grid spacing per level
226  real(dp) :: dr(3, mg_lvl_lo:mg_lvl_hi)
227  !> Minimum coordinates
228  real(dp) :: r_min(3)
229  !> List of all levels
230  type(mg_lvl_t) :: lvls(mg_lvl_lo:mg_lvl_hi)
231  !> Array with all boxes in the tree. Only boxes owned by this task are
232  !> actually allocated
233  type(mg_box_t), allocatable :: boxes(:)
234  !> Buffer for communication with each other task
235  type(mg_buf_t), allocatable :: buf(:)
236 
237  !> Communication info for restriction
238  type(mg_comm_t) :: comm_restrict
239  !> Communication info for prolongation
240  type(mg_comm_t) :: comm_prolong
241  !> Communication info for ghost cell filling
242  type(mg_comm_t) :: comm_ghostcell
243 
244  !> Whether boundary condition data has been stored for mg solution
245  logical :: phi_bc_data_stored = .false.
246 
247  !> Whether a dimension is periodic
248  logical :: periodic(3) = .false.
249 
250  !> To store pre-defined boundary conditions per direction per variable
252 
253  !> Type of operator
254  integer :: operator_type = mg_laplacian
255  !> Type of grid geometry
256  integer :: geometry_type = mg_cartesian
257 
258  !> Whether the mean has to be subtracted from the multigrid solution
259  logical :: subtract_mean = .false.
260  !> Type of multigrid smoother
261  integer :: smoother_type = mg_smoother_gs
262  !> Number of substeps for the smoother (for GSRB this is 2)
263  integer :: n_smoother_substeps = 1
264  !> Number of cycles when doing downwards relaxation
265  integer :: n_cycle_down = 2
266  !> Number of cycles when doing upwards relaxation
267  integer :: n_cycle_up = 2
268  !> Maximum number of cycles on the coarse grid
269  integer :: max_coarse_cycles = 1000
270  integer :: coarsest_grid(3) = 2
271  !> Stop coarse grid when max. residual is smaller than this
272  real(dp) :: residual_coarse_abs = 1e-8_dp
273  !> Stop coarse grid when residual has been reduced by this factor
274  real(dp) :: residual_coarse_rel = 1e-8_dp
275 
276  !> Multigrid operator (e.g., Laplacian)
277  procedure(mg_box_op), pointer, nopass :: box_op => null()
278 
279  !> Multigrid smoother
280  procedure(mg_box_gsrb), pointer, nopass :: box_smoother => null()
281 
282  !> Multigrid prolongation method
283  procedure(mg_box_prolong), pointer, nopass :: box_prolong => null()
284 
285  !> Number of timers
286  integer :: n_timers = 0
287  !> Values for the timers
288  type(mg_timer_t) :: timers(mg_max_timers)
289  end type mg_t
290 
291  interface
292  !> To fill ghost cells near physical boundaries
293  subroutine mg_subr_bc(box, nc, iv, nb, bc_type, bc)
294  import
295  type(mg_box_t), intent(in) :: box
296  integer, intent(in) :: nc
297  integer, intent(in) :: iv !< Index of variable
298  integer, intent(in) :: nb !< Direction
299  integer, intent(out) :: bc_type !< Type of b.c.
300  !> Boundary values
301  real(dp), intent(out) :: bc(nc, nc)
302  end subroutine mg_subr_bc
303 
304  !> To fill ghost cells near refinement boundaries
305  subroutine mg_subr_rb(box, nc, iv, nb, cgc)
306  import
307  type(mg_box_t), intent(inout) :: box
308  integer, intent(in) :: nc
309  integer, intent(in) :: iv !< Index of variable
310  integer, intent(in) :: nb !< Direction
311  !> Coarse data
312  real(dp), intent(in) :: cgc(nc, nc)
313  end subroutine mg_subr_rb
314 
315  !> Subroutine that performs A * cc(..., i_in) = cc(..., i_out)
316  subroutine mg_box_op(mg, id, nc, i_out)
317  import
318  type(mg_t), intent(inout) :: mg
319  integer, intent(in) :: id
320  integer, intent(in) :: nc
321  integer, intent(in) :: i_out
322  end subroutine mg_box_op
323 
324  !> Subroutine that performs Gauss-Seidel relaxation
325  subroutine mg_box_gsrb(mg, id, nc, redblack_cntr)
326  import
327  type(mg_t), intent(inout) :: mg
328  integer, intent(in) :: id
329  integer, intent(in) :: nc
330  integer, intent(in) :: redblack_cntr
331  end subroutine mg_box_gsrb
332 
333  !> Subroutine that performs prolongation to a single child
334  subroutine mg_box_prolong(mg, p_id, dix, nc, iv, fine)
335  import
336  type(mg_t), intent(inout) :: mg
337  integer, intent(in) :: p_id !< Id of parent
338  integer, intent(in) :: dix(3) !< Offset of child in parent grid
339  integer, intent(in) :: nc !< Child grid size
340  integer, intent(in) :: iv !< Prolong from this variable
341  real(dp), intent(out) :: fine(nc, nc, nc) !< Prolonged values
342  end subroutine mg_box_prolong
343  end interface
344 
345  ! Public methods
346  public :: mg_subr_bc
347  public :: mg_subr_rb
348  public :: mg_box_op
349  public :: mg_box_gsrb
350  public :: mg_box_prolong
351  public :: mg_has_children
352  public :: mg_ix_to_ichild
353  public :: mg_get_child_offset
354  public :: mg_highest_uniform_lvl
355  public :: mg_number_of_unknowns
356  public :: mg_get_face_coords
357  public :: mg_add_timer
358  public :: mg_timer_start
359  public :: mg_timer_end
360  public :: mg_timers_show
361 
362  !! File ../src/m_allocate_storage.f90
363 
364  public :: mg_allocate_storage
365  public :: mg_deallocate_storage
366 
367  !! File ../src/m_diffusion.f90
368 
369 !> Module for implicitly solving diffusion equations
370 
371  public :: diffusion_solve
372  public :: diffusion_solve_vcoeff
373 
374  !! File ../src/m_laplacian.f90
375 
376 !> Module which contains multigrid procedures for a Laplacian operator
377 
378  public :: laplacian_set_methods
379 
380  !! File ../src/m_vhelmholtz.f90
381 
382 !> Module which contains multigrid procedures for a Helmholtz operator of the
383 !> form: div(D grad(phi)) - lambda*phi = f, where D has a smooth spatial
384 !> variation
385 
386  !> The lambda used for the Helmholtz equation (should be >= 0)
387  real(dp), public, protected :: vhelmholtz_lambda = 0.0_dp
388 
389  public :: vhelmholtz_set_methods
390  public :: vhelmholtz_set_lambda
391 
392  !! File ../src/m_build_tree.f90
393 
394  ! Public methods
395  public :: mg_build_rectangle
396  public :: mg_add_children
397  public :: mg_set_leaves_parents
398  public :: mg_set_next_level_ids
400  public :: mg_set_neighbors_lvl
401 
402  !! File ../src/m_load_balance.f90
403 !> Module for load balancing a tree (that already has been constructed). The
404 !> load balancing determines which ranks (MPI processes) allocated physical
405 !> storage for boxes. The tree structure itself is present on all processes.
406 
407  ! Public methods
408  public :: mg_load_balance_simple
409  public :: mg_load_balance
410  public :: mg_load_balance_parents
411 
412  !! File ../src/m_vlaplacian.f90
413 
414 !> Module which contains multigrid procedures for a variable-coefficient
415 !> Laplacian operator, assuming the variation is smooth
416 
417  public :: vlaplacian_set_methods
418 
419  !! File ../src/m_communication.f90
420 
421  ! Public methods
422  public :: mg_comm_init
423  public :: sort_and_transfer_buffers
424 
425  !! File ../src/m_ghost_cells.f90
426 
427  ! Public methods
428  public :: mg_ghost_cell_buffer_size
429  public :: mg_fill_ghost_cells
430  public :: mg_fill_ghost_cells_lvl
431  public :: mg_phi_bc_store
432 
433  !! File ../src/m_prolong.f90
434 
435  ! Public methods
436  public :: mg_prolong
437  public :: mg_prolong_buffer_size
438  public :: mg_prolong_sparse
439 
440  !! File ../src/m_helmholtz.f90
441 
442 !> Module which contains multigrid procedures for a Helmholtz operator of the
443 !> form: laplacian(phi) - lambda*phi = f
444 
445  !> The lambda used for the Helmholtz equation (should be >= 0)
446  real(dp), public, protected :: helmholtz_lambda = 0.0_dp
447 
448  public :: helmholtz_set_methods
449  public :: helmholtz_set_lambda
450 
451  !! File ../src/m_multigrid.f90
452 
453  integer :: timer_total_vcycle = -1
454  integer :: timer_total_fmg = -1
455  integer :: timer_smoother = -1
456  integer :: timer_smoother_gc = -1
457  integer :: timer_coarse = -1
458  integer :: timer_correct = -1
459  integer :: timer_update_coarse = -1
460 
461  ! Public methods
462  public :: mg_fas_vcycle
463  public :: mg_fas_fmg
464  public :: mg_set_methods
465  public :: mg_apply_op
466 
467  !! File ../src/m_restrict.f90
468 
469  ! Public methods
470  public :: mg_restrict
471  public :: mg_restrict_lvl
472  public :: mg_restrict_buffer_size
473 
474  !! File ../src/m_mrgrnk.f90
475 
476  Integer, Parameter :: kdp = selected_real_kind(15)
477  public :: mrgrnk
478  private :: kdp
479  private :: i_mrgrnk
480  interface mrgrnk
481  module procedure i_mrgrnk
482  end interface mrgrnk
483 contains
484 
485  !! File ../src/m_data_structures.f90
486 
487  !> Return .true. if a box has children
488  elemental logical function mg_has_children(box)
489  type(mg_box_t), intent(in) :: box
490 
491  ! Boxes are either fully refined or not, so we only need to check one of the
492  ! children
493  mg_has_children = (box%children(1) /= mg_no_box)
494  end function mg_has_children
495 
496  !> Compute the child index for a box with spatial index ix. With child index
497  !> we mean the index in the children(:) array of its parent.
498  integer function mg_ix_to_ichild(ix)
499  integer, intent(in) :: ix(3) !< Spatial index of the box
500  ! The index can range from 1 (all ix odd) and 2**$D (all ix even)
501  mg_ix_to_ichild = 8 - 4 * iand(ix(3), 1) - &
502  2 * iand(ix(2), 1) - iand(ix(1), 1)
503  end function mg_ix_to_ichild
504 
505  !> Get the offset of a box with respect to its parent (e.g. in 2d, there can
506  !> be a child at offset 0,0, one at n_cell/2,0, one at 0,n_cell/2 and one at
507  !> n_cell/2, n_cell/2)
508  pure function mg_get_child_offset(mg, id) result(ix_offset)
509  type(mg_t), intent(in) :: mg
510  integer, intent(in) :: id
511  integer :: ix_offset(3)
512 
513  if (mg%boxes(id)%lvl <= mg%first_normal_lvl) then
514  ix_offset(:) = 0
515  else
516  ix_offset = iand(mg%boxes(id)%ix-1, 1) * &
517  ishft(mg%box_size, -1) ! * n_cell / 2
518  end if
519  end function mg_get_child_offset
520 
521  pure function mg_highest_uniform_lvl(mg) result(lvl)
522  type(mg_t), intent(in) :: mg
523  integer :: lvl
524 
525  do lvl = mg%first_normal_lvl, mg%highest_lvl-1
526  ! Exit if a grid is partially refined
527  if (size(mg%lvls(lvl)%leaves) /= 0 .and. &
528  size(mg%lvls(lvl)%parents) /= 0) exit
529  end do
530  ! If the loop did not exit, we get lvl equals mg%highest_lvl
531  end function mg_highest_uniform_lvl
532 
533  !> Determine total number of unknowns (on leaves)
534  function mg_number_of_unknowns(mg) result(n_unknowns)
535  type(mg_t), intent(in) :: mg
536  integer :: lvl
537  integer(i8) :: n_unknowns
538 
539  n_unknowns = 0
540  do lvl = mg%first_normal_lvl, mg%highest_lvl
541  n_unknowns = n_unknowns + size(mg%lvls(lvl)%leaves)
542  end do
543  n_unknowns = n_unknowns * int(mg%box_size**3, i8)
544  end function mg_number_of_unknowns
545 
546  !> Get coordinates at the face of a box
547  subroutine mg_get_face_coords(box, nb, nc, x)
548  type(mg_box_t), intent(in) :: box
549  integer, intent(in) :: nb
550  integer, intent(in) :: nc
551  real(dp), intent(out) :: x(nc, nc, 3)
552  integer :: i, j, ixs(3-1)
553  integer :: nb_dim
554  real(dp) :: rmin(3)
555 
556  nb_dim = mg_neighb_dim(nb)
557 
558  ! Determine directions perpendicular to neighbor
559  ixs = [(i, i = 1, 3-1)]
560  ixs(nb_dim:) = ixs(nb_dim:) + 1
561 
562  rmin = box%r_min
563  if (.not. mg_neighb_low(nb)) then
564  rmin(nb_dim) = rmin(nb_dim) + box%dr(nb_dim) * nc
565  end if
566 
567  do j = 1, nc
568  do i = 1, nc
569  x(i, j, :) = rmin
570  x(i, j, ixs) = x(i, j, ixs) + ([i, j] - 0.5d0) * box%dr(ixs)
571  end do
572  end do
573  end subroutine mg_get_face_coords
574 
575  integer function mg_add_timer(mg, name)
576  type(mg_t), intent(inout) :: mg
577  character(len=*), intent(in) :: name
578 
579  mg%n_timers = mg%n_timers + 1
580  mg_add_timer = mg%n_timers
581  mg%timers(mg_add_timer)%name = name
582  end function mg_add_timer
583 
584  subroutine mg_timer_start(timer)
585  use mpi
586  type(mg_timer_t), intent(inout) :: timer
587  timer%t0 = mpi_wtime()
588  end subroutine mg_timer_start
589 
590  subroutine mg_timer_end(timer)
591  use mpi
592  type(mg_timer_t), intent(inout) :: timer
593  timer%t = timer%t + mpi_wtime() - timer%t0
594  end subroutine mg_timer_end
595 
596  subroutine mg_timers_show(mg)
597  use mpi
598  type(mg_t), intent(in) :: mg
599  integer :: n, ierr
600  real(dp) :: tmin(mg%n_timers)
601  real(dp) :: tmax(mg%n_timers)
602 
603  call mpi_reduce(mg%timers(1:mg%n_timers)%t, tmin, mg%n_timers, &
604  mpi_double, mpi_min, 0, mg%comm, ierr)
605  call mpi_reduce(mg%timers(1:mg%n_timers)%t, tmax, mg%n_timers, &
606  mpi_double, mpi_max, 0, mg%comm, ierr)
607 
608  if (mg%my_rank == 0) then
609  write(*, "(A20,2A16)") "name ", "min(s)", "max(s)"
610  do n = 1, mg%n_timers
611  write(*, "(A20,2F16.6)") mg%timers(n)%name, &
612  tmin(n), tmax(n)
613  end do
614  end if
615  end subroutine mg_timers_show
616 
617  !! File ../src/m_allocate_storage.f90
618 
619  !> Deallocate all allocatable arrays
620  subroutine mg_deallocate_storage(mg)
621  type(mg_t), intent(inout) :: mg
622  integer :: lvl
623 
624  if (.not. mg%is_allocated) &
625  error stop "deallocate_storage: tree is not allocated"
626 
627  deallocate(mg%boxes)
628  deallocate(mg%buf)
629 
630  deallocate(mg%comm_restrict%n_send)
631  deallocate(mg%comm_restrict%n_recv)
632 
633  deallocate(mg%comm_prolong%n_send)
634  deallocate(mg%comm_prolong%n_recv)
635 
636  deallocate(mg%comm_ghostcell%n_send)
637  deallocate(mg%comm_ghostcell%n_recv)
638 
639  do lvl = mg%lowest_lvl, mg%highest_lvl
640  deallocate(mg%lvls(lvl)%ids)
641  deallocate(mg%lvls(lvl)%leaves)
642  deallocate(mg%lvls(lvl)%parents)
643  deallocate(mg%lvls(lvl)%ref_bnds)
644  deallocate(mg%lvls(lvl)%my_ids)
645  deallocate(mg%lvls(lvl)%my_leaves)
646  deallocate(mg%lvls(lvl)%my_parents)
647  deallocate(mg%lvls(lvl)%my_ref_bnds)
648  end do
649 
650  mg%is_allocated = .false.
651  mg%n_boxes = 0
652  mg%phi_bc_data_stored = .false.
653  end subroutine mg_deallocate_storage
654 
655  !> Allocate communication buffers and local boxes for a tree that has already
656  !> been created
657  subroutine mg_allocate_storage(mg)
659  type(mg_t), intent(inout) :: mg
660  integer :: i, id, lvl, nc
661  integer :: n_send(0:mg%n_cpu-1, 3)
662  integer :: n_recv(0:mg%n_cpu-1, 3)
663  integer :: dsize(3)
664  integer :: n_in, n_out, n_id
665 
666  if (.not. mg%tree_created) &
667  error stop "allocate_storage: tree is not yet created"
668 
669  if (mg%is_allocated) &
670  error stop "allocate_storage: tree is already allocated"
671 
672  do lvl = mg%lowest_lvl, mg%highest_lvl
673  nc = mg%box_size_lvl(lvl)
674  do i = 1, size(mg%lvls(lvl)%my_ids)
675  id = mg%lvls(lvl)%my_ids(i)
676  allocate(mg%boxes(id)%cc(0:nc+1, 0:nc+1, 0:nc+1, &
677  mg_num_vars + mg%n_extra_vars))
678 
679  ! Set all initial values to zero
680  mg%boxes(id)%cc(:, :, :, :) = 0.0_dp
681  end do
682  end do
683 
684  allocate(mg%buf(0:mg%n_cpu-1))
685 
686  call mg_ghost_cell_buffer_size(mg, n_send(:, 1), &
687  n_recv(:, 1), dsize(1))
688  call mg_restrict_buffer_size(mg, n_send(:, 2), &
689  n_recv(:, 2), dsize(2))
690  call mg_prolong_buffer_size(mg, n_send(:, 3), &
691  n_recv(:, 3), dsize(3))
692 
693  do i = 0, mg%n_cpu-1
694  n_out = maxval(n_send(i, :) * dsize(:))
695  n_in = maxval(n_recv(i, :) * dsize(:))
696  n_id = maxval(n_send(i, :))
697  allocate(mg%buf(i)%send(n_out))
698  allocate(mg%buf(i)%recv(n_in))
699  allocate(mg%buf(i)%ix(n_id))
700  end do
701 
702  mg%is_allocated = .true.
703  end subroutine mg_allocate_storage
704 
705  !! File ../src/m_diffusion.f90
706 
707  !> Solve a diffusion equation implicitly, assuming a constant diffusion
708  !> coefficient. The solution at time t should be stored in mg_iphi, which is
709  !> on output replaced by the solution at time t+dt.
710  subroutine diffusion_solve(mg, dt, diffusion_coeff, order, max_res)
712  type(mg_t), intent(inout) :: mg
713  real(dp), intent(in) :: dt
714  real(dp), intent(in) :: diffusion_coeff
715  integer, intent(in) :: order
716  real(dp), intent(in) :: max_res
717  integer, parameter :: max_its = 10
718  integer :: n
719  real(dp) :: res
720 
721  mg%operator_type = mg_helmholtz
722  call mg_set_methods(mg)
723 
724  select case (order)
725  case (1)
726  call helmholtz_set_lambda(1/(dt * diffusion_coeff))
727  call set_rhs(mg, -1/(dt * diffusion_coeff), 0.0_dp)
728  case (2)
729  call helmholtz_set_lambda(0.0d0)
730  call mg_apply_op(mg, mg_irhs)
731  call helmholtz_set_lambda(2/(dt * diffusion_coeff))
732  call set_rhs(mg, -2/(dt * diffusion_coeff), -1.0_dp)
733  case default
734  error stop "diffusion_solve: order should be 1 or 2"
735  end select
736 
737  ! Start with an FMG cycle
738  call mg_fas_fmg(mg, .true., max_res=res)
739 
740  ! Add V-cycles if necessary
741  do n = 1, max_its
742  if (res <= max_res) exit
743  call mg_fas_vcycle(mg, max_res=res)
744  end do
745 
746  if (n == max_its + 1) then
747  if (mg%my_rank == 0) &
748  print *, "Did you specify boundary conditions correctly?"
749  error stop "diffusion_solve: no convergence"
750  end if
751  end subroutine diffusion_solve
752 
753  !> Solve a diffusion equation implicitly, assuming a variable diffusion
754  !> coefficient which has been stored in mg_iveps (also on coarse grids). The
755  !> solution at time t should be stored in mg_iphi, which is on output replaced
756  !> by the solution at time t+dt.
757  subroutine diffusion_solve_vcoeff(mg, dt, order, max_res)
759  type(mg_t), intent(inout) :: mg
760  real(dp), intent(in) :: dt
761  integer, intent(in) :: order
762  real(dp), intent(in) :: max_res
763  integer, parameter :: max_its = 10
764  integer :: n
765  real(dp) :: res
766 
767  mg%operator_type = mg_vhelmholtz
768  call mg_set_methods(mg)
769 
770  select case (order)
771  case (1)
772  call vhelmholtz_set_lambda(1/dt)
773  call set_rhs(mg, -1/dt, 0.0_dp)
774  case (2)
775  call vhelmholtz_set_lambda(0.0d0)
776  call mg_apply_op(mg, mg_irhs)
777  call vhelmholtz_set_lambda(2/dt)
778  call set_rhs(mg, -2/dt, -1.0_dp)
779  case default
780  error stop "diffusion_solve: order should be 1 or 2"
781  end select
782 
783  ! Start with an FMG cycle
784  call mg_fas_fmg(mg, .true., max_res=res)
785 
786  ! Add V-cycles if necessary
787  do n = 1, max_its
788  if (res <= max_res) exit
789  call mg_fas_vcycle(mg, max_res=res)
790  end do
791 
792  if (n == max_its + 1) then
793  if (mg%my_rank == 0) then
794  print *, "Did you specify boundary conditions correctly?"
795  print *, "Or is the variation in diffusion too large?"
796  end if
797  error stop "diffusion_solve: no convergence"
798  end if
799  end subroutine diffusion_solve_vcoeff
800 
801  subroutine set_rhs(mg, f1, f2)
802  type(mg_t), intent(inout) :: mg
803  real(dp), intent(in) :: f1, f2
804  integer :: lvl, i, id, nc
805 
806  do lvl = 1, mg%highest_lvl
807  nc = mg%box_size_lvl(lvl)
808  do i = 1, size(mg%lvls(lvl)%my_leaves)
809  id = mg%lvls(lvl)%my_leaves(i)
810  mg%boxes(id)%cc(1:nc, 1:nc, 1:nc, mg_irhs) = &
811  f1 * mg%boxes(id)%cc(1:nc, 1:nc, 1:nc, mg_iphi) + &
812  f2 * mg%boxes(id)%cc(1:nc, 1:nc, 1:nc, mg_irhs)
813  end do
814  end do
815  end subroutine set_rhs
816 
817  !! File ../src/m_laplacian.f90
818 
819  subroutine laplacian_set_methods(mg)
820  type(mg_t), intent(inout) :: mg
821 
822  if (all(mg%periodic)) then
823  ! For a fully periodic Laplacian, remove the mean from the rhs and phi so
824  ! that a unique and periodic solution can be found
825  mg%subtract_mean = .true.
826  end if
827 
828  select case (mg%geometry_type)
829  case (mg_cartesian)
830  mg%box_op => box_lpl
831 
832  select case (mg%smoother_type)
833  ! case (smoother_jacobi)
834  ! mg%box_smoother => box_jacobi_lpl
836  mg%box_smoother => box_gs_lpl
837  case default
838  error stop "laplacian_set_methods: unsupported smoother type"
839  end select
840  case default
841  error stop "laplacian_set_methods: unsupported geometry"
842  end select
843 
844  end subroutine laplacian_set_methods
845 
846  !> Perform Gauss-Seidel relaxation on box for a Laplacian operator
847  subroutine box_gs_lpl(mg, id, nc, redblack_cntr)
848  type(mg_t), intent(inout) :: mg
849  integer, intent(in) :: id
850  integer, intent(in) :: nc
851  integer, intent(in) :: redblack_cntr !< Iteration counter
852  integer :: i, j, k, i0, di
853  real(dp) :: idr2(3), fac
854  logical :: redblack
855  real(dp), parameter :: sixth = 1/6.0_dp
856 
857  idr2 = 1/mg%dr(:, mg%boxes(id)%lvl)**2
858  fac = 0.5_dp / sum(idr2)
859  i0 = 1
860 
861  redblack = (mg%smoother_type == mg_smoother_gsrb)
862  if (redblack) then
863  di = 2
864  else
865  di = 1
866  end if
867 
868  ! The parity of redblack_cntr determines which cells we use. If
869  ! redblack_cntr is even, we use the even cells and vice versa.
870  associate(cc => mg%boxes(id)%cc, n => mg_iphi)
871  do k = 1, nc
872  do j = 1, nc
873  if (redblack) &
874  i0 = 2 - iand(ieor(redblack_cntr, k+j), 1)
875  do i = i0, nc, di
876  cc(i, j, k, n) = fac * ( &
877  idr2(1) * (cc(i+1, j, k, n) + cc(i-1, j, k, n)) + &
878  idr2(2) * (cc(i, j+1, k, n) + cc(i, j-1, k, n)) + &
879  idr2(3) * (cc(i, j, k+1, n) + cc(i, j, k-1, n)) - &
880  cc(i, j, k, mg_irhs))
881  end do
882  end do
883  end do
884  end associate
885  end subroutine box_gs_lpl
886 
887 ! !> Perform Jacobi relaxation on box for a Laplacian operator
888 ! subroutine box_jacobi_lpl(mg, id, nc, cntr)
889 ! type(mg_t), intent(inout) :: mg
890 ! integer, intent(in) :: id
891 ! integer, intent(in) :: nc
892 ! integer, intent(in) :: cntr !< Not used
893 ! integer :: i, j, k
894 ! real(dp), parameter :: w = 2.0_dp / 3
895 ! real(dp) :: tmp(0:nc+1, 0:nc+1, 0:nc+1)
896 ! real(dp) :: dr2
897 ! #if 3 == 3
898 ! real(dp), parameter :: sixth = 1/6.0_dp
899 ! #endif
900 
901 ! dr2 = mg%dr(mg%boxes(id)%lvl)**2
902 
903 ! associate (box => mg%boxes(id))
904 ! tmp = box%cc(:, :, :, mg_iphi)
905 ! do j=1, nc; do i=1, nc; do k=1, nc
906 ! #if 3 == 2
907 ! box%cc(i, j, mg_iphi) = (1-w) * box%cc(i, j, mg_iphi) + &
908 ! 0.25_dp * w * ( &
909 ! tmp(i+1, j) + tmp(i-1, j) + &
910 ! tmp(i, j+1) + tmp(i, j-1) - &
911 ! dr2 * box%cc(i, j, mg_irhs))
912 ! #elif 3 == 3
913 ! box%cc(i, j, k, mg_iphi) = (1-w) * &
914 ! box%cc(i, j, k, mg_iphi) + &
915 ! sixth * w * ( &
916 ! tmp(i+1, j, k) + tmp(i-1, j, k) + &
917 ! tmp(i, j+1, k) + tmp(i, j-1, k) + &
918 ! tmp(i, j, k+1) + tmp(i, j, k-1) - &
919 ! dr2 * box%cc(i, j, k, mg_irhs))
920 ! #endif
921 ! end do; end do; end do
922 ! end associate
923 ! end subroutine box_jacobi_lpl
924 
925  !> Perform Laplacian operator on a box
926  subroutine box_lpl(mg, id, nc, i_out)
927  type(mg_t), intent(inout) :: mg
928  integer, intent(in) :: id
929  integer, intent(in) :: nc
930  integer, intent(in) :: i_out !< Index of variable to store Laplacian in
931  integer :: i, j, k
932  real(dp) :: idr2(3)
933 
934  idr2 = 1 / mg%dr(:, mg%boxes(id)%lvl)**2
935 
936  associate(cc => mg%boxes(id)%cc, n => mg_iphi)
937  do k = 1, nc
938  do j = 1, nc
939  do i = 1, nc
940  cc(i, j, k, i_out) = &
941  idr2(1) * (cc(i-1, j, k, n) + cc(i+1, j, k, n) &
942  - 2 * cc(i, j, k, n)) &
943  + idr2(2) * (cc(i, j-1, k, n) + cc(i, j+1, k, n) &
944  - 2 * cc(i, j, k, n)) &
945  + idr2(3) * (cc(i, j, k-1, n) + cc(i, j, k+1, n) &
946  - 2 * cc(i, j, k, n))
947  end do
948  end do
949  end do
950  end associate
951  end subroutine box_lpl
952 
953 
954  !! File ../src/m_vhelmholtz.f90
955 
956  subroutine vhelmholtz_set_methods(mg)
957  type(mg_t), intent(inout) :: mg
958 
959  if (mg%n_extra_vars == 0 .and. mg%is_allocated) then
960  error stop "vhelmholtz_set_methods: mg%n_extra_vars == 0"
961  else
962  mg%n_extra_vars = max(1, mg%n_extra_vars)
963  end if
964 
965  mg%subtract_mean = .false.
966 
967  ! Use Neumann zero boundary conditions for the variable coefficient, since
968  ! it is needed in ghost cells.
969  mg%bc(:, mg_iveps)%bc_type = mg_bc_neumann
970  mg%bc(:, mg_iveps)%bc_value = 0.0_dp
971 
972  select case (mg%geometry_type)
973  case (mg_cartesian)
974  mg%box_op => box_vhelmh
975 
976  select case (mg%smoother_type)
978  mg%box_smoother => box_gs_vhelmh
979  case default
980  error stop "vhelmholtz_set_methods: unsupported smoother type"
981  end select
982  case default
983  error stop "vhelmholtz_set_methods: unsupported geometry"
984  end select
985 
986  end subroutine vhelmholtz_set_methods
987 
988  subroutine vhelmholtz_set_lambda(lambda)
989  real(dp), intent(in) :: lambda
990 
991  if (lambda < 0) &
992  error stop "vhelmholtz_set_lambda: lambda < 0 not allowed"
993 
994  vhelmholtz_lambda = lambda
995  end subroutine vhelmholtz_set_lambda
996 
997  !> Perform Gauss-Seidel relaxation on box for a Helmholtz operator
998  subroutine box_gs_vhelmh(mg, id, nc, redblack_cntr)
999  type(mg_t), intent(inout) :: mg
1000  integer, intent(in) :: id
1001  integer, intent(in) :: nc
1002  integer, intent(in) :: redblack_cntr !< Iteration counter
1003  integer :: i, j, k, i0, di
1004  logical :: redblack
1005  real(dp) :: idr2(2*3), u(2*3)
1006  real(dp) :: a0, a(2*3), c(2*3)
1007 
1008  ! Duplicate 1/dr^2 array to multiply neighbor values
1009  idr2(1:2*3:2) = 1/mg%dr(:, mg%boxes(id)%lvl)**2
1010  idr2(2:2*3:2) = idr2(1:2*3:2)
1011  i0 = 1
1012 
1013  redblack = (mg%smoother_type == mg_smoother_gsrb)
1014  if (redblack) then
1015  di = 2
1016  else
1017  di = 1
1018  end if
1019 
1020  ! The parity of redblack_cntr determines which cells we use. If
1021  ! redblack_cntr is even, we use the even cells and vice versa.
1022  associate(cc => mg%boxes(id)%cc, n => mg_iphi, &
1023  i_eps => mg_iveps)
1024  do k = 1, nc
1025  do j = 1, nc
1026  if (redblack) &
1027  i0 = 2 - iand(ieor(redblack_cntr, k+j), 1)
1028  do i = i0, nc, di
1029  a0 = cc(i, j, k, i_eps)
1030  u(1:2) = cc(i-1:i+1:2, j, k, n)
1031  a(1:2) = cc(i-1:i+1:2, j, k, i_eps)
1032  u(3:4) = cc(i, j-1:j+1:2, k, n)
1033  a(3:4) = cc(i, j-1:j+1:2, k, i_eps)
1034  u(5:6) = cc(i, j, k-1:k+1:2, n)
1035  a(5:6) = cc(i, j, k-1:k+1:2, i_eps)
1036  c(:) = 2 * a0 * a(:) / (a0 + a(:)) * idr2
1037 
1038  cc(i, j, k, n) = (sum(c(:) * u(:)) - &
1039  cc(i, j, k, mg_irhs)) / &
1040  (sum(c(:)) + vhelmholtz_lambda)
1041  end do
1042  end do
1043  end do
1044  end associate
1045  end subroutine box_gs_vhelmh
1046 
1047  !> Perform Helmholtz operator on a box
1048  subroutine box_vhelmh(mg, id, nc, i_out)
1049  type(mg_t), intent(inout) :: mg
1050  integer, intent(in) :: id
1051  integer, intent(in) :: nc
1052  integer, intent(in) :: i_out !< Index of variable to store Helmholtz in
1053  integer :: i, j, k
1054  real(dp) :: idr2(2*3), a0, u0, u(2*3), a(2*3)
1055 
1056  ! Duplicate 1/dr^2 array to multiply neighbor values
1057  idr2(1:2*3:2) = 1/mg%dr(:, mg%boxes(id)%lvl)**2
1058  idr2(2:2*3:2) = idr2(1:2*3:2)
1059 
1060  associate(cc => mg%boxes(id)%cc, n => mg_iphi, &
1061  i_eps => mg_iveps)
1062  do k = 1, nc
1063  do j = 1, nc
1064  do i = 1, nc
1065  u0 = cc(i, j, k, n)
1066  a0 = cc(i, j, k, i_eps)
1067  u(1:2) = cc(i-1:i+1:2, j, k, n)
1068  u(3:4) = cc(i, j-1:j+1:2, k, n)
1069  u(5:6) = cc(i, j, k-1:k+1:2, n)
1070  a(1:2) = cc(i-1:i+1:2, j, k, i_eps)
1071  a(3:4) = cc(i, j-1:j+1:2, k, i_eps)
1072  a(5:6) = cc(i, j, k-1:k+1:2, i_eps)
1073 
1074  cc(i, j, k, i_out) = sum(2 * idr2 * &
1075  a0*a(:)/(a0 + a(:)) * (u(:) - u0)) - &
1076  vhelmholtz_lambda * u0
1077  end do
1078  end do
1079  end do
1080  end associate
1081  end subroutine box_vhelmh
1082 
1083  !! File ../src/m_build_tree.f90
1084 
1085  subroutine mg_build_rectangle(mg, domain_size, box_size, dx, r_min, &
1086  periodic, n_finer)
1087  type(mg_t), intent(inout) :: mg
1088  integer, intent(in) :: domain_size(3)
1089  integer, intent(in) :: box_size
1090  real(dp), intent(in) :: dx(3)
1091  real(dp), intent(in) :: r_min(3)
1092  logical, intent(in) :: periodic(3)
1093  integer, intent(in) :: n_finer
1094  integer :: i, j, k, lvl, n, id, nx(3)
1095  integer :: boxes_per_dim(3, mg_lvl_lo:1)
1096  integer :: periodic_offset(3)
1097 
1098  if (modulo(box_size, 2) /= 0) &
1099  error stop "box_size should be even"
1100  if (any(modulo(domain_size, box_size) /= 0)) &
1101  error stop "box_size does not divide domain_size"
1102 
1103  nx = domain_size
1104  mg%box_size = box_size
1105  mg%box_size_lvl(1) = box_size
1106  mg%domain_size_lvl(:, 1) = domain_size
1107  mg%first_normal_lvl = 1
1108  mg%dr(:, 1) = dx
1109  mg%r_min(:) = r_min
1110  mg%periodic = periodic
1111  boxes_per_dim(:, :) = 0
1112  boxes_per_dim(:, 1) = domain_size / box_size
1113 
1114  do lvl = 1, mg_lvl_lo+1, -1
1115  ! For a Gauss-Seidel (non red-black) smoother, we should avoid boxes
1116  ! containing a single cell
1117  if (any(modulo(nx, 2) == 1 .or. nx == mg%coarsest_grid .or. &
1118  (mg%box_size_lvl(lvl) == mg%coarsest_grid .and. &
1119  mg%smoother_type == mg_smoother_gs))) exit
1120 
1121  if (all(modulo(nx/mg%box_size_lvl(lvl), 2) == 0)) then
1122  mg%box_size_lvl(lvl-1) = mg%box_size_lvl(lvl)
1123  boxes_per_dim(:, lvl-1) = boxes_per_dim(:, lvl)/2
1124  mg%first_normal_lvl = lvl-1
1125  else
1126  mg%box_size_lvl(lvl-1) = mg%box_size_lvl(lvl)/2
1127  boxes_per_dim(:, lvl-1) = boxes_per_dim(:, lvl)
1128  end if
1129 
1130  mg%dr(:, lvl-1) = mg%dr(:, lvl) * 2
1131  nx = nx / 2
1132  mg%domain_size_lvl(:, lvl-1) = nx
1133  end do
1134 
1135  mg%lowest_lvl = lvl
1136  mg%highest_lvl = 1
1137 
1138  do lvl = 2, mg_lvl_hi
1139  mg%dr(:, lvl) = mg%dr(:, lvl-1) * 0.5_dp
1140  mg%box_size_lvl(lvl) = box_size
1141  mg%domain_size_lvl(:, lvl) = 2 * mg%domain_size_lvl(:, lvl-1)
1142  end do
1143 
1144  n = sum(product(boxes_per_dim, dim=1)) + n_finer
1145  allocate(mg%boxes(n))
1146 
1147  ! Create lowest level
1148  nx = boxes_per_dim(:, mg%lowest_lvl)
1149  periodic_offset = [nx(1)-1, (nx(2)-1)*nx(1), &
1150  (nx(3)-1) * nx(2) * nx(1)]
1151 
1152  do k=1,nx(3); do j=1,nx(2); do i=1,nx(1)
1153  mg%n_boxes = mg%n_boxes + 1
1154  n = mg%n_boxes
1155 
1156  mg%boxes(n)%rank = 0
1157  mg%boxes(n)%id = n
1158  mg%boxes(n)%lvl = mg%lowest_lvl
1159  mg%boxes(n)%ix(:) = [i, j, k]
1160  mg%boxes(n)%r_min(:) = r_min + (mg%boxes(n)%ix(:) - 1) * &
1161  mg%box_size_lvl(mg%lowest_lvl) * mg%dr(:, mg%lowest_lvl)
1162  mg%boxes(n)%dr(:) = mg%dr(:, mg%lowest_lvl)
1163  mg%boxes(n)%parent = mg_no_box
1164  mg%boxes(n)%children(:) = mg_no_box
1165 
1166  ! Set default neighbors
1167  mg%boxes(n)%neighbors(:) = [n-1, n+1, n-nx(1), n+nx(1), &
1168  n-nx(1)*nx(2), n+nx(1)*nx(2)]
1169 
1170  ! Handle boundaries
1171  where ([i, j, k] == 1 .and. .not. periodic)
1172  mg%boxes(n)%neighbors(1:mg_num_neighbors:2) = &
1174  end where
1175  where ([i, j, k] == 1 .and. periodic)
1176  mg%boxes(n)%neighbors(1:mg_num_neighbors:2) = &
1177  n + periodic_offset
1178  end where
1179 
1180  where ([i, j, k] == nx .and. .not. periodic)
1181  mg%boxes(n)%neighbors(2:mg_num_neighbors:2) = &
1183  end where
1184  where ([i, j, k] == nx .and. periodic)
1185  mg%boxes(n)%neighbors(2:mg_num_neighbors:2) = &
1186  n - periodic_offset
1187  end where
1188  end do; end do; end do
1189 
1190  mg%lvls(mg%lowest_lvl)%ids = [(n, n=1, mg%n_boxes)]
1191 
1192  ! Add higher levels
1193  do lvl = mg%lowest_lvl, 0
1194  if (mg%box_size_lvl(lvl+1) == mg%box_size_lvl(lvl)) then
1195  do i = 1, size(mg%lvls(lvl)%ids)
1196  id = mg%lvls(lvl)%ids(i)
1197  call mg_add_children(mg, id)
1198  end do
1199 
1200  call mg_set_leaves_parents(mg%boxes, mg%lvls(lvl))
1201  call mg_set_next_level_ids(mg, lvl)
1202  call mg_set_neighbors_lvl(mg, lvl+1)
1203  else
1204  do i = 1, size(mg%lvls(lvl)%ids)
1205  id = mg%lvls(lvl)%ids(i)
1206  call add_single_child(mg, id, size(mg%lvls(lvl)%ids))
1207  end do
1208 
1209  call mg_set_leaves_parents(mg%boxes, mg%lvls(lvl))
1210  call mg_set_next_level_ids(mg, lvl)
1211  end if
1212  end do
1213 
1214  call mg_set_leaves_parents(mg%boxes, mg%lvls(1))
1215 
1216  ! No refinement boundaries
1217  do lvl = mg%lowest_lvl, 1
1218  if (allocated(mg%lvls(lvl)%ref_bnds)) &
1219  deallocate(mg%lvls(lvl)%ref_bnds)
1220  allocate(mg%lvls(lvl)%ref_bnds(0))
1221  end do
1222 
1223  mg%tree_created = .true.
1224  end subroutine mg_build_rectangle
1225 
1226  subroutine mg_set_neighbors_lvl(mg, lvl)
1227  type(mg_t), intent(inout) :: mg
1228  integer, intent(in) :: lvl
1229  integer :: i, id
1230 
1231  do i = 1, size(mg%lvls(lvl)%ids)
1232  id = mg%lvls(lvl)%ids(i)
1233  call set_neighbs(mg%boxes, id)
1234  end do
1235  end subroutine mg_set_neighbors_lvl
1236 
1237  subroutine mg_set_next_level_ids(mg, lvl)
1238  type(mg_t), intent(inout) :: mg
1239  integer, intent(in) :: lvl
1240  integer :: n, i, id
1241 
1242  if (allocated(mg%lvls(lvl+1)%ids)) &
1243  deallocate(mg%lvls(lvl+1)%ids)
1244 
1245  ! Set next level ids to children of this level
1246  if (mg%box_size_lvl(lvl+1) == mg%box_size_lvl(lvl)) then
1247  n = mg_num_children * size(mg%lvls(lvl)%parents)
1248  allocate(mg%lvls(lvl+1)%ids(n))
1249 
1250  n = mg_num_children
1251  do i = 1, size(mg%lvls(lvl)%parents)
1252  id = mg%lvls(lvl)%parents(i)
1253  mg%lvls(lvl+1)%ids(n*(i-1)+1:n*i) = mg%boxes(id)%children
1254  end do
1255  else
1256  n = size(mg%lvls(lvl)%parents)
1257  allocate(mg%lvls(lvl+1)%ids(n))
1258 
1259  n = 1
1260  do i = 1, size(mg%lvls(lvl)%parents)
1261  id = mg%lvls(lvl)%parents(i)
1262  mg%lvls(lvl+1)%ids(i) = mg%boxes(id)%children(1)
1263  end do
1264  end if
1265 
1266  end subroutine mg_set_next_level_ids
1267 
1268  ! Set the neighbors of id (using their parent)
1269  subroutine set_neighbs(boxes, id)
1270  type(mg_box_t), intent(inout) :: boxes(:)
1271  integer, intent(in) :: id
1272  integer :: nb, nb_id
1273 
1274  do nb = 1, mg_num_neighbors
1275  if (boxes(id)%neighbors(nb) == mg_no_box) then
1276  nb_id = find_neighb(boxes, id, nb)
1277  if (nb_id > mg_no_box) then
1278  boxes(id)%neighbors(nb) = nb_id
1279  boxes(nb_id)%neighbors(mg_neighb_rev(nb)) = id
1280  end if
1281  end if
1282  end do
1283  end subroutine set_neighbs
1284 
1285  !> Get the id of neighbor nb of boxes(id), through its parent
1286  function find_neighb(boxes, id, nb) result(nb_id)
1287  type(mg_box_t), intent(in) :: boxes(:) !< List with all the boxes
1288  integer, intent(in) :: id !< Box whose neighbor we are looking for
1289  integer, intent(in) :: nb !< Neighbor index
1290  integer :: nb_id, p_id, c_ix, d, old_pid
1291 
1292  p_id = boxes(id)%parent
1293  old_pid = p_id
1294  c_ix = mg_ix_to_ichild(boxes(id)%ix)
1295  d = mg_neighb_dim(nb)
1296 
1297  ! Check if neighbor is in same direction as ix is (low/high). If so,
1298  ! use neighbor of parent
1299  if (mg_child_low(d, c_ix) .eqv. mg_neighb_low(nb)) then
1300  p_id = boxes(p_id)%neighbors(nb)
1301  end if
1302 
1303  ! The child ix of the neighbor is reversed in direction d
1304  nb_id = boxes(p_id)%children(mg_child_rev(c_ix, d))
1305  end function find_neighb
1306 
1307  !> Create a list of leaves and a list of parents for a level
1308  subroutine mg_set_leaves_parents(boxes, level)
1309  type(mg_box_t), intent(in) :: boxes(:) !< List of boxes
1310  type(mg_lvl_t), intent(inout) :: level !< Level type which contains the indices of boxes
1311  integer :: i, id, i_leaf, i_parent
1312  integer :: n_parents, n_leaves
1313 
1314  n_parents = count(mg_has_children(boxes(level%ids)))
1315  n_leaves = size(level%ids) - n_parents
1316 
1317  if (.not. allocated(level%parents)) then
1318  allocate(level%parents(n_parents))
1319  else if (n_parents /= size(level%parents)) then
1320  deallocate(level%parents)
1321  allocate(level%parents(n_parents))
1322  end if
1323 
1324  if (.not. allocated(level%leaves)) then
1325  allocate(level%leaves(n_leaves))
1326  else if (n_leaves /= size(level%leaves)) then
1327  deallocate(level%leaves)
1328  allocate(level%leaves(n_leaves))
1329  end if
1330 
1331  i_leaf = 0
1332  i_parent = 0
1333  do i = 1, size(level%ids)
1334  id = level%ids(i)
1335  if (mg_has_children(boxes(id))) then
1336  i_parent = i_parent + 1
1337  level%parents(i_parent) = id
1338  else
1339  i_leaf = i_leaf + 1
1340  level%leaves(i_leaf) = id
1341  end if
1342  end do
1343  end subroutine mg_set_leaves_parents
1344 
1345  !> Create a list of refinement boundaries (from the coarse side)
1346  subroutine mg_set_refinement_boundaries(boxes, level)
1347  type(mg_box_t), intent(in) :: boxes(:)
1348  type(mg_lvl_t), intent(inout) :: level
1349  integer, allocatable :: tmp(:)
1350  integer :: i, id, nb, nb_id, ix
1351 
1352  if (allocated(level%ref_bnds)) deallocate(level%ref_bnds)
1353 
1354  if (size(level%parents) == 0) then
1355  ! There are no refinement boundaries
1356  allocate(level%ref_bnds(0))
1357  else
1358  allocate(tmp(size(level%leaves)))
1359  ix = 0
1360  do i = 1, size(level%leaves)
1361  id = level%leaves(i)
1362 
1363  do nb = 1, mg_num_neighbors
1364  nb_id = boxes(id)%neighbors(nb)
1365  if (nb_id > mg_no_box) then
1366  if (mg_has_children(boxes(nb_id))) then
1367  ix = ix + 1
1368  tmp(ix) = id
1369  exit
1370  end if
1371  end if
1372  end do
1373  end do
1374 
1375  allocate(level%ref_bnds(ix))
1376  level%ref_bnds(:) = tmp(1:ix)
1377  end if
1378  end subroutine mg_set_refinement_boundaries
1379 
1380  subroutine mg_add_children(mg, id)
1381  type(mg_t), intent(inout) :: mg
1382  integer, intent(in) :: id !< Id of box that gets children
1383  integer :: lvl, i, nb, child_nb(2**(3-1))
1384  integer :: c_ids(mg_num_children)
1385  integer :: c_id, c_ix_base(3)
1386 
1387  if (mg%n_boxes + mg_num_children > size(mg%boxes)) then
1388  error stop "mg_add_children: not enough space"
1389  end if
1390 
1391  c_ids = [(mg%n_boxes+i, i=1,mg_num_children)]
1392  mg%n_boxes = mg%n_boxes + mg_num_children
1393  mg%boxes(id)%children = c_ids
1394  c_ix_base = 2 * mg%boxes(id)%ix - 1
1395  lvl = mg%boxes(id)%lvl+1
1396 
1397  do i = 1, mg_num_children
1398  c_id = c_ids(i)
1399  mg%boxes(c_id)%rank = mg%boxes(id)%rank
1400  mg%boxes(c_id)%ix = c_ix_base + mg_child_dix(:, i)
1401  mg%boxes(c_id)%lvl = lvl
1402  mg%boxes(c_id)%parent = id
1403  mg%boxes(c_id)%children = mg_no_box
1404  mg%boxes(c_id)%neighbors = mg_no_box
1405  mg%boxes(c_id)%r_min = mg%boxes(id)%r_min + &
1406  mg%dr(:, lvl) * mg_child_dix(:, i) * mg%box_size
1407  mg%boxes(c_id)%dr(:) = mg%dr(:, lvl)
1408  end do
1409 
1410  ! Set boundary conditions at children
1411  do nb = 1, mg_num_neighbors
1412  if (mg%boxes(id)%neighbors(nb) < mg_no_box) then
1413  child_nb = c_ids(mg_child_adj_nb(:, nb)) ! Neighboring children
1414  mg%boxes(child_nb)%neighbors(nb) = mg%boxes(id)%neighbors(nb)
1415  end if
1416  end do
1417  end subroutine mg_add_children
1418 
1419  subroutine add_single_child(mg, id, n_boxes_lvl)
1420  type(mg_t), intent(inout) :: mg
1421  integer, intent(in) :: id !< Id of box that gets children
1422  integer, intent(in) :: n_boxes_lvl
1423  integer :: lvl, c_id
1424 
1425  c_id = mg%n_boxes + 1
1426  mg%n_boxes = mg%n_boxes + 1
1427  mg%boxes(id)%children(1) = c_id
1428  lvl = mg%boxes(id)%lvl+1
1429 
1430  mg%boxes(c_id)%rank = mg%boxes(id)%rank
1431  mg%boxes(c_id)%ix = mg%boxes(id)%ix
1432  mg%boxes(c_id)%lvl = lvl
1433  mg%boxes(c_id)%parent = id
1434  mg%boxes(c_id)%children = mg_no_box
1435  where (mg%boxes(id)%neighbors == mg_physical_boundary)
1436  mg%boxes(c_id)%neighbors = mg%boxes(id)%neighbors
1437  elsewhere
1438  mg%boxes(c_id)%neighbors = mg%boxes(id)%neighbors + n_boxes_lvl
1439  end where
1440  mg%boxes(c_id)%r_min = mg%boxes(id)%r_min
1441  mg%boxes(c_id)%dr(:) = mg%dr(:, lvl)
1442 
1443  end subroutine add_single_child
1444 
1445  !! File ../src/m_load_balance.f90
1446 
1447  !> Load balance all boxes in the multigrid tree, by simply distributing the
1448  !> load per grid level. This method will only work well for uniform grids.
1449  !>
1450  !> Note that in a typical application the load balancing of the leaves is
1451  !> already determined, then mg_load_balance_parents can be used.
1452  subroutine mg_load_balance_simple(mg)
1453  type(mg_t), intent(inout) :: mg
1454  integer :: i, id, lvl, single_cpu_lvl
1455  integer :: work_left, my_work, i_cpu
1456 
1457  ! Up to this level, all boxes have to be on a single processor because they
1458  ! have a different size and the communication routines do not support this
1459  single_cpu_lvl = max(mg%first_normal_lvl-1, mg%lowest_lvl)
1460 
1461  do lvl = mg%lowest_lvl, single_cpu_lvl
1462  do i = 1, size(mg%lvls(lvl)%ids)
1463  id = mg%lvls(lvl)%ids(i)
1464  mg%boxes(id)%rank = 0
1465  end do
1466  end do
1467 
1468  ! Distribute the boxes equally. Due to the way the mesh is constructed, the
1469  ! mg%lvls(lvl)%ids array already contains a Morton-like ordering.
1470  do lvl = single_cpu_lvl+1, mg%highest_lvl
1471  work_left = size(mg%lvls(lvl)%ids)
1472  my_work = 0
1473  i_cpu = 0
1474 
1475  do i = 1, size(mg%lvls(lvl)%ids)
1476  if ((mg%n_cpu - i_cpu - 1) * my_work >= work_left) then
1477  i_cpu = i_cpu + 1
1478  my_work = 0
1479  end if
1480 
1481  my_work = my_work + 1
1482  work_left = work_left - 1
1483 
1484  id = mg%lvls(lvl)%ids(i)
1485  mg%boxes(id)%rank = i_cpu
1486  end do
1487  end do
1488 
1489  do lvl = mg%lowest_lvl, mg%highest_lvl
1490  call update_lvl_info(mg, mg%lvls(lvl))
1491  end do
1492 
1493  end subroutine mg_load_balance_simple
1494 
1495  !> Load balance all boxes in the multigrid tree. Compared to
1496  !> mg_load_balance_simple, this method does a better job of setting the ranks
1497  !> of parent boxes
1498  !>
1499  !> Note that in a typical application the load balancing of the leaves is
1500  !> already determined, then mg_load_balance_parents can be used.
1501  subroutine mg_load_balance(mg)
1502  type(mg_t), intent(inout) :: mg
1503  integer :: i, id, lvl, single_cpu_lvl
1504  integer :: work_left, my_work(0:mg%n_cpu), i_cpu
1505  integer :: c_ids(mg_num_children)
1506  integer :: c_ranks(mg_num_children)
1507  integer :: coarse_rank
1508 
1509  ! Up to this level, all boxes have to be on a single processor because they
1510  ! have a different size and the communication routines do not support this
1511  single_cpu_lvl = max(mg%first_normal_lvl-1, mg%lowest_lvl)
1512 
1513  ! Distribute the boxes equally. Due to the way the mesh is constructed, the
1514  ! mg%lvls(lvl)%ids array already contains a Morton-like ordering.
1515  do lvl = mg%highest_lvl, single_cpu_lvl+1, -1
1516  ! For parents determine the rank based on their child ranks
1517  my_work(:) = 0
1518 
1519  do i = 1, size(mg%lvls(lvl)%parents)
1520  id = mg%lvls(lvl)%parents(i)
1521 
1522  c_ids = mg%boxes(id)%children
1523  c_ranks = mg%boxes(c_ids)%rank
1524  i_cpu = most_popular(c_ranks, my_work, mg%n_cpu)
1525  mg%boxes(id)%rank = i_cpu
1526  my_work(i_cpu) = my_work(i_cpu) + 1
1527  end do
1528 
1529  work_left = size(mg%lvls(lvl)%leaves)
1530  i_cpu = 0
1531 
1532  do i = 1, size(mg%lvls(lvl)%leaves)
1533  ! Skip this CPU if it already has enough work
1534  if ((mg%n_cpu - i_cpu - 1) * my_work(i_cpu) >= &
1535  work_left + sum(my_work(i_cpu+1:))) then
1536  i_cpu = i_cpu + 1
1537  end if
1538 
1539  my_work(i_cpu) = my_work(i_cpu) + 1
1540  work_left = work_left - 1
1541 
1542  id = mg%lvls(lvl)%leaves(i)
1543  mg%boxes(id)%rank = i_cpu
1544  end do
1545  end do
1546 
1547  ! Determine most popular CPU for coarse grids
1548  if (single_cpu_lvl < mg%highest_lvl) then
1549  coarse_rank = most_popular(mg%boxes(&
1550  mg%lvls(single_cpu_lvl+1)%ids)%rank, my_work, mg%n_cpu)
1551  else
1552  coarse_rank = 0
1553  end if
1554 
1555  do lvl = mg%lowest_lvl, single_cpu_lvl
1556  do i = 1, size(mg%lvls(lvl)%ids)
1557  id = mg%lvls(lvl)%ids(i)
1558  mg%boxes(id)%rank = coarse_rank
1559  end do
1560  end do
1561 
1562  do lvl = mg%lowest_lvl, mg%highest_lvl
1563  call update_lvl_info(mg, mg%lvls(lvl))
1564  end do
1565 
1566  end subroutine mg_load_balance
1567 
1568  !> Load balance the parents (non-leafs). Assign them to the rank that has most
1569  !> children.
1570  subroutine mg_load_balance_parents(mg)
1571  type(mg_t), intent(inout) :: mg
1572  integer :: i, id, lvl
1573  integer :: c_ids(mg_num_children)
1574  integer :: c_ranks(mg_num_children)
1575  integer :: single_cpu_lvl, coarse_rank
1576  integer :: my_work(0:mg%n_cpu), i_cpu
1577 
1578  ! Up to this level, all boxes have to be on a single processor because they
1579  ! have a different size and the communication routines do not support this
1580  single_cpu_lvl = max(mg%first_normal_lvl-1, mg%lowest_lvl)
1581 
1582  do lvl = mg%highest_lvl-1, single_cpu_lvl+1, -1
1583  my_work(:) = 0
1584 
1585  ! Determine amount of work for the leaves
1586  do i = 1, size(mg%lvls(lvl)%leaves)
1587  id = mg%lvls(lvl)%leaves(i)
1588  i_cpu = mg%boxes(id)%rank
1589  my_work(i_cpu) = my_work(i_cpu) + 1
1590  end do
1591 
1592  do i = 1, size(mg%lvls(lvl)%parents)
1593  id = mg%lvls(lvl)%parents(i)
1594 
1595  c_ids = mg%boxes(id)%children
1596  c_ranks = mg%boxes(c_ids)%rank
1597  i_cpu = most_popular(c_ranks, my_work, mg%n_cpu)
1598  mg%boxes(id)%rank = i_cpu
1599  my_work(i_cpu) = my_work(i_cpu) + 1
1600  end do
1601 
1602  end do
1603 
1604  ! Determine most popular CPU for coarse grids
1605  if (single_cpu_lvl < mg%highest_lvl) then
1606  coarse_rank = most_popular(mg%boxes(&
1607  mg%lvls(single_cpu_lvl+1)%ids)%rank, my_work, mg%n_cpu)
1608  else
1609  coarse_rank = 0
1610  end if
1611 
1612  do lvl = mg%lowest_lvl, single_cpu_lvl
1613  do i = 1, size(mg%lvls(lvl)%ids)
1614  id = mg%lvls(lvl)%ids(i)
1615  mg%boxes(id)%rank = coarse_rank
1616  end do
1617  end do
1618 
1619  do lvl = mg%lowest_lvl, mg%highest_lvl
1620  call update_lvl_info(mg, mg%lvls(lvl))
1621  end do
1622 
1623  end subroutine mg_load_balance_parents
1624 
1625  !> Determine most popular rank in the list. In case of ties, assign the rank
1626  !> with the least work.
1627  pure integer function most_popular(list, work, n_cpu)
1628  integer, intent(in) :: list(:) !< List of MPI ranks
1629  integer, intent(in) :: n_cpu
1630  integer, intent(in) :: work(0:n_cpu-1) !< Existing work per rank
1631  integer :: i, best_count, current_count
1632  integer :: my_work, best_work
1633 
1634  best_count = 0
1635  best_work = 0
1636  most_popular = -1
1637 
1638  do i = 1, size(list)
1639  current_count = count(list == list(i))
1640  my_work = work(list(i))
1641 
1642  ! In case of ties, select task with lowest work
1643  if (current_count > best_count .or. &
1644  current_count == best_count .and. my_work < best_work) then
1645  best_count = current_count
1646  best_work = my_work
1647  most_popular = list(i)
1648  end if
1649  end do
1650 
1651  end function most_popular
1652 
1653  subroutine update_lvl_info(mg, lvl)
1654  type(mg_t), intent(inout) :: mg
1655  type(mg_lvl_t), intent(inout) :: lvl
1656 
1657  lvl%my_ids = pack(lvl%ids, &
1658  mg%boxes(lvl%ids)%rank == mg%my_rank)
1659  lvl%my_leaves = pack(lvl%leaves, &
1660  mg%boxes(lvl%leaves)%rank == mg%my_rank)
1661  lvl%my_parents = pack(lvl%parents, &
1662  mg%boxes(lvl%parents)%rank == mg%my_rank)
1663  lvl%my_ref_bnds = pack(lvl%ref_bnds, &
1664  mg%boxes(lvl%ref_bnds)%rank == mg%my_rank)
1665  end subroutine update_lvl_info
1666 
1667  !! File ../src/m_vlaplacian.f90
1668 
1669  subroutine vlaplacian_set_methods(mg)
1670  type(mg_t), intent(inout) :: mg
1671 
1672  if (mg%n_extra_vars == 0 .and. mg%is_allocated) then
1673  error stop "vlaplacian_set_methods: mg%n_extra_vars == 0"
1674  else
1675  mg%n_extra_vars = max(1, mg%n_extra_vars)
1676  end if
1677 
1678  mg%subtract_mean = .false.
1679 
1680  ! Use Neumann zero boundary conditions for the variable coefficient, since
1681  ! it is needed in ghost cells.
1682  mg%bc(:, mg_iveps)%bc_type = mg_bc_neumann
1683  mg%bc(:, mg_iveps)%bc_value = 0.0_dp
1684 
1685  select case (mg%geometry_type)
1686  case (mg_cartesian)
1687  mg%box_op => box_vlpl
1688 
1689  ! TODO: test whether a custom prolongation works better
1690  ! mg%box_prolong => vlpl_prolong
1691 
1692  select case (mg%smoother_type)
1694  mg%box_smoother => box_gs_vlpl
1695  case default
1696  error stop "vlaplacian_set_methods: unsupported smoother type"
1697  end select
1698 
1699  case default
1700  error stop "vlaplacian_set_methods: unsupported geometry"
1701  end select
1702 
1703  end subroutine vlaplacian_set_methods
1704 
1705  !> Perform Gauss-Seidel relaxation on box for a Laplacian operator
1706  subroutine box_gs_vlpl(mg, id, nc, redblack_cntr)
1707  type(mg_t), intent(inout) :: mg
1708  integer, intent(in) :: id
1709  integer, intent(in) :: nc
1710  integer, intent(in) :: redblack_cntr !< Iteration counter
1711  integer :: i, j, k, i0, di
1712  logical :: redblack
1713  real(dp) :: idr2(2*3), u(2*3)
1714  real(dp) :: a0, a(2*3), c(2*3)
1715 
1716  ! Duplicate 1/dr^2 array to multiply neighbor values
1717  idr2(1:2*3:2) = 1/mg%dr(:, mg%boxes(id)%lvl)**2
1718  idr2(2:2*3:2) = idr2(1:2*3:2)
1719  i0 = 1
1720 
1721  redblack = (mg%smoother_type == mg_smoother_gsrb)
1722  if (redblack) then
1723  di = 2
1724  else
1725  di = 1
1726  end if
1727 
1728  ! The parity of redblack_cntr determines which cells we use. If
1729  ! redblack_cntr is even, we use the even cells and vice versa.
1730  associate(cc => mg%boxes(id)%cc, n => mg_iphi, &
1731  i_eps => mg_iveps)
1732  do k = 1, nc
1733  do j = 1, nc
1734  if (redblack) &
1735  i0 = 2 - iand(ieor(redblack_cntr, k+j), 1)
1736  do i = i0, nc, di
1737  a0 = cc(i, j, k, i_eps)
1738  u(1:2) = cc(i-1:i+1:2, j, k, n)
1739  a(1:2) = cc(i-1:i+1:2, j, k, i_eps)
1740  u(3:4) = cc(i, j-1:j+1:2, k, n)
1741  a(3:4) = cc(i, j-1:j+1:2, k, i_eps)
1742  u(5:6) = cc(i, j, k-1:k+1:2, n)
1743  a(5:6) = cc(i, j, k-1:k+1:2, i_eps)
1744  c(:) = 2 * a0 * a(:) / (a0 + a(:)) * idr2
1745 
1746  cc(i, j, k, n) = (sum(c(:) * u(:)) - &
1747  cc(i, j, k, mg_irhs)) / sum(c(:))
1748  end do
1749  end do
1750  end do
1751  end associate
1752  end subroutine box_gs_vlpl
1753 
1754  !> Perform Laplacian operator on a box
1755  subroutine box_vlpl(mg, id, nc, i_out)
1756  type(mg_t), intent(inout) :: mg
1757  integer, intent(in) :: id
1758  integer, intent(in) :: nc
1759  integer, intent(in) :: i_out !< Index of variable to store Laplacian in
1760  integer :: i, j, k
1761  real(dp) :: idr2(2*3), a0, u0, u(2*3), a(2*3)
1762 
1763  ! Duplicate 1/dr^2 array to multiply neighbor values
1764  idr2(1:2*3:2) = 1/mg%dr(:, mg%boxes(id)%lvl)**2
1765  idr2(2:2*3:2) = idr2(1:2*3:2)
1766 
1767  associate(cc => mg%boxes(id)%cc, n => mg_iphi, &
1768  i_eps => mg_iveps)
1769  do k = 1, nc
1770  do j = 1, nc
1771  do i = 1, nc
1772  u0 = cc(i, j, k, n)
1773  a0 = cc(i, j, k, i_eps)
1774  u(1:2) = cc(i-1:i+1:2, j, k, n)
1775  u(3:4) = cc(i, j-1:j+1:2, k, n)
1776  u(5:6) = cc(i, j, k-1:k+1:2, n)
1777  a(1:2) = cc(i-1:i+1:2, j, k, i_eps)
1778  a(3:4) = cc(i, j-1:j+1:2, k, i_eps)
1779  a(5:6) = cc(i, j, k-1:k+1:2, i_eps)
1780 
1781  cc(i, j, k, i_out) = sum(2 * idr2 * &
1782  a0*a(:)/(a0 + a(:)) * (u(:) - u0))
1783  end do
1784  end do
1785  end do
1786  end associate
1787  end subroutine box_vlpl
1788 
1789  !> Prolong from a parent to a child with index offset dix. This method could
1790  !> sometimes work better than the default prolongation, which does not take
1791  !> the variation in epsilon into account
1792 ! subroutine vlpl_prolong(mg, p_id, dix, nc, iv, fine)
1793 ! type(mg_t), intent(inout) :: mg
1794 ! integer, intent(in) :: p_id !< Id of parent
1795 ! integer, intent(in) :: dix(3) !< Offset of child in parent grid
1796 ! integer, intent(in) :: nc !< Child grid size
1797 ! integer, intent(in) :: iv !< Prolong from this variable
1798 ! real(dp), intent(out) :: fine(nc, nc, nc) !< Prolonged values
1799 
1800 ! integer :: i, j, k, hnc
1801 ! #if 3 == 2
1802 ! integer :: ic, jc
1803 ! real(dp) :: f0, flx, fhx, fly, fhy
1804 ! real(dp) :: c0, clx, chx, cly, chy
1805 ! #elif 3 == 3
1806 ! integer :: ic, jc, kc
1807 ! real(dp) :: f0, flx, fhx, fly, fhy, flz, fhz
1808 ! real(dp) :: c0, clx, chx, cly, chy, clz, chz
1809 ! #endif
1810 
1811 ! hnc = nc/2
1812 
1813 ! associate (crs => mg%boxes(p_id)%cc, &
1814 ! i_eps => mg_iveps)
1815 ! #if 3 == 2
1816 ! do j = 1, hnc
1817 ! jc = j + dix(2)
1818 ! do i = 1, hnc
1819 ! ic = i + dix(1)
1820 
1821 ! f0 = crs(ic, jc, iv)
1822 ! flx = crs(ic-1, jc, iv)
1823 ! fhx = crs(ic+1, jc, iv)
1824 ! fly = crs(ic, jc-1, iv)
1825 ! fhy = crs(ic, jc+1, iv)
1826 
1827 ! c0 = 2 * crs(ic, jc, i_eps)
1828 ! clx = crs(ic-1, jc, i_eps)
1829 ! chx = crs(ic+1, jc, i_eps)
1830 ! cly = crs(ic, jc-1, i_eps)
1831 ! chy = crs(ic, jc+1, i_eps)
1832 
1833 ! fine(2*i-1, 2*j-1) = (f0*c0 + flx*clx + fly*cly) / &
1834 ! (c0 + clx + cly)
1835 ! fine(2*i , 2*j-1) = (f0*c0 + fhx*chx + fly*cly) / &
1836 ! (c0 + chx + cly)
1837 ! fine(2*i-1, 2*j) = (f0*c0 + flx*clx + fhy*chy) / &
1838 ! (c0 + clx + chy)
1839 ! fine(2*i , 2*j) = (f0*c0 + fhx*chx + fhy*chy) / &
1840 ! (c0 + chx + chy)
1841 ! end do
1842 ! end do
1843 ! #elif 3 == 3
1844 ! do k = 1, hnc
1845 ! kc = k + dix(3)
1846 ! do j = 1, hnc
1847 ! jc = j + dix(2)
1848 ! do i = 1, hnc
1849 ! ic = i + dix(1)
1850 
1851 ! f0 = crs(ic, jc, kc, iv)
1852 ! flx = crs(ic-1, jc, kc, iv)
1853 ! fhx = crs(ic+1, jc, kc, iv)
1854 ! fly = crs(ic, jc-1, kc, iv)
1855 ! fhy = crs(ic, jc+1, kc, iv)
1856 ! flz = crs(ic, jc, kc-1, iv)
1857 ! fhz = crs(ic, jc, kc+1, iv)
1858 
1859 ! c0 = crs(ic, jc, kc, i_eps)
1860 ! clx = crs(ic-1, jc, kc, i_eps)
1861 ! chx = crs(ic+1, jc, kc, i_eps)
1862 ! cly = crs(ic, jc-1, kc, i_eps)
1863 ! chy = crs(ic, jc+1, kc, i_eps)
1864 ! clz = crs(ic, jc, kc-1, i_eps)
1865 ! chz = crs(ic, jc, kc+1, i_eps)
1866 
1867 ! fine(2*i-1, 2*j-1, 2*k-1) = (f0*c0 + flx*clx + fly*cly + flz*clz) / &
1868 ! (c0 + clx + cly + clz)
1869 ! fine(2*i, 2*j-1, 2*k-1) = (f0*c0 + fhx*chx + fly*cly + flz*clz) / &
1870 ! (c0 + chx + cly + clz)
1871 ! fine(2*i-1, 2*j, 2*k-1) = (f0*c0 + flx*clx + fhy*chy + flz*clz) / &
1872 ! (c0 + clx + chy + clz)
1873 ! fine(2*i, 2*j, 2*k-1) = (f0*c0 + fhx*chx + fhy*chy + flz*clz) / &
1874 ! (c0 + chx + chy + clz)
1875 ! fine(2*i-1, 2*j-1, 2*k) = (f0*c0 + flx*clx + fly*cly + fhz*chz) / &
1876 ! (c0 + clx + cly + chz)
1877 ! fine(2*i, 2*j-1, 2*k) = (f0*c0 + fhx*chx + fly*cly + fhz*chz) / &
1878 ! (c0 + chx + cly + chz)
1879 ! fine(2*i-1, 2*j, 2*k) = (f0*c0 + flx*clx + fhy*chy + fhz*chz) / &
1880 ! (c0 + clx + chy + chz)
1881 ! fine(2*i, 2*j, 2*k) = (f0*c0 + fhx*chx + fhy*chy + fhz*chz) / &
1882 ! (c0 + chx + chy + chz)
1883 ! end do
1884 ! end do
1885 ! end do
1886 ! #endif
1887 ! end associate
1888 ! end subroutine vlpl_prolong
1889 
1890  !! File ../src/m_communication.f90
1891 
1892  !> Initialize MPI if needed, and store MPI information
1893  subroutine mg_comm_init(mg, comm)
1894  use mpi
1895  type(mg_t), intent(inout) :: mg
1896  !> MPI communicator (default: MPI_COMM_WORLD)
1897  integer, intent(in), optional :: comm
1898  integer :: ierr
1899  logical :: initialized
1900 
1901  call mpi_initialized(initialized, ierr)
1902  if (.not. initialized) then
1903  call mpi_init(ierr)
1904  end if
1905 
1906  if (present(comm)) then
1907  mg%comm = comm
1908  else
1909  mg%comm = mpi_comm_world
1910  end if
1911 
1912  call mpi_comm_rank(mg%comm, mg%my_rank, ierr)
1913  call mpi_comm_size(mg%comm, mg%n_cpu, ierr)
1914  end subroutine mg_comm_init
1915 
1916  subroutine sort_and_transfer_buffers(mg, dsize)
1917  use mpi
1918  type(mg_t), intent(inout) :: mg
1919  integer, intent(in) :: dsize
1920  integer :: i, n_send, n_recv
1921  integer :: send_req(mg%n_cpu)
1922  integer :: recv_req(mg%n_cpu)
1923  integer :: ierr
1924 
1925  n_send = 0
1926  n_recv = 0
1927 
1928  do i = 0, mg%n_cpu - 1
1929  if (mg%buf(i)%i_send > 0) then
1930  n_send = n_send + 1
1931  call sort_sendbuf(mg%buf(i), dsize)
1932  call mpi_isend(mg%buf(i)%send, mg%buf(i)%i_send, mpi_double, &
1933  i, 0, mg%comm, send_req(n_send), ierr)
1934  end if
1935  if (mg%buf(i)%i_recv > 0) then
1936  n_recv = n_recv + 1
1937  call mpi_irecv(mg%buf(i)%recv, mg%buf(i)%i_recv, mpi_double, &
1938  i, 0, mg%comm, recv_req(n_recv), ierr)
1939  end if
1940  end do
1941 
1942  call mpi_waitall(n_recv, recv_req(1:n_recv), mpi_statuses_ignore, ierr)
1943  call mpi_waitall(n_send, send_req(1:n_send), mpi_statuses_ignore, ierr)
1944 
1945  end subroutine sort_and_transfer_buffers
1946 
1947  !> Sort send buffers according to the idbuf array
1948  subroutine sort_sendbuf(gc, dsize)
1950  type(mg_buf_t), intent(inout) :: gc
1951  integer, intent(in) :: dsize !< Size of send buffer elements
1952  integer :: ix_sort(gc%i_ix)
1953  real(dp) :: buf_cpy(gc%i_send)
1954  integer :: i, j, n
1955 
1956  call mrgrnk(gc%ix(1:gc%i_ix), ix_sort)
1957 
1958  buf_cpy = gc%send(1:gc%i_send)
1959 
1960  do n = 1, gc%i_ix
1961  i = (n-1) * dsize
1962  j = (ix_sort(n)-1) * dsize
1963  gc%send(i+1:i+dsize) = buf_cpy(j+1:j+dsize)
1964  end do
1965  gc%ix(1:gc%i_ix) = gc%ix(ix_sort)
1966 
1967  end subroutine sort_sendbuf
1968 
1969  !! File ../src/m_ghost_cells.f90
1970 
1971  !> Specify minimum buffer size (per process) for communication
1972  subroutine mg_ghost_cell_buffer_size(mg, n_send, n_recv, dsize)
1973  type(mg_t), intent(inout) :: mg
1974  integer, intent(out) :: n_send(0:mg%n_cpu-1)
1975  integer, intent(out) :: n_recv(0:mg%n_cpu-1)
1976  integer, intent(out) :: dsize
1977  integer :: i, id, lvl, nc
1978 
1979  allocate(mg%comm_ghostcell%n_send(0:mg%n_cpu-1, &
1980  mg%first_normal_lvl:mg%highest_lvl))
1981  allocate(mg%comm_ghostcell%n_recv(0:mg%n_cpu-1, &
1982  mg%first_normal_lvl:mg%highest_lvl))
1983 
1984  dsize = mg%box_size**(3-1)
1985 
1986  do lvl = mg%first_normal_lvl, mg%highest_lvl
1987  nc = mg%box_size_lvl(lvl)
1988  mg%buf(:)%i_send = 0
1989  mg%buf(:)%i_recv = 0
1990  mg%buf(:)%i_ix = 0
1991 
1992  do i = 1, size(mg%lvls(lvl)%my_ids)
1993  id = mg%lvls(lvl)%my_ids(i)
1994  call buffer_ghost_cells(mg, id, nc, 1, dry_run=.true.)
1995  end do
1996 
1997  if (lvl > 1) then
1998  do i = 1, size(mg%lvls(lvl-1)%my_ref_bnds)
1999  id = mg%lvls(lvl-1)%my_ref_bnds(i)
2000  call buffer_refinement_boundaries(mg, id, nc, 1, dry_run=.true.)
2001  end do
2002  end if
2003 
2004  ! Set ghost cells to received data
2005  mg%buf(:)%i_recv = 0
2006  do i = 1, size(mg%lvls(lvl)%my_ids)
2007  id = mg%lvls(lvl)%my_ids(i)
2008  call set_ghost_cells(mg, id, nc, 1, dry_run=.true.)
2009  end do
2010 
2011  mg%comm_ghostcell%n_send(:, lvl) = mg%buf(:)%i_send/dsize
2012  mg%comm_ghostcell%n_recv(:, lvl) = mg%buf(:)%i_recv/dsize
2013  end do
2014 
2015  n_send = maxval(mg%comm_ghostcell%n_send, dim=2)
2016  n_recv = maxval(mg%comm_ghostcell%n_recv, dim=2)
2017  end subroutine mg_ghost_cell_buffer_size
2018 
2019  !> Store boundary conditions for the solution variable, this can speed up
2020  !> calculations if the same boundary conditions are re-used.
2021  subroutine mg_phi_bc_store(mg)
2022  type(mg_t), intent(inout) :: mg
2023  integer :: lvl, nc
2024 
2025  nc = mg%box_size
2026 
2027  do lvl = mg%lowest_lvl, mg%highest_lvl
2028  nc = mg%box_size_lvl(lvl)
2029  call mg_phi_bc_store_lvl(mg, lvl, nc)
2030  end do
2031 
2032  mg%phi_bc_data_stored = .true.
2033  end subroutine mg_phi_bc_store
2034 
2035  subroutine mg_phi_bc_store_lvl(mg, lvl, nc)
2036  type(mg_t), intent(inout) :: mg
2037  integer, intent(in) :: lvl
2038  integer, intent(in) :: nc
2039  real(dp) :: bc(nc, nc)
2040  integer :: i, id, nb, nb_id, bc_type
2041 
2042  do i = 1, size(mg%lvls(lvl)%my_ids)
2043  id = mg%lvls(lvl)%my_ids(i)
2044  do nb = 1, mg_num_neighbors
2045  nb_id = mg%boxes(id)%neighbors(nb)
2046  if (nb_id < mg_no_box) then
2047  ! Physical boundary
2048  if (associated(mg%bc(nb, mg_iphi)%boundary_cond)) then
2049  call mg%bc(nb, mg_iphi)%boundary_cond(mg%boxes(id), nc, &
2050  mg_iphi, nb, bc_type, bc)
2051  else
2052  bc_type = mg%bc(nb, mg_iphi)%bc_type
2053  bc = mg%bc(nb, mg_iphi)%bc_value
2054  end if
2055 
2056  ! Store the boundary condition type. This is not globally set in
2057  ! the tree, but all negative values are treated the same in
2058  ! other parts of the code
2059  mg%boxes(id)%neighbors(nb) = bc_type
2060 
2061  ! Store ghost cell data in the right-hand side
2062  call box_set_gc(mg%boxes(id), nb, nc, mg_irhs, bc)
2063  end if
2064  end do
2065  end do
2066  end subroutine mg_phi_bc_store_lvl
2067 
2068  !> Fill ghost cells at all grid levels
2069  subroutine mg_fill_ghost_cells(mg, iv)
2070  type(mg_t) :: mg
2071  integer, intent(in) :: iv !< Index of variable
2072  integer :: lvl
2073 
2074  do lvl = mg%lowest_lvl, mg%highest_lvl
2075  call mg_fill_ghost_cells_lvl(mg, lvl, iv)
2076  end do
2077  end subroutine mg_fill_ghost_cells
2078 
2079  !> Fill ghost cells at a grid level
2080  subroutine mg_fill_ghost_cells_lvl(mg, lvl, iv)
2082  type(mg_t) :: mg
2083  integer, intent(in) :: lvl
2084  integer, intent(in) :: iv !< Index of variable
2085  integer :: i, id, dsize, nc
2086 
2087  if (lvl < mg%lowest_lvl) &
2088  error stop "fill_ghost_cells_lvl: lvl < lowest_lvl"
2089  if (lvl > mg%highest_lvl) &
2090  error stop "fill_ghost_cells_lvl: lvl > highest_lvl"
2091 
2092  nc = mg%box_size_lvl(lvl)
2093 
2094  if (lvl >= mg%first_normal_lvl) then
2095  dsize = nc**(3-1)
2096  mg%buf(:)%i_send = 0
2097  mg%buf(:)%i_recv = 0
2098  mg%buf(:)%i_ix = 0
2099 
2100  do i = 1, size(mg%lvls(lvl)%my_ids)
2101  id = mg%lvls(lvl)%my_ids(i)
2102  call buffer_ghost_cells(mg, id, nc, iv, .false.)
2103  end do
2104 
2105  if (lvl > 1) then
2106  do i = 1, size(mg%lvls(lvl-1)%my_ref_bnds)
2107  id = mg%lvls(lvl-1)%my_ref_bnds(i)
2108  call buffer_refinement_boundaries(mg, id, nc, iv, .false.)
2109  end do
2110  end if
2111 
2112  ! Transfer data between processes
2113  mg%buf(:)%i_recv = mg%comm_ghostcell%n_recv(:, lvl) * dsize
2114  call sort_and_transfer_buffers(mg, dsize)
2115 
2116  ! Set ghost cells to received data
2117  mg%buf(:)%i_recv = 0
2118  end if
2119 
2120  do i = 1, size(mg%lvls(lvl)%my_ids)
2121  id = mg%lvls(lvl)%my_ids(i)
2122  call set_ghost_cells(mg, id, nc, iv, .false.)
2123  end do
2124  end subroutine mg_fill_ghost_cells_lvl
2125 
2126  subroutine buffer_ghost_cells(mg, id, nc, iv, dry_run)
2127  type(mg_t), intent(inout) :: mg
2128  integer, intent(in) :: id
2129  integer, intent(in) :: nc
2130  integer, intent(in) :: iv
2131  logical, intent(in) :: dry_run
2132  integer :: nb, nb_id, nb_rank
2133 
2134  do nb = 1, mg_num_neighbors
2135  nb_id = mg%boxes(id)%neighbors(nb)
2136 
2137  if (nb_id > mg_no_box) then
2138  ! There is a neighbor
2139  nb_rank = mg%boxes(nb_id)%rank
2140 
2141  if (nb_rank /= mg%my_rank) then
2142  call buffer_for_nb(mg, mg%boxes(id), nc, iv, nb_id, nb_rank, &
2143  nb, dry_run)
2144  end if
2145  end if
2146  end do
2147  end subroutine buffer_ghost_cells
2148 
2149  subroutine buffer_refinement_boundaries(mg, id, nc, iv, dry_run)
2150  type(mg_t), intent(inout) :: mg
2151  integer, intent(in) :: id
2152  integer, intent(in) :: nc
2153  integer, intent(in) :: iv
2154  logical, intent(in) :: dry_run
2155  integer :: nb, nb_id, c_ids(2**(3-1))
2156  integer :: n, c_id, c_rank
2157 
2158  do nb = 1, mg_num_neighbors
2159  nb_id = mg%boxes(id)%neighbors(nb)
2160  if (nb_id > mg_no_box) then
2161  if (mg_has_children(mg%boxes(nb_id))) then
2162  c_ids = mg%boxes(nb_id)%children(&
2164 
2165  do n = 1, mg_num_children/2
2166  c_id = c_ids(n)
2167  c_rank = mg%boxes(c_id)%rank
2168 
2169  if (c_rank /= mg%my_rank) then
2170  ! Send all coarse ghost cells
2171  call buffer_for_fine_nb(mg, mg%boxes(id), nc, iv, c_id, &
2172  c_rank, nb, dry_run)
2173  end if
2174  end do
2175  end if
2176  end if
2177  end do
2178  end subroutine buffer_refinement_boundaries
2179 
2180  !> The routine that actually fills the ghost cells
2181  subroutine set_ghost_cells(mg, id, nc, iv, dry_run)
2182  type(mg_t), intent(inout) :: mg
2183  integer, intent(in) :: id
2184  integer, intent(in) :: nc
2185  integer, intent(in) :: iv
2186  logical, intent(in) :: dry_run
2187  real(dp) :: bc(nc, nc)
2188  integer :: nb, nb_id, nb_rank, bc_type
2189 
2190  do nb = 1, mg_num_neighbors
2191  nb_id = mg%boxes(id)%neighbors(nb)
2192 
2193  if (nb_id > mg_no_box) then
2194  ! There is a neighbor
2195  nb_rank = mg%boxes(nb_id)%rank
2196 
2197  if (nb_rank /= mg%my_rank) then
2198  call fill_buffered_nb(mg, mg%boxes(id), nb_rank, &
2199  nb, nc, iv, dry_run)
2200  else if (.not. dry_run) then
2201  call copy_from_nb(mg%boxes(id), mg%boxes(nb_id), &
2202  nb, nc, iv)
2203  end if
2204  else if (nb_id == mg_no_box) then
2205  ! Refinement boundary
2206  call fill_refinement_bnd(mg, id, nb, nc, iv, dry_run)
2207  else if (.not. dry_run) then
2208  ! Physical boundary
2209  if (mg%phi_bc_data_stored .and. iv == mg_iphi) then
2210  ! Copy the boundary conditions stored in the ghost cells of the
2211  ! right-hand side
2212  call box_get_gc(mg%boxes(id), nb, nc, mg_irhs, bc)
2213  bc_type = nb_id
2214  else
2215  if (associated(mg%bc(nb, iv)%boundary_cond)) then
2216  call mg%bc(nb, iv)%boundary_cond(mg%boxes(id), nc, iv, &
2217  nb, bc_type, bc)
2218  else
2219  bc_type = mg%bc(nb, iv)%bc_type
2220  bc = mg%bc(nb, iv)%bc_value
2221  end if
2222  end if
2223 
2224  call box_set_gc(mg%boxes(id), nb, nc, iv, bc)
2225  call bc_to_gc(mg, id, nc, iv, nb, bc_type)
2226  end if
2227  end do
2228  end subroutine set_ghost_cells
2229 
2230  subroutine fill_refinement_bnd(mg, id, nb, nc, iv, dry_run)
2231  type(mg_t), intent(inout) :: mg
2232  integer, intent(in) :: id
2233  integer, intent(in) :: nc
2234  integer, intent(in) :: iv
2235  integer, intent(in) :: nb
2236  logical, intent(in) :: dry_run
2237  real(dp) :: gc(nc, nc)
2238  integer :: p_id, p_nb_id, ix_offset(3)
2239  integer :: i, dsize, p_nb_rank
2240 
2241  dsize = nc**(3-1)
2242  p_id = mg%boxes(id)%parent
2243  p_nb_id = mg%boxes(p_id)%neighbors(nb)
2244  p_nb_rank = mg%boxes(p_nb_id)%rank
2245 
2246  if (p_nb_rank /= mg%my_rank) then
2247  i = mg%buf(p_nb_rank)%i_recv
2248  if (.not. dry_run) then
2249  gc = reshape(mg%buf(p_nb_rank)%recv(i+1:i+dsize), shape(gc))
2250  end if
2251  mg%buf(p_nb_rank)%i_recv = mg%buf(p_nb_rank)%i_recv + dsize
2252  else if (.not. dry_run) then
2253  ix_offset = mg_get_child_offset(mg, id)
2254  call box_gc_for_fine_neighbor(mg%boxes(p_nb_id), mg_neighb_rev(nb), &
2255  ix_offset, nc, iv, gc)
2256  end if
2257 
2258  if (.not. dry_run) then
2259  if (associated(mg%bc(nb, iv)%refinement_bnd)) then
2260  call mg%bc(nb, iv)%refinement_bnd(mg%boxes(id), nc, iv, nb, gc)
2261  else
2262  call sides_rb(mg%boxes(id), nc, iv, nb, gc)
2263  end if
2264  end if
2265  end subroutine fill_refinement_bnd
2266 
2267  subroutine copy_from_nb(box, box_nb, nb, nc, iv)
2268  type(mg_box_t), intent(inout) :: box
2269  type(mg_box_t), intent(in) :: box_nb
2270  integer, intent(in) :: nb
2271  integer, intent(in) :: nc
2272  integer, intent(in) :: iv
2273  real(dp) :: gc(nc, nc)
2274 
2275  call box_gc_for_neighbor(box_nb, mg_neighb_rev(nb), nc, iv, gc)
2276  call box_set_gc(box, nb, nc, iv, gc)
2277  end subroutine copy_from_nb
2278 
2279  subroutine buffer_for_nb(mg, box, nc, iv, nb_id, nb_rank, nb, dry_run)
2280  type(mg_t), intent(inout) :: mg
2281  type(mg_box_t), intent(inout) :: box
2282  integer, intent(in) :: nc
2283  integer, intent(in) :: iv
2284  integer, intent(in) :: nb_id
2285  integer, intent(in) :: nb_rank
2286  integer, intent(in) :: nb
2287  logical, intent(in) :: dry_run
2288  integer :: i, dsize
2289  real(dp) :: gc(nc, nc)
2290 
2291  i = mg%buf(nb_rank)%i_send
2292  dsize = nc**(3-1)
2293 
2294  if (.not. dry_run) then
2295  call box_gc_for_neighbor(box, nb, nc, iv, gc)
2296  mg%buf(nb_rank)%send(i+1:i+dsize) = pack(gc, .true.)
2297  end if
2298 
2299  ! Later the buffer is sorted, using the fact that loops go from low to high
2300  ! box id, and we fill ghost cells according to the neighbor order
2301  i = mg%buf(nb_rank)%i_ix
2302  if (.not. dry_run) then
2303  mg%buf(nb_rank)%ix(i+1) = mg_num_neighbors * nb_id + mg_neighb_rev(nb)
2304  end if
2305 
2306  mg%buf(nb_rank)%i_send = mg%buf(nb_rank)%i_send + dsize
2307  mg%buf(nb_rank)%i_ix = mg%buf(nb_rank)%i_ix + 1
2308  end subroutine buffer_for_nb
2309 
2310  subroutine buffer_for_fine_nb(mg, box, nc, iv, fine_id, fine_rank, nb, dry_run)
2311  type(mg_t), intent(inout) :: mg
2312  type(mg_box_t), intent(inout) :: box
2313  integer, intent(in) :: nc
2314  integer, intent(in) :: iv
2315  integer, intent(in) :: fine_id
2316  integer, intent(in) :: fine_rank
2317  integer, intent(in) :: nb
2318  logical, intent(in) :: dry_run
2319  integer :: i, dsize, ix_offset(3)
2320  real(dp) :: gc(nc, nc)
2321 
2322  i = mg%buf(fine_rank)%i_send
2323  dsize = nc**(3-1)
2324 
2325  if (.not. dry_run) then
2326  ix_offset = mg_get_child_offset(mg, fine_id)
2327  call box_gc_for_fine_neighbor(box, nb, ix_offset, nc, iv, gc)
2328  mg%buf(fine_rank)%send(i+1:i+dsize) = pack(gc, .true.)
2329  end if
2330 
2331  ! Later the buffer is sorted, using the fact that loops go from low to high
2332  ! box id, and we fill ghost cells according to the neighbor order
2333  i = mg%buf(fine_rank)%i_ix
2334  if (.not. dry_run) then
2335  mg%buf(fine_rank)%ix(i+1) = mg_num_neighbors * fine_id + &
2336  mg_neighb_rev(nb)
2337  end if
2338 
2339  mg%buf(fine_rank)%i_send = mg%buf(fine_rank)%i_send + dsize
2340  mg%buf(fine_rank)%i_ix = mg%buf(fine_rank)%i_ix + 1
2341  end subroutine buffer_for_fine_nb
2342 
2343  subroutine fill_buffered_nb(mg, box, nb_rank, nb, nc, iv, dry_run)
2344  type(mg_t), intent(inout) :: mg
2345  type(mg_box_t), intent(inout) :: box
2346  integer, intent(in) :: nb_rank
2347  integer, intent(in) :: nb
2348  integer, intent(in) :: nc
2349  integer, intent(in) :: iv
2350  logical, intent(in) :: dry_run
2351  integer :: i, dsize
2352  real(dp) :: gc(nc, nc)
2353 
2354  i = mg%buf(nb_rank)%i_recv
2355  dsize = nc**(3-1)
2356 
2357  if (.not. dry_run) then
2358  gc = reshape(mg%buf(nb_rank)%recv(i+1:i+dsize), shape(gc))
2359  call box_set_gc(box, nb, nc, iv, gc)
2360  end if
2361  mg%buf(nb_rank)%i_recv = mg%buf(nb_rank)%i_recv + dsize
2362 
2363  end subroutine fill_buffered_nb
2364 
2365  subroutine box_gc_for_neighbor(box, nb, nc, iv, gc)
2366  type(mg_box_t), intent(in) :: box
2367  integer, intent(in) :: nb, nc, iv
2368  real(dp), intent(out) :: gc(nc, nc)
2369 
2370  select case (nb)
2371  case (mg_neighb_lowx)
2372  gc = box%cc(1, 1:nc, 1:nc, iv)
2373  case (mg_neighb_highx)
2374  gc = box%cc(nc, 1:nc, 1:nc, iv)
2375  case (mg_neighb_lowy)
2376  gc = box%cc(1:nc, 1, 1:nc, iv)
2377  case (mg_neighb_highy)
2378  gc = box%cc(1:nc, nc, 1:nc, iv)
2379  case (mg_neighb_lowz)
2380  gc = box%cc(1:nc, 1:nc, 1, iv)
2381  case (mg_neighb_highz)
2382  gc = box%cc(1:nc, 1:nc, nc, iv)
2383  end select
2384  end subroutine box_gc_for_neighbor
2385 
2386  !> Get ghost cells for a fine neighbor
2387  subroutine box_gc_for_fine_neighbor(box, nb, di, nc, iv, gc)
2388  type(mg_box_t), intent(in) :: box
2389  integer, intent(in) :: nb !< Direction of fine neighbor
2390  integer, intent(in) :: di(3) !< Index offset of fine neighbor
2391  integer, intent(in) :: nc, iv
2392  real(dp), intent(out) :: gc(nc, nc)
2393  real(dp) :: tmp(0:nc/2+1, 0:nc/2+1), grad(3-1)
2394  integer :: i, j, hnc
2395 
2396  hnc = nc/2
2397 
2398  ! First fill a temporary array with data next to the fine grid
2399  select case (nb)
2400  case (mg_neighb_lowx)
2401  tmp = box%cc(1, di(2):di(2)+hnc+1, di(3):di(3)+hnc+1, iv)
2402  case (mg_neighb_highx)
2403  tmp = box%cc(nc, di(2):di(2)+hnc+1, di(3):di(3)+hnc+1, iv)
2404  case (mg_neighb_lowy)
2405  tmp = box%cc(di(1):di(1)+hnc+1, 1, di(3):di(3)+hnc+1, iv)
2406  case (mg_neighb_highy)
2407  tmp = box%cc(di(1):di(1)+hnc+1, nc, di(3):di(3)+hnc+1, iv)
2408  case (mg_neighb_lowz)
2409  tmp = box%cc(di(1):di(1)+hnc+1, di(2):di(2)+hnc+1, 1, iv)
2410  case (mg_neighb_highz)
2411  tmp = box%cc(di(1):di(1)+hnc+1, di(2):di(2)+hnc+1, nc, iv)
2412  case default
2413  error stop
2414  end select
2415 
2416  ! Now interpolate the coarse grid data to obtain values 'straight' next to
2417  ! the fine grid points
2418  do j = 1, hnc
2419  do i = 1, hnc
2420  grad(1) = 0.125_dp * (tmp(i+1, j) - tmp(i-1, j))
2421  grad(2) = 0.125_dp * (tmp(i, j+1) - tmp(i, j-1))
2422  gc(2*i-1, 2*j-1) = tmp(i, j) - grad(1) - grad(2)
2423  gc(2*i, 2*j-1) = tmp(i, j) + grad(1) - grad(2)
2424  gc(2*i-1, 2*j) = tmp(i, j) - grad(1) + grad(2)
2425  gc(2*i, 2*j) = tmp(i, j) + grad(1) + grad(2)
2426  end do
2427  end do
2428  end subroutine box_gc_for_fine_neighbor
2429 
2430  subroutine box_get_gc(box, nb, nc, iv, gc)
2431  type(mg_box_t), intent(in) :: box
2432  integer, intent(in) :: nb, nc, iv
2433  real(dp), intent(out) :: gc(nc, nc)
2434 
2435  select case (nb)
2436  case (mg_neighb_lowx)
2437  gc = box%cc(0, 1:nc, 1:nc, iv)
2438  case (mg_neighb_highx)
2439  gc = box%cc(nc+1, 1:nc, 1:nc, iv)
2440  case (mg_neighb_lowy)
2441  gc = box%cc(1:nc, 0, 1:nc, iv)
2442  case (mg_neighb_highy)
2443  gc = box%cc(1:nc, nc+1, 1:nc, iv)
2444  case (mg_neighb_lowz)
2445  gc = box%cc(1:nc, 1:nc, 0, iv)
2446  case (mg_neighb_highz)
2447  gc = box%cc(1:nc, 1:nc, nc+1, iv)
2448  end select
2449  end subroutine box_get_gc
2450 
2451  subroutine box_set_gc(box, nb, nc, iv, gc)
2452  type(mg_box_t), intent(inout) :: box
2453  integer, intent(in) :: nb, nc, iv
2454  real(dp), intent(in) :: gc(nc, nc)
2455 
2456  select case (nb)
2457  case (mg_neighb_lowx)
2458  box%cc(0, 1:nc, 1:nc, iv) = gc
2459  case (mg_neighb_highx)
2460  box%cc(nc+1, 1:nc, 1:nc, iv) = gc
2461  case (mg_neighb_lowy)
2462  box%cc(1:nc, 0, 1:nc, iv) = gc
2463  case (mg_neighb_highy)
2464  box%cc(1:nc, nc+1, 1:nc, iv) = gc
2465  case (mg_neighb_lowz)
2466  box%cc(1:nc, 1:nc, 0, iv) = gc
2467  case (mg_neighb_highz)
2468  box%cc(1:nc, 1:nc, nc+1, iv) = gc
2469  end select
2470  end subroutine box_set_gc
2471 
2472  subroutine bc_to_gc(mg, id, nc, iv, nb, bc_type)
2473  type(mg_t), intent(inout) :: mg
2474  integer, intent(in) :: id
2475  integer, intent(in) :: nc
2476  integer, intent(in) :: iv
2477  integer, intent(in) :: nb !< Neighbor direction
2478  integer, intent(in) :: bc_type !< Type of b.c.
2479  real(dp) :: c0, c1, c2, dr
2480 
2481  ! If we call the interior point x1, x2 and the ghost point x0, then a
2482  ! Dirichlet boundary value b can be imposed as:
2483  ! x0 = -x1 + 2*b
2484  ! A Neumann b.c. can be imposed as:
2485  ! x0 = x1 +/- dx * b
2486  ! A continuous boundary (same slope) as:
2487  ! x0 = 2 * x1 - x2
2488  ! Below, we set coefficients to handle these cases
2489  select case (bc_type)
2490  case (mg_bc_dirichlet)
2491  c0 = 2
2492  c1 = -1
2493  c2 = 0
2494  case (mg_bc_neumann)
2495  dr = mg%dr(mg_neighb_dim(nb), mg%boxes(id)%lvl)
2496  c0 = dr * mg_neighb_high_pm(nb) ! This gives a + or - sign
2497  c1 = 1
2498  c2 = 0
2499  case (mg_bc_continuous)
2500  c0 = 0
2501  c1 = 2
2502  c2 = -1
2503  case default
2504  error stop "bc_to_gc: unknown boundary condition"
2505  end select
2506 
2507  select case (nb)
2508  case (mg_neighb_lowx)
2509  mg%boxes(id)%cc(0, 1:nc, 1:nc, iv) = &
2510  c0 * mg%boxes(id)%cc(0, 1:nc, 1:nc, iv) + &
2511  c1 * mg%boxes(id)%cc(1, 1:nc, 1:nc, iv) + &
2512  c2 * mg%boxes(id)%cc(2, 1:nc, 1:nc, iv)
2513  case (mg_neighb_highx)
2514  mg%boxes(id)%cc(nc+1, 1:nc, 1:nc, iv) = &
2515  c0 * mg%boxes(id)%cc(nc+1, 1:nc, 1:nc, iv) + &
2516  c1 * mg%boxes(id)%cc(nc, 1:nc, 1:nc, iv) + &
2517  c2 * mg%boxes(id)%cc(nc-1, 1:nc, 1:nc, iv)
2518  case (mg_neighb_lowy)
2519  mg%boxes(id)%cc(1:nc, 0, 1:nc, iv) = &
2520  c0 * mg%boxes(id)%cc(1:nc, 0, 1:nc, iv) + &
2521  c1 * mg%boxes(id)%cc(1:nc, 1, 1:nc, iv) + &
2522  c2 * mg%boxes(id)%cc(1:nc, 2, 1:nc, iv)
2523  case (mg_neighb_highy)
2524  mg%boxes(id)%cc(1:nc, nc+1, 1:nc, iv) = &
2525  c0 * mg%boxes(id)%cc(1:nc, nc+1, 1:nc, iv) + &
2526  c1 * mg%boxes(id)%cc(1:nc, nc, 1:nc, iv) + &
2527  c2 * mg%boxes(id)%cc(1:nc, nc-1, 1:nc, iv)
2528  case (mg_neighb_lowz)
2529  mg%boxes(id)%cc(1:nc, 1:nc, 0, iv) = &
2530  c0 * mg%boxes(id)%cc(1:nc, 1:nc, 0, iv) + &
2531  c1 * mg%boxes(id)%cc(1:nc, 1:nc, 1, iv) + &
2532  c2 * mg%boxes(id)%cc(1:nc, 1:nc, 2, iv)
2533  case (mg_neighb_highz)
2534  mg%boxes(id)%cc(1:nc, 1:nc, nc+1, iv) = &
2535  c0 * mg%boxes(id)%cc(1:nc, 1:nc, nc+1, iv) + &
2536  c1 * mg%boxes(id)%cc(1:nc, 1:nc, nc, iv) + &
2537  c2 * mg%boxes(id)%cc(1:nc, 1:nc, nc-1, iv)
2538  end select
2539  end subroutine bc_to_gc
2540 
2541  !> Fill ghost cells near refinement boundaries which preserves diffusive fluxes.
2542  subroutine sides_rb(box, nc, iv, nb, gc)
2543  type(mg_box_t), intent(inout) :: box
2544  integer, intent(in) :: nc
2545  integer, intent(in) :: iv
2546  integer, intent(in) :: nb !< Ghost cell direction
2547  !> Interpolated coarse grid ghost cell data (but not yet in the nb direction)
2548  real(dp), intent(in) :: gc(nc, nc)
2549  integer :: di, dj, dk
2550  integer :: i, j, k, ix, dix
2551 
2552  if (mg_neighb_low(nb)) then
2553  ix = 1
2554  dix = 1
2555  else
2556  ix = nc
2557  dix = -1
2558  end if
2559 
2560  select case (mg_neighb_dim(nb))
2561  case (1)
2562  i = ix
2563  di = dix
2564  do k = 1, nc
2565  dk = -1 + 2 * iand(k, 1)
2566  do j = 1, nc
2567  dj = -1 + 2 * iand(j, 1)
2568  box%cc(i-di, j, k, iv) = 0.5_dp * gc(j, k) &
2569  + 0.75_dp * box%cc(i, j, k, iv) &
2570  - 0.25_dp * box%cc(i+di, j, k, iv)
2571  end do
2572  end do
2573  case (2)
2574  j = ix
2575  dj = dix
2576  do k = 1, nc
2577  dk = -1 + 2 * iand(k, 1)
2578  do i = 1, nc
2579  di = -1 + 2 * iand(i, 1)
2580  box%cc(i, j-dj, k, iv) = 0.5_dp * gc(i, k) &
2581  + 0.75_dp * box%cc(i, j, k, iv) &
2582  - 0.25_dp * box%cc(i, j+dj, k, iv)
2583  end do
2584  end do
2585  case (3)
2586  k = ix
2587  dk = dix
2588  do j = 1, nc
2589  dj = -1 + 2 * iand(j, 1)
2590  do i = 1, nc
2591  di = -1 + 2 * iand(i, 1)
2592  box%cc(i, j, k-dk, iv) = 0.5_dp * gc(i, j) &
2593  + 0.75_dp * box%cc(i, j, k, iv) &
2594  - 0.25_dp * box%cc(i, j, k+dk, iv)
2595  end do
2596  end do
2597  end select
2598 
2599  end subroutine sides_rb
2600 
2601  !! File ../src/m_prolong.f90
2602 
2603  !> Specify minimum buffer size (per process) for communication
2604  subroutine mg_prolong_buffer_size(mg, n_send, n_recv, dsize)
2605  type(mg_t), intent(inout) :: mg
2606  integer, intent(out) :: n_send(0:mg%n_cpu-1)
2607  integer, intent(out) :: n_recv(0:mg%n_cpu-1)
2608  integer, intent(out) :: dsize
2609  integer :: lvl, min_lvl
2610 
2611  if (.not. allocated(mg%comm_restrict%n_send)) then
2612  error stop "Call restrict_buffer_size before prolong_buffer_size"
2613  end if
2614 
2615  min_lvl = max(mg%first_normal_lvl-1, mg%lowest_lvl)
2616  allocate(mg%comm_prolong%n_send(0:mg%n_cpu-1, &
2617  min_lvl:mg%highest_lvl))
2618  allocate(mg%comm_prolong%n_recv(0:mg%n_cpu-1, &
2619  min_lvl:mg%highest_lvl))
2620 
2621  mg%comm_prolong%n_recv(:, mg%highest_lvl) = 0
2622  mg%comm_prolong%n_send(:, mg%highest_lvl) = 0
2623 
2624  do lvl = min_lvl, mg%highest_lvl-1
2625  mg%comm_prolong%n_recv(:, lvl) = &
2626  mg%comm_restrict%n_send(:, lvl+1)
2627  mg%comm_prolong%n_send(:, lvl) = &
2628  mg%comm_restrict%n_recv(:, lvl+1)
2629  end do
2630 
2631  ! Send fine grid points, because this is more flexible than sending coarse
2632  ! grid points (e.g., when multiple variables are used for interpolation)
2633  dsize = (mg%box_size)**3
2634  n_send = maxval(mg%comm_prolong%n_send, dim=2)
2635  n_recv = maxval(mg%comm_prolong%n_recv, dim=2)
2636  end subroutine mg_prolong_buffer_size
2637 
2638  !> Prolong variable iv from lvl to variable iv_to at lvl+1
2639  subroutine mg_prolong(mg, lvl, iv, iv_to, method, add)
2641  type(mg_t), intent(inout) :: mg
2642  integer, intent(in) :: lvl !< Level to prolong from
2643  integer, intent(in) :: iv !< Source variable
2644  integer, intent(in) :: iv_to !< Target variable
2645  procedure(mg_box_prolong) :: method !< Prolongation method
2646  logical, intent(in) :: add !< If true, add to current values
2647  integer :: i, id, dsize, nc
2648 
2649  if (lvl == mg%highest_lvl) error stop "cannot prolong highest level"
2650  if (lvl < mg%lowest_lvl) error stop "cannot prolong below lowest level"
2651 
2652  ! Below the first normal level, all boxes are on the same CPU
2653  if (lvl >= mg%first_normal_lvl-1) then
2654  dsize = mg%box_size**3
2655  mg%buf(:)%i_send = 0
2656  mg%buf(:)%i_ix = 0
2657 
2658  do i = 1, size(mg%lvls(lvl)%my_ids)
2659  id = mg%lvls(lvl)%my_ids(i)
2660  call prolong_set_buffer(mg, id, mg%box_size, iv, method)
2661  end do
2662 
2663  mg%buf(:)%i_recv = mg%comm_prolong%n_recv(:, lvl) * dsize
2664  call sort_and_transfer_buffers(mg, dsize)
2665  mg%buf(:)%i_recv = 0
2666  end if
2667 
2668  nc = mg%box_size_lvl(lvl+1)
2669  do i = 1, size(mg%lvls(lvl+1)%my_ids)
2670  id = mg%lvls(lvl+1)%my_ids(i)
2671  call prolong_onto(mg, id, nc, iv, iv_to, add, method)
2672  end do
2673  end subroutine mg_prolong
2674 
2675  !> In case the fine grid is on a different CPU, perform the prolongation and
2676  !> store the fine-grid values in the send buffer.
2677  !>
2678  !> @todo Check whether it's faster to send coarse data and prolong afterwards
2679  subroutine prolong_set_buffer(mg, id, nc, iv, method)
2680  type(mg_t), intent(inout) :: mg
2681  integer, intent(in) :: id
2682  integer, intent(in) :: nc
2683  integer, intent(in) :: iv
2684  procedure(mg_box_prolong) :: method
2685  integer :: i, dix(3)
2686  integer :: i_c, c_id, c_rank, dsize
2687  real(dp) :: tmp(nc, nc, nc)
2688 
2689  dsize = nc**3
2690 
2691  do i_c = 1, mg_num_children
2692  c_id = mg%boxes(id)%children(i_c)
2693  if (c_id > mg_no_box) then
2694  c_rank = mg%boxes(c_id)%rank
2695  if (c_rank /= mg%my_rank) then
2696  dix = mg_get_child_offset(mg, c_id)
2697  call method(mg, id, dix, nc, iv, tmp)
2698 
2699  i = mg%buf(c_rank)%i_send
2700  mg%buf(c_rank)%send(i+1:i+dsize) = pack(tmp, .true.)
2701  mg%buf(c_rank)%i_send = mg%buf(c_rank)%i_send + dsize
2702 
2703  i = mg%buf(c_rank)%i_ix
2704  mg%buf(c_rank)%ix(i+1) = c_id
2705  mg%buf(c_rank)%i_ix = mg%buf(c_rank)%i_ix + 1
2706  end if
2707  end if
2708  end do
2709  end subroutine prolong_set_buffer
2710 
2711  !> Prolong onto a child box
2712  subroutine prolong_onto(mg, id, nc, iv, iv_to, add, method)
2713  type(mg_t), intent(inout) :: mg
2714  integer, intent(in) :: id
2715  integer, intent(in) :: nc
2716  integer, intent(in) :: iv !< Prolong from this variable
2717  integer, intent(in) :: iv_to !< Prolong to this variable
2718  logical, intent(in) :: add !< If true, add to current values
2719  procedure(mg_box_prolong) :: method
2720  integer :: hnc, p_id, p_rank, i, dix(3), dsize
2721  real(dp) :: tmp(nc, nc, nc)
2722 
2723  hnc = nc/2
2724  p_id = mg%boxes(id)%parent
2725  p_rank = mg%boxes(p_id)%rank
2726 
2727  if (p_rank == mg%my_rank) then
2728  dix = mg_get_child_offset(mg, id)
2729  call method(mg, p_id, dix, nc, iv, tmp)
2730  else
2731  dsize = nc**3
2732  i = mg%buf(p_rank)%i_recv
2733  tmp = reshape(mg%buf(p_rank)%recv(i+1:i+dsize), [nc, nc, nc])
2734  mg%buf(p_rank)%i_recv = mg%buf(p_rank)%i_recv + dsize
2735  end if
2736 
2737  if (add) then
2738  mg%boxes(id)%cc(1:nc, 1:nc, 1:nc, iv_to) = &
2739  mg%boxes(id)%cc(1:nc, 1:nc, 1:nc, iv_to) + tmp
2740  else
2741  mg%boxes(id)%cc(1:nc, 1:nc, 1:nc, iv_to) = tmp
2742  end if
2743 
2744  end subroutine prolong_onto
2745 
2746  !> Prolong from a parent to a child with index offset dix
2747  subroutine mg_prolong_sparse(mg, p_id, dix, nc, iv, fine)
2748  type(mg_t), intent(inout) :: mg
2749  integer, intent(in) :: p_id !< Id of parent
2750  integer, intent(in) :: dix(3) !< Offset of child in parent grid
2751  integer, intent(in) :: nc !< Child grid size
2752  integer, intent(in) :: iv !< Prolong from this variable
2753  real(dp), intent(out) :: fine(nc, nc, nc) !< Prolonged values
2754 
2755  integer :: i, j, k, hnc
2756  integer :: ic, jc, kc
2757  real(dp) :: f0, flx, fhx, fly, fhy, flz, fhz
2758 
2759  hnc = nc/2
2760 
2761  associate(crs => mg%boxes(p_id)%cc)
2762  do k = 1, hnc
2763  kc = k + dix(3)
2764  do j = 1, hnc
2765  jc = j + dix(2)
2766  do i = 1, hnc
2767  ic = i + dix(1)
2768 
2769  f0 = 0.25_dp * crs(ic, jc, kc, iv)
2770  flx = 0.25_dp * crs(ic-1, jc, kc, iv)
2771  fhx = 0.25_dp * crs(ic+1, jc, kc, iv)
2772  fly = 0.25_dp * crs(ic, jc-1, kc, iv)
2773  fhy = 0.25_dp * crs(ic, jc+1, kc, iv)
2774  flz = 0.25_dp * crs(ic, jc, kc-1, iv)
2775  fhz = 0.25_dp * crs(ic, jc, kc+1, iv)
2776 
2777  fine(2*i-1, 2*j-1, 2*k-1) = f0 + flx + fly + flz
2778  fine(2*i, 2*j-1, 2*k-1) = f0 + fhx + fly + flz
2779  fine(2*i-1, 2*j, 2*k-1) = f0 + flx + fhy + flz
2780  fine(2*i, 2*j, 2*k-1) = f0 + fhx + fhy + flz
2781  fine(2*i-1, 2*j-1, 2*k) = f0 + flx + fly + fhz
2782  fine(2*i, 2*j-1, 2*k) = f0 + fhx + fly + fhz
2783  fine(2*i-1, 2*j, 2*k) = f0 + flx + fhy + fhz
2784  fine(2*i, 2*j, 2*k) = f0 + fhx + fhy + fhz
2785  end do
2786  end do
2787  end do
2788  end associate
2789  end subroutine mg_prolong_sparse
2790 
2791  !! File ../src/m_helmholtz.f90
2792 
2793  subroutine helmholtz_set_methods(mg)
2794  type(mg_t), intent(inout) :: mg
2795 
2796  mg%subtract_mean = .false.
2797 
2798  select case (mg%geometry_type)
2799  case (mg_cartesian)
2800  mg%box_op => box_helmh
2801 
2802  select case (mg%smoother_type)
2804  mg%box_smoother => box_gs_helmh
2805  case default
2806  error stop "helmholtz_set_methods: unsupported smoother type"
2807  end select
2808  case default
2809  error stop "helmholtz_set_methods: unsupported geometry"
2810  end select
2811 
2812  end subroutine helmholtz_set_methods
2813 
2814  subroutine helmholtz_set_lambda(lambda)
2815  real(dp), intent(in) :: lambda
2816 
2817  if (lambda < 0) &
2818  error stop "helmholtz_set_lambda: lambda < 0 not allowed"
2819 
2820  helmholtz_lambda = lambda
2821  end subroutine helmholtz_set_lambda
2822 
2823  !> Perform Gauss-Seidel relaxation on box for a Helmholtz operator
2824  subroutine box_gs_helmh(mg, id, nc, redblack_cntr)
2825  type(mg_t), intent(inout) :: mg
2826  integer, intent(in) :: id
2827  integer, intent(in) :: nc
2828  integer, intent(in) :: redblack_cntr !< Iteration counter
2829  integer :: i, j, k, i0, di
2830  real(dp) :: idr2(3), fac
2831  logical :: redblack
2832 
2833  idr2 = 1/mg%dr(:, mg%boxes(id)%lvl)**2
2834  fac = 1.0_dp / (2 * sum(idr2) + helmholtz_lambda)
2835  i0 = 1
2836 
2837  redblack = (mg%smoother_type == mg_smoother_gsrb)
2838  if (redblack) then
2839  di = 2
2840  else
2841  di = 1
2842  end if
2843 
2844  ! The parity of redblack_cntr determines which cells we use. If
2845  ! redblack_cntr is even, we use the even cells and vice versa.
2846  associate(cc => mg%boxes(id)%cc, n => mg_iphi)
2847  do k = 1, nc
2848  do j = 1, nc
2849  if (redblack) &
2850  i0 = 2 - iand(ieor(redblack_cntr, k+j), 1)
2851  do i = i0, nc, di
2852  cc(i, j, k, n) = fac * ( &
2853  idr2(1) * (cc(i+1, j, k, n) + cc(i-1, j, k, n)) + &
2854  idr2(2) * (cc(i, j+1, k, n) + cc(i, j-1, k, n)) + &
2855  idr2(3) * (cc(i, j, k+1, n) + cc(i, j, k-1, n)) - &
2856  cc(i, j, k, mg_irhs))
2857  end do
2858  end do
2859  end do
2860  end associate
2861  end subroutine box_gs_helmh
2862 
2863  !> Perform Helmholtz operator on a box
2864  subroutine box_helmh(mg, id, nc, i_out)
2865  type(mg_t), intent(inout) :: mg
2866  integer, intent(in) :: id
2867  integer, intent(in) :: nc
2868  integer, intent(in) :: i_out !< Index of variable to store Helmholtz in
2869  integer :: i, j, k
2870  real(dp) :: idr2(3)
2871 
2872  idr2 = 1 / mg%dr(:, mg%boxes(id)%lvl)**2
2873 
2874  associate(cc => mg%boxes(id)%cc, n => mg_iphi)
2875  do k = 1, nc
2876  do j = 1, nc
2877  do i = 1, nc
2878  cc(i, j, k, i_out) = &
2879  idr2(1) * (cc(i-1, j, k, n) + cc(i+1, j, k, n) &
2880  - 2 * cc(i, j, k, n)) &
2881  + idr2(2) * (cc(i, j-1, k, n) + cc(i, j+1, k, n) &
2882  - 2 * cc(i, j, k, n)) &
2883  + idr2(3) * (cc(i, j, k-1, n) + cc(i, j, k+1, n) &
2884  - 2 * cc(i, j, k, n)) - &
2885  helmholtz_lambda * cc(i, j, k, n)
2886  end do
2887  end do
2888  end do
2889  end associate
2890  end subroutine box_helmh
2891 
2892  !! File ../src/m_multigrid.f90
2893 
2894  subroutine mg_set_methods(mg)
2896  type(mg_t), intent(inout) :: mg
2897 
2898  ! Set default prolongation method (routines below can override this)
2899  mg%box_prolong => mg_prolong_sparse
2900 
2901  select case (mg%operator_type)
2902  case (mg_laplacian)
2903  call laplacian_set_methods(mg)
2904  case (mg_vlaplacian)
2905  call vlaplacian_set_methods(mg)
2906  case (mg_helmholtz)
2907  call helmholtz_set_methods(mg)
2908  case (mg_vhelmholtz)
2909  call vhelmholtz_set_methods(mg)
2910  case default
2911  error stop "mg_set_methods: unknown operator"
2912  end select
2913 
2914  ! For red-black, perform two smoothing sub-steps so that all unknowns are
2915  ! updated per cycle
2916  if (mg%smoother_type == mg_smoother_gsrb) then
2917  mg%n_smoother_substeps = 2
2918  else
2919  mg%n_smoother_substeps = 1
2920  end if
2921  end subroutine mg_set_methods
2922 
2923  subroutine check_methods(mg)
2924  type(mg_t), intent(inout) :: mg
2925 
2926  if (.not. associated(mg%box_op) .or. &
2927  .not. associated(mg%box_smoother)) then
2928  call mg_set_methods(mg)
2929  end if
2930 
2931  end subroutine check_methods
2932 
2933  subroutine mg_add_timers(mg)
2934  type(mg_t), intent(inout) :: mg
2935  timer_total_vcycle = mg_add_timer(mg, "mg total V-cycle")
2936  timer_total_fmg = mg_add_timer(mg, "mg total FMG cycle")
2937  timer_smoother = mg_add_timer(mg, "mg smoother")
2938  timer_smoother_gc = mg_add_timer(mg, "mg smoother g.c.")
2939  timer_coarse = mg_add_timer(mg, "mg coarse")
2940  timer_correct = mg_add_timer(mg, "mg correct")
2941  timer_update_coarse = mg_add_timer(mg, "mg update coarse")
2942  end subroutine mg_add_timers
2943 
2944  !> Perform FAS-FMG cycle (full approximation scheme, full multigrid).
2945  subroutine mg_fas_fmg(mg, have_guess, max_res)
2946  type(mg_t), intent(inout) :: mg
2947  logical, intent(in) :: have_guess !< If false, start from phi = 0
2948  real(dp), intent(out), optional :: max_res !< Store max(abs(residual))
2949  integer :: lvl, i, id
2950 
2951  call check_methods(mg)
2952  if (timer_smoother == -1) call mg_add_timers(mg)
2953 
2954  call mg_timer_start(mg%timers(timer_total_fmg))
2955 
2956  if (.not. have_guess) then
2957  do lvl = mg%highest_lvl, mg%lowest_lvl, -1
2958  do i = 1, size(mg%lvls(lvl)%my_ids)
2959  id = mg%lvls(lvl)%my_ids(i)
2960  mg%boxes(id)%cc(:, :, :, mg_iphi) = 0.0_dp
2961  end do
2962  end do
2963  end if
2964 
2965  ! Ensure ghost cells are filled correctly
2966  call mg_fill_ghost_cells_lvl(mg, mg%highest_lvl, mg_iphi)
2967 
2968  do lvl = mg%highest_lvl, mg%lowest_lvl+1, -1
2969  ! Set rhs on coarse grid and restrict phi
2970  call mg_timer_start(mg%timers(timer_update_coarse))
2971  call update_coarse(mg, lvl)
2972  call mg_timer_end(mg%timers(timer_update_coarse))
2973  end do
2974 
2975  if (mg%subtract_mean) then
2976  ! For fully periodic solutions, the mean source term has to be zero
2977  call subtract_mean(mg, mg_irhs, .false.)
2978  end if
2979 
2980  do lvl = mg%lowest_lvl, mg%highest_lvl
2981  ! Store phi_old
2982  do i = 1, size(mg%lvls(lvl)%my_ids)
2983  id = mg%lvls(lvl)%my_ids(i)
2984  mg%boxes(id)%cc(:, :, :, mg_iold) = &
2985  mg%boxes(id)%cc(:, :, :, mg_iphi)
2986  end do
2987 
2988  if (lvl > mg%lowest_lvl) then
2989  ! Correct solution at this lvl using lvl-1 data
2990  ! phi = phi + prolong(phi_coarse - phi_old_coarse)
2991  call mg_timer_start(mg%timers(timer_correct))
2992  call correct_children(mg, lvl-1)
2993  call mg_timer_end(mg%timers(timer_correct))
2994 
2995  ! Update ghost cells
2996  call mg_fill_ghost_cells_lvl(mg, lvl, mg_iphi)
2997  end if
2998 
2999  ! Perform V-cycle, possibly set residual on last iteration
3000  if (lvl == mg%highest_lvl) then
3001  call mg_fas_vcycle(mg, lvl, max_res, standalone=.false.)
3002  else
3003  call mg_fas_vcycle(mg, lvl, standalone=.false.)
3004  end if
3005  end do
3006 
3007  call mg_timer_end(mg%timers(timer_total_fmg))
3008  end subroutine mg_fas_fmg
3009 
3010  !> Perform FAS V-cycle (full approximation scheme).
3011  subroutine mg_fas_vcycle(mg, highest_lvl, max_res, standalone)
3012  use mpi
3013  type(mg_t), intent(inout) :: mg
3014  integer, intent(in), optional :: highest_lvl !< Maximum level for V-cycle
3015  real(dp), intent(out), optional :: max_res !< Store max(abs(residual))
3016  !> Whether the V-cycle is called by itself (default: true)
3017  logical, intent(in), optional :: standalone
3018  integer :: lvl, min_lvl, i, max_lvl, ierr
3019  real(dp) :: init_res, res
3020  logical :: is_standalone
3021 
3022  is_standalone = .true.
3023  if (present(standalone)) is_standalone = standalone
3024 
3025  call check_methods(mg)
3026  if (timer_smoother == -1) call mg_add_timers(mg)
3027 
3028  if (is_standalone) &
3029  call mg_timer_start(mg%timers(timer_total_vcycle))
3030 
3031  if (mg%subtract_mean .and. .not. present(highest_lvl)) then
3032  ! Assume that this is a stand-alone call. For fully periodic solutions,
3033  ! ensure the mean source term is zero.
3034  call subtract_mean(mg, mg_irhs, .false.)
3035  end if
3036 
3037  min_lvl = mg%lowest_lvl
3038  max_lvl = mg%highest_lvl
3039  if (present(highest_lvl)) max_lvl = highest_lvl
3040 
3041  ! Ensure ghost cells are filled correctly
3042  if (is_standalone) then
3043  call mg_fill_ghost_cells_lvl(mg, max_lvl, mg_iphi)
3044  end if
3045 
3046  do lvl = max_lvl, min_lvl+1, -1
3047  ! Downwards relaxation
3048  call smooth_boxes(mg, lvl, mg%n_cycle_down)
3049 
3050  ! Set rhs on coarse grid, restrict phi, and copy mg_iphi to mg_iold for the
3051  ! correction later
3052  call mg_timer_start(mg%timers(timer_update_coarse))
3053  call update_coarse(mg, lvl)
3054  call mg_timer_end(mg%timers(timer_update_coarse))
3055  end do
3056 
3057  call mg_timer_start(mg%timers(timer_coarse))
3058  if (.not. all(mg%boxes(mg%lvls(min_lvl)%ids)%rank == &
3059  mg%boxes(mg%lvls(min_lvl)%ids(1))%rank)) then
3060  error stop "Multiple CPUs for coarse grid (not implemented yet)"
3061  end if
3062 
3063  init_res = max_residual_lvl(mg, min_lvl)
3064  do i = 1, mg%max_coarse_cycles
3065  call smooth_boxes(mg, min_lvl, mg%n_cycle_up+mg%n_cycle_down)
3066  res = max_residual_lvl(mg, min_lvl)
3067  if (res < mg%residual_coarse_rel * init_res .or. &
3068  res < mg%residual_coarse_abs) exit
3069  end do
3070  call mg_timer_end(mg%timers(timer_coarse))
3071 
3072  ! Do the upwards part of the v-cycle in the tree
3073  do lvl = min_lvl+1, max_lvl
3074  ! Correct solution at this lvl using lvl-1 data
3075  ! phi = phi + prolong(phi_coarse - phi_old_coarse)
3076  call mg_timer_start(mg%timers(timer_correct))
3077  call correct_children(mg, lvl-1)
3078 
3079  ! Have to fill ghost cells after correction
3080  call mg_fill_ghost_cells_lvl(mg, lvl, mg_iphi)
3081  call mg_timer_end(mg%timers(timer_correct))
3082 
3083  ! Upwards relaxation
3084  call smooth_boxes(mg, lvl, mg%n_cycle_up)
3085  end do
3086 
3087  if (present(max_res)) then
3088  init_res = 0.0_dp
3089  do lvl = min_lvl, max_lvl
3090  res = max_residual_lvl(mg, lvl)
3091  init_res = max(res, init_res)
3092  end do
3093  call mpi_allreduce(init_res, max_res, 1, &
3094  mpi_double, mpi_max, mg%comm, ierr)
3095  end if
3096 
3097  ! Subtract mean(phi) from phi
3098  if (mg%subtract_mean) then
3099  call subtract_mean(mg, mg_iphi, .true.)
3100  end if
3101 
3102  if (is_standalone) &
3103  call mg_timer_end(mg%timers(timer_total_vcycle))
3104  end subroutine mg_fas_vcycle
3105 
3106  subroutine subtract_mean(mg, iv, include_ghostcells)
3107  use mpi
3108  type(mg_t), intent(inout) :: mg
3109  integer, intent(in) :: iv
3110  logical, intent(in) :: include_ghostcells
3111  integer :: i, id, lvl, nc, ierr
3112  real(dp) :: sum_iv, mean_iv, volume
3113 
3114  nc = mg%box_size
3115  sum_iv = get_sum(mg, iv)
3116  call mpi_allreduce(sum_iv, mean_iv, 1, &
3117  mpi_double, mpi_sum, mg%comm, ierr)
3118 
3119  ! Divide by total grid volume to get mean
3120  volume = nc**3 * product(mg%dr(:, 1)) * size(mg%lvls(1)%ids)
3121  mean_iv = mean_iv / volume
3122 
3123  do lvl = mg%lowest_lvl, mg%highest_lvl
3124  nc = mg%box_size_lvl(lvl)
3125 
3126  do i = 1, size(mg%lvls(lvl)%my_ids)
3127  id = mg%lvls(lvl)%my_ids(i)
3128  if (include_ghostcells) then
3129  mg%boxes(id)%cc(:, :, :, iv) = &
3130  mg%boxes(id)%cc(:, :, :, iv) - mean_iv
3131  else
3132  mg%boxes(id)%cc(1:nc, 1:nc, 1:nc, iv) = &
3133  mg%boxes(id)%cc(1:nc, 1:nc, 1:nc, iv) - mean_iv
3134  end if
3135  end do
3136  end do
3137  end subroutine subtract_mean
3138 
3139  real(dp) function get_sum(mg, iv)
3140  type(mg_t), intent(in) :: mg
3141  integer, intent(in) :: iv
3142  integer :: lvl, i, id, nc
3143  real(dp) :: w
3144 
3145  get_sum = 0.0_dp
3146  do lvl = 1, mg%highest_lvl
3147  nc = mg%box_size_lvl(lvl)
3148  w = product(mg%dr(:, lvl)) ! Adjust for non-Cartesian cases
3149  do i = 1, size(mg%lvls(lvl)%my_leaves)
3150  id = mg%lvls(lvl)%my_leaves(i)
3151  get_sum = get_sum + w * &
3152  sum(mg%boxes(id)%cc(1:nc, 1:nc, 1:nc, iv))
3153  end do
3154  end do
3155  end function get_sum
3156 
3157  real(dp) function max_residual_lvl(mg, lvl)
3158  type(mg_t), intent(inout) :: mg
3159  integer, intent(in) :: lvl
3160  integer :: i, id, nc
3161  real(dp) :: res
3162 
3163  nc = mg%box_size_lvl(lvl)
3164  max_residual_lvl = 0.0_dp
3165 
3166  do i = 1, size(mg%lvls(lvl)%my_ids)
3167  id = mg%lvls(lvl)%my_ids(i)
3168  call residual_box(mg, id, nc)
3169  res = maxval(abs(mg%boxes(id)%cc(1:nc, 1:nc, 1:nc, mg_ires)))
3170  max_residual_lvl = max(max_residual_lvl, res)
3171  end do
3172  end function max_residual_lvl
3173 
3174  ! subroutine solve_coarse_grid(mg)
3175  ! use m_fishpack
3176  ! type(mg_t), intent(inout) :: mg
3177 
3178  ! real(dp) :: rhs(mg%box_size, mg%box_size, mg%box_size)
3179  ! real(dp) :: rmin(3), rmax(3)
3180  ! integer :: nc, nx(3), my_boxes, total_boxes
3181 
3182  ! my_boxes = size(mg%lvls(1)%my_ids)
3183  ! total_boxes = size(mg%lvls(1)%ids)
3184  ! nc = mg%box_size
3185 
3186  ! if (my_boxes == total_boxes) then
3187  ! nx(:) = nc
3188  ! rmin = [0.0_dp, 0.0_dp, 0.0_dp]
3189  ! rmax = mg%dr(1) * [nc, nc, nc]
3190  ! rhs = mg%boxes(1)%cc(1:nc, 1:nc, 1:nc, mg_irhs)
3191 
3192  ! #if 3 == 2
3193  ! call fishpack_2d(nx, rhs, mg%bc, rmin, rmax)
3194  ! #elif 3 == 3
3195  ! call fishpack_3d(nx, rhs, mg%bc, rmin, rmax)
3196  ! #endif
3197 
3198  ! mg%boxes(1)%cc(1:nc, 1:nc, 1:nc, mg_iphi) = rhs
3199  ! else if (my_boxes > 0) then
3200  ! error stop "Boxes at level 1 at different processors"
3201  ! end if
3202 
3203  ! call fill_ghost_cells_lvl(mg, 1)
3204  ! end subroutine solve_coarse_grid
3205 
3206  ! Set rhs on coarse grid, restrict phi, and copy mg_iphi to mg_iold for the
3207  ! correction later
3208  subroutine update_coarse(mg, lvl)
3209  type(mg_t), intent(inout) :: mg !< Tree containing full grid
3210  integer, intent(in) :: lvl !< Update coarse values at lvl-1
3211  integer :: i, id, nc, nc_c
3212 
3213  nc = mg%box_size_lvl(lvl)
3214  nc_c = mg%box_size_lvl(lvl-1)
3215 
3216  ! Compute residual
3217  do i = 1, size(mg%lvls(lvl)%my_ids)
3218  id = mg%lvls(lvl)%my_ids(i)
3219  call residual_box(mg, id, nc)
3220  end do
3221 
3222  ! Restrict phi and the residual
3223  call mg_restrict_lvl(mg, mg_iphi, lvl)
3224  call mg_restrict_lvl(mg, mg_ires, lvl)
3225 
3226  call mg_fill_ghost_cells_lvl(mg, lvl-1, mg_iphi)
3227 
3228  ! Set rhs_c = laplacian(phi_c) + restrict(res) where it is refined, and
3229  ! store current coarse phi in old.
3230  do i = 1, size(mg%lvls(lvl-1)%my_parents)
3231  id = mg%lvls(lvl-1)%my_parents(i)
3232 
3233  ! Set rhs = L phi
3234  call mg%box_op(mg, id, nc_c, mg_irhs)
3235 
3236  ! Add the fine grid residual to rhs
3237  mg%boxes(id)%cc(1:nc_c, 1:nc_c, 1:nc_c, mg_irhs) = &
3238  mg%boxes(id)%cc(1:nc_c, 1:nc_c, 1:nc_c, mg_irhs) + &
3239  mg%boxes(id)%cc(1:nc_c, 1:nc_c, 1:nc_c, mg_ires)
3240 
3241  ! Story a copy of phi
3242  mg%boxes(id)%cc(:, :, :, mg_iold) = &
3243  mg%boxes(id)%cc(:, :, :, mg_iphi)
3244  end do
3245  end subroutine update_coarse
3246 
3247  ! Sets phi = phi + prolong(phi_coarse - phi_old_coarse)
3248  subroutine correct_children(mg, lvl)
3249  type(mg_t), intent(inout) :: mg
3250  integer, intent(in) :: lvl
3251  integer :: i, id
3252 
3253  do i = 1, size(mg%lvls(lvl)%my_parents)
3254  id = mg%lvls(lvl)%my_parents(i)
3255 
3256  ! Store the correction in mg_ires
3257  mg%boxes(id)%cc(:, :, :, mg_ires) = &
3258  mg%boxes(id)%cc(:, :, :, mg_iphi) - &
3259  mg%boxes(id)%cc(:, :, :, mg_iold)
3260  end do
3261 
3262  call mg_prolong(mg, lvl, mg_ires, mg_iphi, mg%box_prolong, add=.true.)
3263  end subroutine correct_children
3264 
3265  subroutine smooth_boxes(mg, lvl, n_cycle)
3266  type(mg_t), intent(inout) :: mg
3267  integer, intent(in) :: lvl
3268  integer, intent(in) :: n_cycle !< Number of cycles to perform
3269  integer :: n, i, id, nc
3270 
3271  nc = mg%box_size_lvl(lvl)
3272 
3273  do n = 1, n_cycle * mg%n_smoother_substeps
3274  call mg_timer_start(mg%timers(timer_smoother))
3275  do i = 1, size(mg%lvls(lvl)%my_ids)
3276  id = mg%lvls(lvl)%my_ids(i)
3277  call mg%box_smoother(mg, id, nc, n)
3278  end do
3279  call mg_timer_end(mg%timers(timer_smoother))
3280 
3281  call mg_timer_start(mg%timers(timer_smoother_gc))
3282  call mg_fill_ghost_cells_lvl(mg, lvl, mg_iphi)
3283  call mg_timer_end(mg%timers(timer_smoother_gc))
3284  end do
3285  end subroutine smooth_boxes
3286 
3287  subroutine residual_box(mg, id, nc)
3288  type(mg_t), intent(inout) :: mg
3289  integer, intent(in) :: id
3290  integer, intent(in) :: nc
3291 
3292  call mg%box_op(mg, id, nc, mg_ires)
3293 
3294  mg%boxes(id)%cc(1:nc, 1:nc, 1:nc, mg_ires) = &
3295  mg%boxes(id)%cc(1:nc, 1:nc, 1:nc, mg_irhs) &
3296  - mg%boxes(id)%cc(1:nc, 1:nc, 1:nc, mg_ires)
3297  end subroutine residual_box
3298 
3299  !> Apply operator to the tree and store in variable i_out
3300  subroutine mg_apply_op(mg, i_out, op)
3301  type(mg_t), intent(inout) :: mg
3302  integer, intent(in) :: i_out
3303  procedure(mg_box_op), optional :: op
3304  integer :: lvl, i, id, nc
3305 
3306  do lvl = mg%lowest_lvl, mg%highest_lvl
3307  nc = mg%box_size_lvl(lvl)
3308  do i = 1, size(mg%lvls(lvl)%my_ids)
3309  id = mg%lvls(lvl)%my_ids(i)
3310  if (present(op)) then
3311  call op(mg, id, nc, i_out)
3312  else
3313  call mg%box_op(mg, id, nc, i_out)
3314  end if
3315  end do
3316  end do
3317  end subroutine mg_apply_op
3318 
3319  !! File ../src/m_restrict.f90
3320 
3321  !> Specify minimum buffer size (per process) for communication
3322  subroutine mg_restrict_buffer_size(mg, n_send, n_recv, dsize)
3323  type(mg_t), intent(inout) :: mg
3324  integer, intent(out) :: n_send(0:mg%n_cpu-1)
3325  integer, intent(out) :: n_recv(0:mg%n_cpu-1)
3326  integer, intent(out) :: dsize
3327  integer :: n_out(0:mg%n_cpu-1, mg%first_normal_lvl:mg%highest_lvl)
3328  integer :: n_in(0:mg%n_cpu-1, mg%first_normal_lvl:mg%highest_lvl)
3329  integer :: lvl, i, id, p_id, p_rank
3330  integer :: i_c, c_id, c_rank, min_lvl
3331 
3332  n_out(:, :) = 0
3333  n_in(:, :) = 0
3334  min_lvl = max(mg%lowest_lvl+1, mg%first_normal_lvl)
3335 
3336  do lvl = min_lvl, mg%highest_lvl
3337  ! Number of messages to receive (at lvl-1)
3338  do i = 1, size(mg%lvls(lvl-1)%my_parents)
3339  id = mg%lvls(lvl-1)%my_parents(i)
3340  do i_c = 1, mg_num_children
3341  c_id = mg%boxes(id)%children(i_c)
3342 
3343  if (c_id > mg_no_box) then
3344  c_rank = mg%boxes(c_id)%rank
3345  if (c_rank /= mg%my_rank) then
3346  n_in(c_rank, lvl) = n_in(c_rank, lvl) + 1
3347  end if
3348  end if
3349  end do
3350  end do
3351 
3352  ! Number of messages to send
3353  do i = 1, size(mg%lvls(lvl)%my_ids)
3354  id = mg%lvls(lvl)%my_ids(i)
3355 
3356  p_id = mg%boxes(id)%parent
3357  p_rank = mg%boxes(p_id)%rank
3358  if (p_rank /= mg%my_rank) then
3359  n_out(p_rank, lvl) = n_out(p_rank, lvl) + 1
3360  end if
3361 
3362  end do
3363  end do
3364 
3365  allocate(mg%comm_restrict%n_send(0:mg%n_cpu-1, &
3366  mg%first_normal_lvl:mg%highest_lvl))
3367  allocate(mg%comm_restrict%n_recv(0:mg%n_cpu-1, &
3368  mg%first_normal_lvl:mg%highest_lvl))
3369  mg%comm_restrict%n_send = n_out
3370  mg%comm_restrict%n_recv = n_in
3371 
3372  dsize = (mg%box_size/2)**3
3373  n_send = maxval(n_out, dim=2)
3374  n_recv = maxval(n_in, dim=2)
3375  end subroutine mg_restrict_buffer_size
3376 
3377  !> Restrict all levels
3378  subroutine mg_restrict(mg, iv)
3379  type(mg_t), intent(inout) :: mg
3380  integer, intent(in) :: iv
3381  integer :: lvl
3382 
3383  do lvl = mg%highest_lvl, mg%lowest_lvl+1, -1
3384  call mg_restrict_lvl(mg, iv, lvl)
3385  end do
3386  end subroutine mg_restrict
3387 
3388  !> Restrict from lvl to lvl-1
3389  subroutine mg_restrict_lvl(mg, iv, lvl)
3391  type(mg_t), intent(inout) :: mg
3392  integer, intent(in) :: iv
3393  integer, intent(in) :: lvl
3394  integer :: i, id, dsize, nc
3395 
3396  if (lvl <= mg%lowest_lvl) error stop "cannot restrict lvl <= lowest_lvl"
3397 
3398  nc = mg%box_size_lvl(lvl)
3399 
3400  if (lvl >= mg%first_normal_lvl) then
3401  dsize = (nc/2)**3
3402 
3403  mg%buf(:)%i_send = 0
3404  mg%buf(:)%i_ix = 0
3405 
3406  do i = 1, size(mg%lvls(lvl)%my_ids)
3407  id = mg%lvls(lvl)%my_ids(i)
3408  call restrict_set_buffer(mg, id, iv)
3409  end do
3410 
3411  mg%buf(:)%i_recv = mg%comm_restrict%n_recv(:, lvl) * dsize
3412  call sort_and_transfer_buffers(mg, dsize)
3413  mg%buf(:)%i_recv = 0
3414  end if
3415 
3416  do i = 1, size(mg%lvls(lvl-1)%my_parents)
3417  id = mg%lvls(lvl-1)%my_parents(i)
3418  call restrict_onto(mg, id, nc, iv)
3419  end do
3420  end subroutine mg_restrict_lvl
3421 
3422  subroutine restrict_set_buffer(mg, id, iv)
3423  type(mg_t), intent(inout) :: mg
3424  integer, intent(in) :: id
3425  integer, intent(in) :: iv
3426  integer :: i, j, k, n, hnc, p_id, p_rank
3427  real(dp) :: tmp(mg%box_size/2, mg%box_size/2, mg%box_size/2)
3428 
3429  hnc = mg%box_size/2
3430  p_id = mg%boxes(id)%parent
3431  p_rank = mg%boxes(p_id)%rank
3432 
3433  if (p_rank /= mg%my_rank) then
3434  do k = 1, hnc
3435  do j = 1, hnc
3436  do i = 1, hnc
3437  tmp(i, j, k) = 0.125_dp * sum(mg%boxes(id)%cc(2*i-1:2*i, &
3438  2*j-1:2*j, 2*k-1:2*k, iv))
3439  end do
3440  end do
3441  end do
3442 
3443  ! Buffer
3444  n = size(tmp)
3445  i = mg%buf(p_rank)%i_send
3446  mg%buf(p_rank)%send(i+1:i+n) = pack(tmp, .true.)
3447  mg%buf(p_rank)%i_send = mg%buf(p_rank)%i_send + n
3448 
3449  ! To later sort the send buffer according to parent order
3450  i = mg%buf(p_rank)%i_ix
3451  n = mg_ix_to_ichild(mg%boxes(id)%ix)
3452  mg%buf(p_rank)%ix(i+1) = mg_num_children * p_id + n
3453  mg%buf(p_rank)%i_ix = mg%buf(p_rank)%i_ix + 1
3454  end if
3455  end subroutine restrict_set_buffer
3456 
3457  subroutine restrict_onto(mg, id, nc, iv)
3458  type(mg_t), intent(inout) :: mg
3459  integer, intent(in) :: id
3460  integer, intent(in) :: nc
3461  integer, intent(in) :: iv
3462  integer :: i, j, k, hnc, dsize, i_c, c_id
3463  integer :: c_rank, dix(3)
3464 
3465  hnc = nc/2
3466  dsize = hnc**3
3467 
3468  do i_c = 1, mg_num_children
3469  c_id = mg%boxes(id)%children(i_c)
3470  if (c_id == mg_no_box) cycle ! For coarsened grid
3471  c_rank = mg%boxes(c_id)%rank
3472  dix = mg_get_child_offset(mg, c_id)
3473 
3474  if (c_rank == mg%my_rank) then
3475  do j=1, hnc; do i=1, hnc; do k=1, hnc
3476  mg%boxes(id)%cc(dix(1)+i, dix(2)+j, dix(3)+k, iv) = &
3477  0.125_dp * sum(mg%boxes(c_id)%cc(2*i-1:2*i, &
3478  2*j-1:2*j, 2*k-1:2*k, iv))
3479  end do; end do; end do
3480  else
3481  i = mg%buf(c_rank)%i_recv
3482  mg%boxes(id)%cc(dix(1)+1:dix(1)+hnc, &
3483  dix(2)+1:dix(2)+hnc, dix(3)+1:dix(3)+hnc, iv) = &
3484  reshape(mg%buf(c_rank)%recv(i+1:i+dsize), [hnc, hnc, hnc])
3485  mg%buf(c_rank)%i_recv = mg%buf(c_rank)%i_recv + dsize
3486  end if
3487  end do
3488 
3489  end subroutine restrict_onto
3490 
3491  !! File ../src/m_mrgrnk.f90
3492 
3493  Subroutine i_mrgrnk (XDONT, IRNGT)
3494  ! __________________________________________________________
3495  ! MRGRNK = Merge-sort ranking of an array
3496  ! For performance reasons, the first 2 passes are taken
3497  ! out of the standard loop, and use dedicated coding.
3498  ! __________________________________________________________
3499  ! __________________________________________________________
3500  Integer, Dimension (:), Intent (In) :: XDONT
3501  Integer, Dimension (:), Intent (Out) :: IRNGT
3502  ! __________________________________________________________
3503  Integer :: XVALA, XVALB
3504  !
3505  Integer, Dimension (SIZE(IRNGT)) :: JWRKT
3506  Integer :: LMTNA, LMTNC, IRNG1, IRNG2
3507  Integer :: NVAL, IIND, IWRKD, IWRK, IWRKF, JINDA, IINDA, IINDB
3508  !
3509  nval = min(SIZE(xdont), SIZE(irngt))
3510  Select Case (nval)
3511  Case (:0)
3512  Return
3513  Case (1)
3514  irngt(1) = 1
3515  Return
3516  Case Default
3517  Continue
3518  End Select
3519  !
3520  ! Fill-in the index array, creating ordered couples
3521  !
3522  Do iind = 2, nval, 2
3523  If (xdont(iind-1) <= xdont(iind)) Then
3524  irngt(iind-1) = iind - 1
3525  irngt(iind) = iind
3526  Else
3527  irngt(iind-1) = iind
3528  irngt(iind) = iind - 1
3529  End If
3530  End Do
3531  If (modulo(nval, 2) /= 0) Then
3532  irngt(nval) = nval
3533  End If
3534  !
3535  ! We will now have ordered subsets A - B - A - B - ...
3536  ! and merge A and B couples into C - C - ...
3537  !
3538  lmtna = 2
3539  lmtnc = 4
3540  !
3541  ! First iteration. The length of the ordered subsets goes from 2 to 4
3542  !
3543  Do
3544  If (nval <= 2) Exit
3545  !
3546  ! Loop on merges of A and B into C
3547  !
3548  Do iwrkd = 0, nval - 1, 4
3549  If ((iwrkd+4) > nval) Then
3550  If ((iwrkd+2) >= nval) Exit
3551  !
3552  ! 1 2 3
3553  !
3554  If (xdont(irngt(iwrkd+2)) <= xdont(irngt(iwrkd+3))) Exit
3555  !
3556  ! 1 3 2
3557  !
3558  If (xdont(irngt(iwrkd+1)) <= xdont(irngt(iwrkd+3))) Then
3559  irng2 = irngt(iwrkd+2)
3560  irngt(iwrkd+2) = irngt(iwrkd+3)
3561  irngt(iwrkd+3) = irng2
3562  !
3563  ! 3 1 2
3564  !
3565  Else
3566  irng1 = irngt(iwrkd+1)
3567  irngt(iwrkd+1) = irngt(iwrkd+3)
3568  irngt(iwrkd+3) = irngt(iwrkd+2)
3569  irngt(iwrkd+2) = irng1
3570  End If
3571  Exit
3572  End If
3573  !
3574  ! 1 2 3 4
3575  !
3576  If (xdont(irngt(iwrkd+2)) <= xdont(irngt(iwrkd+3))) cycle
3577  !
3578  ! 1 3 x x
3579  !
3580  If (xdont(irngt(iwrkd+1)) <= xdont(irngt(iwrkd+3))) Then
3581  irng2 = irngt(iwrkd+2)
3582  irngt(iwrkd+2) = irngt(iwrkd+3)
3583  If (xdont(irng2) <= xdont(irngt(iwrkd+4))) Then
3584  ! 1 3 2 4
3585  irngt(iwrkd+3) = irng2
3586  Else
3587  ! 1 3 4 2
3588  irngt(iwrkd+3) = irngt(iwrkd+4)
3589  irngt(iwrkd+4) = irng2
3590  End If
3591  !
3592  ! 3 x x x
3593  !
3594  Else
3595  irng1 = irngt(iwrkd+1)
3596  irng2 = irngt(iwrkd+2)
3597  irngt(iwrkd+1) = irngt(iwrkd+3)
3598  If (xdont(irng1) <= xdont(irngt(iwrkd+4))) Then
3599  irngt(iwrkd+2) = irng1
3600  If (xdont(irng2) <= xdont(irngt(iwrkd+4))) Then
3601  ! 3 1 2 4
3602  irngt(iwrkd+3) = irng2
3603  Else
3604  ! 3 1 4 2
3605  irngt(iwrkd+3) = irngt(iwrkd+4)
3606  irngt(iwrkd+4) = irng2
3607  End If
3608  Else
3609  ! 3 4 1 2
3610  irngt(iwrkd+2) = irngt(iwrkd+4)
3611  irngt(iwrkd+3) = irng1
3612  irngt(iwrkd+4) = irng2
3613  End If
3614  End If
3615  End Do
3616  !
3617  ! The Cs become As and Bs
3618  !
3619  lmtna = 4
3620  Exit
3621  End Do
3622  !
3623  ! Iteration loop. Each time, the length of the ordered subsets
3624  ! is doubled.
3625  !
3626  Do
3627  If (lmtna >= nval) Exit
3628  iwrkf = 0
3629  lmtnc = 2 * lmtnc
3630  !
3631  ! Loop on merges of A and B into C
3632  !
3633  Do
3634  iwrk = iwrkf
3635  iwrkd = iwrkf + 1
3636  jinda = iwrkf + lmtna
3637  iwrkf = iwrkf + lmtnc
3638  If (iwrkf >= nval) Then
3639  If (jinda >= nval) Exit
3640  iwrkf = nval
3641  End If
3642  iinda = 1
3643  iindb = jinda + 1
3644  !
3645  ! Shortcut for the case when the max of A is smaller
3646  ! than the min of B. This line may be activated when the
3647  ! initial set is already close to sorted.
3648  !
3649  ! IF (XDONT(IRNGT(JINDA)) <= XDONT(IRNGT(IINDB))) CYCLE
3650  !
3651  ! One steps in the C subset, that we build in the final rank array
3652  !
3653  ! Make a copy of the rank array for the merge iteration
3654  !
3655  jwrkt(1:lmtna) = irngt(iwrkd:jinda)
3656  !
3657  xvala = xdont(jwrkt(iinda))
3658  xvalb = xdont(irngt(iindb))
3659  !
3660  Do
3661  iwrk = iwrk + 1
3662  !
3663  ! We still have unprocessed values in both A and B
3664  !
3665  If (xvala > xvalb) Then
3666  irngt(iwrk) = irngt(iindb)
3667  iindb = iindb + 1
3668  If (iindb > iwrkf) Then
3669  ! Only A still with unprocessed values
3670  irngt(iwrk+1:iwrkf) = jwrkt(iinda:lmtna)
3671  Exit
3672  End If
3673  xvalb = xdont(irngt(iindb))
3674  Else
3675  irngt(iwrk) = jwrkt(iinda)
3676  iinda = iinda + 1
3677  If (iinda > lmtna) exit! Only B still with unprocessed values
3678  xvala = xdont(jwrkt(iinda))
3679  End If
3680  !
3681  End Do
3682  End Do
3683  !
3684  ! The Cs become As and Bs
3685  !
3686  lmtna = 2 * lmtna
3687  End Do
3688  !
3689  Return
3690  !
3691  End Subroutine i_mrgrnk
3692 
3693 end module m_octree_mg_3d
subroutine, public mg_comm_init(mg, comm)
Prolong from a parent to a child with index offset dix. This method could sometimes work better than ...
integer, parameter, public mg_num_children
integer, parameter, public mg_max_timers
Maximum number of timers to use.
subroutine, public mg_timers_show(mg)
elemental logical function, public mg_has_children(box)
Return .true. if a box has children.
subroutine, public mg_load_balance_simple(mg)
Load balance all boxes in the multigrid tree, by simply distributing the load per grid level...
integer, parameter, public i8
Type for 64-bit integers.
integer, parameter, public mg_vlaplacian
Indicates a variable-coefficient Laplacian.
subroutine, public mg_build_rectangle(mg, domain_size, box_size, dx, r_min, periodic, n_finer)
integer, parameter, public mg_no_box
Special value that indicates there is no box.
integer, parameter, public mg_vhelmholtz
Indicates a variable-coefficient Helmholtz equation.
integer, parameter, public mg_max_num_vars
Maximum number of variables.
subroutine, public sort_and_transfer_buffers(mg, dsize)
pure integer function, dimension(3), public mg_get_child_offset(mg, id)
Get the offset of a box with respect to its parent (e.g. in 2d, there can be a child at offset 0...
Lists of blocks per refinement level.
subroutine, public mg_prolong(mg, lvl, iv, iv_to, method, add)
Prolong variable iv from lvl to variable iv_to at lvl+1.
integer, parameter, public mg_ndim
Problem dimension.
integer, parameter, public mg_iold
Index of previous solution (used for correction)
subroutine, public mg_prolong_buffer_size(mg, n_send, n_recv, dsize)
Specify minimum buffer size (per process) for communication.
integer, parameter, public mg_neighb_highx
subroutine, public mg_ghost_cell_buffer_size(mg, n_send, n_recv, dsize)
Specify minimum buffer size (per process) for communication.
integer, parameter, public mg_num_vars
Number of predefined multigrid variables.
integer, parameter, public mg_bc_neumann
Value to indicate a Neumann boundary condition.
integer, parameter, public mg_lvl_hi
Maximum allowed grid level.
integer, parameter, public mg_smoother_gsrb
logical, dimension(3, 8), parameter, public mg_child_low
Box data structure.
integer, parameter, public mg_bc_continuous
Value to indicate a continuous boundary condition.
integer, parameter, public mg_iveps
Index of the variable coefficient (at cell centers)
integer, parameter, public mg_iphi
Index of solution.
integer, parameter, public mg_cylindrical
Cylindrical coordinate system.
integer, parameter, public mg_laplacian
Indicates a standard Laplacian.
integer, dimension(3, 6), parameter, public mg_neighb_dix
integer, parameter, public mg_bc_dirichlet
Value to indicate a Dirichlet boundary condition.
subroutine, public mg_load_balance(mg)
Load balance all boxes in the multigrid tree. Compared to mg_load_balance_simple, this method does a ...
subroutine, public mg_load_balance_parents(mg)
Load balance the parents (non-leafs). Assign them to the rank that has most children.
logical, dimension(6), parameter, public mg_neighb_low
subroutine, public mg_add_children(mg, id)
subroutine, public diffusion_solve_vcoeff(mg, dt, order, max_res)
Solve a diffusion equation implicitly, assuming a variable diffusion coefficient which has been store...
integer(i8) function, public mg_number_of_unknowns(mg)
Determine total number of unknowns (on leaves)
integer, parameter, public mg_lvl_lo
Minimum allowed grid level.
subroutine, public mg_timer_start(timer)
integer, parameter, public mg_helmholtz
Indicates a constant-coefficient Helmholtz equation.
real(dp) function get_sum(mg, iv)
subroutine, public mg_restrict(mg, iv)
Restrict all levels.
integer, dimension(6), parameter, public mg_neighb_high_pm
integer, parameter, public mg_ires
Index of residual.
subroutine, public mg_fill_ghost_cells(mg, iv)
Fill ghost cells at all grid levels.
subroutine, public mg_restrict_lvl(mg, iv, lvl)
Restrict from lvl to lvl-1.
subroutine, public mg_prolong_sparse(mg, p_id, dix, nc, iv, fine)
Prolong from a parent to a child with index offset dix.
integer, parameter, public mg_cartesian
Cartesian coordinate system.
integer, dimension(6), parameter, public mg_neighb_rev
subroutine subtract_mean(mg, iv, include_ghostcells)
subroutine, public vhelmholtz_set_methods(mg)
subroutine, public mg_set_refinement_boundaries(boxes, level)
Create a list of refinement boundaries (from the coarse side)
subroutine, public diffusion_solve(mg, dt, diffusion_coeff, order, max_res)
Solve a diffusion equation implicitly, assuming a constant diffusion coefficient. The solution at tim...
subroutine, public mg_set_neighbors_lvl(mg, lvl)
subroutine, public mg_set_next_level_ids(mg, lvl)
Buffer type (one is used for each pair of communicating processes)
subroutine, public helmholtz_set_lambda(lambda)
subroutine, public vlaplacian_set_methods(mg)
To fill ghost cells near physical boundaries.
subroutine, public mg_set_methods(mg)
integer, parameter, public mg_num_neighbors
integer, parameter, public mg_neighb_highz
subroutine, public mg_fill_ghost_cells_lvl(mg, lvl, iv)
Fill ghost cells at a grid level.
subroutine, public mg_apply_op(mg, i_out, op)
Apply operator to the tree and store in variable i_out.
subroutine, public helmholtz_set_methods(mg)
subroutine, public mg_fas_fmg(mg, have_guess, max_res)
Perform FAS-FMG cycle (full approximation scheme, full multigrid).
subroutine, public mg_fas_vcycle(mg, highest_lvl, max_res, standalone)
Perform FAS V-cycle (full approximation scheme).
integer, parameter, public mg_smoother_gs
integer, parameter, public mg_neighb_lowy
subroutine sort_sendbuf(gc, dsize)
Sort send buffers according to the idbuf array.
integer function, public mg_add_timer(mg, name)
subroutine, public mg_allocate_storage(mg)
Allocate communication buffers and local boxes for a tree that has already been created.
subroutine, public mg_set_leaves_parents(boxes, level)
Create a list of leaves and a list of parents for a level.
integer, parameter, public mg_physical_boundary
Special value that indicates there is a physical boundary.
integer, parameter, public mg_neighb_lowx
subroutine, public mg_restrict_buffer_size(mg, n_send, n_recv, dsize)
Specify minimum buffer size (per process) for communication.
subroutine, public mg_get_face_coords(box, nb, nc, x)
Get coordinates at the face of a box.
integer, parameter, public dp
Type of reals.
subroutine, public mg_deallocate_storage(mg)
Deallocate all allocatable arrays.
integer, parameter, public mg_neighb_highy
real(dp), public, protected vhelmholtz_lambda
Module for implicitly solving diffusion equations.
integer, parameter, public mg_irhs
Index of right-hand side.
subroutine, public laplacian_set_methods(mg)
integer, dimension(8, 3), parameter, public mg_child_rev
integer, parameter, public mg_neighb_lowz
integer, parameter, public mg_spherical
Spherical coordinate system.
subroutine, public mg_timer_end(timer)
pure integer function, public mg_highest_uniform_lvl(mg)
integer, dimension(3, 8), parameter, public mg_child_dix
subroutine, public mg_phi_bc_store(mg)
Store boundary conditions for the solution variable, this can speed up calculations if the same bound...
integer, dimension(6), parameter, public mg_neighb_dim
integer, dimension(4, 6), parameter, public mg_child_adj_nb
subroutine, public vhelmholtz_set_lambda(lambda)
real(dp), public, protected helmholtz_lambda
Module for load balancing a tree (that already has been constructed). The load balancing determines w...
integer function, public mg_ix_to_ichild(ix)
Compute the child index for a box with spatial index ix. With child index we mean the index in the ch...
integer, parameter, public mg_smoother_jacobi