!===============================================!
! Partial derivative schemes for each direction !
!===============================================!
! SUBROUTINE init_derivs() !
! FUNCTION deriv_dx(f) !
! FUNCTION deriv_dy(f) !
! FUNCTION deriv_dz(f) !
!===============================================!
MODULE fff_derivs
  IMPLICIT NONE

  CONTAINS

  !--------------------------------------------------------------!
  ! SUBROUTINE init_derivs !
  !--------------------------------------------------------------!
  ! PURPOSE: !
  ! Calculates coefficients for derivative schemes, based !
  ! on {x,y,z} grids !
  ! NOTE: This routine must be called before using derivatives !
  ! INPUT: !
  ! none !
  ! OUTPUT: !
  ! none; indirect (xder_coeff,yder_coeff,zder_coeff) !
  ! REQUIRES: !
  ! x,y,z,nx,ny,nz,xder_coeff,yder_coeff,zder_coeff: from vars !
  ! HISTORY: !
  ! 17-May-2005: jmm, jimm@ssl.berkeley.edu, from deriv_dx.pro !
  ! Feb-2008: djb, modularized, based on earlier versions !
  !--------------------------------------------------------------!
  SUBROUTINE init_derivs
    USE params, ONLY: r8
    USE vars, ONLY: x, y, z, nx, ny, nz, &
                    xder_coeff, yder_coeff, zder_coeff
    IMPLICIT NONE

    ! Local variable declarations
    REAL(KIND=r8), DIMENSION(:), ALLOCATABLE :: hm, hp
    !------------------------------------------------------------------!
    ! Locally 3rd-order non-uniform finite derivative schemes have a !
    ! three-point stencil of the form !
    ! df(i) = A(i)*f(i-1) + B(i)*f(i) + C(i)*f(i+1) !
    ! !
    ! The coefficients can be determined using Lagrange interpolants !
    ! or Taylor expansions and are given by !
    ! A(i) = -hp(i-1)/(hm(i-1)*(hm(i-1) + hp(i-1))) !
    ! B(i) = -(hm(i-1) - hp(i-1))/(hm(i-1)*hp(i-1)) !
    ! C(i) = hm(i-1)/(hp(i-1)*(hm(i-1) + hp(i-1))) !
    ! where !
    ! hm(i-1) = x(i) - x(i-1) !
    ! hp(i-1) = x(i+1) - x(i) !
    ! !
    ! Boundary points are treated using a one-sided, three-point !
    ! stencil !
    ! !
    ! Lower boundary: !
    ! df(1) = A(1)*f(1) + B(1)*f(2) + C(1)*f(3) !
    ! A(1) = - (2*hm(1) + hp(1))/(hm(1)*(hm(1) + hp(1))) !
    ! B(1) = (hm(1) + hp(1))/(hm(1)*hp(1)) !
    ! C(1) = - hm(1)/(hp(1)*(hm(1) + hp(1))) !
    ! !
    ! Upper boundary: !
    ! df(nx) = A(nx)*f(nx-2) + B(nx)*f(nx-1) + C(nx)*f(nx) !
    ! A(nx) = hp(nx-2)/(hm(nx-2)*(hm(nx-2) + hp(nx-2))) !
    ! B(nx) = - (hm(nx-2) + hp(nx-2))/(hm(nx-2)*hp(nx-2)) !
    ! C(nx) = (2*hp(nx-2) + hm(nx-2))/(hp(nx-2)*(hm(nx-2) + hp(nx-2))) !
    !------------------------------------------------------------------!
    ! x-dimension
    ALLOCATE(hm(nx-2), hp(nx-2))
    hm = x(2:nx-1) - x(1:nx-2)
    hp = x(3:nx) - x(2:nx-1)
    xder_coeff(2:nx-1,1) = - hp/(hm*(hm + hp))
    xder_coeff(2:nx-1,2) = (hp - hm)/(hm*hp)
    xder_coeff(2:nx-1,3) = hm/(hp*(hm + hp))
    xder_coeff(1,1) = - (2.0_r8*hm(1) + hp(1))/(hm(1)*(hm(1) + hp(1)))
    xder_coeff(1,2) = (hm(1) + hp(1))/(hm(1)*hp(1))
    xder_coeff(1,3) = - hm(1)/(hp(1)*(hm(1) + hp(1)))
    xder_coeff(nx,1) = hp(nx-2)/(hm(nx-2)*(hm(nx-2) + hp(nx-2)))
    xder_coeff(nx,2) = - (hm(nx-2) + hp(nx-2))/(hm(nx-2)*hp(nx-2))
    xder_coeff(nx,3) = (2.0_r8*hp(nx-2) + hm(nx-2))/ &
                       (hp(nx-2)*(hm(nx-2) + hp(nx-2)))
    DEALLOCATE(hm, hp)
    ! y-dimension
    ALLOCATE(hm(ny-2), hp(ny-2))
    hm = y(2:ny-1) - y(1:ny-2)
    hp = y(3:ny) - y(2:ny-1)
    yder_coeff(2:ny-1,1) = - hp/(hm*(hm + hp))
    yder_coeff(2:ny-1,2) = (hp - hm)/(hm*hp)
    yder_coeff(2:ny-1,3) = hm/(hp*(hm + hp))
    yder_coeff(1,1) = - (2.0_r8*hm(1) + hp(1))/(hm(1)*(hm(1) + hp(1)))
    yder_coeff(1,2) = (hm(1) + hp(1))/(hm(1)*hp(1))
    yder_coeff(1,3) = - hm(1)/(hp(1)*(hm(1) + hp(1)))
    yder_coeff(ny,1) = hp(ny-2)/(hm(ny-2)*(hm(ny-2) + hp(ny-2)))
    yder_coeff(ny,2) = - (hm(ny-2) + hp(ny-2))/(hm(ny-2)*hp(ny-2))
    yder_coeff(ny,3) = (2.0_r8*hp(ny-2) + hm(ny-2))/ &
                       (hp(ny-2)*(hm(ny-2) + hp(ny-2)))
    DEALLOCATE(hm, hp)
    ! z-dimension
    ALLOCATE(hm(nz-2), hp(nz-2))
    hm = z(2:nz-1) - z(1:nz-2)
    hp = z(3:nz) - z(2:nz-1)
    zder_coeff(2:nz-1,1) = - hp/(hm*(hm + hp))
    zder_coeff(2:nz-1,2) = (hp - hm)/(hm*hp)
    zder_coeff(2:nz-1,3) = hm/(hp*(hm + hp))
    zder_coeff(1,1) = - (2.0_r8*hm(1) + hp(1))/(hm(1)*(hm(1) + hp(1)))
    zder_coeff(1,2) = (hm(1) + hp(1))/(hm(1)*hp(1))
    zder_coeff(1,3) = - hm(1)/(hp(1)*(hm(1) + hp(1)))
    zder_coeff(nz,1) = hp(nz-2)/(hm(nz-2)*(hm(nz-2) + hp(nz-2)))
    zder_coeff(nz,2) = - (hm(nz-2) + hp(nz-2))/(hm(nz-2)*hp(nz-2))
    zder_coeff(nz,3) = (2.0_r8*hp(nz-2) + hm(nz-2))/ &
                       (hp(nz-2)*(hm(nz-2) + hp(nz-2)))
    DEALLOCATE(hm, hp)
  END SUBROUTINE init_derivs

  !-----------------------------------------------------------------!
  ! FUNCTION deriv_dx !
  !-----------------------------------------------------------------!
  ! PURPOSE: !
  ! Calculates partial derivative of a 3-d array w.r.t. first (x) !
  ! index !
  ! INPUT: !
  ! f: function to be differentiated; dimensions (nx,ny,nz) !
  ! OUTPUT: !
  ! deriv_dx: partial derivative, df/dx; dimensions (nx,ny,nz) !
  ! REQUIRES: !
  ! nx,ny,nz,xder_coeff: from vars !
  ! HISTORY: !
  ! 17-May-2005: jmm, jimm@ssl.berkeley.edu, from deriv_dx.pro !
  ! Feb-2008: djb, modularized !
  !-----------------------------------------------------------------!
  FUNCTION deriv_dx(f)
    USE params, ONLY: i4, r8
    USE vars, ONLY: nx, ny, nz, xder_coeff
    IMPLICIT NONE

    ! Dummy variable declarations
    REAL(KIND=r8), DIMENSION(:,:,:), INTENT(IN) :: f
    REAL(KIND=r8), DIMENSION(SIZE(f,1),SIZE(f,2),SIZE(f,3)) :: deriv_dx
    ! Local variable declarations
    INTEGER(KIND=i4) :: i, j, k

    !---------------------------------------------------------!
    ! df/dx(i,j,k) = A*f(i-1,j,k) + B*f(i,j,k) + C*f(i+1,j,k) !
    ! !
    ! A = -hp/(hm*(hm+hp)) !
    ! B = (hp-hm)/(hm*hp) !
    ! C = hm/(hp*(hm+hp)) !
    ! hm = x(i) - x(i-1) !
    ! hp = x(i+1) - x(i) !
    !---------------------------------------------------------!
    ! Compute derivative on interior points
    FORALL(i=2:nx-1, j=1:ny, k=1:nz)
      deriv_dx(i ,j, k) = xder_coeff(i, 1)*f(i-1, j, k) &
                          + xder_coeff(i, 2)*f(i , j, k) &
                          + xder_coeff(i, 3)*f(i+1, j, k)
    END FORALL
    ! Set derivative on boundary points
    deriv_dx(1, :, :) = xder_coeff(1, 1)*f(1, :, :) &
                        + xder_coeff(1, 2)*f(2, :, :) &
                        + xder_coeff(1, 3)*f(3, :, :)
    deriv_dx(nx, :, :) = xder_coeff(nx, 1)*f(nx-2, :, :) &
                         + xder_coeff(nx, 2)*f(nx-1, :, :) &
                         + xder_coeff(nx, 3)*f(nx, :, :)
  END FUNCTION deriv_dx

  !------------------------------------------------------------------!
  ! FUNCTION deriv_dy !
  !------------------------------------------------------------------!
  ! PURPOSE: !
  ! Calculates partial derivative of a 3-d array w.r.t. second (y) !
  ! index !
  ! INPUT: !
  ! f: function to be differentiated; dimensions (nx,ny,nz) !
  ! OUTPUT: !
  ! deriv_dy: partial derivative, df/dy; dimensions (nx,ny,nz) !
  ! REQUIRES: !
  ! nx,ny,nz,yder_coeff: from vars !
  ! HISTORY: !
  ! 17-May-2005: jmm, jimm@ssl.berkeley.edu, from deriv_dx.pro !
  ! Feb-2008: djb, modularized !
  !------------------------------------------------------------------!
  FUNCTION deriv_dy(f)
    USE params, ONLY: i4, r8
    USE vars, ONLY: nx, ny, nz, yder_coeff
    IMPLICIT NONE

    ! Dummy variable declarations
    REAL(KIND=r8), DIMENSION(:,:,:), INTENT(IN) :: f
    REAL(KIND=r8), DIMENSION(SIZE(f,1),SIZE(f,2),SIZE(f,3)) :: deriv_dy
    ! Local variable declarations
    INTEGER(KIND=i4) :: i, j, k

    !---------------------------------------------------------!
    ! df/dy(i,j,k) = A*f(i,j-1,k) + B*f(i,j,k) + C*f(i,j+1,k) !
    ! !
    ! A = -hp/(hm*(hm+hp)) !
    ! B = (hp-hm)/(hm*hp) !
    ! C = hm/(hp*(hm+hp)) !
    ! hm = y(i) - y(i-1) !
    ! hp = y(i+1) - y(i) !
    !---------------------------------------------------------!
    ! Compute derivative on interior points
    FORALL(i=1:nx, j=2:ny-1, k=1:nz)
      deriv_dy(i, j, k) = yder_coeff(j, 1)*f(i, j-1, k) &
                          + yder_coeff(j, 2)*f(i, j , k) &
                          + yder_coeff(j, 3)*f(i, j+1, k)
    END FORALL
    ! Set derivative on boundary points
    deriv_dy(:, 1, :) = yder_coeff(1, 1)*f(:, 1, :) &
                        + yder_coeff(1, 2)*f(:, 2, :) &
                        + yder_coeff(1, 3)*f(:, 3, :)
    deriv_dy(:, ny, :) = yder_coeff(ny, 1)*f(:, ny-2, :) &
                         + yder_coeff(ny, 2)*f(:, ny-1, :) &
                         + yder_coeff(ny, 3)*f(:, ny , :)
  END FUNCTION deriv_dy

  !-----------------------------------------------------------------!
  ! FUNCTION deriv_dz !
  !-----------------------------------------------------------------!
  ! PURPOSE: !
  ! Calculates partial derivative of a 3-d array w.r.t. third (z) !
  ! index !
  ! INPUT: !
  ! f: function to be differentiated; dimensions (nx,ny,nz) !
  ! OUTPUT: !
  ! deriv_dz: partial derivative, df/dz; dimensions (nx,ny,nz) !
  ! REQUIRES: !
  ! nx,ny,nz,zder_coeff: from vars !
  ! HISTORY: !
  ! 17-May-2005: jmm, jimm@ssl.berkeley.edu, from deriv_dx.pro !
  ! Feb-2008: djb, modularized !
  !-----------------------------------------------------------------!
  FUNCTION deriv_dz(f)
    USE params, ONLY: i4, r8
    USE vars, ONLY: nx, ny, nz, zder_coeff
    IMPLICIT NONE

    ! Dummy variable declarations
    REAL(KIND=r8), DIMENSION(:,:,:), INTENT(IN) :: f
    REAL(KIND=r8), DIMENSION(SIZE(f,1),SIZE(f,2),SIZE(f,3)) :: deriv_dz
    ! Local variable declarations
    INTEGER(KIND=i4) :: i, j, k

    !---------------------------------------------------------!
    ! df/dz(i,j,k) = A*f(i,j,k-1) + B*f(i,j,k) + C*f(i,j,k+1) !
    ! !
    ! A = -hp/(hm*(hm+hp)) !
    ! B = (hp-hm)/(hm*hp) !
    ! C = hm/(hp*(hm+hp)) !
    ! hm = z(i) - z(i-1) !
    ! hp = z(i+1) - z(i) !
    !---------------------------------------------------------!
    ! Compute derivative on interior points
    FORALL(i=1:nx, j=1:ny, k=2:nz-1)
      deriv_dz(i, j, k) = zder_coeff(k, 1)*f(i, j, k-1) &
                          + zder_coeff(k, 2)*f(i, j, k) &
                          + zder_coeff(k, 3)*f(i, j, k+1)
    END FORALL
    ! Set derivative on boundary points
    deriv_dz(:, :, 1) = zder_coeff(1, 1)*f(:, :, 1) &
                        + zder_coeff(1, 2)*f(:, :, 2) &
                        + zder_coeff(1, 3)*f(:, :, 3)
    deriv_dz(:, :, nz) = zder_coeff(nz, 1)*f(:, :, nz-2) &
                         + zder_coeff(nz, 2)*f(:, :, nz-1) &
                         + zder_coeff(nz, 3)*f(:, :, nz)
  END FUNCTION deriv_dz

END MODULE fff_derivs
