tree_amg.f90

Go to the documentation of this file.

! Copyright (c) 2024, The Neko Authors
! All rights reserved.
!
! Redistribution and use in source and binary forms, with or without
! modification, are permitted provided that the following conditions
! are met:
!
!   * Redistributions of source code must retain the above copyright
!     notice, this list of conditions and the following disclaimer.
!
!   * Redistributions in binary form must reproduce the above
!     copyright notice, this list of conditions and the following
!     disclaimer in the documentation and/or other materials provided
!     with the distribution.
!
!   * Neither the name of the authors nor the names of its
!     contributors may be used to endorse or promote products derived
!     from this software without specific prior written permission.
!
! THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
! "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
! LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
! FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
! COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
! INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
! BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
! LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
! CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
! LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
! ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
! POSSIBILITY OF SUCH DAMAGE.
!
module tree_amg
  use num_types
  use utils
  use math
  use coefs, only : coef_t
  use mesh, only : mesh_t
  use space, only : space_t
  use ax_product, only: ax_t
  use bc_list, only: bc_list_t
  use gather_scatter, only : gs_t, gs_op_add
  implicit none
  private
 
  type, private :: tamg_node_t
    logical :: isleaf = .true. 
    integer :: gid = -1
    integer :: lvl = -1
    integer :: ndofs = 0
    integer, allocatable :: dofs(:)
    real(kind=rp) :: xyz(3) 
    real(kind=rp), allocatable :: interp_r(:) 
    real(kind=rp), allocatable :: interp_p(:) 
  end type tamg_node_t
 
  type, private :: tamg_lvl_t
    integer :: lvl = -1
    integer :: nnodes = 0
    type(tamg_node_t), allocatable :: nodes(:)
    integer :: fine_lvl_dofs = 0
    real(kind=rp), allocatable :: wrk_in(:) 
    real(kind=rp), allocatable :: wrk_out(:) 
    integer, allocatable :: map_f2c_dof(:)
    integer, allocatable :: map_c2f_dof(:)
    !--!
    integer, allocatable :: nodes_ptr(:)
    integer, allocatable :: nodes_gid(:)
    integer, allocatable :: nodes_dofs(:)
    integer, allocatable :: nodes_gids(:)
    ! could make another array of the same size of nodes_dofs
    ! that stores the parent node gid information
    ! (similar to nodes_gid that stores the gid of each node)
    ! then some loops can be simplified to a single loop
    ! of len(nodes_dofs) instead of going through each node
    ! and looping through nodes_ptr(i) to nodes_ptr(i+1)-1
  end type tamg_lvl_t
 
  type, public :: tamg_hierarchy_t
    integer :: nlvls
    type(tamg_lvl_t), allocatable :: lvl(:)
 
    class(ax_t), pointer :: ax
    type(mesh_t), pointer :: msh
    type(space_t), pointer :: xh
    type(coef_t), pointer :: coef
    type(gs_t), pointer :: gs_h
    type(bc_list_t), pointer :: blst
 
  contains
    procedure, pass(this) :: init => tamg_init
    procedure, pass(this) :: matvec => tamg_matvec
    procedure, pass(this) :: matvec_impl => tamg_matvec_impl
    procedure, pass(this) :: interp_f2c => tamg_restriction_operator
    procedure, pass(this) :: interp_c2f => tamg_prolongation_operator
  end type tamg_hierarchy_t
 
  public tamg_lvl_init, tamg_node_init
 
contains
 
  subroutine tamg_init(this, ax, Xh, coef, msh, gs_h, nlvls, blst)
    class(tamg_hierarchy_t), target, intent(inout) :: this
    class(ax_t), target, intent(in) :: ax
    type(space_t),target, intent(in) :: Xh
    type(coef_t), target, intent(in) :: coef
    type(mesh_t), target, intent(in) :: msh
    type(gs_t), target, intent(in) :: gs_h
    integer, intent(in) :: nlvls
    type(bc_list_t), target, intent(in) :: blst
    integer :: i, n
 
    this%ax => ax
    this%msh => msh
    this%Xh => xh
    this%coef => coef
    this%gs_h => gs_h
    this%blst => blst
 
    if (nlvls .lt. 2) then
      call neko_error("Need to request at least two multigrid levels.")
    end if
 
    this%nlvls = nlvls
    allocate( this%lvl(this%nlvls) )
 
    do i = 1, nlvls
      allocate( this%lvl(i)%map_f2c_dof( coef%dof%size() ))
      allocate( this%lvl(i)%map_c2f_dof( coef%dof%size() ))
    end do
 
  end subroutine tamg_init
 
  subroutine tamg_lvl_init(tamg_lvl, lvl, nnodes, ndofs)
    type(tamg_lvl_t), intent(inout) :: tamg_lvl
    integer, intent(in) :: lvl
    integer, intent(in) :: nnodes
    integer, intent(in) :: ndofs
 
    tamg_lvl%lvl = lvl
    tamg_lvl%nnodes = nnodes
    allocate( tamg_lvl%nodes(tamg_lvl%nnodes) )
    allocate( tamg_lvl%nodes_ptr(tamg_lvl%nnodes+1) )
    allocate( tamg_lvl%nodes_gid(tamg_lvl%nnodes) )
    allocate( tamg_lvl%nodes_dofs(ndofs) )
    allocate( tamg_lvl%nodes_gids(ndofs) )
 
    tamg_lvl%fine_lvl_dofs = ndofs
    allocate( tamg_lvl%wrk_in( ndofs ) )
    allocate( tamg_lvl%wrk_out( ndofs ) )
  end subroutine tamg_lvl_init
 
  subroutine tamg_node_init(node, gid, ndofs)
    type(tamg_node_t), intent(inout) :: node
    integer, intent(in) :: gid
    integer, intent(in) :: ndofs
 
    node%gid = gid
    node%ndofs = ndofs
    allocate( node%dofs( node%ndofs) )
    node%dofs = -1
    allocate( node%interp_r( node%ndofs) )
    allocate( node%interp_p( node%ndofs) )
    node%interp_r = 1.0
    node%interp_p = 1.0
  end subroutine tamg_node_init
 
  recursive subroutine tamg_matvec(this, vec_out, vec_in, lvl_out)
    class(tamg_hierarchy_t), intent(inout) :: this
    real(kind=rp), intent(inout) :: vec_out(:)
    real(kind=rp), intent(inout) :: vec_in(:)
    integer, intent(in) :: lvl_out
    integer :: i, n, e
    call this%matvec_impl(vec_out, vec_in, this%nlvls, lvl_out)
    !call tamg_matvec_flat_impl(this, vec_out, vec_in, this%nlvls, lvl_out)
  end subroutine tamg_matvec
 
  recursive subroutine tamg_matvec_impl(this, vec_out, vec_in, lvl, lvl_out)
    class(tamg_hierarchy_t), intent(inout) :: this
    real(kind=rp), intent(inout) :: vec_out(:)
    real(kind=rp), intent(inout) :: vec_in(:)
    integer, intent(in) :: lvl
    integer, intent(in) :: lvl_out
    integer :: i, n, e
 
    vec_out = 0d0
 
    if (lvl .eq. 0) then
      n = size(vec_in)
      call this%gs_h%op(vec_in, n, gs_op_add)
      call col2( vec_in, this%coef%mult(1,1,1,1), n)
      call this%ax%compute(vec_out, vec_in, this%coef, this%msh, this%Xh)
      call this%gs_h%op(vec_out, n, gs_op_add)
      call this%blst%apply(vec_out, n)
    else
      if (lvl_out .ge. lvl) then
        associate( wrk_in => this%lvl(lvl)%wrk_in, wrk_out => this%lvl(lvl)%wrk_out)
        wrk_in = 0d0
        wrk_out = 0d0
        do n = 1, this%lvl(lvl)%nnodes
          associate(node => this%lvl(lvl)%nodes(n))
          do i = 1, node%ndofs
            wrk_in( node%dofs(i) ) = wrk_in( node%dofs(i) ) + vec_in( node%gid ) * node%interp_p( i )
          end do
          end associate
        end do
 
        call this%matvec_impl(wrk_out, wrk_in, lvl-1, lvl_out)
 
        do n = 1, this%lvl(lvl)%nnodes
          associate(node => this%lvl(lvl)%nodes(n))
          do i = 1, node%ndofs
            vec_out( node%gid ) = vec_out(node%gid ) + wrk_out( node%dofs(i) ) * node%interp_r( i )
          end do
          end associate
        end do
        end associate
      else if (lvl_out .lt. lvl) then
        call this%matvec_impl(vec_out, vec_in, lvl-1, lvl_out)
      else
        call neko_error("TAMG: matvec level numbering problem.")
      end if
    end if
  end subroutine tamg_matvec_impl
 
 
  recursive subroutine tamg_matvec_flat_impl(this, vec_out, vec_in, lvl_blah, lvl_out)
    class(tamg_hierarchy_t), intent(inout) :: this
    real(kind=rp), intent(inout) :: vec_out(:)
    real(kind=rp), intent(inout) :: vec_in(:)
    integer, intent(in) :: lvl_blah
    integer, intent(in) :: lvl_out
    integer :: i, n, cdof, lvl
 
    lvl = lvl_out
    if (lvl .eq. 0) then
      n = size(vec_in)
      call this%gs_h%op(vec_in, n, gs_op_add)
      call col2( vec_in, this%coef%mult(1,1,1,1), n)
      call this%ax%compute(vec_out, vec_in, this%coef, this%msh, this%Xh)
      call this%gs_h%op(vec_out, n, gs_op_add)
      call this%blst%apply(vec_out, n)
    else
 
      associate( wrk_in => this%lvl(1)%wrk_in, wrk_out => this%lvl(1)%wrk_out)
      n = size(wrk_in)
      wrk_out = 0d0
      vec_out = 0d0
 
      do i = 1, n
        cdof = this%lvl(lvl)%map_f2c_dof(i)
        wrk_in(i) = vec_in( cdof )
      end do
 
      call this%gs_h%op(wrk_in, n, gs_op_add)
      call col2( wrk_in, this%coef%mult(1,1,1,1), n)
      call this%ax%compute(wrk_out, wrk_in, this%coef, this%msh, this%Xh)
      call this%gs_h%op(wrk_out, n, gs_op_add)
      call this%blst%apply(wrk_out, n)
 
      do i = 1, n
        cdof = this%lvl(lvl)%map_f2c_dof(i)
        vec_out(cdof) = vec_out(cdof) + wrk_out( i )
      end do
      end associate
 
    end if
  end subroutine tamg_matvec_flat_impl
 
 
  subroutine tamg_restriction_operator(this, vec_out, vec_in, lvl)
    class(tamg_hierarchy_t), intent(inout) :: this
    real(kind=rp), intent(inout) :: vec_out(:)
    real(kind=rp), intent(inout) :: vec_in(:)
    integer, intent(in) :: lvl
    integer :: i, n, node_start, node_end, node_id
 
    vec_out = 0d0
    do n = 1, this%lvl(lvl)%nnodes
      associate(node => this%lvl(lvl)%nodes(n))
      do i = 1, node%ndofs
        vec_out( node%gid ) = vec_out( node%gid ) + vec_in( node%dofs(i) ) * node%interp_r( i )
      end do
      end associate
    end do
    !do n = 1, this%lvl(lvl)%nnodes
    !  node_start = this%lvl(lvl)%nodes_ptr(n)
    !  node_end   = this%lvl(lvl)%nodes_ptr(n+1)-1
    !  node_id    = this%lvl(lvl)%nodes_gid(n)
    !  do i = node_start, node_end
    !    vec_out( node_id ) = vec_out( node_id ) + &
    !      vec_in( this%lvl(lvl)%nodes_dofs(i) )
    !  end do
    !end do
  end subroutine tamg_restriction_operator
 
  subroutine tamg_prolongation_operator(this, vec_out, vec_in, lvl)
    class(tamg_hierarchy_t), intent(inout) :: this
    real(kind=rp), intent(inout) :: vec_out(:)
    real(kind=rp), intent(inout) :: vec_in(:)
    integer, intent(in) :: lvl
    integer :: i, n, node_start, node_end, node_id
 
    vec_out = 0d0
    do n = 1, this%lvl(lvl)%nnodes
      associate(node => this%lvl(lvl)%nodes(n))
      do i = 1, node%ndofs
        vec_out( node%dofs(i) ) = vec_out( node%dofs(i) ) + vec_in( node%gid ) * node%interp_p( i )
      end do
      end associate
    end do
    !do n = 1, this%lvl(lvl)%nnodes
    !  node_start = this%lvl(lvl)%nodes_ptr(n)
    !  node_end   = this%lvl(lvl)%nodes_ptr(n+1)-1
    !  node_id    = this%lvl(lvl)%nodes_gid(n)
    !  do i = node_start, node_end
    !    vec_out( this%lvl(lvl)%nodes_dofs(i) ) = vec_out( this%lvl(lvl)%nodes_dofs(i) ) + &
    !      vec_in(node_id)
    !  end do
    !end do
  end subroutine tamg_prolongation_operator
 
end module tree_amg