source: codes/icosagcm/devel/Python/src/kernels_caldyn_NH.jin @ 615

Last change on this file since 615 was 615, checked in by dubos, 6 years ago

devel/unstructured : import Python layer from HEAT
File size: 10.3 KB
Line 
1{% macro compute_p_and_Aik() %}
2{% set p_and_c2_from_theta=caller %}
3    SEQUENCE_EXT
4      BODY('1,llm')
5          rho_ij = (g*m_ik(CELL))/(Phi_il(UP(CELL))-Phi_il(CELL))
6          {{ p_and_c2_from_theta() }}
7          A_ik(CELL) = c2_mik*(tau/g*rho_ij)**2
8      END_BLOCK
9    END_BLOCK
10{%- endmacro %}
11
12KERNEL(compute_NH_geopot)
13  tau2_g=tau*tau/g
14  g2=g*g
15  gm2 = 1./g2
16  gamma = 1./(1.-kappa)
17   
18  BARRIER
19
20  ! compute Phi_star
21  SEQUENCE_EXT
22    BODY('1,llm+1')
23      Phi_star_il(CELL) = Phi_il(CELL) + tau*g2*(W_il(CELL)/m_il(CELL)-tau)
24    END_BLOCK
25  END_BLOCK
26
27  ! Newton-Raphson iteration : Phi_il contains current guess value
28  DO iter=1,2 ! 2 iterations should be enough
29    ! Compute pressure, A_ik
30    SELECT CASE(caldyn_thermo)
31    CASE(thermo_theta)
32      {% call() compute_p_and_Aik() %}
33        X_ij = (cpp/preff)*kappa*theta(CELL)*rho_ij
34        p_ik(CELL) = preff*(X_ij**gamma)
35        c2_mik = gamma*p_ik(CELL)/(rho_ij*m_ik(CELL)) ! c^2 = gamma*R*T = gamma*p/rho
36      {% endcall %}
37    CASE(thermo_entropy)
38      {% call() compute_p_and_Aik() %}
39        X_ij = Treff*exp(theta(CELL)/cpp) ! theta = Tref.exp(s/Cp)
40        X_ij = (cpp/preff)*kappa*X_ij*rho_ij
41        p_ik(CELL) = preff*(X_ij**gamma)
42        c2_mik = gamma*p_ik(CELL)/(rho_ij*m_ik(CELL)) ! c^2 = gamma*R*T = gamma*p/rho
43      {% endcall %}
44    CASE DEFAULT
45        PRINT *, 'caldyn_thermo not supported by compute_NH_geopot', caldyn_thermo
46        STOP
47    END SELECT
48   
49    ! NB : A(1), A(llm), R(1), R(llm+1) = 0 => x(l)=0 at l=1,llm+1 => flat, rigid top and bottom
50    ! Solve -A(l-1)x(l-1) + B(l)x(l) - A(l)x(l+1) = R(l) using Thomas algorithm
51
52    SEQUENCE_EXT
53      ! Compute residual R_il and B_il
54      PROLOGUE(1)
55        ! bottom interface l=1
56        ml_g2 = gm2*m_il(CELL)
57        B_il(CELL) = A_ik(CELL) + ml_g2 + tau2_g*rho_bot
58        R_il(CELL) = ml_g2*( Phi_il(CELL)-Phi_star_il(CELL))  &
59                  + tau2_g*( p_ik(CELL)-pbot+rho_bot*(Phi_il(CELL)-PHI_BOT(HIDX(CELL))) )
60      END_BLOCK
61      BODY('2,llm')
62        ! inner interfaces
63        ml_g2 = gm2*m_il(CELL)
64        B_il(CELL) = A_ik(CELL)+A_ik(DOWN(CELL)) + ml_g2
65        R_il(CELL) = ml_g2*( Phi_il(CELL)-Phi_star_il(CELL))  &
66                  + tau2_g*(p_ik(CELL)-p_ik(DOWN(CELL)))
67        ! consistency check : if Wil=0 and initial state is in hydrostatic balance
68        ! then Phi_star_il(CELL) = Phi_il(CELL) - tau^2*g^2
69        ! and residual = tau^2*(ml+(1/g)dl_pi)=0
70      END_BLOCK
71      EPILOGUE('llm+1')
72        ! top interface l=llm+1
73        ml_g2 = gm2*m_il(CELL)
74        B_il(CELL) = A_ik(DOWN(CELL)) + ml_g2
75        R_il(CELL) = ml_g2*( Phi_il(CELL)-Phi_star_il(CELL))  &
76                  + tau2_g*( ptop-p_ik(DOWN(CELL)) )
77      END_BLOCK
78      !
79      ! Forward sweep :
80      ! C(0)=0, C(l) = -A(l) / (B(l)+A(l-1)C(l-1)),
81      ! D(0)=0, D(l) = (R(l)+A(l-1)D(l-1)) / (B(l)+A(l-1)C(l-1))
82      PROLOGUE(1)
83        X_ij = 1./B_il(CELL)
84        C_ik(CELL) = -A_ik(CELL) * X_ij
85        D_il(CELL) =  R_il(CELL) * X_ij
86      END_BLOCK
87      BODY('2,llm')
88        X_ij = 1./( B_il(CELL) + A_ik(DOWN(CELL))*C_ik(DOWN(CELL)) )
89        C_ik(CELL) = -A_ik(CELL) * X_ij
90        D_il(CELL) =  (R_il(CELL)+A_ik(DOWN(CELL))*D_il(DOWN(CELL))) * X_ij
91      END_BLOCK
92      EPILOGUE('llm+1')
93        X_ij = 1./( B_il(CELL) + A_ik(DOWN(CELL))*C_ik(DOWN(CELL)) )
94        D_il(CELL) =  (R_il(CELL)+A_ik(DOWN(CELL))*D_il(DOWN(CELL))) * X_ij
95        ! Back substitution :
96        ! x(i) = D(i)-C(i)x(i+1), x(llm+1)=0
97        ! + Newton-Raphson update
98        ! top interface l=llm+1
99        x_il(CELL) = D_il(CELL)
100        Phi_il(CELL) = Phi_il(CELL) - x_il(CELL)
101      END_BLOCK
102      BODY('llm,1,-1')
103        ! Back substitution at lower interfaces
104        x_il(CELL) = D_il(CELL) - C_ik(CELL)*x_il(UP(CELL))
105        Phi_il(CELL) = Phi_il(CELL) - x_il(CELL)
106      END_BLOCK
107    END_BLOCK
108
109    IF(debug_hevi_solver) THEN
110       PRINT *, '[hevi_solver] A,B', iter, MAXVAL(ABS(A_ik)),MAXVAL(ABS(B_il))
111       PRINT *, '[hevi_solver] C,D', iter, MAXVAL(ABS(C_ik)),MAXVAL(ABS(D_il))
112       DO l=1,llm+1
113          WRITE(*,'(A,I2.1,I3.2,E9.2)'), '[hevi_solver] x_il', iter,l, MAXVAL(ABS(x_il(l,:)))
114       END DO
115       DO l=1,llm+1
116          WRITE(*,'(A,I2.1,I3.2,E9.2)'), '[hevi_solver] R_il', iter,l, MAXVAL(ABS(R_il(l,:)))
117       END DO
118    END IF
119
120  END DO ! Newton-Raphson
121
122  BARRIER
123
124  debug_hevi_solver=.FALSE.
125END_BLOCK
126
127KERNEL(caldyn_solver)
128  FORALL_CELLS_EXT('1', 'llm+1')
129    ON_PRIMAL
130       CST_IF(IS_BOTTOM_LEVEL,    m_il(CELL) = .5*rhodz(CELL)                     )
131       CST_IF(IS_INNER_INTERFACE, m_il(CELL) = .5*(rhodz(CELL)+rhodz(DOWN(CELL))) )
132       CST_IF(IS_TOP_INTERFACE,   m_il(CELL) = .5*rhodz(DOWN(CELL))               )
133    END_BLOCK
134  END_BLOCK
135  !
136  IF(tau>0) THEN ! solve implicit problem for geopotential
137    CALL compute_NH_geopot(tau,PHI_BOT_VAR, rhodz, m_il, theta, W, geopot)
138  END IF
139  !
140  ! Compute pressure (pres) and Exner function (pk)
141  ! kappa = R/Cp
142  ! 1-kappa = Cv/Cp
143  ! Cp/Cv = 1/(1-kappa)
144  gamma = 1./(1.-kappa)
145  vreff    = Rd*Treff/preff ! reference specific volume
146  Cvd   = cpp-Rd
147  FORALL_CELLS_EXT()
148    SELECT CASE(caldyn_thermo)
149    CASE(thermo_theta)
150      ON_PRIMAL
151        rho_ij = g*rhodz(CELL)/(geopot(UP(CELL))-geopot(CELL))
152        X_ij = (cpp/preff)*kappa*theta(CELL,1)*rho_ij
153        ! kappa.theta.rho = p/exner
154        ! => X = (p/p0)/(exner/Cp)
155        !      = (p/p0)^(1-kappa)
156        pres(CELL) = preff*(X_ij**gamma)  ! pressure
157        ! Compute Exner function (needed by compute_caldyn_fast)
158        ! other formulae possible if exponentiation is slow
159        pk(CELL)   = cpp*((pres(CELL)/preff)**kappa) ! Exner
160      END_BLOCK
161    CASE(thermo_entropy)
162      ON_PRIMAL
163        rho_ij = g*rhodz(CELL)/(geopot(UP(CELL))-geopot(CELL))
164        T_ij = Treff*exp( (theta(CELL,1)+Rd*log(vreff*rho_ij))/Cvd )
165        pres(CELL) = rho_ij*Rd*T_ij
166        pk(CELL)   = T_ij
167      END_BLOCK
168    CASE DEFAULT
169      STOP
170    END SELECT
171  END_BLOCK
172
173  ! We need a barrier here because we compute pres above and do a vertical difference below
174  BARRIER
175
176  FORALL_CELLS_EXT('1', 'llm+1')
177    ON_PRIMAL
178      CST_IFTHEN(IS_BOTTOM_LEVEL)
179        ! Lower BC
180        dW(CELL)   = (1./g)*(pbot-rho_bot*(geopot(CELL)-PHI_BOT(HIDX(CELL)))-pres(CELL)) - m_il(CELL)
181      CST_ELSEIF(IS_TOP_INTERFACE)
182        ! Top BC
183        dW(CELL)   = (1./g)*(pres(DOWN(CELL))-ptop) - m_il(CELL)
184      CST_ELSE
185        dW(CELL)   = (1./g)*(pres(DOWN(CELL))-pres(CELL)) - m_il(CELL)
186      CST_ENDIF
187      W(CELL)    = W(CELL)+tau*dW(CELL) ! update W
188      dPhi(CELL) = g*g*W(CELL)/m_il(CELL)
189    END_BLOCK
190  END_BLOCK
191
192  ! We need a barrier here because we update W above and do a vertical average below
193  BARRIER
194
195  FORALL_CELLS_EXT()
196    ON_PRIMAL
197      ! compute du = -0.5*g^2.grad(w^2)
198      berni(CELL) = (-.25*g*g)*((W(CELL)/m_il(CELL))**2 + (W(UP(CELL))/m_il(UP(CELL)))**2 )
199    END_BLOCK
200  END_BLOCK
201  FORALL_CELLS_EXT()
202    ON_EDGES
203      du(EDGE) = SIGN*(berni(CELL1)-berni(CELL2))
204    END_BLOCK
205  END_BLOCK
206END_BLOCK
207
208KERNEL(caldyn_vert_NH)
209  FORALL_CELLS_EXT()
210    ON_PRIMAL
211      CST_IF(IS_BOTTOM_LEVEL,     w_ij = ( W(CELL)+.5*W(UP(CELL)) )/mass(CELL) )
212      CST_IF(IS_INNER_LAYER , w_ij = .5*( W(CELL)+W(UP(CELL)) )/mass(CELL) )
213      CST_IF(IS_TOP_LAYER,        w_ij = ( .5*W(CELL)+W(UP(CELL)) )/mass(CELL) )
214      wflux_ij = .5*(wflux(CELL)+wflux(UP(CELL)))
215      W_etadot(CELL) = wflux_ij*w_ij
216      eta_dot(CELL) = wflux_ij / mass(CELL)
217      wcov(CELL) = w_ij*(geopot(UP(CELL))-geopot(CELL))
218    END_BLOCK
219  END_BLOCK
220  ! add NH term to du
221  FORALL_CELLS()
222    ON_EDGES
223      du(EDGE) = du(EDGE) - .5*(wcov(CELL1)+wcov(CELL2))*SIGN*(eta_dot(CELL2)-eta_dot(CELL1))
224    END_BLOCK
225  END_BLOCK
226  ! add NH terms to dW, dPhi
227  ! FIXME : TODO top and bottom
228  FORALL_CELLS('2','llm')
229    ON_PRIMAL
230      dPhi(CELL)=dPhi(CELL)-wflux(CELL)*(geopot(UP(CELL))-geopot(DOWN(CELL)))/(mass(DOWN(CELL))+mass(CELL))
231    END_BLOCK
232  END_BLOCK
233
234  ! We need a barrier here because we compute W_etadot above and do a vertical difference below
235  BARRIER
236
237  FORALL_CELLS('1','llm+1')
238    ON_PRIMAL
239      CST_IFTHEN(IS_BOTTOM_LEVEL)
240         dW(CELL) = dW(CELL) - W_etadot(CELL)
241      CST_ELSEIF(IS_TOP_INTERFACE)
242         dW(CELL) = dW(CELL) + W_etadot(DOWN(CELL))
243      CST_ELSE
244         dW(CELL) = dW(CELL) + W_etadot(DOWN(CELL)) - W_etadot(CELL)
245      CST_ENDIF
246    END_BLOCK
247  END_BLOCK
248END_BLOCK
249
250KERNEL(caldyn_slow_NH)
251  FORALL_CELLS_EXT('1', 'llm+1')
252     ON_PRIMAL
253        CST_IF(IS_INNER_INTERFACE, w_il(CELL) = 2.*W(CELL)/(rhodz(KDOWN(CELL))+rhodz(KUP(CELL))) )
254        CST_IF(IS_BOTTOM_LEVEL,    w_il(CELL) = 2.*W(CELL)/rhodz(KUP(CELL)) )
255        CST_IF(IS_TOP_INTERFACE,   w_il(CELL) = 2.*W(CELL)/rhodz(KDOWN(CELL)) )
256     END_BLOCK
257  END_BLOCK
258  FORALL_CELLS_EXT('1', 'llm+1')
259     ON_EDGES
260       ! compute DePhi, v_el, G_el, F_el
261       ! v_el, W2_el and therefore G_el incorporate metric factor le_de
262       ! while DePhil, W_el and F_el do not
263       W_el         = .5*( W(CELL2)+W(CELL1) )
264       DePhil(EDGE) = SIGN*(Phi(CELL2)-Phi(CELL1))
265       F_el(EDGE)   = DePhil(EDGE)*W_el
266       W2_el        = .5*LE_DE * ( W(CELL1)*w_il(CELL1) + W(CELL2)*w_il(CELL2) )
267       v_el(EDGE)   = .5*LE_DE*(u(KUP(EDGE))+u(KDOWN(EDGE))) ! checked
268       G_el(EDGE)   = v_el(EDGE)*W_el - DePhil(EDGE)*W2_el
269     END_BLOCK
270  END_BLOCK
271
272  FORALL_CELLS_EXT('1', 'llm+1')
273    ! compute GradPhi2, dPhi, dW
274    ON_PRIMAL
275       gPhi2=0.
276       dP=0.
277       divG=0
278       FORALL_EDGES       
279         gPhi2 = gPhi2 + LE_DE*DePhil(EDGE)**2
280         dP = dP + LE_DE*DePhil(EDGE)*v_el(EDGE)
281         divG = divG + SIGN*G_el(EDGE) ! -div(G_el), G_el already has le_de
282       END_BLOCK
283       gradPhi2(CELL) = 1./(2.*AI) * gPhi2
284       dPhi(CELL) = gradPhi2(CELL)*w_il(CELL) - 1./(2.*AI)*dP
285       dW(CELL) = (-1./AI)*divG
286    END_BLOCK
287  END_BLOCK
288
289  ! We need a barrier here because we compute gradPhi2, F_el and w_il above and do a vertical average below
290  BARRIER
291
292  FORALL_CELLS_EXT()
293    ! Compute berni at scalar points
294    ON_PRIMAL
295       u2=0.
296       FORALL_EDGES       
297         u2 = u2 + LE_DE*u(EDGE)**2
298       END_BLOCK
299       berni(CELL) = 1./(4.*AI) * u2 - .25*( gradPhi2(CELL)*w_il(CELL)**2 + gradPhi2(UP(CELL))*w_il(UP(CELL))**2 )
300    END_BLOCK
301  END_BLOCK
302
303  FORALL_CELLS_EXT()
304    ON_EDGES
305    ! Compute mass flux and grad(berni)
306      uu = .5*(rhodz(CELL1)+rhodz(CELL2))*u(EDGE) - .5*( F_el(EDGE)+F_el(UP(EDGE)) )
307      hflux(EDGE) = LE_DE*uu
308      du(EDGE) = SIGN*(berni(CELL1)-berni(CELL2))
309    END_BLOCK
310  END_BLOCK
311
312END_BLOCK
Note: See TracBrowser for help on using the repository browser.