C $Header: /home/ubuntu/mnt/e9_copy/MITgcm/model/src/solve_pentadiagonal.F,v 1.4 2010/08/10 17:58:30 gforget Exp $
C $Name:  $

#include "CPP_OPTIONS.h"

C o Switch to code that has the k-loop inside the 
C   ij-loops, which matters in adjoint mode.
#ifdef ALLOW_AUTODIFF 
#define ALLOW_SOLVERS_KLOOPINSIDE
#endif

CBOP
C     !ROUTINE: SOLVE_PENTADIAGONAL
C     !INTERFACE:
      SUBROUTINE SOLVE_PENTADIAGONAL( 
     I                     iMin,iMax, jMin,jMax,
     U                     a5d, b5d, c5d, d5d, e5d,
     U                     y5d,
     O                     errCode,
     I                     bi, bj, myThid )
C     !DESCRIPTION: \bv
C     *==========================================================*
C     | S/R SOLVE_PENTADIAGONAL                                              
C     | o Solve a penta-diagonal system A*X=Y (dimension Nr)
C     *==========================================================*
C     | o Used to solve implicitly vertical advection & diffusion
C     *==========================================================*
C     \ev

C     !USES:
      IMPLICIT NONE
C     == Global data ==
#include "SIZE.h"
#include "EEPARAMS.h"

C     !INPUT/OUTPUT PARAMETERS:
C     == Routine Arguments ==
C     INPUT:
C     iMin,iMax,jMin,jMax  :: computational domain
C     a5d    :: 2nd  lower diagonal of the pentadiagonal matrix
C     b5d    :: 1rst lower diagonal of the pentadiagonal matrix
C     c5d    :: main diagonal       of the pentadiagonal matrix
C     d5d    :: 1rst upper diagonal of the pentadiagonal matrix
C     e5d    :: 2nd  upper diagonal of the pentadiagonal matrix
C     y5d    :: Y vector (R.H.S.); 
C     bi,bj  :: tile indices
C     myThid :: thread number
C     OUTPUT:
C     y5d    :: X = solution of A*X=Y
C     a5d,b5d,c5d,d5d,e5d :: modified to enable to find Xp solution of
C                        A*Xp=Yp without solving the full system again
C     errCode :: > 0 if singular matrix
      INTEGER iMin,iMax,jMin,jMax
      _RL a5d(1-Olx:sNx+Olx,1-Oly:sNy+Oly,Nr)
      _RL b5d(1-Olx:sNx+Olx,1-Oly:sNy+Oly,Nr)
      _RL c5d(1-Olx:sNx+Olx,1-Oly:sNy+Oly,Nr)
      _RL d5d(1-Olx:sNx+Olx,1-Oly:sNy+Oly,Nr)
      _RL e5d(1-Olx:sNx+Olx,1-Oly:sNy+Oly,Nr)
      _RL y5d(1-Olx:sNx+Olx,1-Oly:sNy+Oly,Nr,nSx,nSy)
      INTEGER errCode
      INTEGER bi, bj, myThid

C     !LOCAL VARIABLES:
C     == Local variables ==
      INTEGER i,j,k
#ifdef ALLOW_SOLVERS_KLOOPINSIDE
      _RL y5d_m1(1-Olx:sNx+Olx,1-Oly:sNy+Oly,Nr)
      _RL a5d_prime(Nr), b5d_prime(Nr)
      _RL c5d_prime(Nr), d5d_prime(Nr), e5d_prime(Nr)
      _RL y5d_prime(Nr), y5d_update(Nr), tmpval
#endif
CEOP

#ifndef ALLOW_SOLVERS_KLOOPINSIDE

      errCode = 0

      DO k=1,Nr
C--   forward sweep (starting from top)
       IF (k.EQ.1) THEN
        DO j=jMin,jMax
         DO i=iMin,iMax
          IF ( c5d(i,j,k).NE.0. _d 0 ) THEN 
           c5d(i,j,k) = 1. _d 0 / c5d(i,j,k)
          ELSE
           c5d(i,j,k) = 0. _d 0
           errCode = 1
          ENDIF
         ENDDO
        ENDDO

       ELSEIF (k.EQ.2) THEN
        DO j=jMin,jMax
         DO i=iMin,iMax
C--       [k] <- [k] - b_k/c_k-1 * [k-1]
          b5d(i,j,k) = b5d(i,j,k)*c5d(i,j,k-1)
          c5d(i,j,k) = c5d(i,j,k) - b5d(i,j,k)*d5d(i,j,k-1)
          d5d(i,j,k) = d5d(i,j,k) - b5d(i,j,k)*e5d(i,j,k-1)
          y5d(i,j,k,bi,bj) = y5d(i,j,k,bi,bj)
     &                     - b5d(i,j,k)*y5d(i,j,k-1,bi,bj)
          IF ( c5d(i,j,k).NE.0. _d 0 ) THEN 
           c5d(i,j,k) = 1. _d 0 / c5d(i,j,k)
          ELSE
           c5d(i,j,k) = 0. _d 0
           errCode = 1
          ENDIF
         ENDDO
        ENDDO

       ELSE
C--   Middle of forward sweep
        DO j=jMin,jMax
         DO i=iMin,iMax
C--       [k] <- [k] - a_k/c_k-2 * [k-2]
          a5d(i,j,k) = a5d(i,j,k)*c5d(i,j,k-2) 
          b5d(i,j,k) = b5d(i,j,k) - a5d(i,j,k)*d5d(i,j,k-2)
          c5d(i,j,k) = c5d(i,j,k) - a5d(i,j,k)*e5d(i,j,k-2)
C--       [k] <- [k] - b_k/c_k-1 * [k-1]
          b5d(i,j,k) = b5d(i,j,k)*c5d(i,j,k-1)
          c5d(i,j,k) = c5d(i,j,k) - b5d(i,j,k)*d5d(i,j,k-1)
          d5d(i,j,k) = d5d(i,j,k) - b5d(i,j,k)*e5d(i,j,k-1)
          y5d(i,j,k,bi,bj) = y5d(i,j,k,bi,bj)
     &                     - b5d(i,j,k)*y5d(i,j,k-1,bi,bj)
     &                     - a5d(i,j,k)*y5d(i,j,k-2,bi,bj)
          IF ( c5d(i,j,k).NE.0. _d 0 ) THEN 
           c5d(i,j,k) = 1. _d 0 / c5d(i,j,k)
          ELSE
           c5d(i,j,k) = 0. _d 0
           errCode = 1
          ENDIF
         ENDDO
        ENDDO
C-      end if k= .. ; end of k loop
       ENDIF
      ENDDO

C--   Backward sweep (starting from bottom)
      DO k=Nr,1,-1
       IF (k.EQ.Nr) THEN
        DO j=jMin,jMax
         DO i=iMin,iMax
          y5d(i,j,k,bi,bj) =   y5d(i,j,k,bi,bj)*c5d(i,j,k)
         ENDDO
        ENDDO
       ELSEIF (k.EQ.Nr-1) THEN
        DO j=jMin,jMax
         DO i=iMin,iMax
          y5d(i,j,k,bi,bj) = ( y5d(i,j,k,bi,bj)
     &                       - d5d(i,j,k)*y5d(i,j,k+1,bi,bj)
     &                       )*c5d(i,j,k)
         ENDDO
        ENDDO
       ELSE
        DO j=jMin,jMax
         DO i=iMin,iMax
          y5d(i,j,k,bi,bj) = ( y5d(i,j,k,bi,bj)
     &                       - d5d(i,j,k)*y5d(i,j,k+1,bi,bj)
     &                       - e5d(i,j,k)*y5d(i,j,k+2,bi,bj)
     &                       )*c5d(i,j,k)
         ENDDO
        ENDDO
C-      end if k= .. ; end of k loop
       ENDIF
      ENDDO

#else  /* ALLOW_SOLVERS_KLOOPINSIDE */

      errCode = 0

C--   Temporary array
      DO j=jMin,jMax
      DO i=iMin,iMax
      DO k=1,Nr
         y5d_m1(i,j,k) = y5d(i,j,k,bi,bj)
      ENDDO
      ENDDO
      ENDDO

C--   Main loop
      DO j=jMin,jMax
      DO i=iMin,iMax

      DO k=1,Nr
        a5d_prime(k) = 0. _d 0
        b5d_prime(k) = 0. _d 0
        c5d_prime(k) = 0. _d 0
        d5d_prime(k) = 0. _d 0
        e5d_prime(k) = 0. _d 0
        y5d_prime(k) = 0. _d 0
        y5d_update(k) = 0. _d 0
      ENDDO

      DO k=1,Nr
C--   forward sweep (starting from top)

          IF (k.EQ.1) THEN
c just copy terms
           a5d_prime(k) = 0. _d 0
           b5d_prime(k) = 0. _d 0
           c5d_prime(k) = c5d(i,j,k)
           d5d_prime(k) = d5d(i,j,k)
           e5d_prime(k) = e5d(i,j,k)
           y5d_prime(k) = y5d_m1(i,j,k)
          ELSEIF (k.EQ.2) THEN
c subtract one term
           a5d_prime(k) = 0. _d 0
           b5d_prime(k) = 0. _d 0
           c5d_prime(k) = c5d(i,j,k)
     &      -b5d(i,j,k)*d5d_prime(k-1)
           d5d_prime(k) = d5d(i,j,k)
     &      -b5d(i,j,k)*e5d_prime(k-1)
           e5d_prime(k) = e5d(i,j,k)
           y5d_prime(k) = y5d_m1(i,j,k)
     &      -b5d(i,j,k)*y5d_prime(k-1)
          ELSE
c subtract two terms
           a5d_prime(k) = 0. _d 0
           b5d_prime(k) = 0. _d 0
           c5d_prime(k) = c5d(i,j,k)
     &      -a5d(i,j,k)*e5d_prime(k-2)
     &      -(b5d(i,j,k)-a5d(i,j,k)*d5d_prime(k-2))*d5d_prime(k-1)
           d5d_prime(k) = d5d(i,j,k)
     &      -(b5d(i,j,k)-a5d(i,j,k)*d5d_prime(k-2))*e5d_prime(k-1)
           e5d_prime(k) = e5d(i,j,k)
           y5d_prime(k) = y5d_m1(i,j,k)
     &      -a5d(i,j,k)*y5d_prime(k-2)
     &      -(b5d(i,j,k)-a5d(i,j,k)*d5d_prime(k-2))*y5d_prime(k-1)
          ENDIF

c normalization
          tmpval=c5d_prime(k)
          IF ( tmpval.NE.0. _d 0 ) THEN
           a5d_prime(k) = a5d_prime(k) / tmpval
           b5d_prime(k) = b5d_prime(k) / tmpval
           c5d_prime(k) = 1. _d 0
           d5d_prime(k) = d5d_prime(k) / tmpval
           e5d_prime(k) = e5d_prime(k) / tmpval
           y5d_prime(k) = y5d_prime(k) / tmpval
          ELSE
           a5d_prime(k) = 0. _d 0
           b5d_prime(k) = 0. _d 0
           c5d_prime(k) = 0. _d 0
           d5d_prime(k) = 0. _d 0
           e5d_prime(k) = 0. _d 0
           y5d_prime(k) = 0. _d 0
           errCode = 1
          ENDIF

      ENDDO

C--   Backward sweep (starting from bottom)
      DO k=Nr,1,-1
       IF (k.EQ.Nr) THEN
          y5d_update(k) =   y5d_prime(k)
       ELSEIF (k.EQ.Nr-1) THEN
          y5d_update(k) =   y5d_prime(k)
     &     - y5d_update(k+1)*d5d_prime(k)
       ELSE
          y5d_update(k) =   y5d_prime(k)
     &     - y5d_update(k+1)*d5d_prime(k)
     &     - y5d_update(k+2)*e5d_prime(k)
       ENDIF
      ENDDO

C--   Update array
      DO k=1,Nr
         y5d(i,j,k,bi,bj)=y5d_update(k)
      ENDDO

      ENDDO
      ENDDO

#endif  /* ALLOW_SOLVERS_KLOOPINSIDE */

      RETURN
      END