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