C $Header: /u/gcmpack/MITgcm/model/src/solve_tridiagonal.F,v 1.13 2016/10/26 00:50:25 jmc Exp $
C $Name:  $

#include "CPP_OPTIONS.h"

CBOP
C     !ROUTINE: SOLVE_TRIDIAGONAL
C     !INTERFACE:
      SUBROUTINE SOLVE_TRIDIAGONAL(
     I                     iMin,iMax, jMin,jMax,
     I                     a3d, b3d, c3d,
     U                     y3d,
     O                     errCode,
     I                     bi, bj, myThid )
C     !DESCRIPTION: \bv
C     *==========================================================*
C     | S/R SOLVE_TRIDIAGONAL
C     | o Solve a tri-diagonal system A*X=Y (dimension Nr)
C     *==========================================================*
C     | o Used to solve implicitly vertical advection & diffusion
#ifdef SOLVE_DIAGONAL_LOWMEMORY
C     | o Allows 2 types of call:
C     | 1) with INPUT errCode < 0 (first call):
C     |   Solve system A*X=Y by LU decomposition and return
C     |     inverse matrix as output (in a3d,b3d,c3d)
C     | 2) with INPUT errCode = 0 (subsequent calls):
C     |   Solve system A*Xp=Yp by using inverse matrix coeff
C     |     from first call output.
#endif /* SOLVE_DIAGONAL_LOWMEMORY */
C     *==========================================================*
C     \ev

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

C     !INPUT/OUTPUT PARAMETERS:
C     -- INPUT: --
C     iMin,iMax,jMin,jMax  :: computational domain
C     a3d     :: matrix lower diagnonal
C     b3d     :: matrix main  diagnonal
C     c3d     :: matrix upper diagnonal
C     y3d     :: Y vector (R.H.S.);
C     bi,bj   :: tile indices
C     myThid  :: thread number
C     -- OUTPUT: --
C     y3d     :: X = solution of A*X=Y
C     errCode :: > 0 if singular matrix
#ifdef SOLVE_DIAGONAL_LOWMEMORY
C     a3d,b3d,c3d :: inverse matrix coeff to enable to find Xp solution
C                    of A*Xp=Yp with subsequent call to this routine
#endif /* SOLVE_DIAGONAL_LOWMEMORY */
      INTEGER iMin,iMax,jMin,jMax
      _RL a3d(1-OLx:sNx+OLx,1-OLy:sNy+OLy,Nr)
      _RL b3d(1-OLx:sNx+OLx,1-OLy:sNy+OLy,Nr)
      _RL c3d(1-OLx:sNx+OLx,1-OLy:sNy+OLy,Nr)
      _RL y3d(1-OLx:sNx+OLx,1-OLy:sNy+OLy,Nr)
      INTEGER errCode
      INTEGER bi, bj, myThid

C     !LOCAL VARIABLES:
C     == Local variables ==
      INTEGER i,j,k
      _RL tmpVar
#ifndef SOLVE_DIAGONAL_LOWMEMORY
      _RL recVar
      _RL y3d_m1(1-OLx:sNx+OLx,1-OLy:sNy+OLy,Nr)
# ifdef SOLVE_DIAGONAL_KINNER
      _RL c3d_m1(1-OLx:sNx+OLx,1-OLy:sNy+OLy,Nr)
      _RL c3d_prime(Nr)
      _RL y3d_prime(Nr)
      _RL y3d_update(Nr)
# else
      _RL c3d_prime(1-OLx:sNx+OLx,1-OLy:sNy+OLy,Nr)
      _RL y3d_prime(1-OLx:sNx+OLx,1-OLy:sNy+OLy,Nr)
# endif
#endif /* SOLVE_DIAGONAL_LOWMEMORY */
CEOP

#ifdef SOLVE_DIAGONAL_LOWMEMORY

      IF ( errCode .LT. 0 ) THEN
       errCode = 0

C--   forward sweep (starting from top) part-1: matrix L.U decomposition
       DO j=jMin,jMax
        DO i=iMin,iMax
          IF ( b3d(i,j,1).NE.0. _d 0 ) THEN
            b3d(i,j,1) = 1. _d 0 / b3d(i,j,1)
          ELSE
            b3d(i,j,1) = 0. _d 0
            errCode = 1
          ENDIF
        ENDDO
       ENDDO

C-    Middle of forward sweep
       DO k=2,Nr
        DO j=jMin,jMax
         DO i=iMin,iMax
          tmpVar = b3d(i,j,k) - a3d(i,j,k)*c3d(i,j,k-1)*b3d(i,j,k-1)
          IF ( tmpVar .NE. 0. _d 0 ) THEN
            b3d(i,j,k) = 1. _d 0 / tmpVar
          ELSE
            b3d(i,j,k) = 0. _d 0
            errCode = 1
          ENDIF
         ENDDO
        ENDDO
       ENDDO

C--   end if errCode < 0
      ENDIF

C--   forward sweep (starting from top) part-2: process RHS
      DO j=jMin,jMax
       DO i=iMin,iMax
         y3d(i,j,1) = y3d(i,j,1)*b3d(i,j,1)
       ENDDO
      ENDDO
      DO k=2,Nr
       DO j=jMin,jMax
        DO i=iMin,iMax
         y3d(i,j,k) = ( y3d(i,j,k)
     &                - a3d(i,j,k)*y3d(i,j,k-1)
     &                )*b3d(i,j,k)
        ENDDO
       ENDDO
      ENDDO

C--   backward sweep (starting from bottom)
      DO k=Nr-1,1,-1
       DO j=jMin,jMax
        DO i=iMin,iMax
          y3d(i,j,k) = y3d(i,j,k)
     &               - c3d(i,j,k)*b3d(i,j,k)*y3d(i,j,k+1)
        ENDDO
       ENDDO
      ENDDO

#else /* ndef SOLVE_DIAGONAL_LOWMEMORY */

      errCode = 0

#ifdef SOLVE_DIAGONAL_KINNER
C--   Temporary array
      DO k=1,Nr
       DO j=1-OLy,sNy+OLy
        DO i=1-OLx,sNx+OLx
         c3d_m1(i,j,k) = c3d(i,j,k)
         y3d_m1(i,j,k) = y3d(i,j,k)
        ENDDO
       ENDDO
      ENDDO

C--   Main loop
      DO j=1-OLy,sNy+OLy
       DO i=1-OLx,sNx+OLx

        DO k=1,Nr
         c3d_prime(k) = 0. _d 0
         y3d_prime(k) = 0. _d 0
         y3d_update(k) = 0. _d 0
        ENDDO

C--   Forward sweep
        DO k=1,Nr
         IF ( k.EQ.1 ) THEN
           IF ( b3d(i,j,1).NE.0. _d 0 ) THEN
             c3d_prime(1) = c3d_m1(i,j,1) / b3d(i,j,1)
             y3d_prime(1) = y3d_m1(i,j,1) / b3d(i,j,1)
           ELSE
             c3d_prime(1) = 0. _d 0
             y3d_prime(1) = 0. _d 0
             errCode = 1
           ENDIF
         ELSE
           tmpVar = b3d(i,j,k) - a3d(i,j,k)*c3d_prime(k-1)
           IF ( tmpVar .NE. 0. _d 0 ) THEN
             recVar = 1. _d 0 / tmpVar
             c3d_prime(k) = c3d_m1(i,j,k)*recVar
             y3d_prime(k) = (y3d_m1(i,j,k) - y3d_prime(k-1)*a3d(i,j,k))
     &                      *recVar
           ELSE
             c3d_prime(k) = 0. _d 0
             y3d_prime(k) = 0. _d 0
             errCode = 1
           ENDIF
         ENDIF
        ENDDO

C--   Backward sweep
        DO k=Nr,1,-1
         IF ( k.EQ.Nr ) THEN
          y3d_update(k)=y3d_prime(k)
         ELSE
          y3d_update(k)=y3d_prime(k)-c3d_prime(k)*y3d_update(k+1)
         ENDIF
        ENDDO

C--   Update array
        DO k=1,Nr
         y3d(i,j,k) = y3d_update(k)
        ENDDO

       ENDDO
      ENDDO

#else  /* ndef SOLVE_DIAGONAL_KINNER */

C--   Init. + copy to temp. array
      DO k=1,Nr
       DO j=1-OLy,sNy+OLy
        DO i=1-OLx,sNx+OLx
         c3d_prime(i,j,k) = 0. _d 0
         y3d_prime(i,j,k) = 0. _d 0
         y3d_m1(i,j,k) = y3d(i,j,k)
        ENDDO
       ENDDO
      ENDDO

CADJ loop = sequential
C--   Forward sweep
      DO k=1,Nr

       DO j=1-OLy,sNy+OLy
        DO i=1-OLx,sNx+OLx
c      DO j=jMin,jMax
c       DO i=iMin,iMax
         IF ( k.EQ.1 ) THEN
           IF ( b3d(i,j,1).NE.0. _d 0 ) THEN
             recVar = 1. _d 0 / b3d(i,j,1)
             c3d_prime(i,j,1) = c3d(i,j,1)*recVar
             y3d_prime(i,j,1) = y3d_m1(i,j,1)*recVar
           ELSE
             c3d_prime(i,j,1) = 0. _d 0
             y3d_prime(i,j,1) = 0. _d 0
             errCode = 1
           ENDIF
         ELSE
           tmpVar = b3d(i,j,k) - a3d(i,j,k)*c3d_prime(i,j,k-1)
           IF ( tmpVar .NE. 0. _d 0 ) THEN
             recVar = 1. _d 0 / tmpVar
             c3d_prime(i,j,k) = c3d(i,j,k)*recVar
             y3d_prime(i,j,k) = ( y3d_m1(i,j,k)
     &                          - a3d(i,j,k)*y3d_prime(i,j,k-1)
     &                          )*recVar
           ELSE
             c3d_prime(i,j,k) = 0. _d 0
             y3d_prime(i,j,k) = 0. _d 0
             errCode = 1
           ENDIF
         ENDIF
        ENDDO
       ENDDO

C-    end k-loop
      ENDDO

CADJ loop = sequential
C--   Backward sweep
      DO k=Nr,1,-1
       DO j=1-OLy,sNy+OLy
        DO i=1-OLx,sNx+OLx
c      DO j=jMin,jMax
c       DO i=iMin,iMax
         IF ( k.EQ.Nr ) THEN
          y3d(i,j,k) = y3d_prime(i,j,k)
         ELSE
          y3d(i,j,k) = y3d_prime(i,j,k)
     &               - c3d_prime(i,j,k)*y3d(i,j,k+1)
         ENDIF
        ENDDO
       ENDDO
C-    end k-loop
      ENDDO

#endif /* SOLVE_DIAGONAL_KINNER */

#endif /* SOLVE_DIAGONAL_LOWMEMORY */

      RETURN
      END