! poissonmg3d -- Copyright 2010 Paolo Luchini
! http://CPLcode.net/Applications/Numerical/Multigrid/
!
! solves the constant-coefficient 3D Poisson problem using a multigrid iteration

!( Permission is hereby granted, free of charge, to any person obtaining a copy
  of this software and associated documentation files (the "Software"), to deal
  in the Software without restriction, including without limitation the rights
  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
  copies of the Software, and to permit persons to whom the Software is
  furnished to do so, subject to the following conditions:
  
  The above copyright notice and this permission notice shall be included in all
  copies or substantial portions of the Software.
  
  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE. !)


!( Use this to impose Dirichlet boundary conditions.
  The array of unknowns has to be declared in the main program as
  var(0..nx,0..ny,0..nz), with no ghost points.
  This routine does nothing.
  !)
SUBROUTINE dirichlet(POINTER TO ARRAY(∗,∗,∗) OF REAL var)
END dirichlet

!( This routine imposes Neumann boundary conditions.
  The array of unknowns has to be declared in the main program as
  var(-1..nx+1,-1..ny+1,-1..nz+1), with one ghost plane on each of the 6 sides.
  !)
SUBROUTINE neumann(POINTER TO ARRAY(∗,∗,∗) OF REAL var)
  nx=var.HI-1; ny=var(0).HI-1; nz=var(0,0).HI-1
  DO var(ix,-1,iz)=var(ix,1,iz); var(ix,ny+1,iz)=var(ix,ny-1,iz) FOR ix=0 TO nx AND iz=0 TO nz
  DO var(-1,iy,iz)=var(1,iy,iz); var(nx+1,iy,iz)=var(nx-1,iy,iz) FOR iy=0 TO ny AND iz=0 TO nz
  DO var(ix,iy,-1)=var(ix,iy,1); var(ix,iy,nz+1)=var(ix,iy,nz-1) FOR ix=0 TO nx AND iy=0 TO ny
END neumann

!( This routine imposes periodic boundary conditions in the x and y directions
  and Neumann boundary conditions in the z direction.
  The array of unknowns has to be declared in the main program as
  var(-1..nx,-1..ny,-1..nz+1), with one ghost plane in the periodic directions
  and two in the Neumann direction.
  !)
SUBROUTINE biperiod(POINTER TO ARRAY(∗,∗,∗) OF REAL var)
  ! In a distributed-memory parallel implementation, this is the only
  ! routine where communication will need to appear.
  DO var(ix,iy,LO)=var(ix,iy,LO+2); var(ix,iy,HI)=var(ix,iy,HI-2) FOR ix=LO+1 TO HI-1 AND iy=LO+1 TO HI-1
  DO var(ix,LO,iz)=var(ix,HI-1,iz); var(ix,HI,iz)=var(ix,LO+1,iz) FOR ix=LO+1 TO HI-1 AND ALL iz
  DO var(LO,iy,iz)=var(HI-1,iy,iz); var(HI,iy,iz)=var(LO+1,iy,iz) FOR ALL iy,iz
END biperiod

!( When no Dirichlet boundary conditions are used, the solution exists subject
  to the compatibility condition that the sum of the r.h.s. should be zero.
  This routine enforces this by subtracting from the r.h.s. its average.
  !)
SUBROUTINE compat(POINTER TO ARRAY(∗,∗,∗) OF REAL rhs)
  REAL s=SUM rhs(ix,iy,iz) FOR ix=LO+1 TO HI-1 AND iy=LO+1 TO HI-1 AND iz=LO+2 TO HI-2
  DO s=~+0.5∗[rhs(ix,iy,LO+1)+rhs(ix,iy,HI-1)] FOR ix=LO+1 TO HI-1 AND iy=LO+1 TO HI-1
  s=s/(rhs.HI1-rhs.LO1-1)/(rhs.HI2-rhs.LO2-1)/(rhs.HI3-rhs.LO3-2)
  DO rhs(ix,iy,iz)=~-s FOR ALL ix,iy,iz
END compat

!( Basic red-black Jacobi iteration used as the smoother. The parameter
  makeghosts can be any one of the above routines used to impose boundary
  conditions, or you can construct your own hybrid.
  !)
SUBROUTINE poissonrb(POINTER TO ARRAY(∗,∗,∗) OF REAL var, rhs; SUBROUTINE(POINTER TO ARRAY(∗,∗,∗) OF REAL v) makeghosts)
  LOOP FOR parity=0 TO 1 
    LOOP FOR ix=var.LO+1 TO var.HI-1 AND iy=var(0).LO+1 TO var(0).HI-1 AND iz=var(0,0).LO+1+[(ix+iy+var(0,0).LO+1+parity) MOD 2] TO var(0,0).HI-1 BY 2
      var(ix,iy,iz)=(var(ix+1,iy,iz)+var(ix-1,iy,iz)+var(ix,iy+1,iz)+var(ix,iy-1,iz)+var(ix,iy,iz+1)+var(ix,iy,iz-1))/6-rhs(ix,iy,iz)
    REPEAT LOOP
    makeghosts(var)
  REPEAT LOOP
END poissonrb

!( Current residual of the difference equation, used both in the multigrid
  routine and to write out the residual itself.
  !)
INLINE FUNCTION rsd(REAL var(∗,∗,∗); INTEGER ix,iy,iz)=(var(ix,iy+1,iz)+var(ix,iy-1,iz)+var(ix+1,iy,iz)+var(ix-1,iy,iz)+var(ix,iy,iz+1)+var(ix,iy,iz-1))/6-var(ix,iy,iz)

! Main multigrid routine. Same calling prototype as poissonrb.
SUBROUTINE poissonmg(POINTER TO ARRAY(∗,∗,∗) OF REAL fvar, frhs; SUBROUTINE(POINTER TO ARRAY(∗,∗,∗) OF REAL var) makeghosts)
  ! apply smoother 1 time
  DO poissonrb(fvar,frhs,makeghosts) FOR 1 TIMES
  ! allocate variables for the next coarser grid
  ARRAY((fvar.LO-1) DIV 2..(fvar.HI+1) DIV 2, (fvar.LO2-1) DIV 2..(fvar.HI2+1) DIV 2, (fvar.LO3-1) DIV 2..(fvar.HI3+1) DIV 2) OF REAL cvar=0, crhs=0
  ! restriction
  LOOP FOR ixc=cvar.LO+1 TO cvar.HI-1 AND iyc=cvar(0).LO+1 TO cvar(0).HI-1 AND izc=cvar(0,0).LO+1 TO cvar(0,0).HI-1
    ixf=2∗ixc; iyf=2∗iyc; izf=2∗izc
    crhs(ixc,iyc,izc)=2∗[frhs(ixf,iyf,izf)-rsd(fvar,ixf,iyf,izf)]
    cvar(ixc,iyc)=-crhs(ixc,iyc)
  REPEAT LOOP
  makeghosts(cvar)
  !  compat(crhs)
  IF fvar.HI DIV 2 MOD 4 = 0 AND fvar(0).HI DIV 2 MOD 4 = 0 AND fvar(0,0).HI DIV 2 MOD 2 = 0 THEN
    ! recursively call itself on the coarser grid in a W-cycle
    DO poissonmg(cvar,crhs,makeghosts) FOR 2 TIMES
    ! or, on the coarsest grid, call redblack fifty times.
  ELSE DO poissonrb(cvar,crhs,makeghosts) FOR 50 TIMES
  ! prolongation
  LOOP FOR ixc=cvar.LO1+1 TO fvar.HI1 DIV 2 AND iyc=cvar.LO2+1 TO fvar.HI2 DIV 2 AND izc=cvar.LO3+1 TO fvar.HI3 DIV 2
    ixf=2∗ixc; iyf=2∗iyc; izf=2∗izc
    fvar(ixf-1,iyf,izf)=~+(cvar(ixc-1,iyc,izc)+cvar(ixc,iyc,izc))∗0.5
    fvar(ixf,iyf-1,izf)=~+(cvar(ixc,iyc-1,izc)+cvar(ixc,iyc,izc))∗0.5
    fvar(ixf,iyf,izf-1)=~+(cvar(ixc,iyc,izc-1)+cvar(ixc,iyc,izc))∗0.5
    fvar(ixf-1,iyf-1,izf-1)=~+[cvar(ixc-1,iyc-1,izc-1)+cvar(ixc,iyc-1,izc-1)+
      cvar(ixc-1,iyc,izc-1)+cvar(ixc,iyc,izc-1)+cvar(ixc-1,iyc-1,izc)+
      cvar(ixc,iyc-1,izc)+cvar(ixc-1,iyc,izc)+cvar(ixc,iyc,izc)]/8
  REPEAT LOOP
  makeghosts(fvar)
  ! apply smoother 2 more times
  DO poissonrb(fvar,frhs,makeghosts) FOR 2 TIMES
END poissonmg

! calculate maximum residual, for convergence test only
REAL FUNCTION mrsd(POINTER TO ARRAY(∗,∗,∗) OF REAL var, rhs)
  RESULT=0
  LOOP FOR ix=var.LO+1 TO var.HI-1 AND iy=var(0).LO+1 TO var(0).HI-1 AND iz=var(0,0).LO+1 TO var(0,0).HI-1
    RESULT=MAX(RESULT,ABS[rsd(var,ix,iy,iz)-rhs(ix,iy,iz)])
  REPEAT LOOP
END mrsd

! Usage example and test
! Uncomment what follows to test the routine.
! For best performance nx,ny,nz should be (possibly different) powers of 2.
! This constraint can be removed with a slightly more complicated bookkeeping.
!(
ASK INTEGER nx,ny,nz
ARRAY(-1..nx,-1..ny,-1..nz+1) OF REAL p,rhs=0
DO p(ix,iy,iz)=RAND() FOR ALL ix,iy,iz
DO rhs(ix,iy,iz)=RAND() FOR ix=LO+1 TO HI-1 AND iy=LO+1 TO HI-1 AND iz=LO+1 TO HI-1
compat(rhs)
REAL oldr, r=0
LOOP main
  oldr=r
  r=mrsd(p,rhs)
!  compat(rhs)
  DO poissonrb(p,rhs,biperiod) FOR 4 TIMES
  poissonmg(p,rhs,biperiod)
  WRITE r, oldr/r
REPEAT main FOR 12 TIMES
!)