1 |
C $Header: /u/gcmpack/MITgcm/model/src/solve_tridiagonal.F,v 1.2 2004/12/05 17:29:36 jmc Exp $ |
2 |
C $Name: $ |
3 |
|
4 |
#include "CPP_OPTIONS.h" |
5 |
|
6 |
C o Switch to code that has the k-loop inside the |
7 |
C ij-loops, which matters in adjoint mode. |
8 |
#ifdef ALLOW_AUTODIFF |
9 |
#define ALLOW_SOLVERS_KLOOPINSIDE |
10 |
#endif |
11 |
|
12 |
CBOP |
13 |
C !ROUTINE: SOLVE_TRIDIAGONAL |
14 |
C !INTERFACE: |
15 |
SUBROUTINE SOLVE_TRIDIAGONAL( |
16 |
I iMin,iMax, jMin,jMax, |
17 |
I a3d, b3d, c3d, |
18 |
U y3d, |
19 |
O errCode, |
20 |
I bi, bj, myThid ) |
21 |
C !DESCRIPTION: \bv |
22 |
C *==========================================================* |
23 |
C | S/R SOLVE_TRIDIAGONAL |
24 |
C | o Solve a tri-diagonal system A*X=Y (dimension Nr) |
25 |
C *==========================================================* |
26 |
C | o Used to solve implicitly vertical advection & diffusion |
27 |
C *==========================================================* |
28 |
C \ev |
29 |
|
30 |
C !USES: |
31 |
IMPLICIT NONE |
32 |
C == Global data == |
33 |
#include "SIZE.h" |
34 |
#include "EEPARAMS.h" |
35 |
|
36 |
C !INPUT/OUTPUT PARAMETERS: |
37 |
C == Routine Arguments == |
38 |
C a3d :: matrix lower diagnonal |
39 |
C b3d :: matrix main diagnonal |
40 |
C c3d :: matrix upper diagnonal |
41 |
C y3d :: Input = Y vector ; Output = X = solution of A*X=Y |
42 |
C errCode :: > 0 if singular matrix |
43 |
INTEGER iMin,iMax,jMin,jMax |
44 |
_RL a3d(1-Olx:sNx+Olx,1-Oly:sNy+Oly,Nr) |
45 |
_RL b3d(1-Olx:sNx+Olx,1-Oly:sNy+Oly,Nr) |
46 |
_RL c3d(1-Olx:sNx+Olx,1-Oly:sNy+Oly,Nr) |
47 |
_RL y3d(1-Olx:sNx+Olx,1-Oly:sNy+Oly,Nr,nSx,nSy) |
48 |
INTEGER errCode |
49 |
INTEGER bi, bj, myThid |
50 |
|
51 |
C !LOCAL VARIABLES: |
52 |
C == Local variables == |
53 |
INTEGER i,j,k |
54 |
_RL tmpvar |
55 |
#ifndef ALLOW_SOLVERS_KLOOPINSIDE |
56 |
_RL bet(1-Olx:sNx+Olx,1-Oly:sNy+Oly,Nr) |
57 |
#else |
58 |
_RL c3d_prime(Nr), y3d_prime(Nr),y3d_update(Nr) |
59 |
_RL c3d_m1(1-Olx:sNx+Olx,1-Oly:sNy+Oly,Nr) |
60 |
_RL y3d_m1(1-Olx:sNx+Olx,1-Oly:sNy+Oly,Nr) |
61 |
#endif |
62 |
CEOP |
63 |
|
64 |
#ifndef ALLOW_SOLVERS_KLOOPINSIDE |
65 |
|
66 |
errCode = 0 |
67 |
|
68 |
C-- Beginning of forward sweep (top level) |
69 |
DO j=jMin,jMax |
70 |
DO i=iMin,iMax |
71 |
IF ( b3d(i,j,1).NE.0. _d 0 ) THEN |
72 |
bet(i,j,1) = 1. _d 0 / b3d(i,j,1) |
73 |
ELSE |
74 |
bet(i,j,1) = 0. _d 0 |
75 |
errCode = 1 |
76 |
ENDIF |
77 |
ENDDO |
78 |
ENDDO |
79 |
|
80 |
C-- Middle of forward sweep |
81 |
DO k=2,Nr |
82 |
DO j=jMin,jMax |
83 |
DO i=iMin,iMax |
84 |
tmpvar = b3d(i,j,k) - a3d(i,j,k)*c3d(i,j,k-1)*bet(i,j,k-1) |
85 |
IF ( tmpvar .NE. 0. _d 0 ) THEN |
86 |
bet(i,j,k) = 1. _d 0 / tmpvar |
87 |
ELSE |
88 |
bet(i,j,k) = 0. _d 0 |
89 |
errCode = 1 |
90 |
ENDIF |
91 |
ENDDO |
92 |
ENDDO |
93 |
ENDDO |
94 |
|
95 |
DO j=jMin,jMax |
96 |
DO i=iMin,iMax |
97 |
y3d(i,j,1,bi,bj) = y3d(i,j,1,bi,bj)*bet(i,j,1) |
98 |
ENDDO |
99 |
ENDDO |
100 |
DO k=2,Nr |
101 |
DO j=jMin,jMax |
102 |
DO i=iMin,iMax |
103 |
y3d(i,j,k,bi,bj) = ( y3d(i,j,k,bi,bj) |
104 |
& - a3d(i,j,k)*y3d(i,j,k-1,bi,bj) |
105 |
& )*bet(i,j,k) |
106 |
ENDDO |
107 |
ENDDO |
108 |
ENDDO |
109 |
|
110 |
C-- Backward sweep |
111 |
CADJ loop = sequential |
112 |
DO k=Nr-1,1,-1 |
113 |
DO j=jMin,jMax |
114 |
DO i=iMin,iMax |
115 |
y3d(i,j,k,bi,bj) = y3d(i,j,k,bi,bj) |
116 |
& - c3d(i,j,k)*bet(i,j,k)*y3d(i,j,k+1,bi,bj) |
117 |
ENDDO |
118 |
ENDDO |
119 |
ENDDO |
120 |
|
121 |
#else /* ALLOW_SOLVERS_KLOOPINSIDE */ |
122 |
|
123 |
errCode = 0 |
124 |
|
125 |
C-- Temporary array |
126 |
DO j=jMin,jMax |
127 |
DO i=iMin,iMax |
128 |
DO k=1,Nr |
129 |
c3d_m1(i,j,k) = c3d(i,j,k) |
130 |
y3d_m1(i,j,k) = y3d(i,j,k,bi,bj) |
131 |
ENDDO |
132 |
ENDDO |
133 |
ENDDO |
134 |
|
135 |
C-- Main loop |
136 |
DO j=jMin,jMax |
137 |
DO i=iMin,iMax |
138 |
|
139 |
DO k=1,Nr |
140 |
c3d_prime(k) = 0. _d 0 |
141 |
y3d_prime(k) = 0. _d 0 |
142 |
y3d_update(k) = 0. _d 0 |
143 |
ENDDO |
144 |
|
145 |
C-- Forward sweep |
146 |
DO k=1,Nr |
147 |
IF ( k.EQ.1 ) THEN |
148 |
IF ( b3d(i,j,1).NE.0. _d 0 ) THEN |
149 |
c3d_prime(1) = c3d_m1(i,j,1) / b3d(i,j,1) |
150 |
y3d_prime(1) = y3d_m1(i,j,1) / b3d(i,j,1) |
151 |
ELSE |
152 |
c3d_prime(1) = 0. _d 0 |
153 |
y3d_prime(1) = 0. _d 0 |
154 |
errCode = 1 |
155 |
ENDIF |
156 |
ELSE |
157 |
tmpvar = b3d(i,j,k) - a3d(i,j,k)*c3d_prime(k-1) |
158 |
IF ( tmpvar .NE. 0. _d 0 ) THEN |
159 |
c3d_prime(k) = c3d_m1(i,j,k) / tmpvar |
160 |
y3d_prime(k) = (y3d_m1(i,j,k) - y3d_prime(k-1)*a3d(i,j,k)) |
161 |
& / tmpvar |
162 |
ELSE |
163 |
c3d_prime(k) = 0. _d 0 |
164 |
y3d_prime(k) = 0. _d 0 |
165 |
errCode = 1 |
166 |
ENDIF |
167 |
ENDIF |
168 |
ENDDO |
169 |
|
170 |
C-- Backward sweep |
171 |
DO k=Nr,1,-1 |
172 |
IF ( k.EQ.Nr ) THEN |
173 |
y3d_update(k)=y3d_prime(k) |
174 |
ELSE |
175 |
y3d_update(k)=y3d_prime(k)-c3d_prime(k)*y3d_update(k+1) |
176 |
ENDIF |
177 |
ENDDO |
178 |
|
179 |
C-- Update array |
180 |
DO k=1,Nr |
181 |
y3d(i,j,k,bi,bj)=y3d_update(k) |
182 |
ENDDO |
183 |
|
184 |
ENDDO |
185 |
ENDDO |
186 |
|
187 |
#endif /* ALLOW_SOLVERS_KLOOPINSIDE */ |
188 |
|
189 |
RETURN |
190 |
END |