model/src/solve_tridiagonal.F

7448700841 Mart*0001 #include "PACKAGES_CONFIG.h"
b83c015289 Jean*0002 #include "CPP_OPTIONS.h"
7448700841 Mart*0003 #ifdef ALLOW_AUTODIFF
                0004 # include "AUTODIFF_OPTIONS.h"
                0005 #endif
b83c015289 Jean*0006
                0007 CBOP
                0008 C     !ROUTINE: SOLVE_TRIDIAGONAL
                0009 C     !INTERFACE:
5cc28932de Jean*0010       SUBROUTINE SOLVE_TRIDIAGONAL(
b83c015289 Jean*0011      I                     iMin,iMax, jMin,jMax,
                0012      I                     a3d, b3d, c3d,
                0013      U                     y3d,
                0014      O                     errCode,
                0015      I                     bi, bj, myThid )
                0016 C     !DESCRIPTION: \bv
                0017 C     *==========================================================*
5cc28932de Jean*0018 C     | S/R SOLVE_TRIDIAGONAL
b83c015289 Jean*0019 C     | o Solve a tri-diagonal system A*X=Y (dimension Nr)
                0020 C     *==========================================================*
                0021 C     | o Used to solve implicitly vertical advection & diffusion
c3aa667b91 Jean*0022 #ifdef SOLVE_DIAGONAL_LOWMEMORY
                0023 C     | o Allows 2 types of call:
                0024 C     | 1) with INPUT errCode < 0 (first call):
                0025 C     |   Solve system A*X=Y by LU decomposition and return
                0026 C     |     inverse matrix as output (in a3d,b3d,c3d)
                0027 C     | 2) with INPUT errCode = 0 (subsequent calls):
                0028 C     |   Solve system A*Xp=Yp by using inverse matrix coeff
                0029 C     |     from first call output.
                0030 #endif /* SOLVE_DIAGONAL_LOWMEMORY */
b83c015289 Jean*0031 C     *==========================================================*
                0032 C     \ev
                0033
                0034 C     !USES:
                0035       IMPLICIT NONE
                0036 C     == Global data ==
                0037 #include "SIZE.h"
                0038 #include "EEPARAMS.h"
                0039
                0040 C     !INPUT/OUTPUT PARAMETERS:
c3aa667b91 Jean*0041 C     -- INPUT: --
                0042 C     iMin,iMax,jMin,jMax  :: computational domain
                0043 C     a3d     :: matrix lower diagnonal
                0044 C     b3d     :: matrix main  diagnonal
                0045 C     c3d     :: matrix upper diagnonal
                0046 C     y3d     :: Y vector (R.H.S.);
                0047 C     bi,bj   :: tile indices
                0048 C     myThid  :: thread number
                0049 C     -- OUTPUT: --
                0050 C     y3d     :: X = solution of A*X=Y
b83c015289 Jean*0051 C     errCode :: > 0 if singular matrix
c3aa667b91 Jean*0052 #ifdef SOLVE_DIAGONAL_LOWMEMORY
                0053 C     a3d,b3d,c3d :: inverse matrix coeff to enable to find Xp solution
                0054 C                    of A*Xp=Yp with subsequent call to this routine
                0055 #endif /* SOLVE_DIAGONAL_LOWMEMORY */
b83c015289 Jean*0056       INTEGER iMin,iMax,jMin,jMax
5cc28932de Jean*0057       _RL a3d(1-OLx:sNx+OLx,1-OLy:sNy+OLy,Nr)
                0058       _RL b3d(1-OLx:sNx+OLx,1-OLy:sNy+OLy,Nr)
                0059       _RL c3d(1-OLx:sNx+OLx,1-OLy:sNy+OLy,Nr)
23a7f3050f Jean*0060       _RL y3d(1-OLx:sNx+OLx,1-OLy:sNy+OLy,Nr)
b83c015289 Jean*0061       INTEGER errCode
                0062       INTEGER bi, bj, myThid
                0063
                0064 C     !LOCAL VARIABLES:
                0065 C     == Local variables ==
                0066       INTEGER i,j,k
5cc28932de Jean*0067       _RL tmpVar
351f151dc9 Patr*0068 #ifndef SOLVE_DIAGONAL_LOWMEMORY
5cc28932de Jean*0069       _RL recVar
                0070       _RL y3d_m1(1-OLx:sNx+OLx,1-OLy:sNy+OLy,Nr)
6b36f7748a Patr*0071 # ifdef SOLVE_DIAGONAL_KINNER
1dd6de01dd Jean*0072       _RL c3d_m1(1-OLx:sNx+OLx,1-OLy:sNy+OLy,Nr)
6b36f7748a Patr*0073       _RL c3d_prime(Nr)
                0074       _RL y3d_prime(Nr)
                0075       _RL y3d_update(Nr)
7448700841 Mart*0076 #  ifdef ALLOW_AUTODIFF_TAMC
                0077 C     kkey :: tape key, depends on i,j,k
                0078       INTEGER kkey
                0079 #  endif
6b36f7748a Patr*0080 # else
5cc28932de Jean*0081       _RL c3d_prime(1-OLx:sNx+OLx,1-OLy:sNy+OLy,Nr)
                0082       _RL y3d_prime(1-OLx:sNx+OLx,1-OLy:sNy+OLy,Nr)
6b36f7748a Patr*0083 # endif
351f151dc9 Patr*0084 #endif /* SOLVE_DIAGONAL_LOWMEMORY */
b83c015289 Jean*0085 CEOP
                0086
351f151dc9 Patr*0087 #ifdef SOLVE_DIAGONAL_LOWMEMORY
6b36f7748a Patr*0088
c3aa667b91 Jean*0089       IF ( errCode .LT. 0 ) THEN
                0090        errCode = 0
b83c015289 Jean*0091
c3aa667b91 Jean*0092 C--   forward sweep (starting from top) part-1: matrix L.U decomposition
b83c015289 Jean*0093        DO j=jMin,jMax
                0094         DO i=iMin,iMax
c3aa667b91 Jean*0095           IF ( b3d(i,j,1).NE.0. _d 0 ) THEN
                0096             b3d(i,j,1) = 1. _d 0 / b3d(i,j,1)
                0097           ELSE
                0098             b3d(i,j,1) = 0. _d 0
                0099             errCode = 1
                0100           ENDIF
b83c015289 Jean*0101         ENDDO
                0102        ENDDO
                0103
c3aa667b91 Jean*0104 C-    Middle of forward sweep
                0105        DO k=2,Nr
                0106         DO j=jMin,jMax
                0107          DO i=iMin,iMax
                0108           tmpVar = b3d(i,j,k) - a3d(i,j,k)*c3d(i,j,k-1)*b3d(i,j,k-1)
                0109           IF ( tmpVar .NE. 0. _d 0 ) THEN
                0110             b3d(i,j,k) = 1. _d 0 / tmpVar
                0111           ELSE
                0112             b3d(i,j,k) = 0. _d 0
                0113             errCode = 1
                0114           ENDIF
                0115          ENDDO
                0116         ENDDO
                0117        ENDDO
                0118
                0119 C--   end if errCode < 0
                0120       ENDIF
                0121
                0122 C--   forward sweep (starting from top) part-2: process RHS
b83c015289 Jean*0123       DO j=jMin,jMax
                0124        DO i=iMin,iMax
23a7f3050f Jean*0125          y3d(i,j,1) = y3d(i,j,1)*b3d(i,j,1)
b83c015289 Jean*0126        ENDDO
                0127       ENDDO
                0128       DO k=2,Nr
                0129        DO j=jMin,jMax
                0130         DO i=iMin,iMax
23a7f3050f Jean*0131          y3d(i,j,k) = ( y3d(i,j,k)
                0132      &                - a3d(i,j,k)*y3d(i,j,k-1)
                0133      &                )*b3d(i,j,k)
b83c015289 Jean*0134         ENDDO
                0135        ENDDO
                0136       ENDDO
                0137
c3aa667b91 Jean*0138 C--   backward sweep (starting from bottom)
b83c015289 Jean*0139       DO k=Nr-1,1,-1
                0140        DO j=jMin,jMax
                0141         DO i=iMin,iMax
23a7f3050f Jean*0142           y3d(i,j,k) = y3d(i,j,k)
                0143      &               - c3d(i,j,k)*b3d(i,j,k)*y3d(i,j,k+1)
b83c015289 Jean*0144         ENDDO
                0145        ENDDO
                0146       ENDDO
                0147
351f151dc9 Patr*0148 #else /* ndef SOLVE_DIAGONAL_LOWMEMORY */
d666da8b50 Gael*0149
c3aa667b91 Jean*0150       errCode = 0
                0151
1dd6de01dd Jean*0152 #ifdef SOLVE_DIAGONAL_KINNER
7448700841 Mart*0153 #ifdef ALLOW_AUTODIFF_TAMC
                0154 CADJ INIT loctape_solvetri = COMMON, Nr*(sNy+2*OLy)*(sNx+2*OLx)
                0155 #endif
d666da8b50 Gael*0156 C--   Temporary array
68073fc7c8 Patr*0157       DO k=1,Nr
5cc28932de Jean*0158        DO j=1-OLy,sNy+OLy
                0159         DO i=1-OLx,sNx+OLx
d666da8b50 Gael*0160          c3d_m1(i,j,k) = c3d(i,j,k)
23a7f3050f Jean*0161          y3d_m1(i,j,k) = y3d(i,j,k)
6b36f7748a Patr*0162         ENDDO
                0163        ENDDO
d666da8b50 Gael*0164       ENDDO
6b36f7748a Patr*0165
d666da8b50 Gael*0166 C--   Main loop
5cc28932de Jean*0167       DO j=1-OLy,sNy+OLy
                0168        DO i=1-OLx,sNx+OLx
d666da8b50 Gael*0169
6b36f7748a Patr*0170         DO k=1,Nr
                0171          c3d_prime(k) = 0. _d 0
                0172          y3d_prime(k) = 0. _d 0
                0173          y3d_update(k) = 0. _d 0
                0174         ENDDO
d666da8b50 Gael*0175
                0176 C--   Forward sweep
6b36f7748a Patr*0177         DO k=1,Nr
7448700841 Mart*0178 #ifdef ALLOW_AUTODIFF_TAMC
                0179           kkey = k + Nr*(i-1+OLx) + Nr*(sNx+2*OLx)*(j-1+OLy)
                0180 CADJ STORE c3d_prime(k) = loctape_solvetri, key = kkey
                0181 CADJ STORE y3d_prime(k) = loctape_solvetri, key = kkey
                0182 #endif
d666da8b50 Gael*0183          IF ( k.EQ.1 ) THEN
                0184            IF ( b3d(i,j,1).NE.0. _d 0 ) THEN
                0185              c3d_prime(1) = c3d_m1(i,j,1) / b3d(i,j,1)
                0186              y3d_prime(1) = y3d_m1(i,j,1) / b3d(i,j,1)
                0187            ELSE
                0188              c3d_prime(1) = 0. _d 0
                0189              y3d_prime(1) = 0. _d 0
                0190              errCode = 1
                0191            ENDIF
                0192          ELSE
5cc28932de Jean*0193            tmpVar = b3d(i,j,k) - a3d(i,j,k)*c3d_prime(k-1)
                0194            IF ( tmpVar .NE. 0. _d 0 ) THEN
                0195              recVar = 1. _d 0 / tmpVar
                0196              c3d_prime(k) = c3d_m1(i,j,k)*recVar
d666da8b50 Gael*0197              y3d_prime(k) = (y3d_m1(i,j,k) - y3d_prime(k-1)*a3d(i,j,k))
5cc28932de Jean*0198      &                      *recVar
d666da8b50 Gael*0199            ELSE
                0200              c3d_prime(k) = 0. _d 0
                0201              y3d_prime(k) = 0. _d 0
                0202              errCode = 1
                0203            ENDIF
                0204          ENDIF
6b36f7748a Patr*0205         ENDDO
d666da8b50 Gael*0206
                0207 C--   Backward sweep
6b36f7748a Patr*0208         DO k=Nr,1,-1
d666da8b50 Gael*0209          IF ( k.EQ.Nr ) THEN
                0210           y3d_update(k)=y3d_prime(k)
                0211          ELSE
                0212           y3d_update(k)=y3d_prime(k)-c3d_prime(k)*y3d_update(k+1)
                0213          ENDIF
6b36f7748a Patr*0214         ENDDO
d666da8b50 Gael*0215
                0216 C--   Update array
6b36f7748a Patr*0217         DO k=1,Nr
23a7f3050f Jean*0218          y3d(i,j,k) = y3d_update(k)
6b36f7748a Patr*0219         ENDDO
                0220
                0221        ENDDO
                0222       ENDDO
                0223
                0224 #else  /* ndef SOLVE_DIAGONAL_KINNER */
                0225
1dd6de01dd Jean*0226 C--   Init. + copy to temp. array
6b36f7748a Patr*0227       DO k=1,Nr
5cc28932de Jean*0228        DO j=1-OLy,sNy+OLy
                0229         DO i=1-OLx,sNx+OLx
6b36f7748a Patr*0230          c3d_prime(i,j,k) = 0. _d 0
                0231          y3d_prime(i,j,k) = 0. _d 0
1dd6de01dd Jean*0232          y3d_m1(i,j,k) = y3d(i,j,k)
6b36f7748a Patr*0233         ENDDO
                0234        ENDDO
                0235       ENDDO
                0236
                0237 CADJ loop = sequential
                0238 C--   Forward sweep
                0239       DO k=1,Nr
                0240
e0fb0cb0b7 Mart*0241        IF ( k .EQ. 1 ) THEN
                0242         DO j=1-OLy,sNy+OLy
                0243          DO i=1-OLx,sNx+OLx
                0244           IF ( b3d(i,j,1).NE.0. _d 0 ) THEN
                0245            recVar = 1. _d 0 / b3d(i,j,1)
                0246            c3d_prime(i,j,1) = c3d(i,j,1)*recVar
                0247            y3d_prime(i,j,1) = y3d_m1(i,j,1)*recVar
                0248           ELSE
                0249            c3d_prime(i,j,1) = 0. _d 0
                0250            y3d_prime(i,j,1) = 0. _d 0
                0251            errCode = 1
                0252           ENDIF
                0253          ENDDO
6b36f7748a Patr*0254         ENDDO
e0fb0cb0b7 Mart*0255        ELSE
                0256         DO j=1-OLy,sNy+OLy
                0257          DO i=1-OLx,sNx+OLx
                0258           tmpVar = b3d(i,j,k) - a3d(i,j,k)*c3d_prime(i,j,k-1)
                0259           IF ( tmpVar .NE. 0. _d 0 ) THEN
                0260            recVar = 1. _d 0 / tmpVar
                0261            c3d_prime(i,j,k) = c3d(i,j,k)*recVar
                0262            y3d_prime(i,j,k) = ( y3d_m1(i,j,k)
                0263      &                        - a3d(i,j,k)*y3d_prime(i,j,k-1)
                0264      &                        )*recVar
                0265           ELSE
                0266            c3d_prime(i,j,k) = 0. _d 0
                0267            y3d_prime(i,j,k) = 0. _d 0
                0268            errCode = 1
                0269           ENDIF
                0270          ENDDO
                0271         ENDDO
                0272        ENDIF
1dd6de01dd Jean*0273
                0274 C-    end k-loop
d666da8b50 Gael*0275       ENDDO
                0276
6b36f7748a Patr*0277 CADJ loop = sequential
                0278 C--   Backward sweep
                0279       DO k=Nr,1,-1
e0fb0cb0b7 Mart*0280        IF ( k.EQ.Nr ) THEN
                0281         DO j=1-OLy,sNy+OLy
                0282          DO i=1-OLx,sNx+OLx
1dd6de01dd Jean*0283           y3d(i,j,k) = y3d_prime(i,j,k)
e0fb0cb0b7 Mart*0284          ENDDO
                0285         ENDDO
                0286        ELSE
                0287         DO j=1-OLy,sNy+OLy
                0288          DO i=1-OLx,sNx+OLx
1dd6de01dd Jean*0289           y3d(i,j,k) = y3d_prime(i,j,k)
                0290      &               - c3d_prime(i,j,k)*y3d(i,j,k+1)
e0fb0cb0b7 Mart*0291          ENDDO
6b36f7748a Patr*0292         ENDDO
e0fb0cb0b7 Mart*0293        ENDIF
1dd6de01dd Jean*0294 C-    end k-loop
d666da8b50 Gael*0295       ENDDO
                0296
6b36f7748a Patr*0297 #endif /* SOLVE_DIAGONAL_KINNER */
                0298
351f151dc9 Patr*0299 #endif /* SOLVE_DIAGONAL_LOWMEMORY */
d666da8b50 Gael*0300
b83c015289 Jean*0301       RETURN
                0302       END