! pde3b.f90  3D non equal h boundary value PDE, iterative solution  
!  u(x,y,z) = x^3  + y^3  +z^3  + 2x^2y + 3xy^2 + 4z^2x + 5y + 6  
!     du/dx = 3x^2 + 0    + 0   + 4xy   + 3y^2  + 4z^2  + 0  + 0  
! d^2u/dx^2 = 6x   + 0    + 0   + 4y    + 0     + 0     + 0  + 0  
!     du/dy = 0    + 3y^2 + 0   + 2x^2  + 6xy   + 0     + 5  + 0  
! d^2u/dy^2 = 0    + 6y   + 0   + 0     + 6x    + 0     + 0  + 0  
!     du/dz = 0    + 0    +3z^2 + 0     + 0     + 8zx   + 0  + 0  
! d^2u/dz^2 = 0    + 0    +6z   + 0     + 0     + 8x    + 0  + 0  
!                                                        
! solve d^2u/dx^2 + d^2u/dy^2 + d^2u/dz^2 = c(x,y,z)     
! by second order method on cube -1,-1,-1 to 1.2,1,1     

! second order numerical difference                      
! d^2u/dx^2 + d^2u/dy^2 + d^2u/dz^2 =                    
! 1/hx^2(u(x-hx,y,z)+u(x+hx,y,z)-2u(x,y,z)) +            
! 1/hy^2(u(x,y-hy,z)+u(x,y+hy,z)-2u(x,y,z))              
! 1/hz^2(u(x,y,z-hz)+u(x,y,z+hz)-2u(x,y,z)) + O(h^2)     
!  = c(x,y,z) = 20x + 10y + 6z                           
! solve for new u(x,y,z) =                               
! ((u(x-hx,y,z)+u(x+hx,y,z))/hx^2 +                      
!  (u(x,y-hy,z)+u(x,y+hy,z))/hy^2 +                      
!  (u(x,y,z-hz)+u(x,y,z+hz))/hz^2 - c(x,y,z)) /          
!  (2/hx^2 + 2/hy^2 + 2/hz^2)                            

module global  
  implicit none
  integer :: nx = 10 ! one less than 'C' when Fortran subscript starts at zero
  integer :: ny = 10
  integer :: nz = 8
  double precision, dimension(0:40,0:40,0:40) :: us ! solution being iterated        
  double precision, dimension(0:40,0:40,0:40) :: ut ! temporary for partial results 
                 ! nx-1 by ny-1 by ny-1 internal, non boundary cells 
  double precision :: xmin = -1.0
  double precision :: ymin = -1.0
  double precision :: zmin = -1.0
  double precision :: xmax =  1.2
  double precision :: ymax =  1.0
  double precision :: zmax =  1.0
  double precision :: hx   ! x step = (xmax-xmin)/nx 
  double precision :: hy   ! y step = (ymax-ymin)/ny 
  double precision :: hz   ! z step = (zmax-zmin)/nz
  double precision :: diff=1.0d100  ! sum abs difference form last iteration 
  double precision :: error         ! sum of absolute exact-computed         
  double precision :: avgerror      ! average error                          
  double precision :: maxerror      ! maximum cell error                     
  interface
    function u(x, y, z) result(value) ! solution for boundary and check 
      implicit none
      double precision, intent(in) :: x, y, z
      double precision :: value
    end function u
  end interface
  interface
    function c(x, y, z) result(term) ! from solve
      implicit none
      double precision, intent(in) :: x, y, z
      double precision :: term
    end function c
  end interface
end module global

program pde3b
  use global
  implicit none
  integer :: iter=0
  
  print *, "pde3b.f90 running"
  hx = (xmax-xmin)/nx ! step size in x 
  hy = (ymax-ymin)/ny ! step size in y 
  hz = (zmax-zmin)/nz ! step size in z 
  print *, "xmin=",xmin, "xmax=",xmax, "nx=",nx, "hx=",hx
  print *, "ymin=",ymin, "ymax=",ymax, "ny=",ny, "hy=",hy
  print *, "zmin=",zmin, "zmax=",zmax, "nz=",nz, "hz=",hz
  
  call init_boundary()
  print *, "initial conditions set"
  
  do while(diff>0.0001 .and. iter<149)
    call iterate()
    call check()
    avgerror = error/((nx-1)*(ny-1)*(nz-1))
    print "('iter=',i4,' diff=',f8.4,' error=',f9.4,' avg=',f8.4,' max=',f8.4)", &
          iter, diff, error, avgerror, maxerror
    iter = iter+1
    if(mod(iter,50).eq.0) then
      call printus()
      call printdiff()
    end if
  end do
  call printexact()
  call printus()
  call printdiff()
end program pde3b 

function u(x, y, z) result(value) ! solution for boundary and check 
  implicit none
  double precision, intent(in) :: x, y, z
  double precision :: value
  value = x*x*x + y*y*y + z*z*z + 2.0*x*x*y + 3.0*x*y*y + &
          4.0*z*z*x + 5.0*y +6.0
end function u

function c(x, y, z) result(term) ! from solve
  implicit none
  double precision, intent(in) :: x, y, z
  double precision :: term
  term = 20.0*x + 10.0*y + 6.0*z
end function c

function u000(i, j, k) result(value) ! new iterated value 
  use global
  implicit none
  integer, intent(in) :: i, j, k
  double precision :: value
  value = ((us(i-1,j,k) + us(i+1,j,k))/(hx*hx) + &
           (us(i,j-1,k) + us(i,j+1,k))/(hy*hy) + &
           (us(i,j,k-1) + us(i,j,k+1))/(hz*hz) - &
           c(xmin+i*hx,ymin+j*hy,zmin+k*hz))   / &
           (2.0/(hx*hx) + 2.0/(hy*hy) + 2.0/(hz*hz))
end function u000

subroutine init_boundary()
  use global
  implicit none
  integer i, j, k
  
  ! set boundary on six sides 
  do i=0,nx
    do j=0,ny
      us(i,j,0)  = u(xmin+i*hx, ymin+j*hy, zmin)
      us(i,j,nz) = u(xmin+i*hx, ymin+j*hy, zmax)
    end do
  end do
  do j=0,ny
    do k=0,nz
      us(0,j,k)  = u(xmin, ymin+j*hy, zmin+k*hz)
      us(nx,j,k) = u(xmax, ymin+j*hy, zmin+k*hz)
    end do
  end do
  do i=0,nx
    do k=0,nz
      us(i,0,k)  = u(xmin+i*hx, ymin, zmin+k*hz)
      us(i,ny,k) = u(xmin+i*hx, ymax, zmin+k*hz)
    end do
  end do
  ! zero internal cells 
  do i=1,nx-1
    do j=1,ny-1
      do k=1,nz-1
        us(i,j,k) = 0.0
      end do
    end do
  end do
end subroutine init_boundary 

subroutine iterate()
  use global
  implicit none
  integer :: i, j, k
  
  interface
    function u000(i, j, k) result(value) ! new iterated value 
      implicit none
      integer, intent(in) :: i, j, k
      double precision :: value
    end function u000
  end interface

  diff = 0.0
  do i=1,nx-1
    do j=1,ny-1
      do k=1,nz-1
        ut(i,j,k) = u000(i,j,k) 
        diff = diff + abs(ut(i,j,k)-us(i,j,k))
      end do
    end do
  end do
  ! copy temp to solution 
  do i=1,nx-1
    do j=1,ny-1
      do k=1,nz-1
        us(i,j,k) = ut(i,j,k)
      end do
    end do
  end do
end subroutine iterate 

subroutine check() ! typically not known, instrumentation 
  use global
  implicit none
  integer :: i, j, k
  double precision :: zexact, err
  
  error = 0.0
  maxerror = 0.0
  do i=1,nx-1
    do j=1,ny-1
      do k=1,nz-1
        zexact = u(xmin+i*hx,ymin+j*hy,zmin+k*hz)
        err = abs(us(i,j,k)-zexact)
        error = error + err
        if(err>maxerror) maxerror=err
      end do
    end do
  end do
end subroutine check 

subroutine printus()
  use global
  implicit none
  integer :: i, j, k
  
  print "(/a)", "computed solution"
  do k=0,nz
    print *, "u(x,y,",k,") plane"
    do i=0,nx
      print "(11f7.3)", (us(i,j,k), j=0,min(ny,10))
    end do
  end do
end subroutine printus 

subroutine printexact()
  use global
  implicit none
  integer :: i, j, k
  
  print "(/a)", "exact solution"
  do k=0,nz
    print *, "u(x,y,",k,") plane"
    do i=0,nx
      print "(11f7.3)", (u(xmin+i*hx,ymin+j*hy,zmin+k*hz),j=0,min(ny,10))
    end do
  end do
end subroutine printexact 

subroutine printdiff()
  use global
  implicit none
  integer :: i, j, k
  
  print "(/a)", "difference exact-computed solution" 
  do k=0,nz
    print *, "u(x,y,",k,") plane"
    do i=1,nx-1
      print "(7x, 10f7.4)", (u(xmin+i*hx,ymin+j*hy,zmin+k*hz)- &
                             us(i,j,k), j=1,min(ny-1,10))
    end do
  end do
end subroutine printdiff