Homec4science
Diffusion SPClibs (feature/fft-split) Browse

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

Last Change
Raw File

stencil_mod.f90
View Options

!>
!> @file stencil_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 stencil
!
! stencil_2d: Implement 2D stencil for matrix-less operations
!
! T.M. Tran, CRPP-EPFL
! August 2013
!
USE iso_fortran_env, ONLY : rkind => real64
IMPLICIT NONE
!
PRIVATE
PUBLIC :: stencil_2d, init, vmx, putmat, laplacian, &
& OPERATOR(*)
!
TYPE stencil_2d
LOGICAL :: nluni
INTEGER, DIMENSION(2) :: ldim, gdim
INTEGER, DIMENSION(2) :: s0, e0, s, e
INTEGER :: npoints
INTEGER, ALLOCATABLE :: id(:,:)
REAL(rkind), ALLOCATABLE :: val(:,:,:)
END TYPE stencil_2d
!
INTERFACE init
MODULE PROCEDURE init_stencil_2d
END INTERFACE init
INTERFACE vmx
MODULE PROCEDURE vmx_stencila_2d
MODULE PROCEDURE vmx_stencilg_2d
END INTERFACE vmx
INTERFACE putmat
module procedure putmat_stencil
END INTERFACE putmat
!
INTERFACE OPERATOR(*)
MODULE PROCEDURE vmx_stencila_2d
MODULE PROCEDURE vmx_stencilg_2d
END INTERFACE OPERATOR(*)
!
CONTAINS
!================================================================================
SUBROUTINE init_stencil_2d(s, e, id, nluni, mat, comm)
!
! stencil_2d constructor
!
USE mpi
INTEGER, INTENT(in) :: s(2), e(2) ! Bounds in each dim.
INTEGER, INTENT(in) :: id(:,:) ! Structure of stencil
LOGICAL, INTENT(in) :: nluni ! Uniform stencil
TYPE(stencil_2d), INTENT(out) :: mat
INTEGER, INTENT(in) :: comm
INTEGER :: me, ndim=2, ierr
INTEGER :: npoints ! Size of the stencil
!
CALL mpi_comm_rank(comm, me, ierr)
!
IF(id(1,1).NE.0 .AND. id(1,2).NE.0) THEN
IF(me.EQ.0) THEN
WRITE(*,*) 'INIT_STENCIL: id(1,:) should be (0,0)!'
CALL mpi_abort(comm, -1, ierr)
END IF
END IF
!
npoints = SIZE(id,1)
mat%npoints = npoints
mat%s = s
mat%e = e
mat%nluni = nluni
IF(nluni) THEN
ALLOCATE(mat%val(1,1, 0:npoints-1))
ELSE
ALLOCATE(mat%val(s(1):e(1), s(2):e(2), 0:npoints-1))
END IF
ALLOCATE(mat%id(0:npoints-1, ndim))
mat%id(:,:) = id(:,:)
mat%val(:,:,:) = 0.0
!
mat%ldim = e-s+1
CALL mpi_allreduce(mat%s, mat%s0, ndim, MPI_INTEGER, MPI_MIN, comm, ierr)
CALL mpi_allreduce(mat%e, mat%e0, ndim, MPI_INTEGER, MPI_MAX, comm, ierr)
mat%gdim = mat%e0 - mat%s0 + 1
!
END SUBROUTINE init_stencil_2d
!================================================================================
FUNCTION vmx_stencila_2d(mat, xarr) RESULT(res)
!
! Return product res = mat*x, where x and res are simple arrays
!
TYPE(stencil_2d), INTENT(in) :: mat
REAL(rkind), ALLOCATABLE, INTENT(in) :: xarr(:,:)
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 = 0.0
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 vmx_stencila_2d
!================================================================================
FUNCTION vmx_stencilg_2d(mat, xarr) RESULT(res)
!
! Return product res= mat*x, where x and res are gvectors
!
USE gvector, ONLY : gvector_2d
TYPE(stencil_2d), INTENT(in) :: mat
TYPE(gvector_2d), INTENT(in) :: xarr
TYPE(gvector_2d) :: res
INTEGER :: k, i, j
INTEGER, DIMENSION(2) :: d, s, e
!
s = xarr%s
e = xarr%e
res = gvector_2d(xarr%s, xarr%e, xarr%g)
!
! Diagonal contributions: d(0) = (0,0)
DO j=s(2),e(2)
DO i=s(1),e(1)
res%val(i,j) = mat%val(i,j,0)*xarr%val(i,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)
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 vmx_stencilg_2d
!================================================================================
SUBROUTINE putmat_stencil(fid, label, mat, str)
USE futils
INTEGER, INTENT(in) :: fid
CHARACTER(len=*), INTENT(in) :: label
TYPE(stencil_2d), INTENT(in) :: mat
CHARACTER(len=*), INTENT(in), OPTIONAL :: str
!
IF(PRESENT(str)) THEN
CALL creatg(fid, label, str)
ELSE
CALL creatg(fid, label)
END IF
!
CALL putarr(fid, TRIM(label)//'/dists', mat%id, ionode=0)
CALL putarrnd(fid, TRIM(label)//'/val', mat%val, (/1,2/))
END SUBROUTINE putmat_stencil
!=======================================================================
SUBROUTINE laplacian(dx, dy, mat)
!
! Construct a Laplacian using 5-point FD discretization
! Assume homegeneous Dirichlet BC on all 4 faces.
!
REAL(rkind), INTENT(in) :: dx, dy
TYPE(stencil_2d), INTENT(inout) :: mat
!
INTEGER :: i, j, k
INTEGER :: ieast, iwest, jsouth, jnorth
INTEGER, DIMENSION(2) :: d
REAL(rkind) :: dx2inv, dy2inv
!
! Assemble the stencil
!
dx2inv = 1.0d0/dx**2
dy2inv = 1.0d0/dy**2
!
mat%val(:,:,0) = -2.0d0*(dx2inv+dy2inv) ! Diagonal
!
DO k=1,mat%npoints-1 ! Off diagonal
d = mat%id(k,:)
DO j=mat%s(2),mat%e(2)
DO i=mat%s(1),mat%e(1)
IF(d(1).EQ.0) THEN ! north and south
mat%val(i,j,k) = dy2inv
ELSE IF(d(2).EQ.0) THEN ! east and west
mat%val(i,j,k) = dx2inv
END IF
END DO
END DO
END DO
!
! Impose Dirichlet BC on all 4 boundaries
!
ieast = mat%s0(1)
IF(ieast .EQ. mat%s(1)) THEN ! East boundary
mat%val(ieast, mat%s(2):mat%e(2), :) = 0.0
mat%val(ieast, mat%s(2):mat%e(2), 0) = 1.0
END IF
iwest = mat%e0(1)
IF(iwest .EQ. mat%e(1)) THEN ! West boundary
mat%val(iwest, mat%s(2):mat%e(2), :) = 0.0
mat%val(iwest, mat%s(2):mat%e(2), 0) = 1.0
END IF
jsouth = mat%s0(2)
IF(jsouth .EQ. mat%s(2)) THEN ! South boundary
mat%val(mat%s(1):mat%e(1), jsouth, :) = 0.0
mat%val(mat%s(1):mat%e(1), jsouth, 0) = 1.0
END IF
jnorth = mat%e0(2)
IF(jnorth .EQ. mat%e(2)) THEN ! North boundary
mat%val(mat%s(1):mat%e(1), jnorth, :) = 0.0
mat%val(mat%s(1):mat%e(1), jnorth, 0) = 1.0
END IF
!
END SUBROUTINE laplacian
!================================================================================
END MODULE stencil
c4science · Help