C $Header: /u/gcmpack/MITgcm/model/src/solve_pentadiagonal.F,v 1.13 2016/10/26 00:50:25 jmc Exp $
C $Name: $
#include "CPP_OPTIONS.h"
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
#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 a5d,b5d,c5d,d5d,e5d)
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 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 errCode :: > 0 if singular matrix
#ifdef SOLVE_DIAGONAL_LOWMEMORY
C a5d,b5d,c5d,d5d,e5d :: 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 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)
INTEGER errCode
INTEGER bi, bj, myThid
C !LOCAL VARIABLES:
C == Local variables ==
INTEGER i,j,k
#ifndef SOLVE_DIAGONAL_LOWMEMORY
_RL tmpVar, recVar
_RL y5d_m1(1-OLx:sNx+OLx,1-OLy:sNy+OLy,Nr)
# ifdef SOLVE_DIAGONAL_KINNER
_RL c5d_prime(Nr)
_RL d5d_prime(Nr)
_RL e5d_prime(Nr)
_RL y5d_prime(Nr)
_RL y5d_update(Nr)
else
_RL c5d_prime(1-OLx:sNx+OLx,1-OLy:sNy+OLy,Nr)
_RL d5d_prime(1-OLx:sNx+OLx,1-OLy:sNy+OLy,Nr)
_RL e5d_prime(1-OLx:sNx+OLx,1-OLy:sNy+OLy,Nr)
_RL y5d_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
DO k=1,Nr
C-- forward sweep (starting from top) part-1: matrix L.U decomposition
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)
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)
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-- end if errCode < 0
ENDIF
DO k=2,Nr
C-- forward sweep (starting from top) part-2: process RHS
IF (k.EQ.2) THEN
DO j=jMin,jMax
DO i=iMin,iMax
y5d(i,j,k) = y5d(i,j,k) - b5d(i,j,k)*y5d(i,j,k-1)
ENDDO
ENDDO
ELSE
DO j=jMin,jMax
DO i=iMin,iMax
y5d(i,j,k) = y5d(i,j,k) - b5d(i,j,k)*y5d(i,j,k-1)
& - a5d(i,j,k)*y5d(i,j,k-2)
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) = y5d(i,j,k)*c5d(i,j,k)
ENDDO
ENDDO
ELSEIF (k.EQ.Nr-1) THEN
DO j=jMin,jMax
DO i=iMin,iMax
y5d(i,j,k) = ( y5d(i,j,k)
& - d5d(i,j,k)*y5d(i,j,k+1)
& )*c5d(i,j,k)
ENDDO
ENDDO
ELSE
DO j=jMin,jMax
DO i=iMin,iMax
y5d(i,j,k) = ( y5d(i,j,k)
& - d5d(i,j,k)*y5d(i,j,k+1)
& - e5d(i,j,k)*y5d(i,j,k+2)
& )*c5d(i,j,k)
ENDDO
ENDDO
C- end if k= .. ; end of k loop
ENDIF
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
y5d_m1(i,j,k) = y5d(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
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
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
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
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
tmpVar = c5d_prime(k)
IF ( tmpVar.NE.0. _d 0 ) THEN
recVar = 1. _d 0 / tmpVar
d5d_prime(k) = d5d_prime(k)*recVar
e5d_prime(k) = e5d_prime(k)*recVar
y5d_prime(k) = y5d_prime(k)*recVar
ELSE
d5d_prime(k) = 0. _d 0
e5d_prime(k) = 0. _d 0
y5d_prime(k) = 0. _d 0
errCode = 1
ENDIF
C-- k loop
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) = y5d_update(k)
ENDDO
C- end i,j loop
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
c5d_prime(i,j,k) = 0. _d 0
d5d_prime(i,j,k) = 0. _d 0
e5d_prime(i,j,k) = 0. _d 0
y5d_prime(i,j,k) = 0. _d 0
y5d_m1(i,j,k) = y5d(i,j,k)
ENDDO
ENDDO
ENDDO
CADJ loop = sequential
DO k=1,Nr
C-- Forward sweep (starting from top)
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
C- just copy terms
c5d_prime(i,j,k) = c5d(i,j,k)
d5d_prime(i,j,k) = d5d(i,j,k)
e5d_prime(i,j,k) = e5d(i,j,k)
y5d_prime(i,j,k) = y5d_m1(i,j,k)
ELSEIF (k.EQ.2) THEN
C- subtract one term
c5d_prime(i,j,k) = c5d(i,j,k)
& -b5d(i,j,k)*d5d_prime(i,j,k-1)
d5d_prime(i,j,k) = d5d(i,j,k)
& -b5d(i,j,k)*e5d_prime(i,j,k-1)
e5d_prime(i,j,k) = e5d(i,j,k)
y5d_prime(i,j,k) = y5d_m1(i,j,k)
& -b5d(i,j,k)*y5d_prime(i,j,k-1)
ELSE
C- subtract two terms
c5d_prime(i,j,k) = c5d(i,j,k)
& -a5d(i,j,k)*e5d_prime(i,j,k-2)
& -(b5d(i,j,k)-a5d(i,j,k)*d5d_prime(i,j,k-2))
& *d5d_prime(i,j,k-1)
d5d_prime(i,j,k) = d5d(i,j,k)
& -(b5d(i,j,k)-a5d(i,j,k)*d5d_prime(i,j,k-2))
& *e5d_prime(i,j,k-1)
e5d_prime(i,j,k) = e5d(i,j,k)
y5d_prime(i,j,k) = y5d_m1(i,j,k)
& -a5d(i,j,k)*y5d_prime(i,j,k-2)
& -(b5d(i,j,k)-a5d(i,j,k)*d5d_prime(i,j,k-2))
& *y5d_prime(i,j,k-1)
ENDIF
C- normalization
tmpVar = c5d_prime(i,j,k)
IF ( tmpVar.NE.0. _d 0 ) THEN
recVar = 1. _d 0 / tmpVar
d5d_prime(i,j,k) = d5d_prime(i,j,k)*recVar
e5d_prime(i,j,k) = e5d_prime(i,j,k)*recVar
y5d_prime(i,j,k) = y5d_prime(i,j,k)*recVar
ELSE
d5d_prime(i,j,k) = 0. _d 0
e5d_prime(i,j,k) = 0. _d 0
y5d_prime(i,j,k) = 0. _d 0
errCode = 1
ENDIF
ENDDO
ENDDO
C- end k-loop
ENDDO
CADJ loop = sequential
C-- Backward sweep (starting from bottom)
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
y5d(i,j,k) = y5d_prime(i,j,k)
ELSEIF (k.EQ.Nr-1) THEN
y5d(i,j,k) = y5d_prime(i,j,k)
& - y5d(i,j,k+1)*d5d_prime(i,j,k)
ELSE
y5d(i,j,k) = y5d_prime(i,j,k)
& - y5d(i,j,k+1)*d5d_prime(i,j,k)
& - y5d(i,j,k+2)*e5d_prime(i,j,k)
ENDIF
ENDDO
ENDDO
C- end k-loop
ENDDO
#endif /* SOLVE_DIAGONAL_KINNER */
#endif /* SOLVE_DIAGONAL_LOWMEMORY */
RETURN
END