source: codes/icosagcm/devel/Python/test/py/Baroclinic_3D_ullrich.py

Last change on this file was 943, checked in by dubos, 5 years ago

devel/Python : fix baroclinic test case
File size: 11.9 KB
Line 
1from __future__ import print_function
2from __future__ import division
3
4from dynamico import getargs
5
6getargs.add("--T", type=float, default=3600.,
7                    help="Length of time slice in seconds")
8getargs.add("--perturb", type=float, default=1.,
9                    help="Amplitude of wind perturbation in m/s")
10getargs.add("--N", type=int, default=48,
11                    help="Number of time slices")
12getargs.add("--Davies_N1", type=int, default=3)
13getargs.add("--Davies_N2", type=int, default=3)
14getargs.add("--nx", type=int, default=200)
15getargs.add("--ny", type=int, default=30)
16getargs.add("--betaplane", action='store_true')
17
18getargs.add("--kappa_divgrad",  type=float, default=3.0e15)
19getargs.add("--kappa_curlcurl", type=float, default=3.0e15)
20getargs.add("--ztop", type=float, default=30000.)
21
22getargs.defaults(dt=360., llm=50)
23
24#getargs.config_log(filename='out.log',filemode='w') # must be done before calling Logger()
25#    getargs.config_log(filename='out.log',filemode='w',
26#                    format='%(processName)s:%(name)s:%(filename)s:%(module)s:%(funcName)s:%(lineno)d:%(levelname)s:%(message)s' )
27
28args = getargs.parse()
29log_master, log_world = getargs.getLogger()
30INFO, DEBUG, DEBUG_ALL, ERROR = log_master.info, log_master.debug, log_world.debug, log_world.error
31
32from dynamico.dev import unstructured as unst
33from dynamico.dev import meshes
34from dynamico.dev import xios
35from dynamico.dev import precision as prec
36from dynamico import dyn
37from dynamico import time_step
38from dynamico import DCMIP
39from dynamico import maps
40from dynamico.kernels import grad, curl, div, KE
41from dynamico.LAM import Davies
42
43import math as math
44import numpy as np
45import time
46from numpy import pi, log, exp, sin, cos
47
48# Baroclinic instability test based on Ullrich et al. 2015, QJRMS
49
50def create_pmesh(nx,ny):
51    filename = 'cart_%03d_%03d.nc'%(nx,ny)
52    INFO('Reading Cartesian mesh ...')
53    meshfile = meshes.DYNAMICO_Format(filename)
54    pmesh = meshes.Unstructured_PMesh(comm,meshfile)
55    pmesh.partition_curvilinear(args.mpi_ni,args.mpi_nj)
56    return pmesh
57
58def baroclinic_3D(pmesh, dx,Lx,Ly,llm,ztop):
59    Rd    = 287.0      # Gas constant for dryy air (j kg^-1 K^-1)
60    T0    = 288.0      # Reference temperature (K)
61    lap   = 0.005      # Lapse rate (K m^-1)
62    b     = 2.         # Non dimensional vertical width parameter
63    u0    = 35.        # Reference zonal wind speed (m s^-1)
64    a     = 6.371229e6 # Radius of the Earth (m)
65    y0    = .5*Ly
66    Cpd   = 1004.5
67    p0    = 1e5
68    up    = args.perturb # amplitude (m/s)
69    xc,yc,lp = 0.,Ly/2.,600000.
70
71    omega  = 7.292e-5  # Angular velocity of the Earth (s^-1)   
72    phi0   = 45.*pi/180.0 # Reference latitude North pi/4 (deg)
73    f0     = 2*omega*sin(phi0) 
74    beta0  = 2*omega*cos(phi0)/a if args.betaplane else 0.
75    fb     = f0 - beta0*y0
76
77    def Phi_xy(y):
78        fc = y*y - (Ly*y/pi)*sin(2*pi*y/Ly)
79        fd = Ly*Ly/(2*pi*pi)*cos(2*pi*y/Ly)
80        return .5*u0*( fb*(y-y0-Ly/(2*pi)*sin(2*pi*y/Ly)) + .5*beta0*(fc-fd-(Ly*Ly/3.)- Ly*Ly/(2*pi*pi)) )
81
82    def Phi_xyeta(y,eta): return T0*g/lap*(1-eta**(Rd*lap/g)) + Phi_xy(y)*log(eta)*exp(-((log(eta)/b)**2))
83    def ulon(x,y,eta):
84        u = -u0*(sin(pi*y/Ly)**2)*log(eta)*(eta**(-log(eta)/(b*b))) 
85        # periodize the envelope of the initial perturbation
86        A = exp(-(((x-xc)**2+(y-yc)**2)/lp**2))
87        A = A + exp(-(((x-xc-Lx)**2+(y-yc)**2)/lp**2))
88        A = A + exp(-(((x-xc+Lx)**2+(y-yc)**2)/lp**2))
89        return  u + up*A
90
91    def tmean(eta) : return T0*eta**(Rd*lap/g)
92    def T(y,eta) : return tmean(eta)+(Phi_xy(y)/Rd)*(((2/(b*b))*(log(eta))**2)-1)*exp(-((0.5*log(eta))**2)) 
93    def p(eta): return p0*eta     # eta  = p/p_s
94
95    def eta(alpha) : return (1.-(lap*ztop*alpha/T0))**(g/(Rd*lap)) # roughly equispaced levels
96   
97    betaplane = maps.BetaPlaneMap(dx, .5*Lx, .5*Ly, f0, beta0, .5*Ly)
98    nqdyn = 1
99    mesh = meshes.Local_Mesh(pmesh, llm, nqdyn, betaplane)
100
101    alpha_k = (np.arange(llm)  +.5)/llm
102    alpha_l = (np.arange(llm+1)+ 0.)/llm
103    x_ik, alpha_ik = np.meshgrid(mesh.lon_i, alpha_k, indexing='ij')
104    y_ik, alpha_ik = np.meshgrid(mesh.lat_i, alpha_k, indexing='ij')
105    x_il, alpha_il = np.meshgrid(mesh.lon_i, alpha_l, indexing='ij')
106    y_il, alpha_il = np.meshgrid(mesh.lat_i, alpha_l, indexing='ij')
107    x_ek, alpha_ek = np.meshgrid(mesh.lon_e, alpha_k, indexing='ij')
108    y_ek, alpha_ek = np.meshgrid(mesh.lat_e, alpha_k, indexing='ij')
109
110#    print(np.shape(alpha_k),np.shape(alpha_l))
111    thermo = dyn.Ideal_perfect(Cpd, Rd, p0, T0)
112
113    eta_il = eta(alpha_il)
114    eta_ik = eta(alpha_ik)
115    eta_ek = eta(alpha_ek)
116#    print('min max eta_il', np.min(eta_il),np.max(eta_il))
117
118    Phi_il = Phi_xyeta(y_il, eta_il)
119    Phi_ik = Phi_xyeta(y_ik, eta_ik)
120    p_ik, p_il = p(eta_ik), p(eta_il)
121    T_ik = T(y_ik, eta_ik)  #ik full level(40), il(41)
122
123    gas = thermo.set_pT(p_ik,T_ik)
124    mass_ik = mesh.field_mass()
125    for l in range(llm): 
126        mass_ik[:,l]=(Phi_il[:,l+1]-Phi_il[:,l])/(g*gas.v[:,l])
127#        mass_ik[:,l]=(p_il[:,l]-p_il[:,l+1])/g
128    Sik, Wil  = gas.s*mass_ik, mesh.field_w()
129   
130    u_ek = mesh.ucov3D(ulon(x_ek, y_ek, eta_ek), 0.*eta_ek)
131
132    INFO('ztop (m) = %f %f' % (Phi_il[0,-1]/g, ztop) )
133    ptop = p(eta(1.))
134    INFO( 'ptop (Pa) = %f %f' % (gas.p[0,-1], ptop) )
135    INFO('ztop(ptop) according to Eq. 7: %f' % (T0/lap*(1.-(ptop/p0)**(Rd*lap/g))) ) 
136
137    params=dyn.Struct()
138    params.ptop=ptop
139    params.dx=dx
140    params.dx_g0=dx/g
141    params.g = g
142
143    # define parameters for lower BC
144    pbot = p(eta_il[:,0])
145    INFO( 'min p, T : %f %s' % (pbot.min(), tmean(pbot/p0)) )
146    gas_bot = thermo.set_pT(pbot, tmean(pbot/p0))
147    params.pbot = gas_bot.p
148    params.rho_bot = 1e6/gas_bot.v
149   
150    return thermo, mesh, params, prec.asnum([mass_ik,Sik,u_ek,Phi_il,Wil]), gas
151
152def diagnose(Phi,S,m,W,u):
153    s=S/m
154    curl_vk = curl(mesh, u)
155    abs_vort_vk = mesh.field_z() # absolute vorticity
156    un = mesh.field_u()
157    v = mesh.field_mass() # specific volume
158    w = mesh.field_mass()
159    z = mesh.field_mass()
160    for l in range(llm):
161        v[:,l]=(Phi[:,l+1]-Phi[:,l])/(g*m[:,l])
162        w[:,l]=.5*g*(W[:,l+1]+W[:,l])/m[:,l]
163        z[:,l]=.5*(Phi[:,l+1]+Phi[:,l])/g
164        un[:,l]=u[:,l]/mesh.de
165        abs_vort_vk[:,l]=curl_vk[:,l] + mesh.fv
166    gas = thermo.set_vs(v,s)
167    ps = gas.p[:,0]+ .5*g*m[:,0]
168    return gas, w, z, ps, un, curl_vk, abs_vort_vk
169
170class myDavies(Davies):
171    def mask(self,x,y):
172        DEBUG('davies dy = %f'% dx)
173        return self.mask0(y,Ly,dx)*self.mask0(-y,0.,dx)
174       
175with xios.Client() as client: # setup XIOS which creates the DYNAMICO communicator
176    comm = client.comm
177    mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size()
178
179    INFO('%d/%d starting'%(mpi_rank,mpi_size))
180    if mpi_rank>0 : 
181        getargs.not_master()
182        unst.setvar('is_mpi_master', False)
183           
184    g, Lx, Ly = 9.80616, 4e7, 6e6
185    nx, ny, llm = args.nx, args.ny, args.llm
186    dx = Lx/nx
187   
188    unst.setvar('g',g)
189    unst.setvar('debug_hevi_solver', False)
190
191    pmesh = create_pmesh(nx,ny)
192    thermo, mesh, params, flow0, gas0 = baroclinic_3D(pmesh, dx,Lx,Ly,llm, args.ztop)
193
194    mass_bl,mass_dak,mass_dbk = meshes.compute_hybrid_coefs(flow0[0])
195    unst.ker.dynamico_init_hybrid(mass_bl,mass_dak,mass_dbk)
196
197    T  = args.T
198    dt = args.dt
199    dz = flow0[3].max()/(params.g*llm)
200    nt = int(math.ceil(T/dt))
201    dt = T/nt
202    INFO( 'Time step : %d x %g = %g s' % (nt,dt,nt*dt))
203
204   
205    caldyn_thermo, caldyn_eta = unst.thermo_entropy, unst.eta_mass
206
207    if False: # time stepping in Python
208        caldyn = unst.Caldyn_NH(caldyn_thermo,caldyn_eta, mesh,thermo,params,params.g)
209        scheme = time_step.ARK2(caldyn.bwd_fast_slow, dt)
210        def next_flow(m,S,u,Phi,W):
211            return scheme.advance((m,S,u,Phi,W),nt)
212
213    else: # time stepping in Fortran
214        scheme = time_step.ARK2(None, dt)
215        if args.hydrostatic:
216            caldyn_step = unst.caldyn_step_HPE(mesh,scheme,1, caldyn_thermo,caldyn_eta,
217                                               thermo,params,params.g)
218        else:
219            caldyn_step = unst.caldyn_step_NH(mesh,scheme,1, caldyn_thermo,caldyn_eta,
220                                              thermo,params,params.g)
221        davies = myDavies(args.Davies_N1, args.Davies_N2, 
222                          mesh.lon_i, mesh.lat_i, mesh.lon_e,mesh.lat_e)
223#        print('mask min/max :', davies.beta_i.min(), davies.beta_i.max() )
224        DEBUG('mask min/max :')
225        DEBUG('%f'% davies.beta_i.min())
226        DEBUG('%f'% davies.beta_i.max())
227
228        def next_flow(m,S,u,Phi,W):
229            # junk,fast,slow = caldyn.bwd_fast_slow(flow, 0.)
230            caldyn_step.mass[:,:], caldyn_step.theta_rhodz[:,:], caldyn_step.u[:,:] = m,S,u
231            caldyn_step.geopot[:,:], caldyn_step.W[:,:] = Phi,W
232            for i in range(nt):
233                caldyn_step.next()
234                davies.relax(llm, caldyn_step, flow0)
235                m,S,u = caldyn_step.mass, caldyn_step.theta_rhodz, caldyn_step.u
236                s = S/m
237                laps, bilaps = mesh.field_mass(), mesh.field_mass()
238                lapu, bilapu = mesh.field_u(), mesh.field_u()
239                unst.ker.dynamico_scalar_laplacian(s,laps)
240                unst.ker.dynamico_scalar_laplacian(laps,bilaps)
241                unst.ker.dynamico_curl_laplacian(u,lapu)
242                unst.ker.dynamico_curl_laplacian(lapu,bilapu)
243                caldyn_step.theta_rhodz[:] = S - dt*args.kappa_divgrad*bilaps*m # Euler step
244                caldyn_step.u[:] = u - dt*args.kappa_curlcurl*bilapu # Euler step
245
246            return (caldyn_step.mass.copy(), caldyn_step.theta_rhodz.copy(), caldyn_step.u.copy(),
247                    caldyn_step.geopot.copy(), caldyn_step.W.copy())
248
249    m,S,u,Phi,W = flow0
250    if caldyn_thermo == unst.thermo_theta:
251        s=S/m
252        theta = thermo.T0*exp(s/thermo.Cpd)
253        S=m*theta
254        title_format = 'Potential temperature at t=%g s (K)'
255    else:
256        title_format = 'Specific entropy at t=%g s (J/K/kg)'
257
258    #T, nslice, dt = 3600., 1, 3600.
259    Nslice = args.N
260
261    temp_v = mesh.field_z(),     
262
263    with xios.Context_Curvilinear(mesh,1, 24*3600) as context:
264        # now XIOS knows about the mesh and we can write to disk
265        for i in range(Nslice+1):
266            context.update_calendar(i+1)
267
268            # Diagnose quantities of interest from prognostic variables m,S,u,Phi,W
269            gas, w, z, ps, un, zeta_vk, zeta_abs_vk = diagnose(Phi,S,m,W,u)
270            KE_ik = KE(mesh,u)
271            du_fast, du_slow = caldyn_step.du_fast[0,:,:], caldyn_step.du_slow[0,:,:] 
272            div_fast, div_slow = div(mesh,du_fast), div(mesh,du_slow)
273            drhodz, hflux = caldyn_step.drhodz[0,:,:], caldyn_step.hflux[:,:]
274
275            # write to disk
276            context.send_field_primal('ps', ps)
277            context.send_field_primal('temp', gas.T)
278
279            context.send_field_primal('p', gas.p)
280#            context.send_field_primal('p', KE_ik)
281#            context.send_field_primal('p', drhodz)
282
283            context.send_field_primal('theta', thermo.T0*exp(gas.s/thermo.Cpd))
284
285            context.send_field_primal('uz', w)
286            context.send_field_primal('Phi', Phi)
287            context.send_field_primal('z', z)
288            context.send_field_primal('div_fast', div_fast)
289
290            context.send_field_primal('div_slow', div_slow)
291
292            context.send_field_dual('curl', zeta_vk)
293            context.send_field_dual('curl_abs', zeta_abs_vk)
294
295            INFO( 'ptop, model top (m) : %f %f' % (unst.getvar('ptop'), Phi.max()/unst.getvar('g')) )
296
297            time1, elapsed1 =time.time(), unst.getvar('elapsed')
298            m,S,u,Phi,W = next_flow(m,S,u,Phi,W)
299            time2, elapsed2 =time.time(), unst.getvar('elapsed')
300            factor = 1000./nt
301            INFO( 'ms per full time step : %f %f' %(factor*(time2-time1), factor*(elapsed2-elapsed1)) )
302            factor = 1e9/(4*nt*nx*ny*llm)
303            INFO( 'nanosec per gridpoint per full time step : %f %f' %
304                  (factor*(time2-time1), factor*(elapsed2-elapsed1)) )
305
306INFO('************DONE************')
Note: See TracBrowser for help on using the repository browser.