source: codes/icosagcm/devel/Python/dynamico/dyn.py @ 615

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

devel/unstructured : import Python layer from HEAT
File size: 19.9 KB
Line 
1import numpy as np
2from scipy import sparse as sparse
3import math as math
4import time_step as time_step
5
6class Struct: pass
7
8def Linf(data):
9    return abs(data).max()
10def compare(msg, data1, data2):
11    print msg, Linf(data1-data2)/Linf(data1)   
12
13def col_vec(data):
14    vec=np.zeros((data.size,1))
15    vec[:,0]=data
16    return vec
17def row_vec(data):
18    vec=np.zeros((1,data.size))
19    vec[0,:]=data
20    return vec
21
22zero=np.zeros((1))
23
24####################### Horizontal mesh #########################
25
26class Periodic_FD:
27    # staggering
28    # m0  m1  m2  m3
29    #   u0  u1  u2  u3
30    def __init__(self, Nx):
31        self.Nx=Nx
32        ones=np.ones(Nx);
33
34        right = sparse.spdiags([ones],[0],Nx,Nx, format='csr')
35        left = sparse.spdiags([ones],[-1],Nx,Nx, format='csr')
36        left[0,-1]=1
37        di = right-left
38        mi = .5*(right+left)
39        self.di=lambda f:di.dot(f)
40        self.mi=lambda f:mi.dot(f)
41        self.xi = col_vec(np.linspace(0,Nx-1,Nx))-Nx/2.
42        self.onesi = np.ones((Nx,1))
43        self.zeroi = np.zeros((Nx,1))
44       
45        right = sparse.spdiags([ones],[1],Nx,Nx, format='csr')
46        right[-1,0]=1
47        left = sparse.spdiags([ones],[0],Nx,Nx, format='csr')
48        dj = right-left
49        mj = .5*(right+left)
50        self.dj=lambda f:dj.dot(f)
51        self.mj=lambda f:mj.dot(f)
52        self.xj = col_vec(np.linspace(0.5,Nx-.5,Nx))-Nx/2.
53        self.onesj = np.ones((Nx,1))
54
55############################## Vertical mesh ############################
56
57class Lorenz:
58    def __init__(self, llm):
59        self.llm = llm
60        self.onesk = np.ones((1,llm))
61        self.onesl = np.ones((1,llm+1))
62        self.etak = row_vec(np.linspace(.5/llm,1.-.5/llm,llm))
63        self.etal = row_vec(np.linspace(0.,1.,llm+1))
64    def fieldk(self,Nx):
65        return np.zeros((Nx,self.llm))
66    def fieldl(self,Nx):
67        return np.zeros((Nx,self.llm+1))
68    def dk(self, fl):
69        return fl[:,1:]-fl[:,:-1]
70    def dl(self, fbot, fk, ftop): 
71    # vertical FD from k to l with extra info at top/bottom
72        llm=self.llm
73        gbot = fk[:,0:1]-fbot
74        gl = fk[:,1:]-fk[:,:-1]
75        gtop = ftop-fk[:,llm-1:llm]
76        return np.array(np.hstack((gbot,gl,gtop)))
77    # Vertical averaging, fac=1 or 1/2 is weight of boundary contributions
78    # mk and ml with same fac are adjoint
79    # ml_ext, mk_int use fac=1/2
80    # mk_ext, ml_int use fac=1
81    def ml(self, fk, fac): # fk has size llm, last index llm-1
82        llm=self.llm
83        fbot, ftop = fk[:,0:1], fk[:,llm-1:llm]
84        fl = .5*(fk[:,1:]+fk[:,:-1]) # size llm-1
85        return np.array(np.hstack((fac*fbot,fl,fac*ftop))) # size llm+1
86    def mk(self, fl, fac): # fl has size llm+1, last index llm
87        llm=self.llm
88        fbot = fac*fl[:,0:1] + .5*fl[:,1:2]
89        ftop = fac*fl[:,llm:llm+1] + .5*fl[:,llm-1:llm]
90        fk = .5*(fl[:,2:-1]+fl[:,1:-2]) # size llm-2
91        return np.array(np.hstack((fbot,fk,ftop)))  # size llm
92    # Averaging of extensive fields
93    def ml_ext(self, fk): return self.ml(fk,.5)
94    def mk_ext(self, fl): return self.mk(fl,1.)
95    # Averaging of intensive fields
96    def ml_int(self, fk): return self.ml(fk,1.)
97    def mk_int(self, fl): return self.mk(fl,.5)
98   
99############################## Slice #############################
100
101class Slice:
102    def __init__(self, xmesh, zmesh):
103        self.xmesh, self.zmesh = xmesh, zmesh
104        self.xik = xmesh.xi*zmesh.onesk
105        self.xjk = xmesh.xj*zmesh.onesk
106        self.xil = xmesh.xi*zmesh.onesl
107        self.xjl = xmesh.xj*zmesh.onesl
108        self.etaik = xmesh.onesi*zmesh.etak
109        self.etail = xmesh.onesi*zmesh.etal
110        self.etajk = xmesh.onesj*zmesh.etak
111        self.etajl = xmesh.onesj*zmesh.etal
112        self.ones_ik=xmesh.onesi*zmesh.onesk
113    def field_ik(self):
114        return self.zmesh.fieldk(self.xmesh.Nx)
115    def field_il(self):
116        return self.zmesh.fieldl(self.xmesh.Nx)
117    def field_jk(self):
118        return self.zmesh.fieldk(self.xmesh.Nx)
119    def field_jl(self):
120        return self.zmesh.fieldl(self.xmesh.Nx)
121    def cumsum_il(self, f0, dkf):
122        f_il = np.hstack((f0, dkf))
123        return np.cumsum(f_il,axis=1)
124
125########################## Thermodynamics #########################
126
127class Gas:
128    def dp(self, dvol_ik):
129        """ Variations of p,T,g at constant s """
130        dpik = -self.c2*dvol_ik/(self.v*self.v) # dp/dv = -(1/v2)dp/drho
131        dTik = -self.dpds*dvol_ik # dp/ds = -d2e/dvds = -dT/dv
132        dgik = self.v*dpik-self.s*dTik # dG = vdp - sdT
133        return dpik, dTik, dgik
134   
135class Ideal_perfect:
136    def __init__(self,Cpd,Rd,p0,T0):
137        # make sure everything is floating-point
138        self.Cpd = 1.*Cpd
139        self.Cvd = 1.*Cpd-Rd
140        self.Rd  = 1.*Rd
141        self.T0  = 1.*T0
142        self.p0  = 1.*p0
143        self.v0  = Rd*T0/(1.*p0)
144    def set_ps(self,p,s):
145        Cpd = self.Cpd ; Rd = self.Rd
146        T = self.T0*np.exp( (s + Rd*np.log(p/self.p0))/Cpd )
147        return self.energies(p,Rd*T/p,T,s)
148    def set_vs(self,v,s):
149        Cvd = self.Cvd ; Rd = self.Rd
150        T = self.T0*np.exp( (s - Rd*np.log(v/self.v0))/Cvd )
151        return self.energies(Rd*T/v,v,T,s)
152    def set_pT(self,p,T):
153        Cpd = self.Cpd ; Rd = self.Rd
154        s = Cpd*np.log(T/self.T0) - Rd*np.log(p/self.p0)
155        return self.energies(p,Rd*T/p,T,s)
156    def set_vT(self,v,T):
157        Cvd = self.Cvd ; Rd = self.Rd
158        s = Cvd*np.log(T/self.T0) + Rd*np.log(v/self.v0)
159        return self.energies(Rd*T/v,v,T,s)
160    def energies(self,p,v,T,s):
161        gas = Gas();
162        gas.p, gas.v, gas.T, gas.s = p, v, T, s
163        gas.dpds = p/self.Cvd
164        gas.c2=((self.Cpd/self.Cvd)*self.Rd)*T
165        gas.e = self.Cvd*T
166        gas.h = self.Cpd*T
167        gas.g = gas.h-T*s
168        return gas
169
170################################# Dynamics ##############################
171
172class BoundCond:
173    def __init__(self, ptop, Phi_bot, pbot, rho_bot):
174        self.ptop, self.Phi_bot = ptop, Phi_bot # standard BC
175        self.pbot, self.rho_bot = pbot, rho_bot # soft ground
176        self.rigid_bottom, self.rigid_top = True, True # by default
177        self.sponge = lambda Phi,ujk,dujk,tau : (ujk+tau*dujk,dujk)
178       
179class Dynamics:
180    def __init__(self, thermo, metric, BC, fac=(1,1,1)):
181        self.thermo, self.metric = thermo, metric
182        self.BC, self.fac = BC, fac
183        self.di, self.dj = metric.di, metric.dj
184        self.dk, self.dl = metric.dk, metric.dl
185        self.mi, self.mj = metric.mi, metric.mj
186        self.mk_ext, self.ml_ext = metric.mk_ext, metric.ml_ext
187        self.mk_int, self.ml_int = metric.mk_int, metric.ml_int
188    def kinetic(self,mik,ujk,Wil,Phi_il):
189        djPhil = self.dj(Phi_il)
190#        ujk, Wil, djPhil = self.mult_by_fac((ujk, Wil, djPhil))
191        return self.metric.kinetic(mik, ujk, Wil, djPhil)
192    def time_steps(self,gas,mik) :
193        dx = self.metric.dx
194        cmax, dz = np.sqrt(gas.c2.max()), gas.v*mik/dx
195        print 'Min/max layer thickness (m)', dz.min(), dz.max()
196        dz=dz.min()
197        dt_hydro=.5*dx/cmax
198        dt_NH=.5/cmax/math.sqrt(dx**-2+dz**-2)
199        print 'Horizontal resolution (m)', dx
200        print 'Max temperature (K), sound speed (m/s)', gas.T.max(), cmax
201        print 'Courant number 1 times steps for HPE and NH (s) : ', dt_hydro, dt_NH
202        return dt_NH, dt_hydro
203
204############################## Hydrostatic dynamics ###########################
205
206class Shallow_hydro: # hydrostatic routines of the Shallow class defined at the end
207    def hydrostatic_pressure(self, ptop, mik):
208        M_il=np.cumsum(mik,1)
209        Mtop_il=col_vec(M_il[:,-1])*self.zmesh.onesk
210        M_ik = Mtop_il - M_il + .5*mik
211        return ptop + M_ik*self.g0_dx
212    def geopotential(self, Phi0_i, mik, gas):
213        dkphi = mik*gas.v*self.g0_dx
214        phi_il = self.mesh.cumsum_il(Phi0_i, dkphi)
215        return phi_il
216    def hydrostatic_adjustment(self, thermo, BC, mik, sik):
217        pik=self.hydrostatic_pressure(BC.ptop, mik)
218        gas=thermo.set_ps(pik, sik)
219        Phi_il=self.geopotential(BC.Phi_bot, mik, gas)
220        return Phi_il, gas
221    def dK_hydro(self, mik, ujk):
222        """Partial derivatives of HPE kinetic energy"""
223        Ujk=self.mj(mik)*ujk/self.dx2
224        Kik=(.5/self.dx2)*self.mi(ujk*ujk)
225        return Kik, Ujk
226    def inicond(self, hybrid, thermo, ps, T, u):
227        """Hydrostatic initialization given
228        surface pressure, temperature, horiz vel""" 
229        Mi=(ps(self.xi)-self.ptop)*self.dx_g0
230        mik=hybrid.mik(Mi)
231        pik=self.hydrostatic_pressure(mik)
232        Tik=T(self.xik, pik)
233        gas=thermo.set_pT(pik, Tik)
234        Phi_il=self.geopotential(mik, gas)
235        pjk=self.mj(pik)
236        ujk=self.dx*u(self.xjk,pjk)
237        return Mi, mik, gas, Phi_il, ujk
238
239class Hydrostatic(Dynamics):
240    def __init__(self,  thermo, metric, BC):
241        Dynamics.__init__(self, thermo, metric, BC, (1,0,0))
242    def max_time_step(self,gas,mik):
243         dt_NH, dt_hydro = self.time_steps(gas,mik)
244         return dt_hydro, dt_hydro
245    def diagnose(self,mik,Sik,ujk,Phi_il,Wil):
246        Phi_il, gas_ik = self.metric.hydrostatic_adjustment(
247                                    self.thermo, self.BC, mik, Sik/mik)
248        return mik,gas_ik,ujk,Phi_il,0*Phi_il
249    def dH(self, mik, sik, ujk):
250        Phi_il, gas_ik = self.metric.hydrostatic_adjustment(
251                                    self.thermo, self.BC, mik, sik)
252        Kik, Ujk = self.metric.dK_hydro(mik, ujk)
253        Bik = Kik + self.mk_int(Phi_il) + gas_ik.g
254        return Bik, gas_ik, Ujk, Phi_il
255    def bwd_fast_slow(self, mik, Sik, ujk, tau):
256        sik = Sik/mik
257        Phi_il, gas_ik = \
258            self.metric.hydrostatic_adjustment(self.thermo, self.BC, mik, sik)
259        sjk = self.mj(sik)
260        Bik = self.mk_int(Phi_il) + gas_ik.g
261        dujk_fast = -self.dj(Bik)-sjk*self.dj(gas_ik.T)
262        ujk, dujk_fast = self.BC.sponge(Phi_il, ujk, dujk_fast, tau) # updates ujk, dujk_dt
263        Kik, Fjk = self.metric.dK_hydro(mik, ujk)
264        dSik = -self.di(sjk*Fjk)
265        dujk = -self.dj(Kik)
266#        print 'Linf(du/dt) (fast,slow) : ', Linf(dujk_fast),Linf(dujk)
267        return ujk, dujk_fast, Fjk, dSik, dujk
268
269############################## NH dynamics ###########################
270
271class Shallow_NH: # NH routines of the Shallow class defined at the end
272    def wil(self, mik, Wil):
273        mil=self.ml_ext(mik)
274        return mil, Wil/mil
275    def lambda_il(self,djPhil):
276        return self.mi(djPhil*djPhil)/self.dx2
277    def kinetic(self, mik, ujk, Wil, djPhil):
278        """NH kinetic energy - coincides with HPE if Wil=0"""
279        dx2,mi,mj,ml_int=self.dx2, self.mi, self.mj, self.ml_int
280        mil,wil = self.wil(mik,Wil)
281        lambda_il = (self.lambda_il(djPhil) + self.g2)*wil
282        Kjk = (.5/dx2)*mj(mik)*ujk**2
283        Kjl = (1./dx2)*djPhil*ml_int(ujk)*mj(Wil)
284        Kil = .5*lambda_il*Wil
285        return Kjk.sum()-Kjl.sum()+Kil.sum()
286    def dK_NH_horiz(self, mik, ujk, Wil, djPhil):
287        """Partial derivatives of horizontal kinetic energy"""
288        dx2, mi, mj = self.dx2, self.mi, self.mj
289        mk_ext, mk_int, ml_int = self.mk_int, self.mk_int, self.ml_int
290        mil, wil = self.wil(mik,Wil)
291        lambda_il = wil*self.lambda_il(djPhil)
292        djPhi_W = mk_ext(djPhil*mj(Wil))
293        dK_dujk = (mj(mik)*ujk-djPhi_W)/dx2
294        dK_dmik = (.5/dx2)*mi(ujk*ujk) - .5*mk_int(lambda_il*wil)
295        ujl = ml_int(ujk)
296        dK_dWil = lambda_il - mi(djPhil*ujl)/dx2
297        dK_djPhil = (djPhil*mj(Wil*wil)-mj(Wil)*ujl)/dx2
298        return dK_dmik, dK_dujk, dK_dWil, dK_djPhil
299    def dK_NH_vert(self, mik, Wil):
300        """Partial derivatives of vertical kinetic energy, 
301        also returns mil,wil to be used by HEVI solver"""
302        mil, wil = self.wil(mik,Wil)
303        dK_dWil = self.g2*wil
304        dK_dmik = -.5*self.mk_int(dK_dWil*wil)
305        return dK_dmik, dK_dWil, mil, wil
306    def inicond_NH(self, thermo, Phi, p, T):
307        """ Initialization given Phi(x,eta), p(x,Phi), T(x,p) """
308        Phi_il = Phi(self.xil, self.mesh.etail)
309        Phi_ik = Phi(self.xik, self.mesh.etaik)
310        pik = p(self.xik, Phi_ik)
311        gas = thermo.set_pT(pik, T(self.xik,pik))
312        vol, Jac = self.volume(Phi_il)
313        mik = self.dk(vol) / gas.v
314        return mik, gas, Phi_il
315    def solver_NH(self, BC, dt, mik,Jac_il,gas_ik,mil):
316        dx_g0, g2 = self.dx_g0, self.g2
317        Ak = (dt*dx_g0/gas_ik.v)**2
318        Ak = gas_ik.c2/mik*Ak
319        ml_g2 = mil/g2
320        Bl = 2*self.ml_ext(Ak) + ml_g2
321        Bl[:,0:1] = Bl[:,0:1] + (dx_g0*dt*dt)*BC.rho_bot
322        return time_step.SymTriDiag(Bl,Ak), ml_g2
323
324class Nonhydro(Dynamics):
325    def max_time_step(self,gas,mik):
326         dt_NH, dt_hydro = self.time_steps(gas,mik)
327         return dt_hydro, dt_hydro
328    def diagnose(self,mik,Sik,ujk,Phi_il,Wil):
329        vol_il,dvol_il = self.metric.volume(Phi_il)
330        gas_ik = self.thermo.set_vs(self.dk(vol_il)/mik, Sik/mik)
331        return mik,gas_ik,ujk,Phi_il,Wil
332    def pressure(self, mik,sik,Phi_il):
333        BC, metric = self.BC, self.metric
334        vol_il, Jac_il = metric.volume(Phi_il)
335        gas_ik = self.thermo.set_vs(self.dk(vol_il)/mik, sik)
336        pbot = BC.pbot + BC.rho_bot*(BC.Phi_bot-Phi_il[:,0:1])
337        ptop = BC.ptop*metric.mesh.ones_ik[:,0:1]
338        return gas_ik, Jac_il, self.dl(pbot, gas_ik.p, ptop)
339    def bwd_fast_slow(self, mik,Sik,ujk,Phi_il,Wil, tau, verbose=False):
340        di, dj, dk, dl = self.di, self.dj, self.dk, self.dl
341        BC, metric = self.BC, self.metric
342        sik = Sik/mik
343        mil, wil = metric.wil(mik,Wil)
344        # solve Phi_il - (tau**2)*X.dl(pk) = Phi_star_il
345        if tau>0.:
346            Phi_star_il = Phi_il + tau*metric.g2*(wil - tau)
347            # Newton iteration : Phi_il contains current guess value
348            for ii in range(2):
349                gas_ik, Jac_il, dl_pi = self.pressure(mik,sik,Phi_il)
350                tridiag, ml_g2 = metric.solver_NH(BC, tau, mik,Jac_il,gas_ik,mil)
351                residual = (ml_g2*(Phi_il-Phi_star_il)
352                            +tau*tau*Jac_il*dl_pi)
353                dPhi_il = tridiag.solve(residual)
354                if verbose :
355                    print 'NH solver : Linf(dPhi_il)=', abs(dPhi_il).max(axis=0)
356                    print 'NH solver : Linf(residual)=', abs(residual).max(axis=0)
357                Phi_il = Phi_il - dPhi_il
358        # update Wil
359        gas_ik, Jac_il, dl_pi = self.pressure(mik,sik,Phi_il)
360        dH_Phiil = mil + Jac_il*dl_pi
361        Wil = Wil - tau*dH_Phiil
362        # update ujk
363        dK_mik, dPhiil_fast, mil, wil = metric.dK_NH_vert(mik, Wil)
364        dH_mik = dK_mik + self.mk_int(Phi_il) + gas_ik.g
365        sjk = self.mj(sik)
366        dujk_fast = -dj(dH_mik)-sjk*dj(gas_ik.T)
367        ujk, dujk_fast = self.BC.sponge(Phi_il, ujk, dujk_fast, tau)
368        # slow trends
369        djPhil = dj(Phi_il)
370        Bik, Fjk, dPhiil_slow, dK_djPhil = self.metric.dK_NH_horiz(mik, ujk, Wil, djPhil)
371        dSik_slow = -di(sjk*Fjk)
372        dujk_slow = -dj(Bik)
373        dWil_slow = di(dK_djPhil)
374        return  ( ujk,Phi_il,Wil,                               # updated flow state
375                dujk_fast,dPhiil_fast,-dH_Phiil,                # fast tendencies
376                Fjk,dSik_slow,dujk_slow,dPhiil_slow,dWil_slow ) # slow tendencies
377
378################## Fully equipped shallow-atmosphere metric  ##################
379
380class Shallow(Shallow_hydro, Shallow_NH):
381    def __init__(self, mesh, dx, g0):
382        self.dx, self.g0  = dx, g0
383        self.dx2, self.g2 = dx*dx, g0*g0
384        self.dx_g0, self.g0_dx = dx/g0, g0/dx
385        self.mesh, xmesh, zmesh = mesh, mesh.xmesh, mesh.zmesh
386        self.xmesh, self.zmesh = xmesh, zmesh
387        self.di, self.dj = xmesh.di, xmesh.dj
388        self.mi, self.mj = xmesh.mi, xmesh.mj
389        self.dk, self.dl = zmesh.dk, zmesh.dl
390        self.mk_int, self.ml_int = zmesh.mk_int, zmesh.ml_int
391        self.mk_ext, self.ml_ext = zmesh.mk_ext, zmesh.ml_ext
392        self.xik = mesh.xik*dx
393        self.xjk = mesh.xjk*dx
394        self.xjl = mesh.xjl*dx
395        self.xil = mesh.xil*dx
396        self.xi  = xmesh.xi*dx
397    def volume(self, phi):  # V(Phi), J=dV/dPhi
398        return phi*self.dx_g0, self.dx_g0
399
400######################## Lagrangian vertical coordinate #####################
401
402class Coord:
403    def __init__(self, dyn):
404        self.dyn, self.metric = dyn, dyn.metric
405       
406class LagrangianHydro(Coord):
407    def bwd_fast_slow(self, flow, tau):
408        dyn=self.dyn
409        mik, Sik, ujk, Phiil, Wil, Mjk, Mil =  flow # Phiil, Wil ignored
410        ujk, dujk_fast, Fjk, dSik, dujk = dyn.bwd_fast_slow(mik, Sik, ujk, tau)
411        flow = (mik, Sik, ujk, 0., 0., Mjk, Mil)
412        fast = (0., 0., dujk_fast, 0., 0., 0., 0.)
413        dmik = -dyn.di(Fjk)
414        slow = (dmik, dSik, dujk, 0., 0., Fjk, 0.)
415        return flow, fast, slow
416
417class LagrangianNH(Coord):
418    def bwd_fast_slow(self, flow, tau):
419        dyn=self.dyn
420        mik, Sik, ujk, Phiil, Wil, Mjk, Mil =  flow
421        ujk, Phiil, Wil, \
422            dujk_fast, dPhiil_fast, dWil_fast, \
423            Fjk, dSik, dujk, dPhiil, dWil = dyn.bwd_fast_slow(mik, Sik, ujk, Phiil, Wil, tau)
424        flow = (mik, Sik, ujk, Phiil, Wil, Mjk, Mil)
425        fast = (0., 0., dujk_fast, dPhiil_fast, dWil_fast, 0., 0.)
426        dmik = -dyn.di(Fjk)
427        slow = (dmik, dSik, dujk, dPhiil, dWil, Fjk, 0.)
428        return flow, fast, slow
429
430######################## Mass-based vertical coordinate #####################
431
432class GeneralCoord(Coord):
433    # this class computes the terms due to vertical transport (eta_dot)
434    # given the vertical mass flux Fil
435    def trend_vert_hydro(self, mik,Sik,ujk, Fil):
436        dyn=self.dyn
437        sik=Sik/mik
438        sil=dyn.ml_int(sik) # centered
439        dSik_dt = - dyn.dk(Fil*sil)
440        dluj = dyn.dl(0.,ujk,0.)
441        dujk_dt = - dyn.mk_ext(dluj*dyn.mj(Fil))/dyn.mj(mik)
442        return dSik_dt, dujk_dt
443    def trend_vert_NH(self, mik,Sik,ujk,Phi_il,Wil, Fil):
444        dyn=self.dyn
445        dSik, dujk = self.trend_vert_hydro(mik,Sik,ujk, Fil)
446        Fik = dyn.mk_int(Fil)
447        wik = dyn.mk_ext(Wil)/mik
448        fil = Fil/dyn.ml_ext(mik) # fil = eta_dot
449        dujk = dujk - dyn.mj(wik*dyn.dk(Phi_il))*dyn.dj(Fik/mik)
450        dWil = - dyn.dl(0.,Fik*wik,0.) # centered vertical flux of vertical momentum
451        dPhiil = - fil*dyn.dl(0., dyn.mk_int(Phi_il), 0.)
452        return dSik, dujk, dPhiil, dWil
453
454class MassCoord(GeneralCoord):
455    # computes the vertical mass flux for a mass-based coordinate
456    def massflux(self,mik,Fjk):
457        def sum(fik): return col_vec(fik.sum(1))*onesk
458        mesh=self.metric.mesh
459        onesk=mesh.zmesh.onesk
460        # mik = Aik*Mi
461        Mik=sum(mik)
462        Aik = mik/Mik
463        # dmik_mass = Aik*dMi = Aik*sumk(dmik_lag)
464        dmik_lag = -self.dyn.di(Fjk)
465        dMik = sum(dmik_lag)
466        dmik_mass = Aik*dMik
467        #    dmik_lag  + di(Fjk)            = 0
468        #    dmik_mass + di(Fjk) + dk(Fil)  = 0
469        # => dmik_mass - dmik_lag + dk(Fil) = 0
470        dkFil = dmik_lag - dmik_mass
471        Fil = mesh.cumsum_il(mesh.xmesh.zeroi, dkFil)
472        return Fil, dmik_mass
473
474class MassBasedHydro(MassCoord):
475    def bwd_fast_slow(self, flow, tau):
476        mik, Sik, ujk, Phiil, Wil, Mjk, Mil =  flow # Phiil, Wil ignored
477        ujk, dujk_fast, Fjk, dSik, dujk = self.dyn.bwd_fast_slow(mik, Sik, ujk, tau)
478        flow = (mik, Sik, ujk, 0., 0., Mjk, Mil)
479        fast = (0., 0., dujk_fast, 0., 0., 0., 0.)
480        Fil, dmik = self.massflux(mik,Fjk)
481        dSik_vert, dujk_vert = self.trend_vert_hydro(mik,Sik,ujk, Fil)
482        slow = (dmik, dSik+dSik_vert, dujk+dujk_vert, 0., 0., Fjk, Fil)
483        return flow, fast, slow
484
485class MassBasedNH(MassCoord):
486    def bwd_fast_slow(self, flow, tau, verbose=False):
487        dyn=self.dyn
488        mik, Sik, ujk, Phiil, Wil, Mjk, Mil =  flow
489        ujk,Phiil,Wil, \
490            dujk_fast,dPhiil_fast,dWil_fast, \
491            Fjk, dSik, dujk, dPhiil, dWil = dyn.bwd_fast_slow(mik,Sik,ujk,Phiil,Wil, tau, verbose)
492        flow = (mik, Sik, ujk, Phiil, Wil, Mjk, Mil)
493        fast = (0.,0., dujk_fast,dPhiil_fast,dWil_fast, 0.,0.)
494        Fil, dmik = self.massflux(mik,Fjk)
495        dSik_vert, dujk_vert, dPhiil_vert, dWil_vert = self.trend_vert_NH(mik,Sik,ujk,Phiil,Wil,Fil)
496        slow = (dmik, dSik+dSik_vert, dujk+dujk_vert, dPhiil+dPhiil_vert, dWil+dWil_vert, Fjk, Fil)
497        return flow, fast, slow
Note: See TracBrowser for help on using the repository browser.