mesh_init.F90

Go to the documentation of this file.

!! Copyright (C) 2002-2006 M. Marques, A. Castro, A. Rubio, G. Bertsch
!! Copyright (C) 2021 S. Ohlmann
!!
!! This program is free software; you can redistribute it and/or modify
!! it under the terms of the GNU General Public License as published by
!! the Free Software Foundation; either version 2, or (at your option)
!! any later version.
!!
!! This program is distributed in the hope that it will be useful,
!! but WITHOUT ANY WARRANTY; without even the implied warranty of
!! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
!! GNU General Public License for more details.
!!
!! You should have received a copy of the GNU General Public License
!! along with this program; if not, write to the Free Software
!! Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
!! 02110-1301, USA.
!!
 
#include "global.h"
 
module mesh_init_oct_m
  use accel_oct_m
  use types_oct_m
  use affine_coordinates_oct_m
  use box_oct_m
  use calc_mode_par_oct_m
  use checksum_interface_oct_m
  use coordinate_system_oct_m
  use debug_oct_m
  use global_oct_m
  use iihash_oct_m
  use index_oct_m
  use math_oct_m
  use merge_sorted_oct_m
  use mesh_oct_m
  use mesh_cube_map_oct_m
  use mesh_partition_oct_m
  use messages_oct_m
  use mpi_oct_m
  use mpi_lib_oct_m
  use multicomm_oct_m
  use namespace_oct_m
  use par_vec_oct_m
  use parser_oct_m
  use partition_oct_m
  use profiling_oct_m
  use restart_oct_m
  use sort_oct_m
  use space_oct_m
  use stencil_oct_m
  use utils_oct_m
 
  implicit none
 
  private
  public ::                    &
    mesh_init_stage_1,         &
    mesh_init_stage_2,         &
    mesh_init_stage_3
 
 
  integer, parameter :: INNER_POINT = 1
  integer, parameter :: ENLARGEMENT_POINT = 2
  integer, parameter :: BOUNDARY = -1
 
contains
 
! ---------------------------------------------------------
  subroutine mesh_init_stage_1(mesh, namespace, space, box, coord_system, spacing, enlarge)
    class(mesh_t),               intent(inout) :: mesh
    type(namespace_t),           intent(in)    :: namespace
    class(space_t),              intent(in)    :: space
    class(box_t), target,        intent(in)    :: box
    class(coordinate_system_t), target, intent(in) :: coord_system
    real(real64),                intent(in)    :: spacing(1:space%dim)
    integer,                     intent(in)    :: enlarge(1:space%dim)
 
    integer :: idir, jj, delta
    real(real64)   :: x(space%dim), chi(space%dim), spacing_new(-1:1), box_bounds(2, space%dim)
    logical :: out
 
    push_sub_with_profile(mesh_init_stage_1)
 
    mesh%box => box
 
    safe_allocate(mesh%spacing(1:space%dim))
    mesh%spacing = spacing ! this number can change in the following
 
    mesh%use_curvilinear = coord_system%local_basis
    ! the volume element is per point for curvilinear coordinates
    mesh%use_vol_pp = mesh%use_curvilinear
    mesh%coord_system => coord_system
 
    call index_init(mesh%idx, space%dim)
    mesh%idx%enlarge = enlarge
 
    ! get box bounds along the axes that generate the grid points
    select type (coord_system)
    class is (affine_coordinates_t)
      box_bounds = box%bounds(coord_system%basis)
    class default
      box_bounds = box%bounds()
    end select
 
    ! adjust nr
    do idir = 1, space%dim
      chi = m_zero
      ! the upper border
      jj = 0
      out = .false.
      do while(.not.out)
        jj = jj + 1
        chi(idir) = real(jj, real64)  * mesh%spacing(idir)
        if (mesh%use_curvilinear) then
          x = coord_system%to_cartesian(chi)
          out = x(idir) > maxval(abs(box_bounds(:, idir))) + box_boundary_delta
        else
          ! do the same comparison here as in simul_box_contains_points
          out = chi(idir) > maxval(abs(box_bounds(:, idir))) + box_boundary_delta
        end if
      end do
      mesh%idx%nr(2, idir) = jj - 1
    end do
 
    ! we have a symmetric mesh (for now)
    mesh%idx%nr(1,:) = -mesh%idx%nr(2,:)
 
    ! we have to adjust a couple of things for the periodic directions
    do idir = 1, space%periodic_dim
      if (mesh%idx%nr(2, idir) == 0) then
        ! this happens if Spacing > box size
        mesh%idx%nr(2, idir) =  1
        mesh%idx%nr(1, idir) = -1
      end if
 
      ! We have to adjust the spacing to be commensurate with the box,
      ! for this we scan the possible values of the grid size around the
      ! one we selected. We choose the size that has the spacing closest
      ! to the requested one.
      do delta = -1, 1
        spacing_new(delta) = m_two*maxval(abs(box_bounds(:, idir))) / real(2 * mesh%idx%nr(2, idir) + 1 - delta, real64)
        spacing_new(delta) = abs(spacing_new(delta) - spacing(idir))
      end do
 
      delta = minloc(spacing_new, dim = 1) - 2
 
      assert(delta >= -1)
      assert(delta <=  1)
 
      mesh%spacing(idir) = m_two*maxval(abs(box_bounds(:, idir))) / real(2 * mesh%idx%nr(2, idir) + 1 - delta, real64)
 
      ! we need to adjust the grid by adding or removing one point
      if (delta == -1) then
        mesh%idx%nr(1, idir) = mesh%idx%nr(1, idir) - 1
      else if (delta == 1) then
        mesh%idx%nr(2, idir) = mesh%idx%nr(2, idir) - 1
      end if
 
    end do
 
    if ( any(abs(mesh%spacing - spacing) > 1.e-6_real64) ) then
      call messages_write('The spacing has been modified to make it commensurate with the periodicity of the system.')
      call messages_warning()
    end if
 
    do idir = space%periodic_dim + 1, space%dim
      if (mesh%idx%nr(2, idir) == 0) then
        write(message(1),'(a,i2)') 'Spacing > box size in direction ', idir
        call messages_fatal(1, namespace=namespace)
      end if
    end do
 
    mesh%idx%ll = mesh%idx%nr(2, :) - mesh%idx%nr(1, :) + 1
    ! compute strides for cubic indices
    mesh%idx%stride(:) = 1
    do idir = 2, space%dim+1
      mesh%idx%stride(idir) = mesh%idx%stride(idir-1) *  &
        (mesh%idx%ll(idir-1) + 2*mesh%idx%enlarge(idir-1))
    end do
 
    pop_sub_with_profile(mesh_init_stage_1)
  end subroutine mesh_init_stage_1
 
  ! ---------------------------------------------------------
  !
  subroutine mesh_init_stage_2(mesh, namespace, space, box, stencil, grp, regenerate)
    class(mesh_t),       intent(inout) :: mesh
    type(namespace_t),   intent(in)    :: namespace
    class(space_t),      intent(in)    :: space
    class(box_t),        intent(in)    :: box
    type(stencil_t),     intent(in)    :: stencil
    type(mpi_grp_t),     intent(in)    :: grp
    logical, optional,   intent(in)    :: regenerate
 
    integer :: is
    real(real64)   :: chi(1:space%dim)
    real(real64) :: pos(space%dim)
    integer :: point(space%dim), point_stencil(space%dim), grid_sizes(space%dim)
    integer(int64) :: global_size
    integer(int32) :: local_size
    integer(int64) :: ispatial, ispatialb, istart, iend, spatial_size, ipg
    integer :: ip, ib, ib2, np, np_boundary, ii
    logical :: found
    type(lihash_t) :: spatial_to_boundary
    integer(int64), allocatable :: boundary_to_spatial(:), boundary_to_spatial_reordered(:)
    integer(int64), allocatable :: grid_to_spatial(:), grid_to_spatial_initial(:), grid_to_spatial_reordered(:)
    integer(int64), allocatable :: spatial_to_grid(:)
    integer, allocatable :: sizes(:)
    integer(int64), allocatable :: offsets(:)
    integer :: size_boundary
#ifdef HAVE_MPI
    integer(int64), pointer :: ptr(:)
    type(mpi_grp_t) :: internode_grp, intranode_grp
#endif
 
    push_sub_with_profile(mesh_init_stage_2)
 
    if (.not. optional_default(regenerate, .false.)) then
      ! enlarge mesh for boundary points
      mesh%idx%nr(1, :) = mesh%idx%nr(1, :) - mesh%idx%enlarge(:)
      mesh%idx%nr(2, :) = mesh%idx%nr(2, :) + mesh%idx%enlarge(:)
    end if
 
    !%Variable MeshIndexType
    !%Type integer
    !%Default idx_cubic
    !%Section Mesh
    !%Description
    !% Determine index type. Must be the same for restarting a calculation.
    !%Option idx_cubic 1
    !% Cubic indices are used to map the spatial information to the grid points.
    !%Option idx_hilbert 2
    !% A Hilbert space-filling curve is used to map the spatial information to
    !% the grid points.
    !%End
    call parse_variable(namespace, 'MeshIndexType', idx_cubic, mesh%idx%type)
 
    grid_sizes = mesh%idx%nr(2, :) - mesh%idx%nr(1, :) + 1
    mesh%idx%offset = grid_sizes/2
    if (space%dim > 1 .and. any(grid_sizes > 2**(int(63/space%dim, int64)))) then
      write(message(1), '(A, I10, A, I2, A)') "Error: grid too large, more than ", 2**(int(63/space%dim, int64)), &
        " points in one direction for ", space%dim, " dimensions. This is not supported."
      call messages_fatal(1, namespace=namespace)
    end if
    global_size = product(int(grid_sizes, int64))
    ! compute the bits per dimension: grid_sizes(i) <= 2**bits
    mesh%idx%bits = maxval(ceiling(log(real(grid_sizes, real64) )/log(2.0_real64)))
 
    if (mesh%idx%type == idx_cubic) then
      spatial_size = global_size
    else if (mesh%idx%type == idx_hilbert) then
      spatial_size = 2**(space%dim*mesh%idx%bits)
    end if
 
    ! use block data decomposition of spatial indices
    istart = spatial_size * grp%rank/grp%size
    iend = spatial_size * (grp%rank+1)/grp%size - 1
    if (.not. (iend - istart + 1 < huge(0_int32))) then
      write(message(1), '(A, I10, A, I2, A)') "Error: local grid too large, more than ", &
        huge(0_int32), " points. This is not supported. Maybe use more MPI ranks?"
      call messages_fatal(1, namespace=namespace)
    end if
    local_size = int(iend - istart + 1, int32)
 
    safe_allocate(grid_to_spatial_initial(1:local_size))
 
    ! get inner grid indices
    ip = 1
    do ispatial = istart, iend
      call index_spatial_to_point(mesh%idx, space%dim, ispatial, point)
      ! first check if point is outside bounding box
      if (any(point < mesh%idx%nr(1, :) + mesh%idx%enlarge)) cycle
      if (any(point > mesh%idx%nr(2, :) - mesh%idx%enlarge)) cycle
      ! then check if point is inside simulation box
      chi = real(point, real64)  * mesh%spacing
      pos = mesh%coord_system%to_cartesian(chi)
      if (.not. box%contains_point(pos)) cycle
      grid_to_spatial_initial(ip) = ispatial
      assert(ip + 1 < huge(ip))
      ip = ip + 1
    end do
    np = ip - 1
 
    call rebalance_array(grp,grid_to_spatial_initial(1:np), grid_to_spatial, sizes)
    np = sizes(grp%rank)
 
    safe_deallocate_a(grid_to_spatial_initial)
 
    safe_allocate(spatial_to_grid(grid_to_spatial(1):grid_to_spatial(np)))
    safe_deallocate_a(sizes)
 
    !$omp parallel do
    do ispatial = grid_to_spatial(1), grid_to_spatial(np)
      spatial_to_grid(ispatial) = -1
    end do
    !$omp parallel do
    do ip = 1, np
      spatial_to_grid(grid_to_spatial(ip)) = ip
    end do
 
    ! get local boundary indices
    call lihash_init(spatial_to_boundary)
    size_boundary = np
    safe_allocate(boundary_to_spatial(1:size_boundary))
    ib = 1
    do ip = 1, np
      call index_spatial_to_point(mesh%idx, space%dim, grid_to_spatial(ip), point)
      do is = 1, stencil%size
        if (stencil%center == is) cycle
        point_stencil(1:space%dim) = point(1:space%dim) + stencil%points(1:space%dim, is)
        ! check if point is in inner part
        call index_point_to_spatial(mesh%idx, space%dim, ispatialb, point_stencil)
        assert(ispatialb >= 0)
        if (ispatialb >= lbound(spatial_to_grid, dim=1, kind=int64) .and. &
          ispatialb <= ubound(spatial_to_grid, dim=1, kind=int64)) then
          if (spatial_to_grid(ispatialb) > 0) cycle
        end if
        ! then check if point is inside simulation box
        chi = real(point_stencil, real64)  * mesh%spacing
        pos = mesh%coord_system%to_cartesian(chi)
        if (box%contains_point(pos)) cycle
        ! it has to be a boundary point now
        ! check if already counted
        ib2 = lihash_lookup(spatial_to_boundary, ispatialb, found)
        if (found) cycle
        boundary_to_spatial(ib) = ispatialb
        call lihash_insert(spatial_to_boundary, ispatialb, ib)
        ib = ib + 1
        ! enlarge array
        if (ib >= size_boundary) then
          size_boundary = size_boundary * 2
          call make_array_larger(boundary_to_spatial, size_boundary)
        end if
      end do
    end do
    np_boundary = ib - 1
    call lihash_end(spatial_to_boundary)
    safe_deallocate_a(spatial_to_grid)
 
    ! reorder inner points
    call reorder_points(namespace, space, grp, mesh%idx, grid_to_spatial, grid_to_spatial_reordered)
    safe_deallocate_a(grid_to_spatial)
 
    call rebalance_array(grp,grid_to_spatial_reordered, grid_to_spatial, sizes)
    np = sizes(grp%rank)
    mesh%np_global = sizes(0)
    do ii = 1, grp%size - 1
      mesh%np_global = mesh%np_global + sizes(ii)
    end do
    safe_deallocate_a(sizes)
    safe_deallocate_a(grid_to_spatial_reordered)
 
    ! reorder boundary points
    call make_array_larger(boundary_to_spatial, np_boundary)
    call reorder_points(namespace, space, grp, mesh%idx, boundary_to_spatial, boundary_to_spatial_reordered)
    safe_deallocate_a(boundary_to_spatial)
 
    call rebalance_array(grp,boundary_to_spatial_reordered, boundary_to_spatial, sizes)
    safe_deallocate_a(boundary_to_spatial_reordered)
 
    ! global grid size
    np_boundary = sizes(grp%rank)
    mesh%np_part_global = mesh%np_global + sizes(0)
    do ii = 1, grp%size - 1
      mesh%np_part_global = mesh%np_part_global + sizes(ii)
    end do
    safe_deallocate_a(sizes)
 
 
    ! get global indices
#ifdef HAVE_MPI
    ! create shared memory window and fill it only on root
    call create_intranode_communicator(grp, intranode_grp, internode_grp)
    call lmpi_create_shared_memory_window(mesh%np_part_global, intranode_grp, &
      mesh%idx%window_grid_to_spatial, mesh%idx%grid_to_spatial_global)
#else
    safe_allocate(mesh%idx%grid_to_spatial_global(1:mesh%np_part_global))
#endif
    ! inner grid
    call get_sizes_offsets(np, sizes, offsets, grp)
    call grp%gatherv(grid_to_spatial, np, mpi_integer8, &
      mesh%idx%grid_to_spatial_global, sizes, offsets, mpi_integer8, 0)
 
    ! boundary indices
    call get_sizes_offsets(np_boundary, sizes, offsets, grp)
    call grp%gatherv(boundary_to_spatial, np_boundary, mpi_integer8, &
      mesh%idx%grid_to_spatial_global(mesh%np_global+1:), sizes, offsets, mpi_integer8, 0)
 
    ! fill global hash map
#ifdef HAVE_MPI
    ! create shared memory window and fill it only on root
    call lmpi_create_shared_memory_window(spatial_size, intranode_grp, &
      mesh%idx%window_spatial_to_grid, ptr)
    mesh%idx%spatial_to_grid_global(0:spatial_size-1) => ptr(1:spatial_size)
#else
    safe_allocate(mesh%idx%spatial_to_grid_global(0:spatial_size-1))
#endif
    if (grp%is_root()) then
      ! fill only on root, then broadcast
      !$omp parallel do
      do ispatial = 0, spatial_size-1
        mesh%idx%spatial_to_grid_global(ispatial) = 0
      end do
      !$omp parallel do
      do ipg = 1, mesh%np_part_global
        mesh%idx%spatial_to_grid_global(mesh%idx%grid_to_spatial_global(ipg)) = ipg
      end do
    end if
 
#ifdef HAVE_MPI
    ! now broadcast the global arrays to local rank 0 on each node
    if (intranode_grp%is_root()) then
      call internode_grp%bcast(mesh%idx%grid_to_spatial_global(1), mesh%np_part_global, mpi_integer8, 0)
      call internode_grp%bcast(mesh%idx%spatial_to_grid_global(0), spatial_size, mpi_integer8, 0)
    end if
    call lmpi_sync_shared_memory_window(mesh%idx%window_grid_to_spatial, intranode_grp)
    call lmpi_sync_shared_memory_window(mesh%idx%window_spatial_to_grid, intranode_grp)
#endif
 
    safe_deallocate_a(offsets)
    safe_deallocate_a(sizes)
 
    safe_deallocate_a(boundary_to_spatial)
    safe_deallocate_a(grid_to_spatial)
 
    pop_sub_with_profile(mesh_init_stage_2)
  end subroutine mesh_init_stage_2
 
! ---------------------------------------------------------
! ---------------------------------------------------------
  subroutine mesh_init_stage_3(mesh, namespace, space, stencil, mc, parent, regenerate)
    class(mesh_t),             intent(inout) :: mesh
    type(namespace_t),         intent(in)    :: namespace
    class(space_t),            intent(in)    :: space
    type(stencil_t),           intent(in)    :: stencil
    type(multicomm_t),         intent(in)    :: mc
    type(mesh_t),    optional, intent(in)    :: parent
    logical,         optional, intent(in)    :: regenerate
 
    integer :: ip, jj(space%dim), np
 
    push_sub_with_profile(mesh_init_stage_3)
 
    call mpi_grp_init(mesh%mpi_grp, mc%group_comm(p_strategy_domains))
 
    ! check if we are running in parallel in domains
    mesh%parallel_in_domains = (mesh%mpi_grp%size > 1)
 
    call checksum_calculate(1, mesh%np_part_global, mesh%idx%grid_to_spatial_global(1), &
      mesh%idx%checksum)
 
    if (mesh%parallel_in_domains) then
      call do_partition()
    else
      ! When running serially those two are the same.
      assert(mesh%np_part_global < huge(mesh%np_part))
      mesh%np      = i8_to_i4(mesh%np_global)
      mesh%np_part = i8_to_i4(mesh%np_part_global)
 
      ! These must be initialized for par_vec_gather, par_vec_scatter to work
      ! as copy operations when running without domain parallelization.
      mesh%pv%np_global = mesh%np_global
      mesh%pv%np_ghost = 0
      mesh%pv%np_bndry = mesh%np_part - mesh%np
      mesh%pv%npart = 1
      mesh%pv%xlocal = 1
    end if
 
    ! Compute mesh%x
    safe_allocate(mesh%x(1:space%dim, 1:mesh%np_part))
    safe_allocate(mesh%x_t(1:mesh%np_part, 1:space%dim))
    do ip = 1, mesh%np_part
      mesh%x(1:space%dim, ip) = mesh_x_global(mesh, mesh_local2global(mesh, ip))
      mesh%x_t(ip, 1:space%dim) = mesh%x(1:space%dim, ip)
    end do
 
    ! save chi, i.e. primitive coordinates
    safe_allocate(mesh%chi(1:space%dim, 1:mesh%np_part))
    do ip = 1, mesh%np_part
      call mesh_local_index_to_coords(mesh, ip, jj)
      mesh%chi(:, ip) = jj*mesh%spacing
    end do
 
    call mesh_get_vol_pp()
 
    ! save inverse jacobian
    if (mesh%coord_system%local_basis) then
      np = mesh%np_part
    else
      np = 1
    end if
    safe_allocate(mesh%jacobian_inverse(1:space%dim, 1:space%dim, np))
    do ip = 1, np
      mesh%jacobian_inverse(:, :, ip) = mesh%coord_system%jacobian_inverse(mesh%chi(:, ip))
    end do
 
    pop_sub_with_profile(mesh_init_stage_3)
 
  contains
    ! ---------------------------------------------------------
    subroutine do_partition()
#ifdef HAVE_MPI
      integer :: jj, ipart, jpart
      integer(int64) :: ipg, jpg
      integer, allocatable :: gindex(:), gedges(:)
      logical, allocatable :: nb(:, :)
      integer              :: idx(space%dim), jx(space%dim)
      type(mpi_comm)       :: graph_comm
      integer              :: iedge, nnb
      logical              :: use_topo, reorder, partition_print
      integer              :: ierr
 
      logical :: has_virtual_partition = .false.
      integer :: vsize
      type(restart_t) :: restart_load, restart_dump
      integer, allocatable :: part_vec(:)
 
      push_sub(mesh_init_stage_3.do_partition)
 
      !Try to load the partition from the restart files
      if (.not. optional_default(regenerate, .false.)) then
        call restart_load%init(namespace, restart_partition, restart_type_load, mc, ierr, mesh=mesh, exact=.true.)
        if (ierr == 0) call mesh_partition_load(restart_load, mesh, ierr)
        call restart_load%end()
      else
        ierr = 0
      end if
 
      if (ierr /= 0) then
 
        !%Variable MeshPartitionVirtualSize
        !%Type integer
        !%Default mesh mpi_grp size
        !%Section Execution::Parallelization
        !%Description
        !% Gives the possibility to change the partition nodes.
        !% Afterward, it crashes.
        !%End
        call parse_variable(namespace, 'MeshPartitionVirtualSize', mesh%mpi_grp%size, vsize)
 
        if (vsize /= mesh%mpi_grp%size) then
          write(message(1),'(a,I7)') "Changing the partition size to", vsize
          write(message(2),'(a)') "The execution will crash."
          call messages_warning(2, namespace=namespace)
          has_virtual_partition = .true.
        else
          has_virtual_partition = .false.
        end if
 
        if (.not. present(parent)) then
          call mesh_partition(mesh, namespace, space, stencil, vsize)
        else
          ! if there is a parent grid, use its partition
          call mesh_partition_from_parent(mesh, parent)
        end if
 
        !Now that we have the partitions, we save them
        call restart_dump%init(namespace, restart_partition, restart_type_dump, mc, ierr, mesh=mesh)
        call mesh_partition_dump(restart_dump, mesh, vsize, ierr)
        call restart_dump%end()
      end if
 
      if (has_virtual_partition) then
        call profiling_end(namespace)
        call print_date("Calculation ended on ")
        write(message(1),'(a)') "Execution has ended."
        write(message(2),'(a)') "If you want to run your system, do not use MeshPartitionVirtualSize."
        call messages_warning(2, namespace=namespace)
        call messages_end()
        call global_end()
        stop
      end if
 
      !%Variable MeshUseTopology
      !%Type logical
      !%Default false
      !%Section Execution::Parallelization
      !%Description
      !% (experimental) If enabled, <tt>Octopus</tt> will use an MPI virtual
      !% topology to map the processors. This can improve performance
      !% for certain interconnection systems.
      !%End
      call parse_variable(namespace, 'MeshUseTopology', .false., use_topo)
 
      if (use_topo) then
        ! At the moment we still need the global partition. This will be removed in near future.
        safe_allocate(part_vec(1:mesh%np_global))
        call partition_get_global(mesh%partition, part_vec(1:mesh%np_global))
 
 
        ! generate a table of neighbours
 
        safe_allocate(nb(1:mesh%mpi_grp%size, 1:mesh%mpi_grp%size))
        nb = .false.
 
        do ipg = 1, mesh%np_global
          ipart = part_vec(ipg)
          call mesh_global_index_to_coords(mesh, ipg, idx)
          do jj = 1, stencil%size
            jx = idx + stencil%points(:, jj)
            jpg = mesh_global_index_from_coords(mesh, jx)
            if (jpg > 0 .and. jpg <= mesh%np_global) then
              jpart = part_vec(jpg)
              if (ipart /= jpart ) nb(ipart, jpart) = .true.
            end if
          end do
        end do
        safe_deallocate_a(part_vec)
 
        ! now generate the information of the graph
 
        safe_allocate(gindex(1:mesh%mpi_grp%size))
        safe_allocate(gedges(1:count(nb)))
 
        ! and now generate it
        iedge = 0
        do ipart = 1, mesh%mpi_grp%size
          do jpart = 1, mesh%mpi_grp%size
            if (nb(ipart, jpart)) then
              iedge = iedge + 1
              gedges(iedge) = jpart - 1
            end if
          end do
          gindex(ipart) = iedge
        end do
 
        assert(iedge == count(nb))
 
        reorder = .true.
        call mpi_graph_create(mesh%mpi_grp%comm, mesh%mpi_grp%size, gindex, gedges, reorder, graph_comm)
 
        ! we have a new communicator
        call mpi_grp_init(mesh%mpi_grp, graph_comm)
 
        safe_deallocate_a(nb)
        safe_deallocate_a(gindex)
        safe_deallocate_a(gedges)
 
      end if
 
      if (optional_default(regenerate, .false.)) call par_vec_end(mesh%pv)
      call par_vec_init(mesh%mpi_grp, mesh%np_global, mesh%idx, stencil,&
        space, mesh%partition, mesh%pv, namespace)
 
      ! check the number of ghost neighbours in parallel
      nnb = 0
      jpart =  mesh%pv%partno
      do ipart = 1, mesh%pv%npart
        if (ipart == jpart) cycle
        if (mesh%pv%ghost_scounts(ipart) /= 0) nnb = nnb + 1
      end do
      assert(nnb >= 0 .and. nnb < mesh%pv%npart)
 
      ! Set local point numbers.
      mesh%np      = mesh%pv%np_local
      mesh%np_part = mesh%np + mesh%pv%np_ghost + mesh%pv%np_bndry
 
      !%Variable PartitionPrint
      !%Type logical
      !%Default true
      !%Section Execution::Parallelization
      !%Description
      !% (experimental) If disabled, <tt>Octopus</tt> will not compute
      !% nor print the partition information, such as local points,
      !% no. of neighbours, ghost points and boundary points.
      !%End
      call parse_variable(namespace, 'PartitionPrint', .true., partition_print)
 
      if (partition_print) then
        call mesh_partition_write_info(mesh, namespace=namespace)
        call mesh_partition_messages_debug(mesh, namespace)
      end if
#endif
 
      pop_sub(mesh_init_stage_3.do_partition)
    end subroutine do_partition
 
 
    ! ---------------------------------------------------------
    subroutine mesh_get_vol_pp()
 
      integer :: ip, np
 
      push_sub(mesh_init_stage_3.mesh_get_vol_pp)
 
      np = 1
      if (mesh%use_vol_pp) np = mesh%np_part
      ! If no local point, we should not try to access the arrays
      if (mesh%np_part == 0) np = 0
 
      safe_allocate(mesh%vol_pp(1:np))
 
      do ip = 1, np
        mesh%vol_pp(ip) = product(mesh%spacing)
      end do
 
      do ip = 1, np
        mesh%vol_pp(ip) = mesh%vol_pp(ip)*mesh%coord_system%jacobian_determinant(mesh%chi(:, ip))
      end do
 
      if (mesh%use_vol_pp .or. mesh%np_part == 0) then
        mesh%volume_element = m_one
      else
        mesh%volume_element = mesh%vol_pp(1)
      end if
 
      ! the weights on the device, for weighted reductions of device batches
      if (mesh%use_vol_pp .and. accel_is_enabled()) then
        call accel_create_buffer(mesh%vol_pp_buffer, accel_mem_read_only, type_float, max(np, 1))
        call accel_create_buffer(mesh%sqrt_vol_pp_buffer, accel_mem_read_only, type_float, max(np, 1))
        if (np > 0) then
          call accel_write_buffer(mesh%vol_pp_buffer, np, mesh%vol_pp)
          call accel_write_buffer(mesh%sqrt_vol_pp_buffer, np, sqrt(mesh%vol_pp))
        end if
      end if
 
      pop_sub(mesh_init_stage_3.mesh_get_vol_pp)
    end subroutine mesh_get_vol_pp
 
  end subroutine mesh_init_stage_3
 
  subroutine rebalance_array(grp, data_input, data_output, output_sizes)
    type(mpi_grp_t),                   intent(in)  :: grp
    integer(int64), contiguous,        intent(in)  :: data_input(:)
    integer(int64), allocatable,       intent(out) :: data_output(:)
    integer, allocatable, optional, intent(out) :: output_sizes(:)
 
    integer, allocatable :: initial_sizes(:), final_sizes(:)
    integer(int64), allocatable :: initial_offsets(:), final_offsets(:)
    integer, allocatable :: scounts(:), sdispls(:)
    integer, allocatable :: rcounts(:), rdispls(:)
    integer :: irank
    integer(int64) :: itmp
 
    push_sub(rebalance_array)
 
    ! collect current sizes of distributed array
    safe_allocate(initial_sizes(0:grp%size-1))
    call grp%allgather(size(data_input), 1, mpi_integer, initial_sizes(0), 1, mpi_integer)
    safe_allocate(initial_offsets(0:grp%size))
    initial_offsets(0) = 0
    do irank = 1, grp%size
      initial_offsets(irank) = initial_offsets(irank-1) + initial_sizes(irank-1)
    end do
 
    ! now redistribute the arrays
    ! use block data decomposition of grid indices
    safe_allocate(final_offsets(0:grp%size))
    safe_allocate(final_sizes(0:grp%size-1))
 
    do irank = 0, grp%size
      final_offsets(irank) = sum(int(initial_sizes, int64)) * irank/grp%size
    end do
    do irank = 0, grp%size - 1
      assert(final_offsets(irank + 1) - final_offsets(irank) < huge(0_int32))
      final_sizes(irank) = int(final_offsets(irank + 1) - final_offsets(irank), int32)
    end do
 
    safe_allocate(scounts(0:grp%size-1))
    safe_allocate(sdispls(0:grp%size-1))
    safe_allocate(rcounts(0:grp%size-1))
    safe_allocate(rdispls(0:grp%size-1))
    ! determine communication pattern
    scounts = 0
    do irank = 0, grp%size - 1
      ! get overlap of initial and final distribution
      itmp = min(final_offsets(irank+1), initial_offsets(grp%rank+1)) - &
        max(final_offsets(irank), initial_offsets(grp%rank))
      assert(itmp < huge(0_int32))
      if (itmp < 0) then
        scounts(irank) = 0
      else
        scounts(irank) = int(itmp, int32)
      end if
    end do
    sdispls(0) = 0
    do irank = 1, grp%size - 1
      sdispls(irank) = sdispls(irank - 1) + scounts(irank - 1)
    end do
    assert(sum(int(scounts, int64)) < huge(0_int32))
    assert(sum(scounts) == initial_sizes(grp%rank))
 
    rcounts = 0
    do irank = 0, grp%size - 1
      ! get overlap of initial and final distribution
      itmp = min(final_offsets(grp%rank+1), initial_offsets(irank+1)) - &
        max(final_offsets(grp%rank), initial_offsets(irank))
      assert(itmp < huge(0_int32))
      if (itmp < 0) then
        rcounts(irank) = 0
      else
        rcounts(irank) = int(itmp, int32)
      end if
    end do
    rdispls(0) = 0
    do irank = 1, grp%size - 1
      rdispls(irank) = rdispls(irank - 1) + rcounts(irank - 1)
    end do
    ! check for consistency between sending and receiving
    assert(sum(rcounts) == final_sizes(grp%rank))
 
    safe_allocate(data_output(1:final_sizes(grp%rank)))
    call grp%alltoallv(data_input, scounts, sdispls, mpi_integer8, &
      data_output, rcounts, rdispls, mpi_integer8)
 
    ! save final sizes of array if optional argument present
    if (present(output_sizes)) then
      safe_allocate(output_sizes(0:grp%size-1))
      output_sizes(:) = final_sizes(:)
    end if
 
    safe_deallocate_a(final_offsets)
    safe_deallocate_a(final_sizes)
 
    safe_deallocate_a(scounts)
    safe_deallocate_a(sdispls)
    safe_deallocate_a(rcounts)
    safe_deallocate_a(rdispls)
 
 
    pop_sub(rebalance_array)
  end subroutine rebalance_array
 
  subroutine reorder_points(namespace, space, grp, idx, grid_to_spatial, grid_to_spatial_reordered)
    type(namespace_t), intent(in) :: namespace
    class(space_t), intent(in) :: space
    type(mpi_grp_t), intent(in) :: grp
    type(index_t), intent(in) :: idx
    integer(int64), contiguous, intent(in) :: grid_to_spatial(:)
    integer(int64), allocatable, intent(out) :: grid_to_spatial_reordered(:)
 
    integer :: bsize(space%dim), order, default
    integer :: nn, idir, ipg, ip, number_of_blocks(space%dim)
    type(block_t) :: blk
    integer, parameter :: &
      ORDER_BLOCKS     =  1, &
      order_original    =  2, &
      order_cube       =  3
    integer :: point(1:space%dim)
    integer(int64), allocatable :: reorder_indices(:), reorder_recv(:)
    integer, allocatable :: index_map(:), indices(:)
    integer(int64), allocatable :: grid_to_spatial_recv(:)
    integer, allocatable :: initial_sizes(:)
    integer(int64), allocatable :: initial_offsets(:)
    integer(int64) :: istart, iend, indstart, indend, spatial_size
    integer :: irank, local_size, num_recv
    integer :: iunique, nunique
    integer :: direction
    logical :: increase_with_dimension
 
    integer, allocatable :: scounts(:), sdispls(:), rcounts(:), rdispls(:)
    integer(int64), allocatable :: spatial_cutoff(:)
 
    push_sub(reorder_points)
 
    !%Variable MeshOrder
    !%Type integer
    !%Section Execution::Optimization
    !%Description
    !% This variable controls how the grid points are mapped to a
    !% linear array for global arrays. For runs that are parallel
    !% in domains, the local mesh order may be different (see
    !% <tt>MeshLocalOrder</tt>).
    !% The default is blocks when serial in domains and cube when
    !% parallel in domains with the local mesh order set to blocks.
    !%Option order_blocks 1
    !% The grid is mapped using small parallelepipedic grids. The size
    !% of the blocks is controlled by <tt>MeshBlockSize</tt>.
    !%Option order_original 2
    !% The original order of the indices is used to map the grid.
    !%Option order_cube 3
    !% The grid is mapped using a full cube, i.e. without blocking.
    !%End
    default = order_blocks
    call parse_variable(namespace, 'MeshOrder', default, order)
    ! no reordering in 1D necessary
    if (space%dim == 1) then
      order = order_original
    end if
 
    !%Variable MeshBlockDirection
    !%Type integer
    !%Section Execution::Optimization
    !%Description
    !% Determines the direction in which the dimensions are chosen to compute
    !% the blocked index for sorting the mesh points (see MeshBlockSize).
    !% The default is increase_with_dimensions, corresponding to xyz ordering
    !% in 3D.
    !%Option increase_with_dimension 1
    !% The fastest changing index is in the first dimension, i.e., in 3D this
    !% corresponds to ordering in xyz directions.
    !%Option decrease_with_dimension 2
    !% The fastest changing index is in the last dimension, i.e., in 3D this
    !% corresponds to ordering in zyx directions.
    !%End
    call parse_variable(namespace, 'MeshBlockDirection', 1, direction)
    increase_with_dimension = direction == 1
    if (direction /= 1 .and. direction /= 2) then
      call messages_input_error(namespace, 'MeshBlockDirection')
    end if
 
    select case (order)
    case (order_original)
      ! only copy points, they stay in their original ordering
      safe_allocate(grid_to_spatial_reordered(1:size(grid_to_spatial)))
      grid_to_spatial_reordered(1:size(grid_to_spatial)) = grid_to_spatial(1:size(grid_to_spatial))
    case (order_blocks, order_cube)
      if (order == order_cube) then
        bsize = idx%nr(2, :) - idx%nr(1, :) + 1
      else
        !%Variable MeshBlockSize
        !%Type block
        !%Section Execution::Optimization
        !%Description
        !% To improve memory-access locality when calculating derivatives,
        !% <tt>Octopus</tt> arranges mesh points in blocks. This variable
        !% controls the size of this blocks in the different
        !% directions. The default is selected according to the value of
        !% the <tt>StatesBlockSize</tt> variable. (This variable only affects the
        !% performance of <tt>Octopus</tt> and not the results.)
        !%End
        if (conf%target_states_block_size < 16) then
          bsize(1) = 80 * 4 / abs(conf%target_states_block_size)
          if (space%dim > 1) bsize(2) = 4
          if (space%dim > 2) bsize(3:) = 10
        else
          bsize(1) = max(4 * 16 / abs(conf%target_states_block_size), 1)
          if (space%dim > 1) bsize(2) = 15
          if (space%dim > 2) bsize(3:) = 15
        end if
 
        if (parse_block(namespace, 'MeshBlockSize', blk) == 0) then
          nn = parse_block_cols(blk, 0)
          if (nn /= space%dim) then
            message(1) = "Error: number of entries in MeshBlockSize must match the number of dimensions."
            call messages_fatal(1, namespace=namespace)
          end if
          do idir = 1, nn
            call parse_block_integer(blk, 0, idir - 1, bsize(idir))
          end do
        end if
      end if
 
      number_of_blocks = (idx%nr(2, :) - idx%nr(1, :) + 1) / bsize + 1
 
 
      ! do the global reordering in parallel, use block data decomposition of global indices
      ! reorder indices along blocked parallelepiped curve
 
      ! collect current sizes of distributed array
      safe_allocate(initial_sizes(0:grp%size-1))
      call grp%allgather(size(grid_to_spatial), 1, mpi_integer, initial_sizes(0), 1, mpi_integer)
      safe_allocate(initial_offsets(0:grp%size))
      initial_offsets(0) = 0
      do irank = 1, grp%size
        initial_offsets(irank) = initial_offsets(irank-1) + initial_sizes(irank-1)
      end do
 
      ! get local range and size
      istart = initial_offsets(grp%rank)
      iend = initial_offsets(grp%rank + 1) - 1
      assert(iend - istart + 1 < huge(0_int32))
      local_size = int(iend - istart + 1, int32)
      assert(local_size == initial_sizes(grp%rank))
 
      ! compute new indices locally
      safe_allocate(reorder_indices(1:local_size))
      safe_allocate(indices(1:local_size))
      safe_allocate(grid_to_spatial_reordered(1:local_size))
      !$omp parallel do private(point)
      do ip = 1, local_size
        call index_spatial_to_point(idx, space%dim, grid_to_spatial(ip), point)
        point = point + idx%offset
        reorder_indices(ip) = get_blocked_index(space%dim, point, bsize, number_of_blocks, increase_with_dimension)
      end do
      ! parallel sort according to the new indices
      ! sort the local array
      call sort(reorder_indices, indices)
      ! save reordered indices to send to other processes
      !$omp parallel do
      do ip = 1, local_size
        grid_to_spatial_reordered(ip) = grid_to_spatial(indices(ip))
      end do
 
      ! get minimum and maximum
      if(local_size > 0) then
        indstart = reorder_indices(1)
        indend = reorder_indices(local_size)
      else
        indstart = huge(1_int64)
        indend = 0
      end if
      call grp%allreduce_inplace(indstart, 1, mpi_integer8, mpi_min)
      call grp%allreduce_inplace(indend, 1, mpi_integer8, mpi_max)
      spatial_size = indend - indstart + 1
 
      ! get index ranges for each rank
      safe_allocate(spatial_cutoff(0:grp%size-1))
      do irank = 0, grp%size - 1
        spatial_cutoff(irank) = spatial_size * (irank+1)/grp%size + indstart
      end do
 
      safe_allocate(scounts(0:grp%size-1))
      safe_allocate(sdispls(0:grp%size-1))
      safe_allocate(rcounts(0:grp%size-1))
      safe_allocate(rdispls(0:grp%size-1))
      ! get send counts
      scounts = 0
      irank = 0
      ! the indices are ordered, so we can go through them and increase
      ! the rank to which they are associated to when we cross a cutoff
      do ip = 1, local_size
        if (reorder_indices(ip) >= spatial_cutoff(irank)) then
          ! this do loop is needed in case some ranks do not have any points
          do while (reorder_indices(ip) >= spatial_cutoff(irank))
            irank = irank + 1
          end do
          assert(irank < grp%size)
        end if
        scounts(irank) = scounts(irank) + 1
      end do
      safe_deallocate_a(spatial_cutoff)
      assert(sum(scounts) == local_size)
 
      ! compute communication pattern (sdispls, rcounts, rdispls)
      sdispls(0) = 0
      do irank = 1, grp%size - 1
        sdispls(irank) = sdispls(irank - 1) + scounts(irank - 1)
      end do
 
      call grp%alltoall(scounts, 1, mpi_integer, &
        rcounts, 1, mpi_integer)
 
      rdispls(0) = 0
      do irank = 1, grp%size - 1
        rdispls(irank) = rdispls(irank - 1) + rcounts(irank - 1)
      end do
 
      ! make sure the arrays get allocated also if we do not receive anything
      num_recv = max(sum(rcounts), 1)
      ! communicate the locally sorted indices
      safe_allocate(reorder_recv(1:num_recv))
      call grp%alltoallv(reorder_indices, scounts, sdispls, mpi_integer8, &
        reorder_recv, rcounts, rdispls, mpi_integer8)
      safe_deallocate_a(reorder_indices)
 
      ! communicate the corresponding spatial indices
      safe_allocate(grid_to_spatial_recv(1:num_recv))
      call grp%alltoallv(grid_to_spatial_reordered, scounts, sdispls, mpi_integer8, &
        grid_to_spatial_recv, rcounts, rdispls, mpi_integer8)
      safe_deallocate_a(grid_to_spatial_reordered)
 
      ! do k-way merge of sorted indices
      safe_allocate(reorder_indices(1:num_recv))
      safe_allocate(index_map(1:num_recv))
      if (sum(rcounts) > 0) then
        call merge_sorted_arrays(reorder_recv, rcounts, reorder_indices, index_map)
 
        ! get number of unique indices, needed for boundary
        nunique = 1
        do ipg = 2, sum(rcounts)
          if (reorder_indices(ipg) /= reorder_indices(ipg-1)) then
            nunique = nunique + 1
          end if
        end do
 
        ! reorder according to new order, but remove duplicate entries
        safe_allocate(grid_to_spatial_reordered(1:nunique))
        iunique = 1
        grid_to_spatial_reordered(iunique) = grid_to_spatial_recv(index_map(1))
        do ipg = 2, sum(rcounts)
          if (reorder_indices(ipg) /= reorder_indices(ipg-1)) then
            iunique = iunique + 1
            grid_to_spatial_reordered(iunique) = grid_to_spatial_recv(index_map(ipg))
          end if
        end do
      else
        safe_allocate(grid_to_spatial_reordered(1:0))
      end if
 
      safe_deallocate_a(initial_offsets)
      safe_deallocate_a(initial_sizes)
 
      safe_deallocate_a(reorder_indices)
      safe_deallocate_a(reorder_recv)
 
      safe_deallocate_a(grid_to_spatial_recv)
      safe_deallocate_a(index_map)
      safe_deallocate_a(indices)
 
      safe_deallocate_a(scounts)
      safe_deallocate_a(sdispls)
      safe_deallocate_a(rcounts)
      safe_deallocate_a(rdispls)
 
    end select
    pop_sub(reorder_points)
  end subroutine reorder_points
 
  subroutine get_sizes_offsets(local_size, sizes, offsets, mpi_grp)
    integer,                      intent(in)  :: local_size
    integer, allocatable,         intent(out) :: sizes(:)
    integer(int64), allocatable,  intent(out) :: offsets(:)
    type(mpi_grp_t),              intent(in)  :: mpi_grp
 
    integer :: irank
 
    push_sub(get_sizes_offsets)
 
    safe_allocate(sizes(0:mpi_grp%size-1))
    call mpi_grp%allgather(local_size, 1, mpi_integer, sizes(0), 1, mpi_integer)
    safe_allocate(offsets(0:mpi_grp%size))
    offsets(0) = 0
    do irank = 1, mpi_grp%size
      offsets(irank) = offsets(irank-1) + sizes(irank-1)
    end do
 
    pop_sub(get_sizes_offsets)
  end subroutine get_sizes_offsets
 
end module mesh_init_oct_m
 
!! Local Variables:
!! mode: f90
!! coding: utf-8
!! End: