Homec4science
Diffusion SPClibs (feature/fft-split) Browse

spclibs/multigrid/src/parmg_mod.f90
03540012fedefeature/fft-split

Last Change
Raw File

parmg_mod.f90
View Options

!>
!> @file parmg_mod.f90
!>
!> @brief
!>
!> @copyright
!> Copyright (©) 2021 EPFL (Ecole Polytechnique Fédérale de Lausanne)
!> SPC (Swiss Plasma Center)
!>
!> spclibs is free software: you can redistribute it and/or modify it under
!> the terms of the GNU Lesser General Public License as published by the Free
!> Software Foundation, either version 3 of the License, or (at your option)
!> any later version.
!>
!> spclibs 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 Lesser General Public License
!> along with this program. If not, see <https://www.gnu.org/licenses/>.
!>
!> @authors
!> (in alphabetical order)
!> @author Trach-Minh Tran <trach-minh.tran@epfl.ch>
!>
MODULE parmg
!
! parmg: Utilities for parallel multigrid
!
! T.M. Tran, CRPP-EPFL
! December 2013
!
USE mpi
USE iso_fortran_env, ONLY : rkind => real64
USE gvector, ONLY : gvector_2d
USE stencil, ONLY : stencil_2d
IMPLICIT NONE
!
PRIVATE
PUBLIC :: grid2_type, mg_info, create_grid, mg, get_lmax, coarse, disp, exchange, &
& get_resids, jacobi, prolong, init_restrict, restrict, &
& norm_vec, norm_mat
!
TYPE grid2_type
INTEGER, DIMENSION(2) :: s, e, s0, e0, npt_loc, npt
REAL(rkind), ALLOCATABLE :: x(:)
REAL(rkind), ALLOCATABLE :: y(:)
TYPE(gvector_2d) :: f
TYPE(gvector_2d) :: v
TYPE(stencil_2d) :: fdmat
TYPE(stencil_2d) :: restrict_mat
END TYPE grid2_type
!
TYPE mg_info
INTEGER :: comm ! Communicator
INTEGER :: nu1 ! Relaxation down sweeps
INTEGER :: nu2 ! Relaxation up sweeps
INTEGER :: mu ! mu-cycle number
INTEGER :: nu0 ! Number of FMG cycles
INTEGER :: levels ! Number of mg levels
INTEGER :: direct_solve_nits ! Jacobit nits for direct_solve
CHARACTER(len=4) :: relax ! Type of relation
REAL(rkind) :: omega ! for weighted Jacobi relaxation
END TYPE mg_info
!
INTERFACE create_grid
MODULE PROCEDURE create_grid_2d
END INTERFACE create_grid
INTERFACE mg
MODULE PROCEDURE mg_2d
END INTERFACE mg
INTERFACE coarse
MODULE PROCEDURE coarse_1d, coarse_2d
END INTERFACE coarse
INTERFACE exchange
MODULE PROCEDURE exchange_g_2d, exchange_g_2d_new
END INTERFACE exchange
INTERFACE prolong
MODULE PROCEDURE prolong_g_2d
END INTERFACE prolong
INTERFACE restrict
MODULE PROCEDURE restrict_g_2d
END INTERFACE restrict
INTERFACE jacobi
MODULE PROCEDURE jacobi_stencila_2d
MODULE PROCEDURE jacobi_stencilg_2d
END INTERFACE jacobi
INTERFACE get_resids
MODULE PROCEDURE resids_stencila_2d
MODULE PROCEDURE resids_stencilg_2d
END INTERFACE get_resids
INTERFACE disp
MODULE PROCEDURE dispi_0, dispi_1
END INTERFACE disp
CONTAINS
!
!--------------------------------------------------------------------------------
SUBROUTINE create_grid_2d(x, y, s_in, e_in, id, prb, grids, comm)
!
! Create arrays of partitionned grids
!
USE stencil, ONLY : init
!
REAL(rkind), INTENT(in) :: x(0:), y(0:) ! Global coordinates
INTEGER, INTENT(in) :: s_in(2), e_in(2) ! Partition of finest grid
INTEGER, INTENT(in) :: id(:,:) ! Structure of stencil
CHARACTER(len=*), INTENT(in) :: prb
INTEGER, INTENT(in) :: comm
TYPE(grid2_type) :: grids(:)
!
INTEGER :: levels
INTEGER :: s0(2), e0(2), s(2), e(2)
INTEGER :: npt_loc(2), npt_loc_min(2), npt_glob(2)
INTEGER :: l, ierr
!
levels = SIZE(grids)
s = s_in
e = e_in
!
DO l=1,levels
IF(l.GT.1) THEN
CALL coarse(s,e)
END IF
npt_loc = e-s+1
CALL mpi_allreduce(s, s0, 2, MPI_INTEGER, MPI_MIN, comm, ierr)
CALL mpi_allreduce(e, e0, 2, MPI_INTEGER, MPI_MAX, comm, ierr)
CALL mpi_allreduce(npt_loc, npt_loc_min, 2, MPI_INTEGER, MPI_MIN, comm, ierr)
IF(MINVAL(npt_loc_min) .LT. 2) THEN
PRINT*, 'CREATE_GRID: number intervals too small!'
STOP
END IF
npt_glob = e0+1
grids(l)%s0 = s0
grids(l)%e0 = e0
grids(l)%s = s
grids(l)%e = e
grids(l)%npt_loc = npt_loc
grids(l)%npt = npt_glob
grids(l)%f = gvector_2d(s, e, [1,1]) ! Arrays with ghost cell
grids(l)%v = gvector_2d(s, e, [1,1])
ALLOCATE(grids(l)%x(s0(1):e0(1))) ! Global coords (x,y)
ALLOCATE(grids(l)%y(s0(2):e0(2)))
IF(l.EQ.1) THEN
grids(1)%x = x
grids(1)%y = y
ELSE
grids(l)%x(:) = grids(l-1)%x(0::2)
grids(l)%y(:) = grids(l-1)%y(0::2)
END IF
END DO
!
! Set up FD matrix
!
DO l=1,levels
s = grids(l)%s
e = grids(l)%e
CALL init(s, e, id, .FALSE., grids(l)%fdmat, comm)
END DO
!
! Set up restriction stencil
!
DO l=2,levels
CALL init_restrict(grids(l), prb, comm)
END DO
!
END SUBROUTINE create_grid_2d
!--------------------------------------------------------------------------------
RECURSIVE SUBROUTINE mg_2d(grids, info, l)
!
! Execute a recursive V-cycle
!
USE gvector, ONLY : ASSIGNMENT(=), OPERATOR(+)
TYPE(grid2_type), INTENT(inout) :: grids(:)
TYPE(mg_info), INTENT(in) :: info
INTEGER, INTENT(in) :: l
!
TYPE(gvector_2d) :: resids, v_prolong
INTEGER, DIMENSION(2) :: s0, e0, s, e, g=[1,1]
INTEGER :: comm, levels, k
!
comm = info%comm
levels = info%levels
!
s0 = grids(l)%s0; e0 = grids(l)%e0
s = grids(l)%s; e = grids(l)%e
resids = gvector_2d(s, e, g)
!
IF(l.EQ.levels) THEN
CALL direct_solve(grids(l)%fdmat, grids(l)%v, grids(l)%f)
ELSE
CALL relax(info%nu1)
resids = get_resids(comm, grids(l)%fdmat, grids(l)%v, grids(l)%f)
CALL exchange(comm, resids)
CALL restrict(grids(l+1)%restrict_mat, resids, grids(l+1)%f)
grids(l+1)%v = 0.0d0
!
! Only 1 call to the coarsest level
DO k=1,MERGE(info%mu, 1, (l+1).NE.levels)
CALL mg(grids, info, l+1)
END DO
!
v_prolong = gvector_2d(s, e, g)
CALL exchange(comm, grids(l+1)%v)
CALL prolong(grids(l+1)%v, v_prolong)
grids(l)%v = grids(l)%v + v_prolong
CALL relax(info%nu2)
END IF
!
CONTAINS
SUBROUTINE relax(nu)
INTEGER, INTENT(in) :: nu
SELECT CASE (TRIM(info%relax))
CASE ("jac")
CALL jacobi(comm, grids(l)%fdmat, info%omega, nu, grids(l)%v, grids(l)%f)
CASE default
PRINT*, "relax ", info%relax, " NOT IMPLEMENTED!"
STOP
END SELECT
END SUBROUTINE relax
SUBROUTINE direct_solve(mat, v, f)
TYPE(stencil_2d), INTENT(in) :: mat
TYPE(gvector_2d), INTENT(inout) :: v
TYPE(gvector_2d), INTENT(in) :: f
v = 0.0d0
CALL jacobi(comm, mat, 1.0_rkind, info%direct_solve_nits, v, f)
END SUBROUTINE direct_solve
END SUBROUTINE mg_2d
!--------------------------------------------------------------------------------
FUNCTION get_lmax(s_in, npt_loc, npt_min, comm) RESULT(lmax)
!
! Get max number of levels on all processes
!
INTEGER :: lmax
INTEGER, INTENT(in) :: s_in, npt_loc, npt_min, comm
INTEGER :: me, ierr
INTEGER :: s, e, kpt_loc, kpt, kpt_loc_min
!
CALL mpi_comm_rank(comm, me, ierr)
s = s_in
kpt_loc = npt_loc
e = s+npt_loc-1
lmax = 1
DO
CALL mpi_allreduce(kpt_loc, kpt_loc_min, 1, MPI_INTEGER, MPI_MIN, &
& comm, ierr)
CALL mpi_allreduce(kpt_loc, kpt, 1, MPI_INTEGER, MPI_SUM, &
& comm, ierr)
!
! Stop if npt-1 not even or when minumum local npt is attained
IF(MODULO(kpt-1,2).NE.0 .OR. kpt_loc_min .LE. npt_min) EXIT
!
lmax = lmax+1
CALL coarse(s, e)
kpt_loc = e-s+1
END DO
!
END FUNCTION get_lmax
!--------------------------------------------------------------------------------
SUBROUTINE coarse_1d(s, e)
!
! Compute (s,e) of next coarse grid
!
INTEGER, INTENT(inout) :: s, e
INTEGER :: s0, npt, i
!
! Previous odd indices are discarded
s0 = s
IF( MODULO(s0,2) .NE. 0 ) THEN
s0 = s+1
END IF
!
! Count local number of points
npt = 0
DO i=s0,e,2
npt = npt+1
END DO
!
! Coarse s, e
s = s0/2
e = s + npt - 1
END SUBROUTINE coarse_1d
!--------------------------------------------------------------------------------
SUBROUTINE coarse_2d(s, e)
!
! Compute (s,e) of next coarse grid
!
INTEGER, INTENT(inout) :: s(2), e(2)
!
CALL coarse_1d(s(1), e(1))
CALL coarse_1d(s(2), e(2))
END SUBROUTINE coarse_2d
!--------------------------------------------------------------------------------
SUBROUTINE dispi_0(str, a, comm)
!
! Display integer local scalar
!
INTEGER, INTENT(in) :: a, comm
CHARACTER(len=*), INTENT(in) :: str
INTEGER :: npes, me, ierr
INTEGER, ALLOCATABLE, DIMENSION(:) :: a_gather(:)
!
CALL MPI_COMM_RANK(comm, me, ierr)
CALL MPI_COMM_SIZE(comm, npes, ierr)
ALLOCATE(a_gather(npes))
CALL MPI_GATHER(a, 1, MPI_INTEGER, a_gather, 1, MPI_INTEGER, &
& 0, comm, ierr)
IF( me.EQ.0 ) THEN
WRITE(*,'(a/(20i6))') str, a_gather
END IF
DEALLOCATE(a_gather)
END SUBROUTINE dispi_0
!--------------------------------------------------------------------------------
SUBROUTINE dispi_1(str, a, comm)
!
! Display integer local array
!
INTEGER, INTENT(in) :: a(:), comm
CHARACTER(len=*), INTENT(in) :: str
INTEGER :: npes, me, ierr, n, i
INTEGER, ALLOCATABLE, DIMENSION(:) :: a_gather(:,:)
!
n = SIZE(a,1)
CALL MPI_COMM_RANK(comm, me, ierr)
CALL MPI_COMM_SIZE(comm, npes, ierr)
ALLOCATE(a_gather(n,npes))
CALL MPI_GATHER(a, n, MPI_INTEGER, a_gather, n, MPI_INTEGER, &
& 0, comm, ierr)
IF( me.EQ.0 ) THEN
WRITE(*,'(a)') str
DO i=1,n
WRITE(*,'(20i6)') a_gather(i,:)
END DO
END IF
DEALLOCATE(a_gather)
END SUBROUTINE dispi_1
!--------------------------------------------------------------------------------
SUBROUTINE exchange_g_2d_new(comm, u)
!
! Exchange ghost cells with (west,east,south,north) neighbors.
! Assume same ghost cells on each dimension:
! u%g(1) : number of ghost cells on west and east boundaries
! u%g(2) : number of ghost cells on south and north boundaries
!
INTEGER, INTENT(in) :: comm
TYPE(gvector_2d), INTENT(inout) :: u
INTEGER :: neighs(4), ierr
!
CALL mpi_cart_shift(comm, 0, 1, neighs(1), neighs(2), ierr)
CALL mpi_cart_shift(comm, 1, 1, neighs(3), neighs(4), ierr)
CALL exchange_g_2d(comm, neighs, u)
END SUBROUTINE exchange_g_2d_new
!--------------------------------------------------------------------------------
SUBROUTINE exchange_g_2d(comm, neighs, u)
!
! Exchange ghost cells with (west,east,south,north) neighbors.
! Assume same ghost cells on each dimension:
! u%g(1) : number of ghost cells on west and east boundaries
! u%g(2) : number of ghost cells on south and north boundaries
!
INTEGER, INTENT(in) :: comm
INTEGER, INTENT(in) :: neighs(4)
TYPE(gvector_2d), INTENT(inout) :: u
!
INTEGER :: cols, rows
INTEGER :: ierr
INTEGER, PARAMETER :: ndim=2
INTEGER, DIMENSION(ndim) :: g, lb, ub, s, e, n
!
s = u%s
e = u%e
g = u%g
lb = s - g
ub = e + g
n = ub - lb + 1 ! include ghost cells
!
! g(2) matrix full rows with stride n(1)
CALL mpi_type_vector(n(2), g(2), n(1), MPI_DOUBLE_PRECISION, rows, ierr)
CALL mpi_type_commit(rows, ierr)
!
! g(1) contiguous matrix full columns
CALL mpi_type_contiguous(n(1)*g(1), MPI_DOUBLE_PRECISION, cols, ierr)
CALL mpi_type_commit(cols, ierr)
!
! Exchange along first dimension
CALL mpi_sendrecv(u%val(s(1), lb(2)), 1, rows, neighs(1), 0, &
& u%val(e(1)+1,lb(2)), 1, rows, neighs(2), 0, &
& comm, MPI_STATUS_IGNORE, ierr)
CALL mpi_sendrecv(u%val(e(1)-g(1)+1,lb(2)), 1, rows, neighs(2), 0, &
& u%val(lb(1), lb(2)), 1, rows, neighs(1), 0, &
& comm, MPI_STATUS_IGNORE, ierr)
!
! Exchange along second dimension
CALL mpi_sendrecv(u%val(lb(1),s(2)), 1, cols, neighs(3), 0, &
& u%val(lb(1),e(2)+1), 1, cols, neighs(4), 0, &
& comm, MPI_STATUS_IGNORE, ierr)
CALL mpi_sendrecv(u%val(lb(1),e(2)-g(2)+1), 1, cols, neighs(4), 0, &
& u%val(lb(1),lb(2)), 1, cols, neighs(3), 0, &
& comm, MPI_STATUS_IGNORE, ierr)
END SUBROUTINE exchange_g_2d
!--------------------------------------------------------------------------------
SUBROUTINE prolong_g_2d(vbar, v)
!
! 2D bilinear prolongation
!
TYPE(gvector_2d), INTENT(in) :: vbar
TYPE(gvector_2d), INTENT(inout) :: v
!
INTEGER :: i,j,i1,i2,j1,j2
!
i1 = v%s(1)-MODULO(v%s(1),2); i2 = v%e(1)+MODULO(v%e(1),2)
j1 = v%s(2)-MODULO(v%s(2),2); j2 = v%e(2)+MODULO(v%e(2),2)
!
! Even numbered nodes on fine mesh
!
DO j=j1,j2,2
DO i=i1,i2,2
v%val(i,j) = vbar%val(i/2,j/2)
END DO
END DO
!
! Linear interpolation on x
!
DO j=j1,j2,2
DO i=i1+1,i2-1,2
v%val(i,j) = 0.5d0*(v%val(i-1,j)+v%val(i+1,j))
END DO
END DO
!
! Linear interpolation on y
!
DO j=j1+1,j2-1,2
DO i=i1,i2
v%val(i,j) = 0.5d0*(v%val(i,j-1)+v%val(i,j+1))
END DO
END DO
END SUBROUTINE prolong_g_2d
!--------------------------------------------------------------------------------
SUBROUTINE init_restrict(grid, prb, comm)
!
! Set up restriction stencil
!
USE stencil, ONLY : init
TYPE(grid2_type), INTENT(inout) :: grid
CHARACTER(len=*), INTENT(in) :: prb
INTEGER, INTENT(in) :: comm
!
INTEGER, PARAMETER :: npoints=9, ndim=2
INTEGER :: s(2), e(2), n(2), id(9,2)
INTEGER :: i, j
!
! Stencil structure initialization
!
s = grid%s
e = grid%e
n = grid%npt-1
!
! N
! 6---7---8
! | | |
! W 4---0---5 E Numbering of stencil
! | | |
! 1---2---3
! S
!
id = RESHAPE([0, -1, 0, 1,-1, 1,-1, 0, 1, &
& 0, -1,-1,-1, 0, 0, 1, 1, 1], &
& [npoints, ndim])
CALL init(s, e, id, .FALSE., grid%restrict_mat, comm)
!
! Fill in stencil
!
DO j=s(2),e(2)
DO i=s(1),e(1)
grid%restrict_mat%val(i,j,:) = [4._rkind, &
& 1._rkind, 2._rkind, 1._rkind, &
& 2._rkind, 2._rkind, &
& 1._rkind, 2._rkind, 1._rkind ]&
& / 16._rkind
END DO
END DO
!
! Apply Dirichlet BC
!
IF(s(1).EQ.0 .AND. prb(1:1).EQ.'d') THEN ! West face
grid%restrict_mat%val(s(1),:,3) = 0._rkind
grid%restrict_mat%val(s(1),:,5) = 0._rkind
grid%restrict_mat%val(s(1),:,8) = 0._rkind
END IF
IF(e(1).EQ.n(1) .AND. prb(2:2).EQ.'d') THEN ! East face
grid%restrict_mat%val(e(1),:,1) = 0._rkind
grid%restrict_mat%val(e(1),:,4) = 0._rkind
grid%restrict_mat%val(e(1),:,6) = 0._rkind
END IF
IF(s(2).EQ.0 .AND. prb(3:3).EQ.'d') THEN ! South face
grid%restrict_mat%val(:,s(2),6) = 0._rkind
grid%restrict_mat%val(:,s(2),7) = 0._rkind
grid%restrict_mat%val(:,s(2),8) = 0._rkind
END IF
IF(e(2).EQ.n(2) .AND. prb(4:4).EQ.'d') THEN ! North face
grid%restrict_mat%val(:,e(2),1) = 0._rkind
grid%restrict_mat%val(:,e(2),2) = 0._rkind
grid%restrict_mat%val(:,e(2),3) = 0._rkind
END IF
END SUBROUTINE init_restrict
!--------------------------------------------------------------------------------
SUBROUTINE jacobi_stencila_2d(mat, omega, nu, v, f)
!
! Weighted Jacobi relaxation
!
TYPE(stencil_2d),INTENT(in) :: mat
REAL(rkind), INTENT(in) :: omega
INTEGER, INTENT(in) :: nu
REAL(rkind), ALLOCATABLE, INTENT(inout) :: v(:,:)
REAL(rkind), ALLOCATABLE, INTENT(in) :: f(:,:)
!
REAL(rkind), ALLOCATABLE :: temp(:,:), inv_diag(:,:)
INTEGER, DIMENSION(2) :: smin, emax, s, e, d, lb, ub
INTEGER :: it, k, i, j
!
s(:) = mat%s(:)
e(:) = mat%e(:)
smin(:) = mat%s0(:)
emax(:) = mat%e0(:)
!
ALLOCATE(temp(s(1):e(1),s(2):e(2)))
ALLOCATE(inv_diag(s(1):e(1),s(2):e(2)))
!
inv_diag(:,:) = omega/mat%val(:,:,0)
DO it=1,nu
temp(:,:) = f(s(1):e(1),s(2):e(2))
DO k=1,mat%npoints-1 ! exclude the diagonal term, f - (L+U)*v
d(:) = mat%id(k,:)
lb = MAX(smin, smin-d, mat%s)
ub = MIN(emax, emax-d, mat%e)
DO j=lb(2),ub(2)
DO i=lb(1),ub(1)
temp(i,j) = temp(i,j) - mat%val(i,j,k)*v(i+d(1),j+d(2))
END DO
END DO
END DO
temp = temp * inv_diag
v(s(1):e(1),s(2):e(2)) = (1.d0-omega)*v(s(1):e(1),s(2):e(2)) + temp
END DO
!
DEALLOCATE(temp)
DEALLOCATE(inv_diag)
END SUBROUTINE jacobi_stencila_2d
!--------------------------------------------------------------------------------
SUBROUTINE jacobi_stencilg_2d(comm, mat, omega, nu, v, f)
!
! Weighted Jacobi relaxation
!
USE gvector, ONLY : ASSIGNMENT(=)
INTEGER, INTENT(in) :: comm
TYPE(stencil_2d),INTENT(in) :: mat
REAL(rkind), INTENT(in) :: omega
INTEGER, INTENT(in) :: nu
TYPE(gvector_2d), INTENT(inout) :: v
TYPE(gvector_2d), INTENT(in) :: f
!
REAL(rkind), ALLOCATABLE :: temp(:,:), inv_diag(:,:)
INTEGER, DIMENSION(2) :: s, e, d
INTEGER :: it, k, i, j
!
s(:) = v%s(:)
e(:) = v%e(:)
!
ALLOCATE(temp(s(1):e(1),s(2):e(2)))
ALLOCATE(inv_diag(s(1):e(1),s(2):e(2)))
!
inv_diag(:,:) = omega/mat%val(:,:,0)
DO it=1,nu
CALL exchange(comm, v)
temp(:,:) = f%val(s(1):e(1),s(2):e(2))
DO k=1,mat%npoints-1 ! exclude the diagonal term, f - (L+U)*v
d(:) = mat%id(k,:)
DO j=s(2),e(2)
DO i=s(1),e(1)
temp(i,j) = temp(i,j) - &
& mat%val(i,j,k) * v%val(i+d(1),j+d(2))
END DO
END DO
END DO
temp = temp * inv_diag
v%val(s(1):e(1),s(2):e(2)) = (1.d0-omega) * v%val(s(1):e(1),s(2):e(2)) + &
& temp
END DO
END SUBROUTINE jacobi_stencilg_2d
!--------------------------------------------------------------------------------
FUNCTION resids_stencila_2d(mat, xarr, farr) RESULT(res)
!
! Return residuals res = mat*x, where x, farr and res are simple arrays
!
TYPE(stencil_2d), INTENT(in) :: mat
REAL(rkind), ALLOCATABLE, INTENT(in) :: xarr(:,:)
REAL(rkind), ALLOCATABLE, INTENT(in) :: farr(:,:)
REAL(rkind) :: res(LBOUND(xarr,1):UBOUND(xarr,1), &
& LBOUND(xarr,2):UBOUND(xarr,2))
INTEGER :: k, i, j
INTEGER, DIMENSION(2) :: smin, emax, d, lb, ub
!
smin(:) = mat%s0(:)
emax(:) = mat%e0(:)
res = farr
DO k=0,mat%npoints-1
d(:) = mat%id(k,:)
lb = MAX(smin, smin-d, mat%s)
ub = MIN(emax, emax-d, mat%e)
DO j=lb(2),ub(2)
DO i=lb(1),ub(1)
res(i,j) = res(i,j) - mat%val(i,j,k)*xarr(i+d(1),j+d(2))
END DO
END DO
END DO
END FUNCTION resids_stencila_2d
!--------------------------------------------------------------------------------
FUNCTION resids_stencilg_2d(comm, mat, xarr, farr) RESULT(res)
!
! Return residuals res= f-mat*x, where x, f and res are gvectors
!
INTEGER, INTENT(in) :: comm
TYPE(stencil_2d), INTENT(in) :: mat
TYPE(gvector_2d), INTENT(inout) :: xarr
TYPE(gvector_2d), INTENT(in) :: farr
TYPE(gvector_2d) :: res
INTEGER :: k, i, j
INTEGER, DIMENSION(2) :: s, e, d
!
s = xarr%s
e = xarr%e
res = gvector_2d(xarr%s, xarr%e, xarr%g)
res%val = farr%val
CALL exchange(comm, xarr)
DO k=0,mat%npoints-1
d(:) = mat%id(k,:)
DO j=s(2),e(2)
DO i=s(1),e(1)
res%val(i,j) = res%val(i,j) - mat%val(i,j,k)*xarr%val(i+d(1),j+d(2))
END DO
END DO
END DO
END FUNCTION resids_stencilg_2d
!--------------------------------------------------------------------------------
SUBROUTINE restrict_g_2d(mat, f, fbar)
!
! 2D full weighting restriction
!
TYPE(stencil_2d), INTENT(in) :: mat
TYPE(gvector_2d), INTENT(in) :: f
TYPE(gvector_2d), INTENT(inout) :: fbar
!
INTEGER, DIMENSION(2) :: s, e, d
INTEGER :: k, i, j
!
s = fbar%s
e = fbar%e
!
! Diagonal contributions: d(0) = (0,0)
DO j=s(2),e(2)
DO i=s(1),e(1)
fbar%val(i,j) = mat%val(i,j,0) * f%val(2*i,2*j)
END DO
END DO
!
DO k=1,mat%npoints-1
d(:) = mat%id(k,:)
DO j=s(2),e(2)
DO i=s(1),e(1)
fbar%val(i,j) = fbar%val(i,j) + &
& mat%val(i,j,k) * f%val(2*i+d(1),2*j+d(2))
END DO
END DO
END DO
END SUBROUTINE restrict_g_2d
!--------------------------------------------------------------------------------
REAL(rkind) FUNCTION norm_vec(x, comm, root)
!
! Infinity vector norm
!
TYPE(gvector_2d), INTENT(in) :: x
INTEGER, INTENT(in) :: comm
INTEGER, OPTIONAL, intent(in) :: root
REAL(rkind) :: temp
INTEGER :: ierr
temp = MAXVAL( ABS(x%val(x%s(1):x%e(1),x%s(2):x%e(2))) )
IF(PRESENT(root)) THEN
CALL mpi_reduce(temp, norm_vec, 1, MPI_DOUBLE_PRECISION, MPI_MAX, &
& root, comm, ierr)
ELSE
CALL mpi_allreduce(temp, norm_vec, 1, MPI_DOUBLE_PRECISION, MPI_MAX, &
& comm, ierr)
END IF
END FUNCTION norm_vec
!--------------------------------------------------------------------------------
REAL(rkind) FUNCTION norm_mat(mat, comm, root)
!
! Infinity matrix norm
!
TYPE(stencil_2d), INTENT(in) :: mat
INTEGER, INTENT(in) :: comm
INTEGER, OPTIONAL, intent(in) :: root
REAL(rkind) :: arr_temp(mat%s(1):mat%e(1),mat%s(2):mat%e(2))
REAL(rkind) :: temp
INTEGER :: i, j, s(2), e(2), ierr
s = mat%s; e = mat%e
DO j=s(2),e(2)
DO i=s(1),e(1)
arr_temp(i,j) = SUM(ABS(mat%val(i,j,:)))
END DO
END DO
temp = MAXVAL(arr_temp)
IF(PRESENT(root)) THEN
CALL mpi_reduce(temp, norm_mat, 1, MPI_DOUBLE_PRECISION, MPI_MAX, &
& root, comm, ierr)
ELSE
CALL mpi_allreduce(temp, norm_mat, 1, MPI_DOUBLE_PRECISION, MPI_MAX, &
& comm, ierr)
END IF
END FUNCTION norm_mat
!--------------------------------------------------------------------------------
END MODULE parmg
c4science · Help